hooks: add response hook

Fixes #216

bittorrent/bittorrent.go

@@ -84,8 +84,8 @@ type AnnounceRequest struct {
 // response.
 type AnnounceResponse struct {
 	Compact     bool
-	Complete    int32
-	Incomplete  int32
+	Complete    uint32
+	Incomplete  uint32
 	Interval    time.Duration
 	MinInterval time.Duration
 	IPv4Peers   []Peer

frontend/udp/writer.go

@@ -37,8 +37,8 @@ func WriteAnnounce(w io.Writer, txID []byte, resp *bittorrent.AnnounceResponse,
 		writeHeader(buf, txID, announceActionID)
 	}
 	binary.Write(buf, binary.BigEndian, uint32(resp.Interval/time.Second))
-	binary.Write(buf, binary.BigEndian, uint32(resp.Incomplete))
-	binary.Write(buf, binary.BigEndian, uint32(resp.Complete))
+	binary.Write(buf, binary.BigEndian, resp.Incomplete)
+	binary.Write(buf, binary.BigEndian, resp.Complete)
 
 	peers := resp.IPv4Peers
 	if v6 {

middleware/hooks.go

@@ -2,6 +2,8 @@ package middleware
 
 import (
 	"context"
+	"errors"
+	"net"
 
 	"github.com/chihaya/chihaya/bittorrent"
 	"github.com/chihaya/chihaya/storage"
@@ -64,3 +66,62 @@ func (h *swarmInteractionHook) HandleScrape(ctx context.Context, _ *bittorrent.S
 	// Scrapes have no effect on the swarm.
 	return ctx, nil
 }
+
+// ErrInvalidIP indicates an invalid IP for an Announce.
+var ErrInvalidIP = errors.New("invalid IP")
+
+type skipResponseHook struct{}
+
+// SkipResponseHookKey is a key for the context of an Announce or Scrape to
+// control whether the response middleware should run.
+// Any non-nil value set for this key will cause the response middleware to
+// skip.
+var SkipResponseHookKey = skipResponseHook{}
+
+type scrapeAddressType struct{}
+
+// ScrapeIsIPv6Key is the key under which to store whether or not the
+// address used to request a scrape was an IPv6 address.
+// The value is expected to be of type bool.
+// A missing value or a value that is not a bool for this key is equivalent to
+// it being set to false.
+var ScrapeIsIPv6Key = scrapeAddressType{}
+
+type responseHook struct {
+	store storage.PeerStore
+}
+
+func (h *responseHook) HandleAnnounce(ctx context.Context, req *bittorrent.AnnounceRequest, resp *bittorrent.AnnounceResponse) (_ context.Context, err error) {
+	if ctx.Value(SkipResponseHookKey) != nil {
+		return ctx, nil
+	}
+
+	s := h.store.ScrapeSwarm(req.InfoHash, len(req.IP) == net.IPv6len)
+	resp.Incomplete = s.Incomplete
+	resp.Complete = s.Complete
+
+	switch len(req.IP) {
+	case net.IPv4len:
+		resp.IPv4Peers, err = h.store.AnnouncePeers(req.InfoHash, req.Left == 0, int(req.NumWant), req.Peer)
+	case net.IPv6len:
+		resp.IPv6Peers, err = h.store.AnnouncePeers(req.InfoHash, req.Left == 0, int(req.NumWant), req.Peer)
+	default:
+		return ctx, ErrInvalidIP
+	}
+
+	return ctx, err
+}
+
+func (h *responseHook) HandleScrape(ctx context.Context, req *bittorrent.ScrapeRequest, resp *bittorrent.ScrapeResponse) (context.Context, error) {
+	if ctx.Value(SkipResponseHookKey) != nil {
+		return ctx, nil
+	}
+
+	v6, _ := ctx.Value(ScrapeIsIPv6Key).(bool)
+
+	for _, infoHash := range req.InfoHashes {
+		resp.Files[infoHash] = h.store.ScrapeSwarm(infoHash, v6)
+	}
+
+	return ctx, nil
+}

middleware/middleware.go

@@ -25,7 +25,7 @@ func NewLogic(cfg Config, peerStore storage.PeerStore, preHooks, postHooks []Hoo
 	l := &Logic{
 		announceInterval: cfg.AnnounceInterval,
 		peerStore:        peerStore,
-		preHooks:         preHooks,
+		preHooks:         append(preHooks, &responseHook{store: peerStore}),
 		postHooks:        append(postHooks, &swarmInteractionHook{store: peerStore}),
 	}
 
@@ -44,7 +44,9 @@ type Logic struct {
 // HandleAnnounce generates a response for an Announce.
 func (l *Logic) HandleAnnounce(ctx context.Context, req *bittorrent.AnnounceRequest) (resp *bittorrent.AnnounceResponse, err error) {
 	resp = &bittorrent.AnnounceResponse{
-		Interval: l.announceInterval,
+		Interval:    l.announceInterval,
+		MinInterval: l.announceInterval,
+		Compact:     req.Compact,
 	}
 	for _, h := range l.preHooks {
 		if ctx, err = h.HandleAnnounce(ctx, req, resp); err != nil {