diff --git a/p2p/src/peers.rs b/p2p/src/peers.rs index 8926220936..fcbdea90f7 100644 --- a/p2p/src/peers.rs +++ b/p2p/src/peers.rs @@ -37,7 +37,7 @@ use crate::types::{ pub struct Peers { pub adapter: Arc, store: PeerStore, - peers: RwLock>>, + peers: RwLock>>, dandelion_relay: RwLock>>, 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> { - 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()) + } +} diff --git a/p2p/src/serv.rs b/p2p/src/serv.rs index 45fc383542..cb4b8e5d23 100644 --- a/p2p/src/serv.rs +++ b/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); } diff --git a/p2p/src/store.rs b/p2p/src/store.rs index 921070706a..46d10c9ce6 100644 --- a/p2p/src/store.rs +++ b/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 { - self.db - .iter::(&to_key( - PEER_PREFIX, - &mut format!("{}", peer_addr.ip()).into_bytes(), - )) - .unwrap() - .collect::>() - } - /// List all known peers /// Used for /v1/peers/all api endpoint pub fn all_peers(&self) -> Vec { @@ -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 { - 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()) + } }