From fe991128edc6a7b96b5b4634edb23f7222e536c7 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Fri, 2 Feb 2018 19:07:20 +1100
Subject: [PATCH] Rework connection stat reconciliation with Torrent and
 refactor doppleganger handling

---
 client.go     | 33 +++++++++++++++------------------
 connection.go |  8 ++++++++
 torrent.go    | 17 ++++++++---------
 3 files changed, 31 insertions(+), 27 deletions(-)

diff --git a/client.go b/client.go
index efe5f736..3072dcd1 100644
--- a/client.go
+++ b/client.go
@@ -714,7 +714,7 @@ func (cl *Client) outgoingConnection(t *Torrent, addr string, ps peerSource) {
 	}
 	defer c.Close()
 	c.Discovery = ps
-	cl.runInitiatedHandshookConn(c, t)
+	cl.runHandshookConn(c, t, true)
 }
 
 // The port number for incoming peer connections. 0 if the client isn't
@@ -816,16 +816,6 @@ func (cl *Client) connBTHandshake(c *connection, ih *metainfo.Hash) (ret metainf
 	return
 }
 
-func (cl *Client) runInitiatedHandshookConn(c *connection, t *Torrent) {
-	if c.PeerID == cl.peerID {
-		connsToSelf.Add(1)
-		addr := c.conn.RemoteAddr().String()
-		cl.dopplegangerAddrs[addr] = struct{}{}
-		return
-	}
-	cl.runHandshookConn(c, t, true)
-}
-
 func (cl *Client) runReceivedConn(c *connection) {
 	err := c.conn.SetDeadline(time.Now().Add(cl.config.HandshakesTimeout))
 	if err != nil {
@@ -843,17 +833,24 @@ func (cl *Client) runReceivedConn(c *connection) {
 	}
 	cl.mu.Lock()
 	defer cl.mu.Unlock()
-	if c.PeerID == cl.peerID {
-		// Because the remote address is not necessarily the same as its
-		// client's torrent listen address, we won't record the remote address
-		// as a doppleganger. Instead, the initiator can record *us* as the
-		// doppleganger.
-		return
-	}
 	cl.runHandshookConn(c, t, false)
 }
 
 func (cl *Client) runHandshookConn(c *connection, t *Torrent, outgoing bool) {
+	t.reconcileHandshakeStats(c)
+	if c.PeerID == cl.peerID {
+		if outgoing {
+			connsToSelf.Add(1)
+			addr := c.conn.RemoteAddr().String()
+			cl.dopplegangerAddrs[addr] = struct{}{}
+		} else {
+			// Because the remote address is not necessarily the same as its
+			// client's torrent listen address, we won't record the remote address
+			// as a doppleganger. Instead, the initiator can record *us* as the
+			// doppleganger.
+		}
+		return
+	}
 	c.conn.SetWriteDeadline(time.Time{})
 	c.r = deadlineReader{c.conn, c.r}
 	completedHandshakeConnectionFlags.Add(c.connectionFlags(), 1)
diff --git a/connection.go b/connection.go
index b04c9c14..550d5c10 100644
--- a/connection.go
+++ b/connection.go
@@ -1282,3 +1282,11 @@ func (c *connection) sendChunk(r request, msg func(pp.Message) bool) (more bool,
 	c.lastChunkSent = time.Now()
 	return
 }
+
+func (c *connection) setTorrent(t *Torrent) {
+	if c.t != nil {
+		panic("connection already associated with a torrent")
+	}
+	c.t = t
+	t.conns[c] = struct{}{}
+}
diff --git a/torrent.go b/torrent.go
index d5a45f4c..60119b59 100644
--- a/torrent.go
+++ b/torrent.go
@@ -1410,6 +1410,13 @@ func (t *Torrent) numTotalPeers() int {
 	return len(peers)
 }
 
+// Reconcile bytes transferred before connection was associated with a
+// torrent.
+func (t *Torrent) reconcileHandshakeStats(c *connection) {
+	t.stats.wroteBytes(c.stats.BytesWritten)
+	t.stats.readBytes(c.stats.BytesRead)
+}
+
 // Returns true if the connection is added.
 func (t *Torrent) addConnection(c *connection, outgoing bool) bool {
 	if t.cl.closed.IsSet() {
@@ -1451,15 +1458,7 @@ func (t *Torrent) addConnection(c *connection, outgoing bool) bool {
 	if len(t.conns) >= t.maxEstablishedConns {
 		panic(len(t.conns))
 	}
-	if c.t != nil {
-		panic("connection already associated with a torrent")
-	}
-	// Reconcile bytes transferred before connection was associated with a
-	// torrent.
-	t.stats.wroteBytes(c.stats.BytesWritten)
-	t.stats.readBytes(c.stats.BytesRead)
-	c.t = t
-	t.conns[c] = struct{}{}
+	c.setTorrent(t)
 	return true
 }
 
-- 
2.50.0