transport: Add accept_pending/reject_pending for inbound connections and introduce inbound limits

src/transport/dummy.rs

@@ -77,6 +77,14 @@ impl Transport for DummyTransport {
         Ok(())
     }
 
+    fn accept_pending(&mut self, _connection_id: ConnectionId) -> crate::Result<()> {
+        Ok(())
+    }
+
+    fn reject_pending(&mut self, _connection_id: ConnectionId) -> crate::Result<()> {
+        Ok(())
+    }
+
     fn reject(&mut self, _: ConnectionId) -> crate::Result<()> {
         Ok(())
     }

src/transport/manager/limits.rs

@@ -101,6 +101,17 @@ impl ConnectionLimits {
         Ok(usize::MAX)
     }
 
+    /// Called before accepting a new incoming connection.
+    pub fn on_incoming(&mut self) -> Result<(), ConnectionLimitsError> {
+        if let Some(max_incoming_connections) = self.config.max_incoming_connections {
+            if self.incoming_connections.len() >= max_incoming_connections {
+                return Err(ConnectionLimitsError::MaxIncomingConnectionsExceeded);
+            }
+        }
+
+        Ok(())
+    }
+
     /// Called when a new connection is established.
     pub fn on_connection_established(
         &mut self,
@@ -114,11 +125,9 @@ impl ConnectionLimits {
                     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);
-                }
+        } else if let Some(max_outgoing_connections) = self.config.max_outgoing_connections {
+            if self.outgoing_connections.len() >= max_outgoing_connections {
+                return Err(ConnectionLimitsError::MaxOutgoingConnectionsExceeded);
             }
         }
 
@@ -127,10 +136,8 @@ impl ConnectionLimits {
             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);
-            }
+        } else if self.config.max_outgoing_connections.is_some() {
+            self.outgoing_connections.insert(connection_id);
         }
 
         Ok(())

src/transport/manager/mod.rs

@@ -508,18 +508,14 @@ impl TransportManager {
             record.set_connection_id(connection_id);
 
             #[cfg(feature = "quic")]
-            if address.iter().find(|p| std::matches!(p, Protocol::QuicV1)).is_some() {
+            if address.iter().any(|p| std::matches!(&p, Protocol::QuicV1)) {
                 quic.push(address.clone());
                 transports.insert(SupportedTransport::Quic);
                 continue;
             }
 
             #[cfg(feature = "websocket")]
-            if address
-                .iter()
-                .find(|p| std::matches!(p, Protocol::Ws(_) | Protocol::Wss(_)))
-                .is_some()
-            {
+            if address.iter().any(|p| std::matches!(&p, Protocol::Ws(_) | Protocol::Wss(_))) {
                 websocket.push(address.clone());
                 transports.insert(SupportedTransport::WebSocket);
                 continue;
@@ -839,6 +835,11 @@ impl TransportManager {
         }
     }
 
+    fn on_pending_incoming_connection(&mut self) -> crate::Result<()> {
+        self.connection_limits.on_incoming()?;
+        Ok(())
+    }
+
     /// Handle closed connection.
     fn on_connection_closed(
         &mut self,
@@ -1713,7 +1714,34 @@ impl TransportManager {
                                 }
                                 Ok(None) => {}
                             }
-                        }
+                        },
+                        TransportEvent::PendingInboundConnection { connection_id } => {
+                            if self.on_pending_incoming_connection().is_ok() {
+                                tracing::trace!(
+                                    target: LOG_TARGET,
+                                    ?connection_id,
+                                    "accept pending incoming connection",
+                                );
+
+                                let _ = self
+                                    .transports
+                                    .get_mut(&transport)
+                                    .expect("transport to exist")
+                                    .accept_pending(connection_id);
+                            } else {
+                                tracing::debug!(
+                                    target: LOG_TARGET,
+                                    ?connection_id,
+                                    "reject pending incoming connection",
+                                );
+
+                                let _ = self
+                                    .transports
+                                    .get_mut(&transport)
+                                    .expect("transport to exist")
+                                    .reject_pending(connection_id);
+                            }
+                        },
                         event => panic!("event not supported: {event:?}"),
                     }
                 },
@@ -2563,7 +2591,7 @@ mod tests {
 
             peer_context.state = PeerState::Connected {
                 record,
-                dial_record: dial_record,
+                dial_record,
             };
         }
 

src/transport/mod.rs

@@ -39,6 +39,8 @@ pub mod websocket;
 pub(crate) mod dummy;
 pub(crate) mod manager;
 
+pub use manager::limits::{ConnectionLimitsConfig, ConnectionLimitsError};
+
 /// Timeout for opening a connection.
 pub(crate) const CONNECTION_OPEN_TIMEOUT: Duration = Duration::from_secs(10);
 
@@ -121,6 +123,11 @@ pub(crate) enum TransportEvent {
         endpoint: Endpoint,
     },
 
