transport/manager: Add connection limits for inbound and outbound established connections

src/config.rs

@@ -29,9 +29,9 @@ use crate::{
         notification, request_response, UserProtocol,
     },
     transport::{
-        quic::config::Config as QuicConfig, tcp::config::Config as TcpConfig,
-        webrtc::config::Config as WebRtcConfig, websocket::config::Config as WebSocketConfig,
-        MAX_PARALLEL_DIALS,
+        manager::limits::ConnectionLimitsConfig, quic::config::Config as QuicConfig,
+        tcp::config::Config as TcpConfig, webrtc::config::Config as WebRtcConfig,
+        websocket::config::Config as WebSocketConfig, MAX_PARALLEL_DIALS,
     },
     types::protocol::ProtocolName,
     PeerId,
@@ -109,6 +109,9 @@ pub struct ConfigBuilder {
 
     /// Maximum number of parallel dial attempts.
     max_parallel_dials: usize,
+
+    /// Connection limits config.
+    connection_limits: ConnectionLimitsConfig,
 }
 
 impl Default for ConfigBuilder {
@@ -137,6 +140,7 @@ impl ConfigBuilder {
             notification_protocols: HashMap::new(),
             request_response_protocols: HashMap::new(),
             known_addresses: Vec::new(),
+            connection_limits: ConnectionLimitsConfig::default(),
         }
     }
 
@@ -243,6 +247,12 @@ impl ConfigBuilder {
         self
     }
 
+    /// Set connection limits configuration.
+    pub fn with_connection_limits(mut self, config: ConnectionLimitsConfig) -> Self {
+        self.connection_limits = config;
+        self
+    }
+
     /// Build [`Litep2pConfig`].
     pub fn build(mut self) -> Litep2pConfig {
         let keypair = match self.keypair {
@@ -267,6 +277,7 @@ impl ConfigBuilder {
             notification_protocols: self.notification_protocols,
             request_response_protocols: self.request_response_protocols,
             known_addresses: self.known_addresses,
+            connection_limits: self.connection_limits,
         }
     }
 }
@@ -320,4 +331,7 @@ pub struct Litep2pConfig {
 
     /// Known addresses.
     pub(crate) known_addresses: Vec<(PeerId, Vec<Multiaddr>)>,
+
+    /// Connection limits config.
+    pub(crate) connection_limits: ConnectionLimitsConfig,
 }

src/error.rs

@@ -28,6 +28,7 @@
 
 use crate::{
     protocol::Direction,
+    transport::manager::limits::ConnectionLimitsError,
     types::{protocol::ProtocolName, ConnectionId, SubstreamId},
     PeerId,
 };
@@ -118,6 +119,8 @@ pub enum Error {
     ChannelClogged,
     #[error("Connection doesn't exist: `{0:?}`")]
     ConnectionDoesntExist(ConnectionId),
+    #[error("Exceeded connection limits `{0:?}`")]
+    ConnectionLimit(ConnectionLimitsError),
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -243,6 +246,12 @@ impl From<quinn::ConnectionError> for Error {
     }
 }
 
