From 23df6fa97667e9e419089cbbf65667f94e6d2289 Mon Sep 17 00:00:00 2001
From: Gary Yu <gairy.yu@gmail.com>
Date: Tue, 1 Jan 2019 11:43:05 +0800
Subject: [PATCH] fix: check_banned must care about ip address only, and ignore
 the peer port info (#2267)

---
 p2p/src/peers.rs | 19 ++++++++++++++++---
 p2p/src/store.rs | 11 +++++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/p2p/src/peers.rs b/p2p/src/peers.rs
index 271fa566..b918e26b 100644
--- a/p2p/src/peers.rs
+++ b/p2p/src/peers.rs
@@ -208,9 +208,22 @@ impl Peers {
 	}
 
 	pub fn is_banned(&self, peer_addr: SocketAddr) -> bool {
-		if let Ok(peer_data) = self.store.get_peer(peer_addr) {
-			if peer_data.flags == State::Banned {
-				return true;
+		if global::is_production_mode() {
+			// Ban only cares about ip address, no mather what port.
+			// so, we query all saved peers with one same ip address, and ignore port
+			let peers_data = self.store.find_peers_by_ip(peer_addr);
+			for peer_data in peers_data {
+				if peer_data.flags == State::Banned {
+					return true;
+				}
+			}
+		} else {
+			// For travis-ci test, we need run multiple nodes in one server, with same ip address.
+			// so, just query the ip address and the port
+			if let Ok(peer_data) = self.store.get_peer(peer_addr) {
+				if peer_data.flags == State::Banned {
+					return true;
+				}
 			}
 		}
 		false
diff --git a/p2p/src/store.rs b/p2p/src/store.rs
index 20bee786..92107070 100644
--- a/p2p/src/store.rs
+++ b/p2p/src/store.rs
@@ -162,6 +162,17 @@ impl PeerStore {
 		peers.iter().take(count).cloned().collect()
 	}
 
+	/// Query all peers with same IP address, and ignore the port
+	pub fn find_peers_by_ip(&self, peer_addr: SocketAddr) -> Vec<PeerData> {
+		self.db
+			.iter::<PeerData>(&to_key(
+				PEER_PREFIX,
+				&mut format!("{}", peer_addr.ip()).into_bytes(),
+			))
+			.unwrap()
+			.collect::<Vec<_>>()
+	}
+
 	/// List all known peers
 	/// Used for /v1/peers/all api endpoint
 	pub fn all_peers(&self) -> Vec<PeerData> {