wip

p2p/src/peers.rs

@@ -37,7 +37,7 @@ use crate::types::{
 pub struct Peers {
 	pub adapter: Arc<dyn ChainAdapter>,
 	store: PeerStore,
-	peers: RwLock<HashMap<SocketAddr, Arc<Peer>>>,
+	peers: RwLock<HashMap<String, Arc<Peer>>>,
 	dandelion_relay: RwLock<HashMap<i64, Arc<Peer>>>,
 	config: P2PConfig,
 }
@@ -72,11 +72,8 @@ impl Peers {
 		}
 		debug!("Saving newly connected peer {}.", addr);
 		self.save_peer(&peer_data)?;
+		self.peers.write().insert(peer_key(&addr), peer.clone());
 
-		{
-			let mut peers = self.peers.write();
-			peers.insert(addr, peer.clone());
-		}
 		Ok(())
 	}
 
@@ -131,17 +128,7 @@ impl Peers {
 	}
 
 	pub fn is_known(&self, addr: &SocketAddr) -> bool {
-		self.peers.read().contains_key(addr)
-	}
-
-	/// Check whether an ip address is in the active peers list, ignore the port
-	pub fn is_known_ip(&self, addr: &SocketAddr) -> bool {
-		for socket in self.peers.read().keys() {
-			if addr.ip() == socket.ip() {
-				return true;
-			}
-		}
-		return false;
+		self.peers.read().contains_key(&peer_key(addr))
 	}
 
 	/// Get vec of peers we are currently connected to.
@@ -168,7 +155,7 @@ impl Peers {
 
 	/// Get a peer we're connected to by address.
 	pub fn get_connected_peer(&self, addr: &SocketAddr) -> Option<Arc<Peer>> {
-		self.peers.read().get(addr).map(|p| p.clone())
+		self.peers.read().get(&peer_key(addr)).map(|p| p.clone())
 	}
 
 	/// Number of peers currently connected to.
@@ -259,22 +246,9 @@ impl Peers {
 	}
 
 	pub fn is_banned(&self, peer_addr: SocketAddr) -> bool {
-		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;
-				}
+		if let Ok(peer) = self.store.get_peer(peer_addr) {
+			if peer.flags == State::Banned {
+				return true;
 			}
 		}
 		false
@@ -498,9 +472,10 @@ impl Peers {
 		// now clean up peer map based on the list to remove
 		{
 			let mut peers = self.peers.write();
-			for p in rm {
-				let _ = peers.get(&p).map(|p| p.stop());
-				peers.remove(&p);
+			for ref p in rm {
+				let key = peer_key(p);
+				let _ = peers.get(&key).map(|peer| peer.stop());
+				peers.remove(&key);
 			}
 		}
 	}
@@ -710,3 +685,14 @@ impl NetAdapter for Peers {
 		}
 	}
 }
+
+// TODO - PeerAddr struct for encapsulation and reuse here and in the store.
+// If this is a "loopback" address then we care about the port and construct the key "ip:port".
+// Otherwise we *only* care about the ip and we ignore the port.
+fn peer_key(addr: &SocketAddr) -> String {
+	if addr.ip().is_loopback() {
+		format!("{}:{}", addr.ip(), addr.port())
+	} else {
+		format!("{}", addr.ip())
+	}
+}

p2p/src/serv.rs

@@ -29,7 +29,7 @@ use crate::peer::Peer;
 use crate::peers::Peers;
 use crate::store::PeerStore;
 use crate::types::{
-	Capabilities, ChainAdapter, Error, NetAdapter, P2PConfig, ReasonForBan, Seeding, TxHashSetRead,
+	Capabilities, ChainAdapter, Error, NetAdapter, P2PConfig, ReasonForBan, TxHashSetRead,
 };
 use crate::util::{Mutex, StopState};
 use chrono::prelude::{DateTime, Utc};
@@ -191,13 +191,16 @@ impl Server {
 	/// different sets of peers themselves. In addition, it prevent potential
 	/// duplicate connections, malicious or not.
 	fn check_undesirable(&self, stream: &TcpStream) -> bool {
-		// peer has been banned, go away!
 		if let Ok(peer_addr) = stream.peer_addr() {
-			let banned = self.peers.is_banned(peer_addr);
-			let known_ip =
-				self.peers.is_known_ip(&peer_addr) && self.config.seeding_type == Seeding::DNSSeed;
-			if banned || known_ip {
-				debug!("Peer {} banned or known, refusing connection.", peer_addr);
+			if self.peers.is_banned(peer_addr) {
+				debug!("Peer {} banned, refusing connection.", peer_addr);
+				if let Err(e) = stream.shutdown(Shutdown::Both) {
+					debug!("Error shutting down conn: {:?}", e);
+				}
+				return true;
+			}
+			if self.peers.is_known(&peer_addr) {
+				debug!("Peer {} already known, refusing connection.", peer_addr);
 				if let Err(e) = stream.shutdown(Shutdown::Both) {
 					debug!("Error shutting down conn: {:?}", e);
 				}

p2p/src/store.rs

@@ -29,7 +29,7 @@ use grin_store::{self, option_to_not_found, to_key, Error};
 
 const STORE_SUBPATH: &'static str = "peers";
 
-const PEER_PREFIX: u8 = 'p' as u8;
+const PEER_PREFIX: u8 = 'P' as u8;
 
 /// Types of messages
 enum_from_primitive! {
@@ -162,17 +162,6 @@ 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> {
@@ -226,9 +215,15 @@ impl PeerStore {
 	}
 }
 
+// If this is a "loopback" address then we care about the port and construct the key "ip:port".
+// Otherwise we *only* care about the ip and we ignore the port.
 fn peer_key(peer_addr: SocketAddr) -> Vec<u8> {
-	to_key(
-		PEER_PREFIX,
-		&mut format!("{}:{}", peer_addr.ip(), peer_addr.port()).into_bytes(),
-	)
+	if peer_addr.ip().is_loopback() {
+		to_key(
+			PEER_PREFIX,
+			&mut format!("{}:{}", peer_addr.ip(), peer_addr.port()).into_bytes(),
+		)
+	} else {
+		to_key(PEER_PREFIX, &mut format!("{}", peer_addr.ip()).into_bytes())
+	}
 }