From 85d5feafa3db5a878ed81e3548337792f230e6b8 Mon Sep 17 00:00:00 2001
From: Ignotus Peverell <igno.peverell@protonmail.com>
Date: Tue, 2 Oct 2018 07:17:29 -0700
Subject: [PATCH] Fix concurrency issue around peer add and start. Fixes #1585
 (#1633)

* Fix concurrency issue around peer add and start. Fixes #1585
---
 p2p/src/peers.rs | 32 ++++++++++++++++++--------------
 p2p/src/serv.rs  | 22 ++++++++++++----------
 2 files changed, 30 insertions(+), 24 deletions(-)

diff --git a/p2p/src/peers.rs b/p2p/src/peers.rs
index 8d8dd3275..ffd7e903a 100644
--- a/p2p/src/peers.rs
+++ b/p2p/src/peers.rs
@@ -56,25 +56,29 @@ impl Peers {
 
 	/// Adds the peer to our internal peer mapping. Note that the peer is still
 	/// returned so the server can run it.
-	pub fn add_connected(&self, p: Peer) -> Result<Arc<RwLock<Peer>>, Error> {
-		debug!(LOGGER, "Saving newly connected peer {}.", p.info.addr);
-		let peer_data = PeerData {
-			addr: p.info.addr,
-			capabilities: p.info.capabilities,
-			user_agent: p.info.user_agent.clone(),
-			flags: State::Healthy,
-			last_banned: 0,
-			ban_reason: ReasonForBan::None,
-		};
+	pub fn add_connected(&self, peer: Arc<RwLock<Peer>>) -> Result<(), Error> {
+		let peer_data: PeerData;
+		let addr: SocketAddr;
+		{
+			let p = peer.read().unwrap();
+			peer_data = PeerData {
+				addr: p.info.addr,
+				capabilities: p.info.capabilities,
+				user_agent: p.info.user_agent.clone(),
+				flags: State::Healthy,
+				last_banned: 0,
+				ban_reason: ReasonForBan::None,
+			};
+			addr = p.info.addr.clone();
+		}
+		debug!(LOGGER, "Saving newly connected peer {}.", addr);
 		self.save_peer(&peer_data)?;
 
-		let addr = p.info.addr.clone();
-		let apeer = Arc::new(RwLock::new(p));
 		{
 			let mut peers = self.peers.write().unwrap();
-			peers.insert(addr, apeer.clone());
+			peers.insert(addr, peer.clone());
 		}
-		Ok(apeer)
+		Ok(())
 	}
 
 	// Update the dandelion relay
diff --git a/p2p/src/serv.rs b/p2p/src/serv.rs
index 8474e00d8..a278504c4 100644
--- a/p2p/src/serv.rs
+++ b/p2p/src/serv.rs
@@ -178,20 +178,20 @@ impl Server {
 				let addr = SocketAddr::new(self.config.host, self.config.port);
 				let total_diff = self.peers.total_difficulty();
 
-				let peer = Peer::connect(
+				let peer = Arc::new(RwLock::new(Peer::connect(
 					&mut stream,
 					self.capabilities,
 					total_diff,
 					addr,
 					&self.handshake,
 					self.peers.clone(),
-				)?;
-				let added = self.peers.add_connected(peer)?;
+				)?));
 				{
-					let mut peer = added.write().unwrap();
+					let mut peer = peer.write().unwrap();
 					peer.start(stream);
 				}
-				Ok(added)
+				self.peers.add_connected(peer.clone())?;
+				Ok(peer)
 			}
 			Err(e) => {
 				debug!(
@@ -211,16 +211,18 @@ impl Server {
 		let total_diff = self.peers.total_difficulty();
 
 		// accept the peer and add it to the server map
-		let peer = Peer::accept(
+		let peer = Arc::new(RwLock::new(Peer::accept(
 			&mut stream,
 			self.capabilities,
 			total_diff,
 			&self.handshake,
 			self.peers.clone(),
-		)?;
-		let added = self.peers.add_connected(peer)?;
-		let mut peer = added.write().unwrap();
-		peer.start(stream);
+		)?));
+		{
+			let mut peer = peer.write().unwrap();
+			peer.start(stream);
+		}
+		self.peers.add_connected(peer)?;
 		Ok(())
 	}