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()

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<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(),
         }
     }
 }