From d86b8e14de889f8205426ce2ef4227dfc47e1256 Mon Sep 17 00:00:00 2001
From: Jimmy Zelinskie <jimmy.zelinskie+git@gmail.com>
Date: Sat, 1 Oct 2016 21:35:07 -0400
Subject: [PATCH] middleware: fix scrape data in announces

This change adjusts it such that only empty Swarms return the request
Peer. This also increments the Scrape data properly if a peer is added.

Fixes #241.
---
 middleware/hooks.go | 43 +++++++++++++++++++++++++------------------
 1 file changed, 25 insertions(+), 18 deletions(-)

diff --git a/middleware/hooks.go b/middleware/hooks.go
index 4ba71f7..4fb65dd 100644
--- a/middleware/hooks.go
+++ b/middleware/hooks.go
@@ -96,36 +96,43 @@ func (h *responseHook) HandleAnnounce(ctx context.Context, req *bittorrent.Annou
 		return ctx, nil
 	}
 
+	// Add the Scrape data to the response.
 	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.getPeers(req.InfoHash, req.Left == 0, int(req.NumWant), req.Peer)
-	case net.IPv6len:
-		resp.IPv6Peers, err = h.getPeers(req.InfoHash, req.Left == 0, int(req.NumWant), req.Peer)
-	default:
-		return ctx, ErrInvalidIP
-	}
-
+	err = h.appendPeers(req, resp)
 	return ctx, err
 }
 
-func (h *responseHook) getPeers(ih bittorrent.InfoHash, seeder bool, numWant int, p bittorrent.Peer) ([]bittorrent.Peer, error) {
-	peers, err := h.store.AnnouncePeers(ih, seeder, numWant, p)
+func (h *responseHook) appendPeers(req *bittorrent.AnnounceRequest, resp *bittorrent.AnnounceResponse) error {
+	seeding := req.Left == 0
+	peers, err := h.store.AnnouncePeers(req.InfoHash, seeding, int(req.NumWant), req.Peer)
 	if err != nil && err != storage.ErrResourceDoesNotExist {
-		return nil, err
+		return err
 	}
 
-	// Insert announcing peer.
-	// Some clients expect at least themselves in every announce response.
-	if len(peers) < numWant || len(peers) == 0 {
-		return append(peers, p), nil
+	// Some clients expect a minimum of their own peer representation returned to
+	// them if they are the only peer in a swarm.
+	if len(peers) == 0 {
+		if seeding {
+			resp.Complete++
+		} else {
+			resp.Incomplete++
+		}
+		peers = append(peers, req.Peer)
 	}
 
-	peers[len(peers)-1] = p
-	return peers, nil
+	switch len(req.IP) {
+	case net.IPv4len:
+		resp.IPv4Peers = peers
+	case net.IPv6len:
+		resp.IPv6Peers = peers
+	default:
+		panic("peer IP is not IPv4 or IPv6 length")
+	}
+
+	return nil
 }
 
 func (h *responseHook) HandleScrape(ctx context.Context, req *bittorrent.ScrapeRequest, resp *bittorrent.ScrapeResponse) (context.Context, error) {