diff --git a/core/src/swarm/swarm.rs b/core/src/swarm/swarm.rs index 7d372d83..127daebf 100644 --- a/core/src/swarm/swarm.rs +++ b/core/src/swarm/swarm.rs @@ -34,6 +34,7 @@ use crate::{ use futures::prelude::*; use smallvec::SmallVec; use std::{error, fmt, io, ops::{Deref, DerefMut}}; +use std::collections::HashSet; /// Contains the state of the network, plus the way it should behave. pub type Swarm = ExpandedSwarm< @@ -73,6 +74,9 @@ where /// List of multiaddresses we're listening on, after account for external IP addresses and /// similar mechanisms. external_addrs: SmallVec<[Multiaddr; 8]>, + + /// List of nodes for which we deny any incoming connection. + banned_peers: HashSet, } impl Deref for @@ -210,6 +214,20 @@ where TBehaviour: NetworkBehaviour, me.external_addrs.push(addr); } } + + /// Bans a peer by its peer ID. + /// + /// Any incoming connection and any dialing attempt will immediately be rejected. + /// This function has no effect is the peer is already banned. + pub fn ban_peer_id(&mut self, peer_id: PeerId) { + self.banned_peers.insert(peer_id.clone()); + self.raw_swarm.peer(peer_id).into_connected().map(|c| c.close()); + } + + /// Unbans a peer. + pub fn unban_peer_id(&mut self, peer_id: PeerId) { + self.banned_peers.remove(&peer_id); + } } impl Stream for @@ -258,7 +276,14 @@ where TBehaviour: NetworkBehaviour, self.behaviour.inject_node_event(conn_info.peer_id().clone(), event); }, Async::Ready(RawSwarmEvent::Connected { conn_info, endpoint }) => { - self.behaviour.inject_connected(conn_info.peer_id().clone(), endpoint); + if self.banned_peers.contains(conn_info.peer_id()) { + self.raw_swarm.peer(conn_info.peer_id().clone()) + .into_connected() + .expect("the RawSwarm just notified us that we were connected; QED") + .close(); + } else { + self.behaviour.inject_connected(conn_info.peer_id().clone(), endpoint); + } }, Async::Ready(RawSwarmEvent::NodeClosed { conn_info, endpoint, .. }) => { self.behaviour.inject_disconnected(conn_info.peer_id(), endpoint); @@ -313,7 +338,11 @@ where TBehaviour: NetworkBehaviour, let _ = Swarm::dial_addr(self, address); }, Async::Ready(NetworkBehaviourAction::DialPeer { peer_id }) => { - Swarm::dial(self, peer_id) + if self.banned_peers.contains(&peer_id) { + self.behaviour.inject_dial_failure(&peer_id); + } else { + Swarm::dial(self, peer_id); + } }, Async::Ready(NetworkBehaviourAction::SendEvent { peer_id, event }) => { if let Some(mut peer) = self.raw_swarm.peer(peer_id).into_connected() { @@ -444,6 +473,7 @@ where TBehaviour: NetworkBehaviour, supported_protocols, listened_addrs: SmallVec::new(), external_addrs: SmallVec::new(), + banned_peers: HashSet::new(), } } }