From 01fd102dba892973f786d589194ae96e01b3f4bb Mon Sep 17 00:00:00 2001
From: Alexander Simmerl <a.simmerl@gmail.com>
Date: Wed, 23 May 2018 01:53:37 +0200
Subject: [PATCH] Incoporate review feedback

---
 p2p/peer.go          | 18 +++++++++---------
 p2p/peer_set_test.go |  4 +---
 p2p/switch.go        |  7 +++----
 3 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/p2p/peer.go b/p2p/peer.go
index 39656efc5..447225bfb 100644
--- a/p2p/peer.go
+++ b/p2p/peer.go
@@ -41,7 +41,7 @@ type peerConn struct {
 	persistent bool
 	config     *PeerConfig
 	conn       net.Conn // source connection
-	ips        []net.IP
+	ip         net.IP
 }
 
 // ID only exists for SecretConnection.
@@ -52,16 +52,16 @@ func (pc peerConn) ID() ID {
 
 // Return the IP from the connection RemoteAddr
 func (pc peerConn) RemoteIP() net.IP {
-	if len(pc.ips) > 0 {
-		return pc.ips[0]
+	if pc.ip != nil {
+		return pc.ip
 	}
 
+	// In test cases a conn could not be present at all or be an in-memory
+	// implementation where we want to return a fake ip.
 	if pc.conn == nil || pc.conn.RemoteAddr().String() == "pipe" {
-		pc.ips = []net.IP{
-			net.IP{172, 16, 0, byte(atomic.AddUint32(&testIPSuffix, 1))},
-		}
+		pc.ip = net.IP{172, 16, 0, byte(atomic.AddUint32(&testIPSuffix, 1))}
 
-		return pc.ips[0]
+		return pc.ip
 	}
 
 	host, _, err := net.SplitHostPort(pc.conn.RemoteAddr().String())
@@ -74,9 +74,9 @@ func (pc peerConn) RemoteIP() net.IP {
 		panic(err)
 	}
 
-	pc.ips = ips
+	pc.ip = ips[0]
 
-	return ips[0]
+	return pc.ip
 }
 
 // peer implements Peer.
diff --git a/p2p/peer_set_test.go b/p2p/peer_set_test.go
index b043a36b2..fc3004684 100644
--- a/p2p/peer_set_test.go
+++ b/p2p/peer_set_test.go
@@ -26,9 +26,7 @@ func randPeer(ip net.IP) *peer {
 		},
 	}
 
-	p.ips = []net.IP{
-		ip,
-	}
+	p.ip = ip
 
 	return p
 }
diff --git a/p2p/switch.go b/p2p/switch.go
index eba8679e7..6ea7e408f 100644
--- a/p2p/switch.go
+++ b/p2p/switch.go
@@ -534,8 +534,6 @@ func (sw *Switch) addPeer(pc peerConn) error {
 		return err
 	}
 
-	// dont connect to multiple peers on the same IP
-
 	// NOTE: if AuthEnc==false, we don't have a peerID until after the handshake.
 	// If AuthEnc==true then we already know the ID and could do the checks first before the handshake,
 	// but it's simple to just deal with both cases the same after the handshake.
@@ -578,8 +576,9 @@ func (sw *Switch) addPeer(pc peerConn) error {
 		return ErrSwitchDuplicatePeerID{peerID}
 	}
 
-	// check ips for both the connection addr and the self reported addr
-	if sw.peers.HasIP(pc.RemoteIP()) {
+	// Check for duplicate connection or peer info IP.
+	if sw.peers.HasIP(pc.RemoteIP()) ||
+		sw.peers.HasIP(peerNodeInfo.NetAddress().IP) {
 		return ErrSwitchDuplicatePeerIP{pc.RemoteIP()}
 	}