From 16509ac3dbd41af8189383e5c6b822da5be2806e Mon Sep 17 00:00:00 2001
From: Ethan Buchman <ethan@coinculture.info>
Date: Fri, 5 May 2017 19:19:46 -0400
Subject: [PATCH] p2p: fix race by peer.Start() before peers.Add()

---
 p2p/peer_set.go    |  4 +---
 p2p/switch.go      | 20 ++++++++++++--------
 p2p/switch_test.go |  1 +
 3 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/p2p/peer_set.go b/p2p/peer_set.go
index f3bc1edaf..c5206d2d5 100644
--- a/p2p/peer_set.go
+++ b/p2p/peer_set.go
@@ -16,7 +16,6 @@ type IPeerSet interface {
 
 // PeerSet is a special structure for keeping a table of peers.
 // Iteration over the peers is super fast and thread-safe.
-// We also track how many peers per IP range and avoid too many
 type PeerSet struct {
 	mtx    sync.Mutex
 	lookup map[string]*peerSetItem
@@ -35,8 +34,7 @@ func NewPeerSet() *PeerSet {
 	}
 }
 
-// Returns false if peer with key (PubKeyEd25519) is already in set
-// or if we have too many peers from the peer's IP range
+// Returns false if peer with key (PubKeyEd25519) is already set
 func (ps *PeerSet) Add(peer *Peer) error {
 	ps.mtx.Lock()
 	defer ps.mtx.Unlock()
diff --git a/p2p/switch.go b/p2p/switch.go
index cab446f29..9d778f4bd 100644
--- a/p2p/switch.go
+++ b/p2p/switch.go
@@ -76,8 +76,7 @@ type Switch struct {
 }
 
 var (
-	ErrSwitchDuplicatePeer      = errors.New("Duplicate peer")
-	ErrSwitchMaxPeersPerIPRange = errors.New("IP range has too many peers")
+	ErrSwitchDuplicatePeer = errors.New("Duplicate peer")
 )
 
 func NewSwitch(config *cfg.P2PConfig) *Switch {
@@ -221,12 +220,10 @@ func (sw *Switch) AddPeer(peer *Peer) error {
 		return err
 	}
 
-	// Add the peer to .peers
-	// ignore if duplicate or if we already have too many for that IP range
-	if err := sw.peers.Add(peer); err != nil {
-		sw.Logger.Info("Ignoring peer", "error", err, "peer", peer)
-		peer.Stop()
-		return err
+	// Check for duplicate peer
+	if sw.peers.Has(peer.Key) {
+		return ErrSwitchDuplicatePeer
+
 	}
 
 	// Start peer
@@ -234,6 +231,13 @@ func (sw *Switch) AddPeer(peer *Peer) error {
 		sw.startInitPeer(peer)
 	}
 
+	// Add the peer to .peers.
+	// We start it first so that a peer in the list is safe to Stop.
+	// It should not err since we already checked peers.Has()
+	if err := sw.peers.Add(peer); err != nil {
+		return err
+	}
+
 	sw.Logger.Info("Added peer", "peer", peer)
 	return nil
 }
diff --git a/p2p/switch_test.go b/p2p/switch_test.go
index a6e960cdf..eed7d1fab 100644
--- a/p2p/switch_test.go
+++ b/p2p/switch_test.go
@@ -279,6 +279,7 @@ func TestSwitchReconnectsToPersistentPeer(t *testing.T) {
 	// simulate failure by closing connection
 	peer.CloseConn()
 
+	// TODO: actually detect the disconnection and wait for reconnect
 	time.Sleep(100 * time.Millisecond)
 
 	assert.NotZero(sw.Peers().Size())