manager: Enhance connection lifecycle with custom connection policies

src/codec/identity.rs

@@ -100,9 +100,11 @@ mod tests {
     fn decoding_smaller_payloads() {
         let mut codec = Identity::new(100);
         let bytes = vec![3u8; 64];
+        let copy = bytes.clone();
         let mut bytes = BytesMut::from(&bytes[..]);
 
-        let decoded = codec.decode(&mut bytes);
+        // The smaller payload will not be decoded as the identity code needs 100 bytes.
+        assert!(codec.decode(&mut bytes).unwrap().is_none());
     }
 
     #[test]

src/config.rs

@@ -29,7 +29,7 @@ use crate::{
         notification, request_response, UserProtocol,
     },
     transport::{
-        manager::limits::ConnectionLimitsConfig, tcp::config::Config as TcpConfig,
+        manager::limits::ConnectionMiddleware, tcp::config::Config as TcpConfig,
         KEEP_ALIVE_TIMEOUT, MAX_PARALLEL_DIALS,
     },
     types::protocol::ProtocolName,
@@ -119,11 +119,11 @@ pub struct ConfigBuilder {
     /// Maximum number of parallel dial attempts.
     max_parallel_dials: usize,
 
-    /// Connection limits config.
-    connection_limits: ConnectionLimitsConfig,
-
     /// Close the connection if no substreams are open within this time frame.
     keep_alive_timeout: Duration,
+
+    /// Connection middleware.
+    connection_middleware: Option<Box<dyn ConnectionMiddleware>>,
 }
 
 impl Default for ConfigBuilder {
@@ -155,8 +155,8 @@ impl ConfigBuilder {
             notification_protocols: HashMap::new(),
             request_response_protocols: HashMap::new(),
             known_addresses: Vec::new(),
-            connection_limits: ConnectionLimitsConfig::default(),
             keep_alive_timeout: KEEP_ALIVE_TIMEOUT,
+            connection_middleware: None,
         }
     }
 
@@ -266,9 +266,9 @@ impl ConfigBuilder {
         self
     }
 
-    /// Set connection limits configuration.
-    pub fn with_connection_limits(mut self, config: ConnectionLimitsConfig) -> Self {
-        self.connection_limits = config;
+    /// Set connection middleware.
+    pub fn with_connection_middleware(mut self, middleware: Box<dyn ConnectionMiddleware>) -> Self {
+        self.connection_middleware = Some(middleware);
         self
     }
 
@@ -305,8 +305,8 @@ impl ConfigBuilder {
             notification_protocols: self.notification_protocols,
             request_response_protocols: self.request_response_protocols,
             known_addresses: self.known_addresses,
-            connection_limits: self.connection_limits,
             keep_alive_timeout: self.keep_alive_timeout,
+            connection_middleware: self.connection_middleware.take(),
         }
     }
 }
@@ -364,9 +364,9 @@ pub struct Litep2pConfig {
     /// Known addresses.
     pub(crate) known_addresses: Vec<(PeerId, Vec<Multiaddr>)>,
 
-    /// Connection limits config.
-    pub(crate) connection_limits: ConnectionLimitsConfig,
-
     /// Close the connection if no substreams are open within this time frame.
     pub(crate) keep_alive_timeout: Duration,
+
+    /// Connection middleware.
+    pub(crate) connection_middleware: Option<Box<dyn ConnectionMiddleware>>,
 }

src/lib.rs

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

src/protocol/libp2p/kademlia/mod.rs

@@ -1230,10 +1230,7 @@ mod tests {
     use crate::{
         codec::ProtocolCodec,
         crypto::ed25519::Keypair,
-        transport::{
-            manager::{limits::ConnectionLimitsConfig, TransportManager},
-            KEEP_ALIVE_TIMEOUT,
-        },
+        transport::{manager::TransportManager, ConnectionLimits, KEEP_ALIVE_TIMEOUT},
         types::protocol::ProtocolName,
         BandwidthSink,
     };
@@ -1251,7 +1248,7 @@ mod tests {
             HashSet::new(),
             BandwidthSink::new(),
             8usize,
-            ConnectionLimitsConfig::default(),
+            Some(Box::new(ConnectionLimits::new(Default::default()))),
         );
 
         let peer = PeerId::random();

src/protocol/mdns.rs

@@ -336,7 +336,7 @@ mod tests {
     use super::*;
     use crate::{
         crypto::ed25519::Keypair,
-        transport::manager::{limits::ConnectionLimitsConfig, TransportManager},
+        transport::{manager::TransportManager, ConnectionLimits},
         BandwidthSink,
     };
     use futures::StreamExt;
@@ -354,7 +354,7 @@ mod tests {
             HashSet::new(),
             BandwidthSink::new(),
             8usize,
-            ConnectionLimitsConfig::default(),
+            Some(Box::new(ConnectionLimits::new(Default::default()))),
         );
 
         let mdns1 = Mdns::new(
@@ -377,7 +377,7 @@ mod tests {
             HashSet::new(),
             BandwidthSink::new(),
             8usize,
-            ConnectionLimitsConfig::default(),
+            Some(Box::new(ConnectionLimits::new(Default::default()))),
         );
 
         let mdns2 = Mdns::new(

src/protocol/notification/tests/mod.rs

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

src/protocol/request_response/tests.rs

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