Resolve "Batch unban network command"

massalabs · Oct 6, 2021 · b9d8622 · b9d8622
1 parent a8469df
commit b9d8622
Show file tree

Hide file tree

Showing 6 changed files with 13 additions and 16 deletions.
diff --git a/api-private/src/lib.rs b/api-private/src/lib.rs
@@ -164,12 +164,7 @@ impl MassaPrivate for ApiMassaPrivate {
 
     fn unban(&self, ips: Vec<IpAddr>) -> BoxFuture<Result<(), PrivateApiError>> {
         let network_command_sender = self.network_command_sender.clone();
-        let closure = async move || {
-            for ip in ips {
-                network_command_sender.unban(ip).await?
-            }
-            Ok(())
-        };
+        let closure = async move || Ok(network_command_sender.unban(ips).await?);
         Box::pin(closure())
     }
 }
diff --git a/communication/src/network/network_controller.rs b/communication/src/network/network_controller.rs
@@ -151,9 +151,9 @@ impl NetworkCommandSender {
         Ok(())
     }
 
-    pub async fn unban(&self, ip: IpAddr) -> Result<(), CommunicationError> {
+    pub async fn unban(&self, ips: Vec<IpAddr>) -> Result<(), CommunicationError> {
         self.0
-            .send(NetworkCommand::Unban(ip))
+            .send(NetworkCommand::Unban(ips))
             .await
             .map_err(|_| CommunicationError::ChannelError("could not send Unban command".into()))?;
         Ok(())

diff --git a/communication/src/network/network_worker.rs b/communication/src/network/network_worker.rs
@@ -51,7 +51,7 @@ pub enum NetworkCommand {
     GetPeers(oneshot::Sender<Peers>),
     GetBootstrapPeers(oneshot::Sender<BootstrapPeers>),
     Ban(NodeId),
-    Unban(IpAddr),
+    Unban(Vec<IpAddr>),
     BlockNotFound {
         node: NodeId,
         block_id: BlockId,

diff --git a/communication/src/network/peer_info_database.rs b/communication/src/network/peer_info_database.rs
@@ -310,11 +310,13 @@ impl PeerInfoDatabase {
         res
     }
 
-    pub async fn unban(&mut self, ip: IpAddr) -> Result<(), CommunicationError> {
-        if let Some(peer) = self.peers.get_mut(&ip) {
-            peer.banned = false;
-        } else {
-            return Ok(());
+    pub async fn unban(&mut self, ips: Vec<IpAddr>) -> Result<(), CommunicationError> {
+        for ip in ips.into_iter() {
+            if let Some(peer) = self.peers.get_mut(&ip) {
+                peer.banned = false;
+            } else {
+                return Ok(());
+            }
         }
         cleanup_peers(
             &self.cfg,

diff --git a/communication/src/network/tests/scenarios.rs b/communication/src/network/tests/scenarios.rs
@@ -278,7 +278,7 @@ async fn test_peer_ban() {
 
             //unban connection1.
             network_command_sender
-                .unban(mock_addr.ip())
+                .unban(vec![mock_addr.ip()])
                 .await
                 .expect("error during send unban command.");
 

diff --git a/massa-node/src/main.rs b/massa-node/src/main.rs
@@ -553,7 +553,7 @@ async fn on_api_event(
         ApiEvent::Unban(ip) => {
             massa_trace!("massa-node.main.run.select.api_event.unban", {});
             network_command_sender
-                .unban(ip)
+                .unban(vec![ip])
                 .await
                 .expect("unban failed in api_event_receiver.wait_event")
         }