+    PendingInboundConnection {
+        /// Connection ID.
+        connection_id: ConnectionId,
+    },
+
     /// Connection opened to remote but not yet negotiated.
     ConnectionOpened {
         /// Connection ID.
@@ -176,6 +183,12 @@ pub(crate) trait Transport: Stream + Unpin + Send {
     /// Accept negotiated connection.
     fn accept(&mut self, connection_id: ConnectionId) -> crate::Result<()>;
 
+    /// Accept pending connection.
+    fn accept_pending(&mut self, connection_id: ConnectionId) -> crate::Result<()>;
+
+    /// Reject pending connection.
+    fn reject_pending(&mut self, connection_id: ConnectionId) -> crate::Result<()>;
+
     /// Reject negotiated connection.
     fn reject(&mut self, connection_id: ConnectionId) -> crate::Result<()>;
 

src/transport/quic/mod.rs

@@ -36,7 +36,7 @@ use crate::{
 
 use futures::{future::BoxFuture, stream::FuturesUnordered, Stream, StreamExt};
 use multiaddr::{Multiaddr, Protocol};
-use quinn::{ClientConfig, Connection, Endpoint, IdleTimeout};
+use quinn::{ClientConfig, Connecting, Connection, Endpoint, IdleTimeout};
 
 use std::{
     collections::{HashMap, HashSet},
@@ -80,6 +80,9 @@ pub(crate) struct QuicTransport {
     /// Pending dials.
     pending_dials: HashMap<ConnectionId, Multiaddr>,
 
+    /// Pending inbound connections.
+    pending_inbound_connections: HashMap<ConnectionId, Connecting>,
+
     /// Pending connections.
     pending_connections:
         FuturesUnordered<BoxFuture<'static, (ConnectionId, Result<NegotiatedConnection, Error>)>>,
@@ -110,6 +113,22 @@ impl QuicTransport {
         Some(p2p_cert.peer_id())
     }
 
+    /// Handle inbound accepted connection.
+    fn on_inbound_connection(&mut self, connection_id: ConnectionId, connection: Connecting) {
+        self.pending_connections.push(Box::pin(async move {
+            let connection = match connection.await {
+                Ok(connection) => connection,
+                Err(error) => return (connection_id, Err(error.into())),
+            };
+
+            let Some(peer) = Self::extract_peer_id(&connection) else {
+                return (connection_id, Err(Error::InvalidCertificate));
+            };
+
+            (connection_id, Ok(NegotiatedConnection { peer, connection }))
+        }));
+    }
+
     /// Handle established connection.
     fn on_connection_established(
         &mut self,
@@ -193,6 +212,7 @@ impl TransportBuilder for QuicTransport {
                 opened_raw: HashMap::new(),
                 pending_open: HashMap::new(),
                 pending_dials: HashMap::new(),
+                pending_inbound_connections: HashMap::new(),
                 pending_raw_connections: FuturesUnordered::new(),
                 pending_connections: FuturesUnordered::new(),
             },
@@ -291,6 +311,23 @@ impl Transport for QuicTransport {
             .map_or(Err(Error::ConnectionDoesntExist(connection_id)), |_| Ok(()))
     }
 
+    fn accept_pending(&mut self, connection_id: ConnectionId) -> crate::Result<()> {
+        let connection = self
+            .pending_inbound_connections
+            .remove(&connection_id)
+            .ok_or(Error::ConnectionDoesntExist(connection_id))?;
+
+        self.on_inbound_connection(connection_id, connection);
+
+        Ok(())
+    }
+
+    fn reject_pending(&mut self, connection_id: ConnectionId) -> crate::Result<()> {
+        self.pending_inbound_connections
+            .remove(&connection_id)
+            .map_or(Err(Error::ConnectionDoesntExist(connection_id)), |_| Ok(()))
+    }
+
     fn open(
         &mut self,
         connection_id: ConnectionId,
@@ -406,26 +443,19 @@ impl Stream for QuicTransport {
     type Item = TransportEvent;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        while let Poll::Ready(Some(connection)) = self.listener.poll_next_unpin(cx) {
+        if let Poll::Ready(Some(connection)) = self.listener.poll_next_unpin(cx) {
             let connection_id = self.context.next_connection_id();
 
             tracing::trace!(
                 target: LOG_TARGET,
                 ?connection_id,
-                "accept connection",
+                "pending inbound connection",
             );
 
-            self.pending_connections.push(Box::pin(async move {
-                let connection = match connection.await {
-                    Ok(connection) => connection,
-                    Err(error) => return (connection_id, Err(error.into())),
-                };
+            self.pending_inbound_connections.insert(connection_id, connection);
 
-                let Some(peer) = Self::extract_peer_id(&connection) else {
-                    return (connection_id, Err(Error::InvalidCertificate));
-                };
-
-                (connection_id, Ok(NegotiatedConnection { peer, connection }))
+            return Poll::Ready(Some(TransportEvent::PendingInboundConnection {
+                connection_id,
             }));
         }
 
@@ -545,6 +575,15 @@ mod tests {
         ));
 
         transport2.dial(ConnectionId::new(), listen_address).unwrap();
+
+        let event = transport1.next().await.unwrap();
+        match event {
+            TransportEvent::PendingInboundConnection { connection_id } => {
+                transport1.accept_pending(connection_id).unwrap();
+            }
+            _ => panic!("unexpected event"),
+        }
+
         let (res1, res2) = tokio::join!(transport1.next(), transport2.next());
 
         assert!(std::matches!(

src/transport/tcp/connection.rs

@@ -247,7 +247,10 @@ impl TcpConnection {
         })
         .await
         {
-            Err(_) => Err(Error::Timeout),
+            Err(_) => {
+                tracing::trace!(target: LOG_TARGET, ?connection_id, "connection timed out during negotiation");
+                Err(Error::Timeout)
+            }
             Ok(result) => result,
         }
     }