From 45a5e632aedf061fb885fad4f3ceb57cb3b1c001 Mon Sep 17 00:00:00 2001
From: Leo Balduf <balduf@hm.edu>
Date: Wed, 28 Sep 2016 03:02:37 -0400
Subject: [PATCH] responseHook: always include the announcing peer

Fixes #231
Fixes #232
---
 middleware/hooks.go | 20 ++++++++++++++++++--
 storage/storage.go  |  2 ++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/middleware/hooks.go b/middleware/hooks.go
index 769fb7d..4ba71f7 100644
--- a/middleware/hooks.go
+++ b/middleware/hooks.go
@@ -102,9 +102,9 @@ func (h *responseHook) HandleAnnounce(ctx context.Context, req *bittorrent.Annou
 
 	switch len(req.IP) {
 	case net.IPv4len:
-		resp.IPv4Peers, err = h.store.AnnouncePeers(req.InfoHash, req.Left == 0, int(req.NumWant), req.Peer)
+		resp.IPv4Peers, err = h.getPeers(req.InfoHash, req.Left == 0, int(req.NumWant), req.Peer)
 	case net.IPv6len:
-		resp.IPv6Peers, err = h.store.AnnouncePeers(req.InfoHash, req.Left == 0, int(req.NumWant), req.Peer)
+		resp.IPv6Peers, err = h.getPeers(req.InfoHash, req.Left == 0, int(req.NumWant), req.Peer)
 	default:
 		return ctx, ErrInvalidIP
 	}
@@ -112,6 +112,22 @@ func (h *responseHook) HandleAnnounce(ctx context.Context, req *bittorrent.Annou
 	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)
+	if err != nil && err != storage.ErrResourceDoesNotExist {
+		return nil, 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
+	}
+
+	peers[len(peers)-1] = p
+	return peers, nil
+}
+
 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/storage.go b/storage/storage.go
index bfca330..b2f7020 100644
--- a/storage/storage.go
+++ b/storage/storage.go
@@ -51,6 +51,8 @@ type PeerStore interface {
 	// - if seeder is true, should ideally return more leechers than seeders
 	// - if seeder is false, should ideally return more seeders than
 	//   leechers
+	//
+	// Returns ErrResourceDoesNotExist if the provided infoHash is not tracked.
 	AnnouncePeers(infoHash bittorrent.InfoHash, seeder bool, numWant int, p bittorrent.Peer) (peers []bittorrent.Peer, err error)
 
 	// ScrapeSwarm returns information required to answer a scrape request