middleware: add sanitization hook

middleware/hooks.go

@@ -70,6 +70,48 @@ func (h *swarmInteractionHook) HandleScrape(ctx context.Context, _ *bittorrent.S
 // ErrInvalidIP indicates an invalid IP for an Announce.
 var ErrInvalidIP = errors.New("invalid IP")
 
+// sanitizationHook enforces semantic assumptions about requests that may have
+// not been accounted for in a tracker frontend.
+//
+// The SanitizationHook performs the following checks:
+// - maxNumWant: Checks whether the numWant parameter of an announce is below
+//     a limit. Sets it to the limit if the value is higher.
+// - defaultNumWant: Checks whether the numWant parameter of an announce is
+//     zero. Sets it to the default if it is.
+// - IP sanitization: Checks whether the announcing Peer's IP address is either
+//     IPv4 or IPv6. Returns ErrInvalidIP if the address is neither IPv4 nor
+//     IPv6. Sets the Peer.AddressFamily field accordingly. Truncates IPv4
+//     addresses to have a length of 4 bytes.
+type sanitizationHook struct {
+	maxNumWant     uint32
+	defaultNumWant uint32
+}
+
+func (h *sanitizationHook) HandleAnnounce(ctx context.Context, req *bittorrent.AnnounceRequest, resp *bittorrent.AnnounceResponse) (context.Context, error) {
+	if req.NumWant > h.maxNumWant {
+		req.NumWant = h.maxNumWant
+	}
+
+	if req.NumWant == 0 {
+		req.NumWant = h.defaultNumWant
+	}
+
+	if ip := req.Peer.IP.To4(); ip != nil {
+		req.Peer.IP.IP = ip
+		req.Peer.IP.AddressFamily = bittorrent.IPv4
+	} else if len(req.Peer.IP.IP) == net.IPv6len { // implies req.Peer.IP.To4() == nil
+		req.Peer.IP.AddressFamily = bittorrent.IPv6
+	} else {
+		return ctx, ErrInvalidIP
+	}
+
+	return ctx, nil
+}
+
+func (h *sanitizationHook) HandleScrape(ctx context.Context, req *bittorrent.ScrapeRequest, resp *bittorrent.ScrapeResponse) (context.Context, error) {
+	return ctx, nil
+}
+
 type skipResponseHook struct{}
 
 // SkipResponseHookKey is a key for the context of an Announce or Scrape to
@@ -97,7 +139,7 @@ func (h *responseHook) HandleAnnounce(ctx context.Context, req *bittorrent.Annou
 	}
 
 	// Add the Scrape data to the response.
-	s := h.store.ScrapeSwarm(req.InfoHash, len(req.IP) == net.IPv6len)
+	s := h.store.ScrapeSwarm(req.InfoHash, req.IP.AddressFamily)
 	resp.Incomplete = s.Incomplete
 	resp.Complete = s.Complete
 
@@ -123,13 +165,13 @@ func (h *responseHook) appendPeers(req *bittorrent.AnnounceRequest, resp *bittor
 		peers = append(peers, req.Peer)
 	}
 
-	switch len(req.IP) {
-	case net.IPv4len:
+	switch req.IP.AddressFamily {
+	case bittorrent.IPv4:
 		resp.IPv4Peers = peers
-	case net.IPv6len:
+	case bittorrent.IPv6:
 		resp.IPv6Peers = peers
 	default:
-		panic("peer IP is not IPv4 or IPv6 length")
+		panic("attempted to append peer that is neither IPv4 nor IPv6")
 	}
 
 	return nil
@@ -143,7 +185,11 @@ func (h *responseHook) HandleScrape(ctx context.Context, req *bittorrent.ScrapeR
 	v6, _ := ctx.Value(ScrapeIsIPv6Key).(bool)
 
 	for _, infoHash := range req.InfoHashes {
-		resp.Files[infoHash] = h.store.ScrapeSwarm(infoHash, v6)
+		if v6 {
+			resp.Files[infoHash] = h.store.ScrapeSwarm(infoHash, bittorrent.IPv6)
+		} else {
+			resp.Files[infoHash] = h.store.ScrapeSwarm(infoHash, bittorrent.IPv4)
+		}
 	}
 
 	return ctx, nil

middleware/middleware.go

@@ -16,6 +16,8 @@ import (
 // Config holds the configuration common across all middleware.
 type Config struct {
 	AnnounceInterval time.Duration `yaml:"announce_interval"`
+	MaxNumWant       uint32        `yaml:"max_numwant"`
+	DefaultNumWant   uint32        `yaml:"default_numwant"`
 }
 
 var _ frontend.TrackerLogic = &Logic{}
@@ -26,10 +28,13 @@ func NewLogic(cfg Config, peerStore storage.PeerStore, preHooks, postHooks []Hoo
 	l := &Logic{
 		announceInterval: cfg.AnnounceInterval,
 		peerStore:        peerStore,
-		preHooks:         append(preHooks, &responseHook{store: peerStore}),
+		preHooks:         []Hook{&sanitizationHook{maxNumWant: cfg.MaxNumWant, defaultNumWant: cfg.DefaultNumWant}},
 		postHooks:        append(postHooks, &swarmInteractionHook{store: peerStore}),
 	}
 
+	l.preHooks = append(l.preHooks, preHooks...)
+	l.preHooks = append(l.preHooks, &responseHook{store: peerStore})
+
 	return l
 }
 

middleware/middleware_test.go

@@ -0,0 +1,94 @@
+package middleware
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/chihaya/chihaya/bittorrent"
+)
+
+// nopHook is a Hook to measure the overhead of a no-operation Hook through
+// benchmarks.
+type nopHook struct{}
+
+func (h *nopHook) HandleAnnounce(ctx context.Context, req *bittorrent.AnnounceRequest, resp *bittorrent.AnnounceResponse) (context.Context, error) {
+	return ctx, nil
+}
+
+func (h *nopHook) HandleScrape(ctx context.Context, req *bittorrent.ScrapeRequest, resp *bittorrent.ScrapeResponse) (context.Context, error) {
+	return ctx, nil
+}
+
+type hookList []Hook
+
+func (hooks hookList) handleAnnounce(ctx context.Context, req *bittorrent.AnnounceRequest) (resp *bittorrent.AnnounceResponse, err error) {
+	resp = &bittorrent.AnnounceResponse{
+		Interval:    60,
+		MinInterval: 60,
+		Compact:     true,
+	}
+
+	for _, h := range []Hook(hooks) {
+		if ctx, err = h.HandleAnnounce(ctx, req, resp); err != nil {
+			return nil, err
+		}
+	}
+
+	return resp, nil
+}
+
+func benchHookListV4(b *testing.B, hooks hookList) {
+	req := &bittorrent.AnnounceRequest{Peer: bittorrent.Peer{IP: bittorrent.IP{IP: net.ParseIP("1.2.3.4"), AddressFamily: bittorrent.IPv4}}}
+	benchHookList(b, hooks, req)
+}
+
+func benchHookListV6(b *testing.B, hooks hookList) {
+	req := &bittorrent.AnnounceRequest{Peer: bittorrent.Peer{IP: bittorrent.IP{IP: net.ParseIP("fc00::0001"), AddressFamily: bittorrent.IPv6}}}
+	benchHookList(b, hooks, req)
+}
+
+func benchHookList(b *testing.B, hooks hookList, req *bittorrent.AnnounceRequest) {
+	ctx := context.Background()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		resp, err := hooks.handleAnnounce(ctx, req)
+		require.Nil(b, err)
+		require.NotNil(b, resp)
+	}
+}
+
+func BenchmarkHookOverhead(b *testing.B) {
+	b.Run("none-v4", func(b *testing.B) {
+		benchHookListV4(b, hookList{})
+	})
+
+	b.Run("none-v6", func(b *testing.B) {
+		benchHookListV6(b, hookList{})
+	})
+
+	var nopHooks hookList
+	for i := 1; i < 4; i++ {
+		nopHooks = append(nopHooks, &nopHook{})
+		b.Run(fmt.Sprintf("%dnop-v4", i), func(b *testing.B) {
+			benchHookListV4(b, nopHooks)
+		})
+		b.Run(fmt.Sprintf("%dnop-v6", i), func(b *testing.B) {
+			benchHookListV6(b, nopHooks)
+		})
+	}
+
+	var sanHooks hookList
+	for i := 1; i < 4; i++ {
+		sanHooks = append(sanHooks, &sanitizationHook{maxNumWant: 50})
+		b.Run(fmt.Sprintf("%dsanitation-v4", i), func(b *testing.B) {
+			benchHookListV4(b, sanHooks)
+		})
+		b.Run(fmt.Sprintf("%dsanitation-v6", i), func(b *testing.B) {
+			benchHookListV6(b, sanHooks)
+		})
+	}
+}