Browse Source

net: Implement peer connection rejection

parazyd 2 years ago
parent
commit
59f95dd0fb
3 changed files with 32 additions and 1 deletions
  1. 6 0
      src/net/acceptor.rs
  2. 7 1
      src/net/connector.rs
  3. 19 0
      src/net/hosts.rs

+ 6 - 0
src/net/acceptor.rs

@@ -116,6 +116,12 @@ impl Acceptor {
             // Now we wait for a new connection.
             match listener.next().await {
                 Ok((stream, url)) => {
+                    // Check if we reject this peer
+                    if self.session.upgrade().unwrap().p2p().hosts().is_rejected(&url).await {
+                        debug!(target: "net::acceptor::run_accept_loop()", "Peer {} is rejected", url);
+                        continue
+                    }
+
                     // Create the new Channel.
                     let session = self.session.clone();
                     let channel = Channel::new(stream, url, session).await;

+ 7 - 1
src/net/connector.rs

@@ -18,6 +18,7 @@
 
 use std::time::Duration;
 
+use log::debug;
 use url::Url;
 
 use super::{
@@ -26,7 +27,7 @@ use super::{
     settings::SettingsPtr,
     transport::Dialer,
 };
-use crate::Result;
+use crate::{Error, Result};
 
 /// Create outbound socket connections
 pub struct Connector {
@@ -44,6 +45,11 @@ impl Connector {
 
     /// Establish an outbound connection
     pub async fn connect(&self, url: &Url) -> Result<(Url, ChannelPtr)> {
+        if self.session.upgrade().unwrap().p2p().hosts().is_rejected(url).await {
+            debug!(target: "net::connector::connect", "Peer {} is rejected", url);
+            return Err(Error::ConnectFailed)
+        }
+
         let mut endpoint = url.clone();
 
         let transports = &self.settings.allowed_transports;

+ 19 - 0
src/net/hosts.rs

@@ -47,6 +47,9 @@ pub struct Hosts {
     /// Internet interrupt (goblins unplugging cables)
     quarantine: RwLock<HashMap<Url, usize>>,
 
+    /// Peers we reject from connecting
+    rejected: RwLock<HashSet<Url>>,
+
     /// Subscriber listening for store updates
     store_subscriber: SubscriberPtr<usize>,
 
@@ -60,6 +63,7 @@ impl Hosts {
         Arc::new(Self {
             addrs: RwLock::new(HashSet::new()),
             quarantine: RwLock::new(HashMap::new()),
+            rejected: RwLock::new(HashSet::new()),
             store_subscriber: Subscriber::new(),
             settings,
         })
@@ -210,6 +214,21 @@ impl Hosts {
         }
     }
 
+    /// Check if a given peer should be rejected
+    pub async fn is_rejected(&self, peer: &Url) -> bool {
+        self.rejected.read().await.contains(peer)
+    }
+
+    /// Mark a peer as rejected
+    pub async fn mark_rejected(&self, peer: &Url) {
+        self.rejected.write().await.insert(peer.clone());
+    }
+
+    /// Unmark a rejected peer
+    pub async fn unmark_rejected(&self, peer: &Url) {
+        self.rejected.write().await.remove(peer);
+    }
+
     /// Check if the host list is empty.
     pub async fn is_empty(&self) -> bool {
         self.addrs.read().await.is_empty()