From ea7eea3f84b0abdddb36ecc81b0a8a3fad4c9202 Mon Sep 17 00:00:00 2001
From: Gary Yu <gairy.yu@gmail.com>
Date: Mon, 31 Dec 2018 07:15:58 +0800
Subject: [PATCH] Skip connecting PeerWithSelf by checking ip address (#2253)

* skip the connecting to self
---
 p2p/src/handshake.rs | 16 +++++++++++++++-
 p2p/src/serv.rs      | 18 +++++++++++-------
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/p2p/src/handshake.rs b/p2p/src/handshake.rs
index 0b8b3c2ff..041c81267 100644
--- a/p2p/src/handshake.rs
+++ b/p2p/src/handshake.rs
@@ -14,7 +14,7 @@
 
 use crate::util::RwLock;
 use std::collections::VecDeque;
-use std::net::{SocketAddr, TcpStream};
+use std::net::{Shutdown, SocketAddr, TcpStream};
 use std::sync::Arc;
 
 use chrono::prelude::*;
@@ -29,6 +29,7 @@ use crate::peer::Peer;
 use crate::types::{Capabilities, Direction, Error, P2PConfig, PeerInfo, PeerLiveInfo};
 
 const NONCES_CAP: usize = 100;
+const ADDRS_CAP: usize = 10;
 
 /// Handles the handshake negotiation when two peers connect and decides on
 /// protocol.
@@ -36,6 +37,8 @@ pub struct Handshake {
 	/// Ring buffer of nonces sent to detect self connections without requiring
 	/// a node id.
 	nonces: Arc<RwLock<VecDeque<u64>>>,
+	/// Ring buffer of self addr(s) collected from PeerWithSelf detection (by nonce).
+	pub addrs: Arc<RwLock<VecDeque<SocketAddr>>>,
 	/// The genesis block header of the chain seen by this node.
 	/// We only want to connect to other nodes seeing the same chain (forks are
 	/// ok).
@@ -48,6 +51,7 @@ impl Handshake {
 	pub fn new(genesis: Hash, config: P2PConfig) -> Handshake {
 		Handshake {
 			nonces: Arc::new(RwLock::new(VecDeque::with_capacity(NONCES_CAP))),
+			addrs: Arc::new(RwLock::new(VecDeque::with_capacity(ADDRS_CAP))),
 			genesis,
 			config,
 		}
@@ -145,7 +149,17 @@ impl Handshake {
 		} else {
 			// check the nonce to see if we are trying to connect to ourselves
 			let nonces = self.nonces.read();
+			let addr = extract_ip(&hand.sender_addr.0, &conn);
 			if nonces.contains(&hand.nonce) {
+				// save ip addresses of ourselves
+				let mut addrs = self.addrs.write();
+				addrs.push_back(addr);
+				if addrs.len() >= ADDRS_CAP {
+					addrs.pop_front();
+				}
+				if let Err(e) = conn.shutdown(Shutdown::Both) {
+					debug!("Error shutting down conn: {:?}", e);
+				}
 				return Err(Error::PeerWithSelf);
 			}
 		}
diff --git a/p2p/src/serv.rs b/p2p/src/serv.rs
index 75bb3a701..fb34aaa44 100644
--- a/p2p/src/serv.rs
+++ b/p2p/src/serv.rs
@@ -22,6 +22,7 @@ use crate::lmdb;
 
 use crate::core::core;
 use crate::core::core::hash::Hash;
+use crate::core::global;
 use crate::core::pow::Difficulty;
 use crate::handshake::Handshake;
 use crate::peer::Peer;
@@ -108,13 +109,16 @@ impl Server {
 			return Err(Error::ConnectionClose);
 		}
 
-		// check ip and port to see if we are trying to connect to ourselves
-		// todo: this can't detect all cases of PeerWithSelf, for example config.host is '0.0.0.0'
-		//
-		if self.config.port == addr.port()
-			&& (addr.ip().is_loopback() || addr.ip() == self.config.host)
-		{
-			return Err(Error::PeerWithSelf);
+		if global::is_production_mode() {
+			let hs = self.handshake.clone();
+			let addrs = hs.addrs.read();
+			if addrs.contains(&addr) {
+				debug!(
+					"connect: ignore the connecting to PeerWithSelf, addr: {}",
+					addr
+				);
+				return Err(Error::PeerWithSelf);
+			}
 		}
 
 		if let Some(p) = self.peers.get_connected_peer(addr) {