swarm: ban connections based on peerid (#1065)

* swarm: ban connections based on peerid

* Update core/src/swarm/swarm.rs

Co-Authored-By: montekki <fedor.sakharov@gmail.com>

* Update core/src/swarm/swarm.rs

Co-Authored-By: montekki <fedor.sakharov@gmail.com>

* remove the testing code

* close connections on ban ban_peer_id

* adds code to unban a peer

* simplify connection closing code

* remove blank line

* ignore DialPeer actions and inject_dial_failure

* Update core/src/swarm/swarm.rs

Co-Authored-By: montekki <fedor.sakharov@gmail.com>

* Update core/src/swarm/swarm.rs

Co-Authored-By: montekki <fedor.sakharov@gmail.com>

* bring back the .expect()
This commit is contained in:
Fedor Sakharov 2019-04-18 19:17:14 +03:00 committed by Pierre Krieger
parent 79c7307b0f
commit 7dc95e78f7

View File

@ -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<TTransport, TBehaviour> = 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<PeerId>,
}
impl<TTransport, TBehaviour, TInEvent, TOutEvent, THandler, THandlerErr> Deref for
@ -210,6 +214,20 @@ where TBehaviour: NetworkBehaviour<ProtocolsHandler = THandler>,
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<TTransport, TBehaviour, TMuxer, TInEvent, TOutEvent, THandler, THandlerErr> Stream for
@ -258,7 +276,14 @@ where TBehaviour: NetworkBehaviour<ProtocolsHandler = THandler>,
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<ProtocolsHandler = THandler>,
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(),
}
}
}