From 9d7a72027ccd01ad295fac4082161cf707635e64 Mon Sep 17 00:00:00 2001
From: Pierre Krieger <pierre.krieger1708@gmail.com>
Date: Mon, 29 Apr 2019 12:05:20 +0200
Subject: [PATCH] Fix #2403 (#2404)

* Fix #2403

* Apply suggestions from code review

Co-Authored-By: tomaka <pierre.krieger1708@gmail.com>
---
 .../src/custom_proto/behaviour.rs             | 73 ++++++++++---------
 1 file changed, 40 insertions(+), 33 deletions(-)

diff --git a/substrate/core/network-libp2p/src/custom_proto/behaviour.rs b/substrate/core/network-libp2p/src/custom_proto/behaviour.rs
index ca1363fab1..eb50325dc0 100644
--- a/substrate/core/network-libp2p/src/custom_proto/behaviour.rs
+++ b/substrate/core/network-libp2p/src/custom_proto/behaviour.rs
@@ -612,31 +612,25 @@ where
 	}
 
 	fn inject_connected(&mut self, peer_id: PeerId, connected_point: ConnectedPoint) {
-		match (self.peers.entry(peer_id), connected_point) {
-			(Entry::Occupied(mut entry), connected_point) => {
-				match mem::replace(entry.get_mut(), PeerState::Poisoned) {
-					PeerState::Requested | PeerState::PendingRequest { .. } |
-					PeerState::Banned { .. } => {
-						debug!(target: "sub-libp2p", "Libp2p => Connected({:?}): Connection \
-							requested by PSM (through {:?})", entry.key(), connected_point);
-						debug!(target: "sub-libp2p", "Handler({:?}) <= Enable", entry.key());
-						self.events.push(NetworkBehaviourAction::SendEvent {
-							peer_id: entry.key().clone(),
-							event: CustomProtoHandlerIn::Enable(connected_point.clone().into()),
-						});
-						*entry.into_mut() = PeerState::Enabled { open: false, connected_point };
-					}
-					st @ _ => {
-						// This is a serious bug either in this state machine or in libp2p.
-						error!(target: "sub-libp2p", "Received inject_connected for \
-							already-connected node; state is {:?}", st);
-						*entry.into_mut() = st;
-						return
-					}
-				}
+		match (self.peers.entry(peer_id.clone()).or_insert(PeerState::Poisoned), connected_point) {
+			(st @ &mut PeerState::Requested, connected_point) |
+			(st @ &mut PeerState::PendingRequest { .. }, connected_point) => {
+				debug!(target: "sub-libp2p", "Libp2p => Connected({:?}): Connection \
+					requested by PSM (through {:?})", peer_id, connected_point
+				);
+				debug!(target: "sub-libp2p", "Handler({:?}) <= Enable", peer_id);
+				self.events.push(NetworkBehaviourAction::SendEvent {
+					peer_id: peer_id.clone(),
+					event: CustomProtoHandlerIn::Enable(connected_point.clone().into()),
+				});
+				*st = PeerState::Enabled { open: false, connected_point };
 			}
 
-			(Entry::Vacant(entry), connected_point @ ConnectedPoint::Listener { .. }) => {
+			// Note: it may seem weird that "Banned" nodes get treated as if there were absent.
+			// This is because the word "Banned" means "temporarily prevent outgoing connections to
+			// this node", and not "banned" in the sense that we would refuse the node altogether.
+			(st @ &mut PeerState::Poisoned, connected_point @ ConnectedPoint::Listener { .. }) |
+			(st @ &mut PeerState::Banned { .. }, connected_point @ ConnectedPoint::Listener { .. }) => {
 				let incoming_id = self.next_incoming_index.clone();
 				self.next_incoming_index.0 = match self.next_incoming_index.0.checked_add(1) {
 					Some(v) => v,
@@ -646,27 +640,40 @@ where
 					}
 				};
 				debug!(target: "sub-libp2p", "Libp2p => Connected({:?}): Incoming connection",
-					entry.key());
+					peer_id);
 				debug!(target: "sub-libp2p", "PSM <= Incoming({:?}, {:?}): Through {:?}",
-					incoming_id, entry.key(), connected_point);
-				self.peerset.incoming(entry.key().clone(), incoming_id);
+					incoming_id, peer_id, connected_point);
+				self.peerset.incoming(peer_id.clone(), incoming_id);
 				self.incoming.push(IncomingPeer {
-					peer_id: entry.key().clone(),
+					peer_id: peer_id.clone(),
 					alive: true,
 					incoming_id,
 				});
-				entry.insert(PeerState::Incoming { connected_point });
+				*st = PeerState::Incoming { connected_point };
 			}
 
-			(Entry::Vacant(entry), connected_point) => {
+			(st @ &mut PeerState::Poisoned, connected_point) |
+			(st @ &mut PeerState::Banned { .. }, connected_point) => {
+				let banned_until = if let PeerState::Banned { until } = st {
+					Some(*until)
+				} else {
+					None
+				};
 				debug!(target: "sub-libp2p", "Libp2p => Connected({:?}): Requested by something \
-					else than PSM, disabling", entry.key());
-				debug!(target: "sub-libp2p", "Handler({:?}) <= Disable", entry.key());
+					else than PSM, disabling", peer_id);
+				debug!(target: "sub-libp2p", "Handler({:?}) <= Disable", peer_id);
 				self.events.push(NetworkBehaviourAction::SendEvent {
-					peer_id: entry.key().clone(),
+					peer_id: peer_id.clone(),
 					event: CustomProtoHandlerIn::Disable,
 				});
-				entry.insert(PeerState::Disabled { open: false, connected_point, banned_until: None });
+				*st = PeerState::Disabled { open: false, connected_point, banned_until };
+			}
+
+			st => {
+				// This is a serious bug either in this state machine or in libp2p.
+				error!(target: "sub-libp2p", "Received inject_connected for \
+					already-connected node; state is {:?}", st
+				);
 			}
 		}
 	}