fix reconnect id missing

rickyyangz · Dec 17, 2018 · dbeddbc · dbeddbc
1 parent a1ad253
commit dbeddbc
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 6 deletions.
diff --git a/p2p/netaddress.go b/p2p/netaddress.go
@@ -32,7 +32,16 @@ type NetAddress struct {
 
 // IDAddressString returns id@hostPort.
 func IDAddressString(id ID, hostPort string) string {
-	return fmt.Sprintf("%s@%s", id, hostPort)
+	protocol := ""
+	if strings.Contains(hostPort, "://") {
+		parts := strings.Split(hostPort, "://")
+		protocol, hostPort = parts[0], parts[1]
+	}
+	if protocol != "" {
+		return fmt.Sprintf("%s://%s@%s", protocol, id, hostPort)
+	} else {
+		return fmt.Sprintf("%s@%s", id, hostPort)
+	}
 }
 
 // NewNetAddress returns a new NetAddress using the provided TCP

diff --git a/p2p/netaddress_test.go b/p2p/netaddress_test.go
@@ -20,6 +20,22 @@ func TestNewNetAddress(t *testing.T) {
 	}, "Calling NewNetAddress with UDPAddr should not panic in testing")
 }
 
+func TestIDAddressString(t *testing.T) {
+	testCases := []struct {
+		id       ID
+		hostPort string
+		expect   string
+	}{
+		{"123xxx", "tcp://127.0.0.1:8080", "tcp://123xxx@127.0.0.1:8080"},
+		{"123xxx", "udp://127.0.0.1:8080", "udp://123xxx@127.0.0.1:8080"},
+		{"123xxx", "127.0.0.1:8080", "123xxx@127.0.0.1:8080"},
+	}
+	for _, tc := range testCases {
+		ac := IDAddressString(tc.id, tc.hostPort)
+		assert.Equal(t, tc.expect, ac)
+	}
+}
+
 func TestNewNetAddressStringWithOptionalID(t *testing.T) {
 	testCases := []struct {
 		addr     string

diff --git a/p2p/pex/addrbook_test.go b/p2p/pex/addrbook_test.go
@@ -211,7 +211,7 @@ func randIPv4Address(t *testing.T) *p2p.NetAddress {
 		)
 		port := cmn.RandIntn(65535-1) + 1
 		id := p2p.ID(hex.EncodeToString(cmn.RandBytes(p2p.IDByteLength)))
-		idAddr := p2p.IDAddressString(id, fmt.Sprintf("%v:%v", ip, port))
+		idAddr := p2p.IDAddressString(id, fmt.Sprintf("tcp://%v:%v", ip, port))
 		addr, err := p2p.NewNetAddressString(idAddr)
 		assert.Nil(t, err, "error generating rand network address")
 		if addr.Routable() {

diff --git a/p2p/transport.go b/p2p/transport.go
@@ -395,12 +395,17 @@ func (mt *MultiplexTransport) wrapPeer(
 	ni NodeInfo,
 	cfg peerConfig,
 ) Peer {
+	var originalNetAddr *NetAddress
+	if cfg.outbound {
+		originalNetAddr, _ = NewNetAddressString(IDAddressString(ni.ID, ni.ListenAddr))
+	}
 	p := newPeer(
 		peerConn{
-			conn:       c,
-			config:     &mt.p2pConfig,
-			outbound:   cfg.outbound,
-			persistent: cfg.persistent,
+			conn:         c,
+			config:       &mt.p2pConfig,
+			outbound:     cfg.outbound,
+			persistent:   cfg.persistent,
+			originalAddr: originalNetAddr,
 		},
 		mt.mConfig,
 		ni,