Add get peers limits

diff --git a/tracker/http/server/server.go b/tracker/http/server/server.go
index ed08ac924ef3376cc7bcbe7b43e4e0a57b607967..82c402dec68e7cbe371df504e75dbec13e4def17 100644 (file)
--- 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 dc95bce6841c4dc4adc77def60046ad12b524df3..f7d5ecb8beae4cb139fe0660932cc7af73728bca 100644 (file)
--- 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 95e7a5e7a44d5bc8fe03e8838757578d251a1d45..dfcc6500a0e12776c865b49a0eb1c195c7cb8001 100644 (file)
--- 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
        }