+impl From<ConnectionLimitsError> for Error {
+    fn from(error: ConnectionLimitsError) -> Self {
+        Error::ConnectionLimit(error)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/lib.rs

@@ -143,6 +143,7 @@ impl Litep2p {
             supported_transports,
             bandwidth_sink.clone(),
             litep2p_config.max_parallel_dials,
+            litep2p_config.connection_limits,
         );
 
         // add known addresses to `TransportManager`, if any exist

src/protocol/libp2p/kademlia/mod.rs

@@ -897,8 +897,11 @@ mod tests {
 
     use super::*;
     use crate::{
-        codec::ProtocolCodec, crypto::ed25519::Keypair, transport::manager::TransportManager,
-        types::protocol::ProtocolName, BandwidthSink,
+        codec::ProtocolCodec,
+        crypto::ed25519::Keypair,
+        transport::manager::{limits::ConnectionLimitsConfig, TransportManager},
+        types::protocol::ProtocolName,
+        BandwidthSink,
     };
     use tokio::sync::mpsc::channel;
 
@@ -914,6 +917,7 @@ mod tests {
             HashSet::new(),
             BandwidthSink::new(),
             8usize,
+            ConnectionLimitsConfig::default(),
         );
 
         let peer = PeerId::random();

src/protocol/mdns.rs

@@ -334,7 +334,11 @@ impl Mdns {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{crypto::ed25519::Keypair, transport::manager::TransportManager, BandwidthSink};
+    use crate::{
+        crypto::ed25519::Keypair,
+        transport::manager::{limits::ConnectionLimitsConfig, TransportManager},
+        BandwidthSink,
+    };
     use futures::StreamExt;
     use multiaddr::Protocol;
 
@@ -350,6 +354,7 @@ mod tests {
             HashSet::new(),
             BandwidthSink::new(),
             8usize,
+            ConnectionLimitsConfig::default(),
         );
 
         let mdns1 = Mdns::new(
@@ -372,6 +377,7 @@ mod tests {
             HashSet::new(),
             BandwidthSink::new(),
             8usize,
+            ConnectionLimitsConfig::default(),
         );
 
         let mdns2 = Mdns::new(

src/protocol/notification/tests/mod.rs

@@ -29,7 +29,7 @@ use crate::{
         },
         InnerTransportEvent, ProtocolCommand, TransportService,
     },
-    transport::manager::TransportManager,
+    transport::manager::{limits::ConnectionLimitsConfig, TransportManager},
     types::protocol::ProtocolName,
     BandwidthSink, PeerId,
 };
@@ -53,6 +53,7 @@ fn make_notification_protocol() -> (
         HashSet::new(),
         BandwidthSink::new(),
         8usize,
+        ConnectionLimitsConfig::default(),
     );
 
     let peer = PeerId::random();

src/protocol/request_response/tests.rs

@@ -29,7 +29,7 @@ use crate::{
         InnerTransportEvent, TransportService,
     },
     substream::Substream,
-    transport::manager::TransportManager,
+    transport::manager::{limits::ConnectionLimitsConfig, TransportManager},
     types::{RequestId, SubstreamId},
     BandwidthSink, Error, PeerId, ProtocolName,
 };
@@ -51,6 +51,7 @@ fn protocol() -> (
         HashSet::new(),
         BandwidthSink::new(),
         8usize,
+        ConnectionLimitsConfig::default(),
     );
 
     let peer = PeerId::random();

src/transport/manager/limits.rs

@@ -0,0 +1,204 @@
+// Copyright 2024 litep2p developers
+//
+// Permission is hereby granted, free of charge, to any person obtaining a
+// copy of this software and associated documentation files (the "Software"),
+// to deal in the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+//! Limits for the transport manager.
+
+use crate::types::ConnectionId;
+
+use std::collections::HashSet;
+
+/// Configuration for the connection limits.
+#[derive(Debug, Clone, Default)]
+pub struct ConnectionLimitsConfig {
+    /// Maximum number of incoming connections that can be established.
+    max_incoming_connections: Option<usize>,
+    /// Maximum number of outgoing connections that can be established.
+    max_outgoing_connections: Option<usize>,
+}
+
+impl ConnectionLimitsConfig {
+    /// Configures the maximum number of incoming connections that can be established.
+    pub fn max_incoming_connections(mut self, limit: Option<usize>) -> Self {
+        self.max_incoming_connections = limit;
+        self
+    }
+
+    /// Configures the maximum number of outgoing connections that can be established.
+    pub fn max_outgoing_connections(mut self, limit: Option<usize>) -> Self {
+        self.max_outgoing_connections = limit;
+        self
+    }
+}
+
+/// Error type for connection limits.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ConnectionLimitsError {
+    /// Maximum number of incoming connections exceeded.
+    MaxIncomingConnectionsExceeded,
+    /// Maximum number of outgoing connections exceeded.
+    MaxOutgoingConnectionsExceeded,
+}
+
+/// Connection limits.
+#[derive(Debug, Clone)]
+pub struct ConnectionLimits {
+    /// Configuration for the connection limits.
+    config: ConnectionLimitsConfig,
+
+    /// Established incoming connections.
+    incoming_connections: HashSet<ConnectionId>,
+    /// Established outgoing connections.
+    outgoing_connections: HashSet<ConnectionId>,
+}
+
+impl ConnectionLimits {
+    /// Creates a new connection limits instance.
+    pub fn new(config: ConnectionLimitsConfig) -> Self {
+        let max_incoming_connections = config.max_incoming_connections.unwrap_or(0);
+        let max_outgoing_connections = config.max_outgoing_connections.unwrap_or(0);
+
+        Self {
+            config,
+            incoming_connections: HashSet::with_capacity(max_incoming_connections),
+            outgoing_connections: HashSet::with_capacity(max_outgoing_connections),
+        }
+    }
+
+    /// Called when dialing an address.
+    ///
+    /// Returns the number of outgoing connections permitted to be established.
+    /// It is guaranteed that at least one connection can be established if the method returns `Ok`.
+    /// The number of available outgoing connections can influence the maximum parallel dials to a
+    /// single address.
+    ///
+    /// If the maximum number of outgoing connections is not set, `Ok(usize::MAX)` is returned.
+    pub fn on_dial_address(&mut self) -> Result<usize, ConnectionLimitsError> {
+        if let Some(max_outgoing_connections) = self.config.max_outgoing_connections {
+            if self.outgoing_connections.len() >= max_outgoing_connections {
+                return Err(ConnectionLimitsError::MaxOutgoingConnectionsExceeded);
+            }
+
+            return Ok(max_outgoing_connections - self.outgoing_connections.len());
+        }
+
+        Ok(usize::MAX)
+    }
+
+    /// Called when a new connection is established.
+    pub fn on_connection_established(
+        &mut self,
+        connection_id: ConnectionId,
+        is_listener: bool,
+    ) -> Result<(), ConnectionLimitsError> {
+        // Check connection limits.
+        if is_listener {
+            if let Some(max_incoming_connections) = self.config.max_incoming_connections {
+                if self.incoming_connections.len() >= max_incoming_connections {
+                    return Err(ConnectionLimitsError::MaxIncomingConnectionsExceeded);
+                }
+            }
+        } else {
+            if let Some(max_outgoing_connections) = self.config.max_outgoing_connections {
+                if self.outgoing_connections.len() >= max_outgoing_connections {
+                    return Err(ConnectionLimitsError::MaxOutgoingConnectionsExceeded);
+                }
+            }
+        }
+
+        // Keep track of the connection.
+        if is_listener {
+            if self.config.max_incoming_connections.is_some() {
+                self.incoming_connections.insert(connection_id);
+            }
+        } else {
+            if self.config.max_outgoing_connections.is_some() {
+                self.outgoing_connections.insert(connection_id);
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Called when a connection is closed.
+    pub fn on_connection_closed(&mut self, connection_id: ConnectionId) {
+        self.incoming_connections.remove(&connection_id);
+        self.outgoing_connections.remove(&connection_id);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::types::ConnectionId;
+
+    #[test]
+    fn connection_limits() {
+        let config = ConnectionLimitsConfig::default()
+            .max_incoming_connections(Some(3))
+            .max_outgoing_connections(Some(2));
+        let mut limits = ConnectionLimits::new(config);
+
+        let connection_id_in_1 = ConnectionId::random();
+        let connection_id_in_2 = ConnectionId::random();
+        let connection_id_out_1 = ConnectionId::random();
+        let connection_id_out_2 = ConnectionId::random();
+        let connection_id_in_3 = ConnectionId::random();
+        let connection_id_out_3 = ConnectionId::random();
+
+        // Establish incoming connection.
+        assert!(limits.on_connection_established(connection_id_in_1, true).is_ok());
+        assert_eq!(limits.incoming_connections.len(), 1);
+
+        assert!(limits.on_connection_established(connection_id_in_2, true).is_ok());
+        assert_eq!(limits.incoming_connections.len(), 2);
+
+        assert!(limits.on_connection_established(connection_id_in_3, true).is_ok());
+        assert_eq!(limits.incoming_connections.len(), 3);
+
+        assert_eq!(
+            limits.on_connection_established(ConnectionId::random(), true).unwrap_err(),
+            ConnectionLimitsError::MaxIncomingConnectionsExceeded
+        );
+        assert_eq!(limits.incoming_connections.len(), 3);
+
+        // Establish outgoing connection.
+        assert!(limits.on_connection_established(connection_id_out_1, false).is_ok());
+        assert_eq!(limits.incoming_connections.len(), 3);
+        assert_eq!(limits.outgoing_connections.len(), 1);
+
+        assert!(limits.on_connection_established(connection_id_out_2, false).is_ok());
+        assert_eq!(limits.incoming_connections.len(), 3);
+        assert_eq!(limits.outgoing_connections.len(), 2);
+
+        assert_eq!(
+            limits.on_connection_established(connection_id_out_3, false).unwrap_err(),
+            ConnectionLimitsError::MaxOutgoingConnectionsExceeded
+        );
+
+        // Close connections with peer a.
+        limits.on_connection_closed(connection_id_in_1);
+        assert_eq!(limits.incoming_connections.len(), 2);
+        assert_eq!(limits.outgoing_connections.len(), 2);
+
+        limits.on_connection_closed(connection_id_out_1);
+        assert_eq!(limits.incoming_connections.len(), 2);
+        assert_eq!(limits.outgoing_connections.len(), 1);
+    }
+}