From 6a244e904a6464ec15596e0599e0bd667e4caab3 Mon Sep 17 00:00:00 2001
From: Ignotus Peverell <igno.peverell@protonmail.com>
Date: Thu, 10 Jan 2019 01:22:48 +0000
Subject: [PATCH] Fix add_peer for ban, remove useless disconnect

---
 p2p/src/peers.rs | 16 ++++++++++++++++
 p2p/src/serv.rs  | 18 +++++-------------
 p2p/src/types.rs |  1 +
 3 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/p2p/src/peers.rs b/p2p/src/peers.rs
index 13a12e9f6..1dcd0d90b 100644
--- a/p2p/src/peers.rs
+++ b/p2p/src/peers.rs
@@ -80,6 +80,22 @@ impl Peers {
 		Ok(())
 	}
 
+	/// Add a peer as banned to block future connections, usually due to failed
+	/// handshake
+	pub fn add_banned(&self, addr: SocketAddr, ban_reason: ReasonForBan) -> Result<(), Error> {
+		let peer_data = PeerData {
+			addr,
+			capabilities: Capabilities::UNKNOWN,
+			user_agent: "".to_string(),
+			flags: State::Banned,
+			last_banned: Utc::now().timestamp(),
+			ban_reason,
+			last_connected: Utc::now().timestamp(),
+		};
+		debug!("Banning peer {}.", addr);
+		self.save_peer(&peer_data)
+	}
+
 	// Update the dandelion relay
 	pub fn update_dandelion_relay(&self) {
 		let peers = self.outgoing_connected_peers();
diff --git a/p2p/src/serv.rs b/p2p/src/serv.rs
index 8e8f0ec8b..17644888d 100644
--- a/p2p/src/serv.rs
+++ b/p2p/src/serv.rs
@@ -28,8 +28,8 @@ use crate::core::pow::Difficulty;
 use crate::handshake::Handshake;
 use crate::peer::Peer;
 use crate::peers::Peers;
-use crate::store::{PeerStore, State};
-use crate::types::{Capabilities, ChainAdapter, Error, NetAdapter, P2PConfig, TxHashSetRead};
+use crate::store::PeerStore;
+use crate::types::{Capabilities, ChainAdapter, Error, NetAdapter, P2PConfig, TxHashSetRead, ReasonForBan};
 use crate::util::{Mutex, StopState};
 use chrono::prelude::{DateTime, Utc};
 
@@ -87,7 +87,7 @@ impl Server {
 						let sc = stream.try_clone();
 						if let Err(e) = self.handle_new_peer(stream) {
 							warn!("Error accepting peer {}: {:?}", peer_addr.to_string(), e);
-							let _ = self.peers.update_state(peer_addr, State::Banned);
+							let _ = self.peers.add_banned(peer_addr, ReasonForBan::BadHandshake);
 						} else if let Ok(s) = sc {
 							connected_sockets.insert(peer_addr, s);
 						}
@@ -174,21 +174,13 @@ impl Server {
 		let total_diff = self.peers.total_difficulty();
 
 		// accept the peer and add it to the server map
-		let mut peer = match Peer::accept(
+		let mut peer = Peer::accept(
 			&mut stream,
 			self.capabilities,
 			total_diff,
 			&self.handshake,
 			self.peers.clone(),
-		) {
-			Ok(p) => p,
-			Err(e) => {
-				// in theory, should be shutdown when stream gets dropped,
-				// practically doesn't seem to be happening
-				let _ = stream.shutdown(Shutdown::Both);
-				return Err(e);
-			}
-		};
+		)?;
 		peer.start(stream);
 		self.peers.add_connected(Arc::new(peer))?;
 		Ok(())
diff --git a/p2p/src/types.rs b/p2p/src/types.rs
index 36a674bd3..09bb3d73b 100644
--- a/p2p/src/types.rs
+++ b/p2p/src/types.rs
@@ -244,6 +244,7 @@ enum_from_primitive! {
 		BadTxHashSet = 4,
 		ManualBan = 5,
 		FraudHeight = 6,
+		BadHandshake = 7,
 	}
 }