feat: prevent banning of connected base node in wallet

tari-project · Dec 7, 2021 · 1abe5bb · 1abe5bb
1 parent 53a5174
commit 1abe5bb
Show file tree

Hide file tree

Showing 5 changed files with 52 additions and 6 deletions.
diff --git a/base_layer/wallet/src/wallet.rs b/base_layer/wallet/src/wallet.rs
@@ -293,8 +293,14 @@ where
         );
 
         self.comms.peer_manager().add_peer(peer.clone()).await?;
-        self.wallet_connectivity.set_base_node(peer);
-
+        if let Some(current_node) = self.wallet_connectivity.get_current_base_node_id() {
+            self.comms
+                .connectivity()
+                .remove_peer_from_allow_list(current_node)
+                .await?;
+        }
+        self.wallet_connectivity.set_base_node(peer.clone());
+        self.comms.connectivity().add_peer_to_allow_list(peer.node_id).await?;
         Ok(())
     }
 

diff --git a/common/src/configuration/utils.rs b/common/src/configuration/utils.rs
@@ -170,7 +170,7 @@ pub fn default_config(bootstrap: &ConfigBootstrap) -> Config {
     cfg.set_default("base_node.weatherwax.pruning_horizon", 0).unwrap();
     cfg.set_default("base_node.weatherwax.pruned_mode_cleanup_interval", 50)
         .unwrap();
-    cfg.set_default("base_node.weatherwax.flood_ban_max_msg_count", 1000)
+    cfg.set_default("base_node.weatherwax.flood_ban_max_msg_count", 10000)
         .unwrap();
     cfg.set_default("base_node.weatherwax.peer_seeds", Vec::<String>::new())
         .unwrap();
@@ -215,7 +215,8 @@ pub fn default_config(bootstrap: &ConfigBootstrap) -> Config {
     cfg.set_default("base_node.igor.pruning_horizon", 0).unwrap();
     cfg.set_default("base_node.igor.pruned_mode_cleanup_interval", 50)
         .unwrap();
-    cfg.set_default("base_node.igor.flood_ban_max_msg_count", 1000).unwrap();
+    cfg.set_default("base_node.igor.flood_ban_max_msg_count", 10000)
+        .unwrap();
     cfg.set_default("base_node.igor.grpc_enabled", false).unwrap();
     cfg.set_default("base_node.igor.grpc_base_node_address", "127.0.0.1:18142")
         .unwrap();

diff --git a/comms/src/connectivity/manager.rs b/comms/src/connectivity/manager.rs
@@ -93,6 +93,7 @@ impl ConnectivityManager {
             shutdown_signal: self.shutdown_signal,
             #[cfg(feature = "metrics")]
             uptime: Some(Instant::now()),
+            allow_list: vec![],
         }
         .spawn()
     }
@@ -149,6 +150,7 @@ struct ConnectivityManagerActor {
     shutdown_signal: ShutdownSignal,
     #[cfg(feature = "metrics")]
     uptime: Option<Instant>,
+    allow_list: Vec<NodeId>,
 }
 
 impl ConnectivityManagerActor {
@@ -271,8 +273,25 @@ impl ConnectivityManagerActor {
                 let _ = reply.send(states);
             },
             BanPeer(node_id, duration, reason) => {
-                if let Err(err) = self.ban_peer(&node_id, duration, reason).await {
-                    error!(target: LOG_TARGET, "Error when banning peer: {:?}", err);
+                if !self.allow_list.contains(&node_id) {
+                    if let Err(err) = self.ban_peer(&node_id, duration, reason).await {
+                        error!(target: LOG_TARGET, "Error when banning peer: {:?}", err);
+                    }
+                } else {
+                    info!(
+                        target: LOG_TARGET,
+                        "Peer is excluded from being banned, NodeId: {:?}", node_id
+                    );
+                }
+            },
+            AddPeerToAllowList(node_id) => {
+                if !self.allow_list.contains(&node_id) {
+                    self.allow_list.push(node_id)
+                }
+            },
+            RemovePeerFromAllowList(node_id) => {
+                if let Some(index) = self.allow_list.iter().position(|x| *x == node_id) {
+                    self.allow_list.remove(index);
                 }
             },
             GetActiveConnections(reply) => {

diff --git a/comms/src/connectivity/requester.rs b/comms/src/connectivity/requester.rs
@@ -99,6 +99,8 @@ pub enum ConnectivityRequest {
     GetAllConnectionStates(oneshot::Sender<Vec<PeerConnectionState>>),
     GetActiveConnections(oneshot::Sender<Vec<PeerConnection>>),
     BanPeer(NodeId, Duration, String),
+    AddPeerToAllowList(NodeId),
+    RemovePeerFromAllowList(NodeId),
 }
 
 #[derive(Debug, Clone)]
@@ -243,6 +245,22 @@ impl ConnectivityRequester {
             .await
     }
 
+    pub async fn add_peer_to_allow_list(&mut self, node_id: NodeId) -> Result<(), ConnectivityError> {
+        self.sender
+            .send(ConnectivityRequest::AddPeerToAllowList(node_id))
+            .await
+            .map_err(|_| ConnectivityError::ActorDisconnected)?;
+        Ok(())
+    }
+
+    pub async fn remove_peer_from_allow_list(&mut self, node_id: NodeId) -> Result<(), ConnectivityError> {
+        self.sender
+            .send(ConnectivityRequest::RemovePeerFromAllowList(node_id))
+            .await
+            .map_err(|_| ConnectivityError::ActorDisconnected)?;
+        Ok(())
+    }
+
     pub async fn wait_started(&mut self) -> Result<(), ConnectivityError> {
         let (reply_tx, reply_rx) = oneshot::channel();
         self.sender

diff --git a/comms/src/test_utils/mocks/connectivity_manager.rs b/comms/src/test_utils/mocks/connectivity_manager.rs
@@ -258,6 +258,8 @@ impl ConnectivityManagerMock {
             },
             GetAllConnectionStates(_) => unimplemented!(),
             BanPeer(_, _, _) => {},
+            AddPeerToAllowList(_) => {},
+            RemovePeerFromAllowList(_) => {},
             GetActiveConnections(reply) => {
                 self.state
                     .with_state(|state| reply.send(state.active_conns.values().cloned().collect()).unwrap())