Selective peer connection block (#1581)

Author: sanity

Cargo.lock

@@ -46,6 +46,8 @@ async fn base_node_test_config(
     gateways: Vec<String>,
     public_port: Option<u16>,
     ws_api_port: u16,
+    // New parameter to specify addresses this node should block
+    blocked_addresses: Option<Vec<std::net::SocketAddr>>,
 ) -> anyhow::Result<(ConfigArgs, PresetConfig)> {
     if is_gateway {
         assert!(public_port.is_some());
@@ -72,6 +74,9 @@ async fn base_node_test_config(
             address: Some(Ipv4Addr::LOCALHOST.into()),
             network_port: public_port,
             bandwidth_limit: None,
+            // Assuming the new field 'blocked_addresses' is added to NetworkArgs
+            // and it takes Option<Vec<SocketAddr>>
+            blocked_addresses,
         },
         config_paths: {
             freenet::config::ConfigPathsArgs {
@@ -154,6 +159,7 @@ async fn test_ping_multi_node() -> TestResult {
             vec![],
             Some(network_socket_gw.local_addr()?.port()),
             ws_api_port_socket_gw.local_addr()?.port(),
+            None, // blocked_addresses
         )
         .await?;
         let public_port = cfg.network_api.public_port.unwrap();
@@ -168,6 +174,7 @@ async fn test_ping_multi_node() -> TestResult {
         vec![serde_json::to_string(&config_gw_info)?],
         None,
         ws_api_port_socket_node1.local_addr()?.port(),
+        None, // blocked_addresses
     )
     .await?;
     let ws_api_port_node1 = config_node1.ws_api.ws_api_port.unwrap();
@@ -178,6 +185,7 @@ async fn test_ping_multi_node() -> TestResult {
         vec![serde_json::to_string(&config_gw_info)?],
         None,
         ws_api_port_socket_node2.local_addr()?.port(),
+        None, // blocked_addresses
     )
     .await?;
     let ws_api_port_node2 = config_node2.ws_api.ws_api_port.unwrap();
@@ -667,6 +675,7 @@ async fn test_ping_application_loop() -> TestResult {
             vec![],
             Some(network_socket_gw.local_addr()?.port()),
             ws_api_port_socket_gw.local_addr()?.port(),
+            None, // blocked_addresses
         )
         .await?;
         let public_port = cfg.network_api.public_port.unwrap();
@@ -681,6 +690,7 @@ async fn test_ping_application_loop() -> TestResult {
         vec![serde_json::to_string(&config_gw_info)?],
         None,
         ws_api_port_socket_node1.local_addr()?.port(),
+        None, // blocked_addresses
     )
     .await?;
     let ws_api_port_node1 = config_node1.ws_api.ws_api_port.unwrap();
@@ -691,6 +701,7 @@ async fn test_ping_application_loop() -> TestResult {
         vec![serde_json::to_string(&config_gw_info)?],
         None,
         ws_api_port_socket_node2.local_addr()?.port(),
+        None, // blocked_addresses
     )
     .await?;
     let ws_api_port_node2 = config_node2.ws_api.ws_api_port.unwrap();

apps/freenet-ping/Cargo.lock

@@ -97,6 +97,7 @@ impl Default for ConfigArgs {
                 gateways: None,
                 location: None,
                 bandwidth_limit: None,
+                blocked_addresses: None,
             },
             ws_api: WebsocketApiArgs {
                 address: Some(default_listening_address()),
@@ -320,6 +321,10 @@ impl ConfigArgs {
                 public_port: self.network_api.public_port,
                 ignore_protocol_version: self.network_api.ignore_protocol_checking,
                 bandwidth_limit: self.network_api.bandwidth_limit,
+                blocked_addresses: self
+                    .network_api
+                    .blocked_addresses
+                    .map(|addrs| addrs.into_iter().collect()),
             },
             ws_api: WebsocketApiConfig {
                 // the websocket API is always local
@@ -486,6 +491,10 @@ pub struct NetworkArgs {
     /// Hard limit the bandwidth usage for upstream traffic.
     #[arg(long)]
     pub bandwidth_limit: Option<usize>,
+
+    /// List of IP:port addresses to refuse connections to/from.
+    #[arg(long, num_args = 0..)]
+    pub blocked_addresses: Option<Vec<SocketAddr>>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -518,7 +527,7 @@ impl NetworkArgs {
     }
 }
 
-#[derive(Debug, Copy, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct NetworkApiConfig {
     /// Address to listen to locally
     #[serde(default = "default_listening_address", rename = "network-address")]
@@ -545,6 +554,10 @@ pub struct NetworkApiConfig {
 
     /// Hard limit the bandwidth usage for upstream traffic.
     pub bandwidth_limit: Option<usize>,
+
+    /// List of IP:port addresses to refuse connections to/from.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub blocked_addresses: Option<HashSet<SocketAddr>>,
 }
 
 mod port_allocation;

apps/freenet-ping/app/tests/run_app.rs

@@ -125,6 +125,7 @@ pub struct NodeConfig {
     pub(crate) min_number_conn: Option<usize>,
     pub(crate) max_upstream_bandwidth: Option<Rate>,
     pub(crate) max_downstream_bandwidth: Option<Rate>,
+    pub(crate) blocked_addresses: Option<HashSet<SocketAddr>>,
 }
 
 impl NodeConfig {
@@ -170,13 +171,14 @@ impl NodeConfig {
             network_listener_ip: config.network_api.address,
             network_listener_port: config.network_api.port,
             location: config.location.map(Location::new),
-            config: Arc::new(config),
+            config: Arc::new(config.clone()),
             max_hops_to_live: None,
             rnd_if_htl_above: None,
             max_number_conn: None,
             min_number_conn: None,
             max_upstream_bandwidth: None,
             max_downstream_bandwidth: None,
+            blocked_addresses: config.network_api.blocked_addresses.clone(),
         })
     }
 

crates/core/src/config/mod.rs

@@ -48,6 +48,8 @@ pub(crate) enum ConnectionError {
     FailedConnectOp,
     #[error("unwanted connection")]
     UnwantedConnection,
+    #[error("connection to/from address {0} blocked by local policy")]
+    AddressBlocked(std::net::SocketAddr),
 
     // errors produced while handling the connection:
     #[error("IO error: {0}")]
@@ -80,6 +82,7 @@ impl Clone for ConnectionError {
             Self::TransportError(err) => Self::TransportError(err.clone()),
             Self::FailedConnectOp => Self::FailedConnectOp,
             Self::UnwantedConnection => Self::UnwantedConnection,
+            Self::AddressBlocked(addr) => Self::AddressBlocked(*addr),
         }
     }
 }

crates/core/src/node/mod.rs

@@ -49,6 +49,8 @@ pub(super) enum HandshakeError {
     TransportError(#[from] TransportError),
     #[error("receibed an unexpected message at this point: {0}")]
     UnexpectedMessage(Box<NetMessage>),
+    #[error("connection error: {0}")]
+    ConnectionError(#[from] super::ConnectionError),
 }
 
 #[derive(Debug)]
@@ -124,6 +126,7 @@ pub(super) enum ExternConnection {
     Dropped {
         peer: PeerId,
     },
+    DropConnectionByAddr(SocketAddr),
 }
 
 /// Used for communicating with the HandshakeHandler.
@@ -149,6 +152,14 @@ impl HanshakeHandlerMsg {
             .map_err(|_| HandshakeError::ChannelClosed)?;
         Ok(())
     }
+
+    pub async fn drop_connection_by_addr(&self, remote_addr: SocketAddr) -> Result<()> {
+        self.0
+            .send(ExternConnection::DropConnectionByAddr(remote_addr))
+            .await
+            .map_err(|_| HandshakeError::ChannelClosed)?;
+        Ok(())
+    }
 }
 
 type OutboundMessageSender = mpsc::Sender<NetMessage>;
@@ -526,6 +537,12 @@ impl HandshakeHandler {
                         Some(ExternConnection::Dropped { peer }) => {
                             self.connected.remove(&peer.addr);
                             self.outbound_messages.remove(&peer.addr);
+                            self.connecting.remove(&peer.addr);
+                        }
+                        Some(ExternConnection::DropConnectionByAddr(addr)) => {
+                            self.connected.remove(&addr);
+                            self.outbound_messages.remove(&addr);
+                            self.connecting.remove(&addr);
                         }
                         None => return Err(HandshakeError::ChannelClosed),
                     }

crates/core/src/node/network_bridge.rs

@@ -125,6 +125,7 @@ pub(in crate::node) struct P2pConnManager {
     this_location: Option<Location>,
     check_version: bool,
     bandwidth_limit: Option<usize>,
+    blocked_addresses: Option<HashSet<SocketAddr>>,
 }
 
 impl P2pConnManager {
@@ -154,6 +155,7 @@ impl P2pConnManager {
             this_location: config.location,
             check_version: !config.config.network_api.ignore_protocol_version,
             bandwidth_limit: config.config.network_api.bandwidth_limit,
+            blocked_addresses: config.blocked_addresses.clone(),
         })
     }
 
@@ -245,10 +247,6 @@ impl P2pConnManager {
                                 }
                             }
                         }
-
-                        ConnEvent::HandshakeAction(action) => {
-                            self.handle_handshake_action(action, &mut state).await?;
-                        }
                         ConnEvent::ClosedChannel => {
                             tracing::info!("Notification channel closed");
                             break;
@@ -334,7 +332,7 @@ impl P2pConnManager {
         &mut self,
         state: &mut EventListenerState,
         handshake_handler: &mut HandshakeHandler,
-        handshake_handler_msg: &HanshakeHandlerMsg,
+        handshake_handler_msg: &HanshakeHandlerMsg, // already passed here
         notification_channel: &mut EventLoopNotificationsReceiver,
         node_controller: &mut Receiver<NodeEvent>,
         client_wait_for_transaction: &mut ContractHandlerChannel<WaitingResolution>,
@@ -353,8 +351,18 @@ impl P2pConnManager {
             msg = self.conn_bridge_rx.recv() => {
                 Ok(self.handle_bridge_msg(msg))
             }
-            msg = handshake_handler.wait_for_events() => {
-                Ok(self.handle_handshake_msg(msg))
+            handshake_event_res = handshake_handler.wait_for_events() => {
+                match handshake_event_res {
+                    Ok(event) => {
+                        self.handle_handshake_action(event, state, handshake_handler_msg).await?;
+                        Ok(EventResult::Continue)
+                    }
+                    Err(HandshakeError::ChannelClosed) => Ok(EventResult::Event(ConnEvent::ClosedChannel)),
+                    Err(e) => {
+                        tracing::warn!("Handshake error: {:?}", e);
+                        Ok(EventResult::Continue)
+                    }
+                }
             }
             msg = node_controller.recv() => {
                 Ok(self.handle_node_controller_msg(msg))
@@ -460,13 +468,25 @@ impl P2pConnManager {
     async fn handle_connect_peer(
         &mut self,
         peer: PeerId,
-        callback: Box<dyn ConnectResultSender>,
+        mut callback: Box<dyn ConnectResultSender>,
         tx: Transaction,
         handshake_handler_msg: &HanshakeHandlerMsg,
         state: &mut EventListenerState,
         is_gw: bool,
     ) -> anyhow::Result<()> {
         tracing::info!(tx = %tx, remote = %peer, "Connecting to peer");
+        if let Some(blocked_addrs) = &self.blocked_addresses {
+            if blocked_addrs.contains(&peer.addr) {
+                tracing::info!(tx = %tx, remote = %peer.addr, "Outgoing connection to peer blocked by local policy");
+                // Ensure ConnectionError is correctly namespaced if HandshakeError::ConnectionError expects it directly
+                callback
+                    .send_result(Err(HandshakeError::ConnectionError(
+                        crate::node::network_bridge::ConnectionError::AddressBlocked(peer.addr),
+                    )))
+                    .await?;
+                return Ok(());
+            }
+        }
         state.awaiting_connection.insert(peer.addr, callback);
         let res = timeout(
             Duration::from_secs(10),
@@ -492,6 +512,7 @@ impl P2pConnManager {
         &mut self,
         event: HandshakeEvent,
         state: &mut EventListenerState,
+        handshake_handler_msg: &HanshakeHandlerMsg, // Parameter added
     ) -> anyhow::Result<()> {
         match event {
             HandshakeEvent::InboundConnection {
@@ -502,6 +523,16 @@ impl P2pConnManager {
                 op,
                 forward_info,
             } => {
+                if let Some(blocked_addrs) = &self.blocked_addresses {
+                    if blocked_addrs.contains(&joiner.addr) {
+                        tracing::info!(%id, remote = %joiner.addr, "Inbound connection from peer blocked by local policy");
+                        // Not proceeding with adding connection or processing the operation.
+                        handshake_handler_msg
+                            .drop_connection_by_addr(joiner.addr)
+                            .await?;
+                        return Ok(());
+                    }
+                }
                 let (tx, rx) = mpsc::channel(1);
                 self.connections.insert(joiner.clone(), tx);
                 let was_reserved = {
@@ -733,21 +764,15 @@ impl P2pConnManager {
         }
     }
 
-    fn handle_handshake_msg(&self, msg: Result<HandshakeEvent, HandshakeError>) -> EventResult {
-        match msg {
-            Ok(event) => EventResult::Event(ConnEvent::HandshakeAction(event)),
-            Err(HandshakeError::ChannelClosed) => EventResult::Event(ConnEvent::ClosedChannel),
-            _ => EventResult::Continue,
-        }
-    }
-
     fn handle_node_controller_msg(&self, msg: Option<NodeEvent>) -> EventResult {
         match msg {
             Some(msg) => EventResult::Event(ConnEvent::NodeAction(msg)),
             None => EventResult::Event(ConnEvent::ClosedChannel),
         }
     }
 
+    // Removed handle_handshake_msg as it's integrated into wait_for_event
+
     fn handle_client_transaction_subscription(
         &self,
         event_id: Result<(ClientId, WaitingTransaction), anyhow::Error>,
@@ -875,7 +900,6 @@ enum EventResult {
 enum ConnEvent {
     InboundMessage(NetMessage),
     OutboundMessage(NetMessage),
-    HandshakeAction(HandshakeEvent),
     NodeAction(NodeEvent),
     ClosedChannel,
 }

crates/core/src/node/network_bridge/handshake.rs

@@ -71,6 +71,7 @@ async fn base_node_test_config(
             address: Some(Ipv4Addr::LOCALHOST.into()),
             network_port: public_port,
             bandwidth_limit: None,
+            blocked_addresses: None,
         },
         config_paths: {
             freenet::config::ConfigPathsArgs {