From acb09fcf799f4fd5e3a81323b1e41ddd1247a127 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Thu, 8 Dec 2022 14:04:23 +1100
Subject: [PATCH] Add get peers limits

---
 tracker/http/server/server.go |  4 +++-
 tracker/server.go             | 21 ++++++++++++++++++---
 tracker/udp/server/server.go  |  7 ++++++-
 3 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/tracker/http/server/server.go b/tracker/http/server/server.go
index ed08ac92..82c402de 100644
--- a/tracker/http/server/server.go
+++ b/tracker/http/server/server.go
@@ -9,6 +9,7 @@ import (
 	"strconv"
 
 	"github.com/anacrolix/dht/v2/krpc"
+	"github.com/anacrolix/generics"
 	"github.com/anacrolix/log"
 
 	"github.com/anacrolix/torrent/bencode"
@@ -78,7 +79,8 @@ func (me Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		PeerId:   peerId,
 		Event:    event,
 		Port:     addrPort.Port(),
-	}, addrPort)
+		NumWant:  -1,
+	}, addrPort, tracker.GetPeersOpts{generics.Some[uint](200)})
 	if err != nil {
 		log.Printf("error serving announce: %v", err)
 		http.Error(w, "error handling announce", http.StatusInternalServerError)
diff --git a/tracker/server.go b/tracker/server.go
index dc95bce6..f7d5ecb8 100644
--- a/tracker/server.go
+++ b/tracker/server.go
@@ -14,7 +14,10 @@ import (
 
 // This is reserved for stuff like filtering by IP version, avoiding an announcer's IP or key,
 // limiting return count, etc.
-type GetPeersOpts struct{}
+type GetPeersOpts struct {
+	// Negative numbers are not allowed.
+	MaxCount generics.Option[uint]
+}
 
 type InfoHash = [20]byte
 
@@ -79,7 +82,7 @@ func addMissing(orig []PeerInfo, new peerSet) {
 }
 
 func (me *AnnounceHandler) Serve(
-	ctx context.Context, req AnnounceRequest, addr AnnounceAddr,
+	ctx context.Context, req AnnounceRequest, addr AnnounceAddr, opts GetPeersOpts,
 ) (peers []PeerInfo, err error) {
 	err = me.AnnounceTracker.TrackAnnounce(ctx, req, addr)
 	if err != nil {
@@ -91,7 +94,19 @@ func (me *AnnounceHandler) Serve(
 	me.mu.Lock()
 	op.Value, op.Ok = me.ongoingUpstreamAugmentations[infoHash]
 	me.mu.Unlock()
-	peers, err = me.AnnounceTracker.GetPeers(ctx, infoHash, GetPeersOpts{})
+	// Apply num_want limit to max count. I really can't tell if this is the right place to do it,
+	// but it seems the most flexible.
+	if req.NumWant != -1 {
+		newCount := uint(req.NumWant)
+		if opts.MaxCount.Ok {
+			if newCount < opts.MaxCount.Value {
+				opts.MaxCount.Value = newCount
+			}
+		} else {
+			opts.MaxCount = generics.Some(newCount)
+		}
+	}
+	peers, err = me.AnnounceTracker.GetPeers(ctx, infoHash, opts)
 	if err != nil {
 		return
 	}
diff --git a/tracker/udp/server/server.go b/tracker/udp/server/server.go
index 95e7a5e7..dfcc6500 100644
--- a/tracker/udp/server/server.go
+++ b/tracker/udp/server/server.go
@@ -11,6 +11,7 @@ import (
 	"net/netip"
 
 	"github.com/anacrolix/dht/v2/krpc"
+	"github.com/anacrolix/generics"
 	"github.com/anacrolix/log"
 
 	"github.com/anacrolix/torrent/tracker"
@@ -94,7 +95,11 @@ func (me *Server) handleAnnounce(
 		err = fmt.Errorf("converting source net.Addr to AnnounceAddr: %w", err)
 		return err
 	}
-	peers, err := me.Announce.Serve(ctx, req, announceAddr)
+	opts := tracker.GetPeersOpts{MaxCount: generics.Some[uint](50)}
+	if addrFamily == udp.AddrFamilyIpv4 {
+		opts.MaxCount = generics.Some[uint](150)
+	}
+	peers, err := me.Announce.Serve(ctx, req, announceAddr, opts)
 	if err != nil {
 		return err
 	}
-- 
2.48.1