diff --git a/middleware/hooks.go b/middleware/hooks.go index 2f0574b..2946591 100644 --- a/middleware/hooks.go +++ b/middleware/hooks.go @@ -116,11 +116,22 @@ func (h *responseHook) HandleAnnounce(ctx context.Context, req *bittorrent.Annou func (h *responseHook) appendPeers(req *bittorrent.AnnounceRequest, resp *bittorrent.AnnounceResponse) error { seeding := req.Left == 0 max := int(req.NumWant) - peers, err := h.store.AnnouncePeers(req.InfoHash, seeding, max, req.Peer) + storePeers, err := h.store.AnnouncePeers(req.InfoHash, seeding, max, req.Peer) if err != nil && !errors.Is(err, storage.ErrResourceDoesNotExist) { return err } err = nil + peers := make([]bittorrent.Peer, 0, len(resp.IPv4Peers)+len(resp.IPv6Peers)+len(storePeers)) + + // append peers, which added in middleware + if req.Peer.Addr().Is6() { + peers = append(peers, resp.IPv6Peers...) + peers = append(peers, resp.IPv4Peers...) + } else { + peers = append(peers, resp.IPv4Peers...) + peers = append(peers, resp.IPv6Peers...) + } + peers = append(peers, storePeers...) // Some clients expect a minimum of their own peer representation returned to // them if they are the only peer in a swarm. @@ -133,37 +144,30 @@ func (h *responseHook) appendPeers(req *bittorrent.AnnounceRequest, resp *bittor peers = append(peers, req.Peer) } - switch addr := req.Peer.Addr(); { - case addr.Is4(): - resp.IPv4Peers = mergePeers(resp.IPv4Peers, peers, max) - case addr.Is6(): - resp.IPv6Peers = mergePeers(resp.IPv6Peers, peers, max) - default: - err = bittorrent.ErrInvalidIP + uniquePeers := make(map[bittorrent.Peer]interface{}, len(peers)) + + resp.IPv4Peers = make([]bittorrent.Peer, 0, len(peers)/2) + resp.IPv6Peers = make([]bittorrent.Peer, 0, len(peers)/2) + + for _, p := range peers { + if err != nil || len(uniquePeers) > max { + break + } + if _, found := uniquePeers[p]; !found { + uniquePeers[p] = nil + if p.Addr().Is6() { + resp.IPv6Peers = append(resp.IPv6Peers, p) + } else if p.Addr().Is4() { + resp.IPv4Peers = append(resp.IPv4Peers, p) + } else { + err = bittorrent.ErrInvalidIP + } + } } return err } -func mergePeers(p0, p1 []bittorrent.Peer, max int) (result []bittorrent.Peer) { - peers := make(map[string]bittorrent.Peer, len(p0)+len(p1)) - for _, p := range p0 { - peers[p.RawString()] = p - } - for _, p := range p1 { - peers[p.RawString()] = p - } - result = make([]bittorrent.Peer, 0, len(peers)) - for _, v := range peers { - if len(peers) < max { - result = append(result, v) - } else { - break - } - } - return -} - func (h *responseHook) HandleScrape(ctx context.Context, req *bittorrent.ScrapeRequest, resp *bittorrent.ScrapeResponse) (context.Context, error) { if ctx.Value(SkipResponseHookKey) != nil { return ctx, nil diff --git a/storage/keydb/storage.go b/storage/keydb/storage.go index 3df6ce3..1db0d9b 100644 --- a/storage/keydb/storage.go +++ b/storage/keydb/storage.go @@ -75,6 +75,7 @@ func newStore(cfg r.Config) (*store, error) { logFields: cfg.LogFields(), peerTTL: uint(cfg.PeerLifetime.Seconds()), } + st.logFields["name"] = Name } return st, err diff --git a/storage/memory/storage.go b/storage/memory/storage.go index b35b1c9..5a24e1b 100644 --- a/storage/memory/storage.go +++ b/storage/memory/storage.go @@ -98,8 +98,8 @@ type peerShard struct { type swarm struct { // map serialized peer to mtime - seeders map[string]int64 - leechers map[string]int64 + seeders map[bittorrent.Peer]int64 + leechers map[bittorrent.Peer]int64 } type peerStore struct { @@ -191,25 +191,23 @@ func (ps *peerStore) PutSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) error default: } - pk := p.RawString() - shard := ps.shards[ps.shardIndex(ih, p.Addr().Is6())] shard.Lock() if _, ok := shard.swarms[ih]; !ok { shard.swarms[ih] = swarm{ - seeders: make(map[string]int64), - leechers: make(map[string]int64), + seeders: make(map[bittorrent.Peer]int64), + leechers: make(map[bittorrent.Peer]int64), } } // If this peer isn't already a seeder, update the stats for the swarm. - if _, ok := shard.swarms[ih].seeders[pk]; !ok { + if _, ok := shard.swarms[ih].seeders[p]; !ok { shard.numSeeders++ } // Update the peer in the swarm. - shard.swarms[ih].seeders[pk] = ps.getClock() + shard.swarms[ih].seeders[p] = ps.getClock() shard.Unlock() return nil @@ -222,8 +220,6 @@ func (ps *peerStore) DeleteSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) err default: } - pk := p.RawString() - shard := ps.shards[ps.shardIndex(ih, p.Addr().Is6())] shard.Lock() @@ -232,13 +228,13 @@ func (ps *peerStore) DeleteSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) err return storage.ErrResourceDoesNotExist } - if _, ok := shard.swarms[ih].seeders[pk]; !ok { + if _, ok := shard.swarms[ih].seeders[p]; !ok { shard.Unlock() return storage.ErrResourceDoesNotExist } shard.numSeeders-- - delete(shard.swarms[ih].seeders, pk) + delete(shard.swarms[ih].seeders, p) if len(shard.swarms[ih].seeders)|len(shard.swarms[ih].leechers) == 0 { delete(shard.swarms, ih) @@ -255,25 +251,23 @@ func (ps *peerStore) PutLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error default: } - pk := p.RawString() - shard := ps.shards[ps.shardIndex(ih, p.Addr().Is6())] shard.Lock() if _, ok := shard.swarms[ih]; !ok { shard.swarms[ih] = swarm{ - seeders: make(map[string]int64), - leechers: make(map[string]int64), + seeders: make(map[bittorrent.Peer]int64), + leechers: make(map[bittorrent.Peer]int64), } } // If this peer isn't already a leecher, update the stats for the swarm. - if _, ok := shard.swarms[ih].leechers[pk]; !ok { + if _, ok := shard.swarms[ih].leechers[p]; !ok { shard.numLeechers++ } // Update the peer in the swarm. - shard.swarms[ih].leechers[pk] = ps.getClock() + shard.swarms[ih].leechers[p] = ps.getClock() shard.Unlock() return nil @@ -286,8 +280,6 @@ func (ps *peerStore) DeleteLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) er default: } - pk := p.RawString() - shard := ps.shards[ps.shardIndex(ih, p.Addr().Is6())] shard.Lock() @@ -296,13 +288,13 @@ func (ps *peerStore) DeleteLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) er return storage.ErrResourceDoesNotExist } - if _, ok := shard.swarms[ih].leechers[pk]; !ok { + if _, ok := shard.swarms[ih].leechers[p]; !ok { shard.Unlock() return storage.ErrResourceDoesNotExist } shard.numLeechers-- - delete(shard.swarms[ih].leechers, pk) + delete(shard.swarms[ih].leechers, p) if len(shard.swarms[ih].seeders)|len(shard.swarms[ih].leechers) == 0 { delete(shard.swarms, ih) @@ -319,43 +311,40 @@ func (ps *peerStore) GraduateLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) default: } - pk := p.RawString() - shard := ps.shards[ps.shardIndex(ih, p.Addr().Is6())] shard.Lock() if _, ok := shard.swarms[ih]; !ok { shard.swarms[ih] = swarm{ - seeders: make(map[string]int64), - leechers: make(map[string]int64), + seeders: make(map[bittorrent.Peer]int64), + leechers: make(map[bittorrent.Peer]int64), } } // If this peer is a leecher, update the stats for the swarm and remove them. - if _, ok := shard.swarms[ih].leechers[pk]; ok { + if _, ok := shard.swarms[ih].leechers[p]; ok { shard.numLeechers-- - delete(shard.swarms[ih].leechers, pk) + delete(shard.swarms[ih].leechers, p) } // If this peer isn't already a seeder, update the stats for the swarm. - if _, ok := shard.swarms[ih].seeders[pk]; !ok { + if _, ok := shard.swarms[ih].seeders[p]; !ok { shard.numSeeders++ } // Update the peer in the swarm. - shard.swarms[ih].seeders[pk] = ps.getClock() + shard.swarms[ih].seeders[p] = ps.getClock() shard.Unlock() return nil } -func parsePeers(peersMap map[string]int64, maxCount int, skipPeerID string) (peers []bittorrent.Peer) { - for pk := range peersMap { +func parsePeers(peersMap map[bittorrent.Peer]int64, maxCount int, skipPeer bittorrent.Peer) (peers []bittorrent.Peer) { + for p := range peersMap { if maxCount == 0 { break } - if pk != skipPeerID { - p, _ := bittorrent.NewPeer(pk) + if p != skipPeer { peers = append(peers, p) maxCount-- } @@ -363,16 +352,16 @@ func parsePeers(peersMap map[string]int64, maxCount int, skipPeerID string) (pee return } -func (ps *peerStore) getPeers(shard *peerShard, ih bittorrent.InfoHash, maxCount int, leechersOnly bool, skipPeerID string) (peers []bittorrent.Peer) { +func (ps *peerStore) getPeers(shard *peerShard, ih bittorrent.InfoHash, maxCount int, leechersOnly bool, skipPeer bittorrent.Peer) (peers []bittorrent.Peer) { shard.RLock() defer shard.RUnlock() if swarm, ok := shard.swarms[ih]; ok { if !leechersOnly { - peers = append(peers, parsePeers(swarm.seeders, maxCount, skipPeerID)...) + peers = append(peers, parsePeers(swarm.seeders, maxCount, skipPeer)...) maxCount -= len(peers) } if maxCount > 0 { - peers = append(peers, parsePeers(swarm.leechers, maxCount, skipPeerID)...) + peers = append(peers, parsePeers(swarm.leechers, maxCount, skipPeer)...) } } return @@ -385,19 +374,19 @@ func (ps *peerStore) AnnouncePeers(ih bittorrent.InfoHash, seeder bool, numWant default: } - peerID, isV6 := peer.RawString(), peer.Addr().Is6() + isV6 := peer.Addr().Is6() if seeder { // Append leechers as possible. - peers = ps.getPeers(ps.shards[ps.shardIndex(ih, isV6)], ih, numWant, true, peerID) + peers = ps.getPeers(ps.shards[ps.shardIndex(ih, isV6)], ih, numWant, true, peer) if numWant -= len(peers); numWant > 0 { - peers = append(peers, ps.getPeers(ps.shards[ps.shardIndex(ih, !isV6)], ih, numWant, true, peerID)...) + peers = append(peers, ps.getPeers(ps.shards[ps.shardIndex(ih, !isV6)], ih, numWant, true, peer)...) } } else { // Append as many seeders as possible. - peers = ps.getPeers(ps.shards[ps.shardIndex(ih, isV6)], ih, numWant, false, peerID) + peers = ps.getPeers(ps.shards[ps.shardIndex(ih, isV6)], ih, numWant, false, peer) if numWant -= len(peers); numWant > 0 { - peers = append(peers, ps.getPeers(ps.shards[ps.shardIndex(ih, !isV6)], ih, numWant, false, peerID)...) + peers = append(peers, ps.getPeers(ps.shards[ps.shardIndex(ih, !isV6)], ih, numWant, false, peer)...) } }