From 4a7fbf61703e1e398031151f5674ae229eed3589 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Sun, 27 Nov 2016 00:05:19 +1100
Subject: [PATCH] Add peers received from received announce_peer DHT messages
 to the Client

Addresses #133
---
 client.go     | 17 +++++++++++++++++
 connection.go | 15 +++++++--------
 torrent.go    |  2 +-
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/client.go b/client.go
index 94c02f02..2e6ded8b 100644
--- a/client.go
+++ b/client.go
@@ -309,6 +309,9 @@ func NewClient(cfg *Config) (cl *Client, err error) {
 		if dhtCfg.Conn == nil && cl.utpSock != nil {
 			dhtCfg.Conn = cl.utpSock
 		}
+		if dhtCfg.OnAnnouncePeer == nil {
+			dhtCfg.OnAnnouncePeer = cl.onDHTAnnouncePeer
+		}
 		cl.dHT, err = dht.NewServer(&dhtCfg)
 		if err != nil {
 			return
@@ -1341,3 +1344,17 @@ func (cl *Client) newConnection(nc net.Conn) (c *connection) {
 	c.r = rateLimitedReader{cl.downloadLimit, c.r}
 	return
 }
+
+func (cl *Client) onDHTAnnouncePeer(ih metainfo.Hash, p dht.Peer) {
+	cl.mu.Lock()
+	defer cl.mu.Unlock()
+	t, ok := cl.Torrent(ih)
+	if !ok {
+		return
+	}
+	t.addPeers([]Peer{{
+		IP:     p.IP,
+		Port:   p.Port,
+		Source: peerSourceDHTAnnouncePeer,
+	}})
+}
diff --git a/connection.go b/connection.go
index 727f97f7..1d67e9dc 100644
--- a/connection.go
+++ b/connection.go
@@ -28,13 +28,14 @@ import (
 
 var optimizedCancels = expvar.NewInt("optimizedCancels")
 
-type peerSource byte
+type peerSource string
 
 const (
-	peerSourceTracker  = '\x00' // It's the default.
-	peerSourceIncoming = 'I'
-	peerSourceDHT      = 'H'
-	peerSourcePEX      = 'X'
+	peerSourceTracker         = "T" // It's the default.
+	peerSourceIncoming        = "I"
+	peerSourceDHTGetPeers     = "Hg"
+	peerSourceDHTAnnouncePeer = "Ha"
+	peerSourcePEX             = "X"
 )
 
 // Maintains the state of a connection with a peer.
@@ -155,9 +156,7 @@ func (cn *connection) connectionFlags() (ret string) {
 	if cn.encrypted {
 		c('E')
 	}
-	if cn.Discovery != 0 {
-		c(byte(cn.Discovery))
-	}
+	ret += string(cn.Discovery)
 	if cn.uTP {
 		c('T')
 	}
diff --git a/torrent.go b/torrent.go
index 8610a5ec..3fd3f0b2 100644
--- a/torrent.go
+++ b/torrent.go
@@ -1268,7 +1268,7 @@ func (t *Torrent) consumeDHTAnnounce(pvs <-chan dht.PeersValues) {
 				addPeers = append(addPeers, Peer{
 					IP:     cp.IP[:],
 					Port:   cp.Port,
-					Source: peerSourceDHT,
+					Source: peerSourceDHTGetPeers,
 				})
 				key := (&net.UDPAddr{
 					IP:   cp.IP[:],
-- 
2.48.1