From e4934f101478aa0acf633373ff021d8a56bdbbf0 Mon Sep 17 00:00:00 2001
From: Jeromy <jeromyj@gmail.com>
Date: Tue, 16 Feb 2016 15:30:14 -0800
Subject: [PATCH] add remote multiaddr from active connection to peerstore

---
 p2p/protocol/identify/id.go      | 28 +++++++++++++++-------------
 p2p/protocol/identify/id_test.go | 11 +++++++++--
 2 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/p2p/protocol/identify/id.go b/p2p/protocol/identify/id.go
index 1f1e2c132..a8fe4d4cc 100644
--- a/p2p/protocol/identify/id.go
+++ b/p2p/protocol/identify/id.go
@@ -82,21 +82,21 @@ func (ids *IDService) IdentifyConn(c inet.Conn) {
 		log.Event(context.TODO(), "IdentifyOpenFailed", c.RemotePeer())
 		c.Close()
 		return
-	} else {
-		bwc := ids.Host.GetBandwidthReporter()
-		s = mstream.WrapStream(s, ID, bwc)
-
-		// ok give the response to our handler.
-		if err := msmux.SelectProtoOrFail(ID, s); err != nil {
-			log.Debugf("error writing stream header for %s", ID)
-			log.Event(context.TODO(), "IdentifyOpenFailed", c.RemotePeer())
-			s.Close()
-			return
-		} else {
-			ids.ResponseHandler(s)
-		}
 	}
 
+	bwc := ids.Host.GetBandwidthReporter()
+	s = mstream.WrapStream(s, ID, bwc)
+
+	// ok give the response to our handler.
+	if err := msmux.SelectProtoOrFail(ID, s); err != nil {
+		log.Debugf("error writing stream header for %s", ID)
+		log.Event(context.TODO(), "IdentifyOpenFailed", c.RemotePeer())
+		s.Close()
+		return
+	}
+
+	ids.ResponseHandler(s)
+
 	ids.currmu.Lock()
 	ch, found := ids.currid[c]
 	delete(ids.currid, c)
@@ -190,6 +190,8 @@ func (ids *IDService) consumeMessage(mes *pb.Identify, c inet.Conn) {
 		lmaddrs = append(lmaddrs, maddr)
 	}
 
+	lmaddrs = append(lmaddrs, c.RemoteMultiaddr())
+
 	// update our peerstore with the addresses. here, we SET the addresses, clearing old ones.
 	// We are receiving from the peer itself. this is current address ground truth.
 	ids.Host.Peerstore().SetAddrs(p, lmaddrs, peer.ConnectedAddrTTL)
diff --git a/p2p/protocol/identify/id_test.go b/p2p/protocol/identify/id_test.go
index 53a5729d5..6c1d1e127 100644
--- a/p2p/protocol/identify/id_test.go
+++ b/p2p/protocol/identify/id_test.go
@@ -38,6 +38,7 @@ func subtestIDService(t *testing.T, postDialWait time.Duration) {
 
 	// the IDService should be opened automatically, by the network.
 	// what we should see now is that both peers know about each others listen addresses.
+	t.Log("test peer1 has peer2 addrs correctly")
 	testKnowsAddrs(t, h1, h2p, h2.Peerstore().Addrs(h2p)) // has them
 	testHasProtocolVersions(t, h1, h2p)
 
@@ -49,8 +50,12 @@ func subtestIDService(t *testing.T, postDialWait time.Duration) {
 	}
 	<-h2.IDService().IdentifyWait(c[0])
 
+	addrs := h1.Peerstore().Addrs(h1p)
+	addrs = append(addrs, c[0].RemoteMultiaddr())
+
 	// and the protocol versions.
-	testKnowsAddrs(t, h2, h1p, h1.Peerstore().Addrs(h1p)) // has them
+	t.Log("test peer2 has peer1 addrs correctly")
+	testKnowsAddrs(t, h2, h1p, addrs) // has them
 	testHasProtocolVersions(t, h2, h1p)
 }
 
@@ -58,7 +63,9 @@ func testKnowsAddrs(t *testing.T, h host.Host, p peer.ID, expected []ma.Multiadd
 	actual := h.Peerstore().Addrs(p)
 
 	if len(actual) != len(expected) {
-		t.Error("dont have the same addresses")
+		t.Errorf("expected: %s", expected)
+		t.Errorf("actual: %s", actual)
+		t.Fatal("dont have the same addresses")
 	}
 
 	have := map[string]struct{}{}