From 79cc9a4baeddeb6142023732ab132cf6f3852c3b Mon Sep 17 00:00:00 2001
From: Jakub Sztandera <kubuxu@protocol.ai>
Date: Mon, 7 Jun 2021 18:35:39 +0200
Subject: [PATCH] Fix race in adding connections to connsByPeer

Side A could already be running Identify with side B before B had the
connection in the connsByPeer, this lead to Connectedness returning
NotConnected and Identify failing.

Signed-off-by: Jakub Sztandera <kubuxu@protocol.ai>
---
 p2p/net/mock/mock_peernet.go | 47 +++++++++++++++++++++++-------------
 1 file changed, 30 insertions(+), 17 deletions(-)

diff --git a/p2p/net/mock/mock_peernet.go b/p2p/net/mock/mock_peernet.go
index 426c28fdc1..67409e7dbd 100644
--- a/p2p/net/mock/mock_peernet.go
+++ b/p2p/net/mock/mock_peernet.go
@@ -160,11 +160,41 @@ func (pn *peernet) connect(p peer.ID) (*conn, error) {
 func (pn *peernet) openConn(r peer.ID, l *link) *conn {
 	lc, rc := l.newConnPair(pn)
 	log.Debugf("%s opening connection to %s", pn.LocalPeer(), lc.RemotePeer())
+	addConnPair(pn, rc.net, lc, rc)
+
 	go rc.net.remoteOpenedConn(rc)
 	pn.addConn(lc)
 	return lc
 }
 
+// addConnPair adds connection to both peernets at the same time
+// must be followerd by pn1.addConn(c1) and pn2.addConn(c2)
+func addConnPair(pn1, pn2 *peernet, c1, c2 *conn) {
+	pn1.Lock()
+	pn2.Lock()
+
+	add := func(pn *peernet, c *conn) {
+		_, found := pn.connsByPeer[c.RemotePeer()]
+		if !found {
+			pn.connsByPeer[c.RemotePeer()] = map[*conn]struct{}{}
+		}
+		pn.connsByPeer[c.RemotePeer()][c] = struct{}{}
+
+		_, found = pn.connsByLink[c.link]
+		if !found {
+			pn.connsByLink[c.link] = map[*conn]struct{}{}
+		}
+		pn.connsByLink[c.link][c] = struct{}{}
+	}
+	add(pn1, c1)
+	add(pn2, c2)
+
+	c1.notifLk.Lock()
+	c2.notifLk.Lock()
+	pn2.Unlock()
+	pn1.Unlock()
+}
+
 func (pn *peernet) remoteOpenedConn(c *conn) {
 	log.Debugf("%s accepting connection from %s", pn.LocalPeer(), c.RemotePeer())
 	pn.addConn(c)
@@ -174,24 +204,7 @@ func (pn *peernet) remoteOpenedConn(c *conn) {
 // addConn constructs and adds a connection
 // to given remote peer over given link
 func (pn *peernet) addConn(c *conn) {
-	pn.Lock()
-
-	_, found := pn.connsByPeer[c.RemotePeer()]
-	if !found {
-		pn.connsByPeer[c.RemotePeer()] = map[*conn]struct{}{}
-	}
-	pn.connsByPeer[c.RemotePeer()][c] = struct{}{}
-
-	_, found = pn.connsByLink[c.link]
-	if !found {
-		pn.connsByLink[c.link] = map[*conn]struct{}{}
-	}
-	pn.connsByLink[c.link][c] = struct{}{}
-
-	c.notifLk.Lock()
 	defer c.notifLk.Unlock()
-	pn.Unlock()
-
 	// Call this after unlocking as it might cause us to immediately close
 	// the connection and remove it from the swarm.
 	c.setup()