Enforce a timeout on inbound kad substreams using StreamSet

Author: teor2345

Cargo.lock

@@ -121,7 +121,8 @@ libp2p-yamux = { version = "0.47.0", path = "muxers/yamux" }
 asynchronous-codec = { version = "0.7.0" }
 env_logger = "0.11"
 futures = "0.3.30"
-futures-bounded = { version = "0.3.0", features = ["futures-timer"] }
+# TODO: replace with version = "0.3.1" once released upstream
+futures-bounded = { git = "https://github.com/thomaseizinger/rust-futures-bounded", rev = "012803d343b5c604e65d3c238a8cd7a145616447", features = ["futures-timer"] }
 futures-rustls = { version = "0.26.0", default-features = false }
 getrandom = "0.2"
 if-watch = "3.2.1"

Cargo.toml

@@ -27,8 +27,8 @@ use std::{
 };
 
 use either::Either;
-use futures::{channel::oneshot, prelude::*, stream::SelectAll};
-use futures_bounded::Delay;
+use futures::{channel::oneshot, prelude::*};
+use futures_bounded::{Delay, StreamSet};
 use libp2p_core::{upgrade, ConnectedPoint};
 use libp2p_identity::PeerId;
 use libp2p_swarm::{
@@ -78,7 +78,8 @@ pub struct Handler {
     pending_messages: VecDeque<(KadRequestMsg, QueryId)>,
 
     /// List of active inbound substreams with the state they are in.
-    inbound_substreams: SelectAll<InboundSubstreamState>,
+    /// The streams are typed `InboundSubstreamState`, but the set uses the item type.
+    inbound_substreams: StreamSet<ConnectionHandlerEvent<ProtocolConfig, (), HandlerEvent>>,
 
     /// The connected endpoint of the connection that the handler
     /// is associated with.
@@ -120,8 +121,6 @@ enum InboundSubstreamState {
     PendingFlush(UniqueConnecId, KadInStreamSink<Stream>),
     /// The substream is being closed.
     Closing(KadInStreamSink<Stream>),
-    /// The substream was cancelled in favor of a new one.
-    Cancelled,
 
     Poisoned {
         phantom: PhantomData<QueryId>,
@@ -174,9 +173,6 @@ impl InboundSubstreamState {
             | InboundSubstreamState::Closing(substream) => {
                 *self = InboundSubstreamState::Closing(substream);
             }
-            InboundSubstreamState::Cancelled => {
-                *self = InboundSubstreamState::Cancelled;
-            }
             InboundSubstreamState::Poisoned { .. } => unreachable!(),
         }
     }
@@ -462,7 +458,10 @@ impl Handler {
             endpoint,
             remote_peer_id,
             next_connec_unique_id: UniqueConnecId(0),
-            inbound_substreams: Default::default(),
+            inbound_substreams: StreamSet::new(
+                move || Delay::futures_timer(substreams_timeout),
+                MAX_NUM_STREAMS,
+            ),
             outbound_substreams: futures_bounded::FuturesTupleSet::new(
                 move || Delay::futures_timer(substreams_timeout),
                 MAX_NUM_STREAMS,
@@ -519,38 +518,45 @@ impl Handler {
             });
         }
 
-        if self.inbound_substreams.len() == MAX_NUM_STREAMS {
-            if let Some(s) = self.inbound_substreams.iter_mut().find(|s| {
-                matches!(
-                    s,
-                    // An inbound substream waiting to be reused.
-                    InboundSubstreamState::WaitingMessage { first: false, .. }
-                )
-            }) {
-                *s = InboundSubstreamState::Cancelled;
+        let connec_unique_id = self.next_connec_unique_id;
+        self.next_connec_unique_id.0 += 1;
+        let new_substream = InboundSubstreamState::WaitingMessage {
+            first: true,
+            connection_id: connec_unique_id,
+            substream: protocol,
+        };
+
+        if self.inbound_substreams.len() >= MAX_NUM_STREAMS {
+            if let Some(s) = self
+                .inbound_substreams
+                .iter_mut_of_type::<InboundSubstreamState>()
+                .find(|s| {
+                    matches!(
+                        **s,
+                        // An inbound substream waiting to be reused.
+                        InboundSubstreamState::WaitingMessage { first: false, .. }
+                    )
+                })
+            {
+                *s.get_mut() = new_substream;
                 tracing::debug!(
                     peer=?self.remote_peer_id,
                     "New inbound substream to peer exceeds inbound substream limit. \
-                    Removed older substream waiting to be reused."
+                    Replacing older substream that was waiting to be reused."
                 )
             } else {
                 tracing::warn!(
                     peer=?self.remote_peer_id,
                     "New inbound substream to peer exceeds inbound substream limit. \
                      No older substream waiting to be reused. Dropping new substream."
                 );
-                return;
             }
+        } else {
+            self.inbound_substreams
+                .try_push(new_substream)
+                .map_err(|_| ())
+                .expect("Just checked that stream set is not full; qed");
         }
-
-        let connec_unique_id = self.next_connec_unique_id;
-        self.next_connec_unique_id.0 += 1;
-        self.inbound_substreams
-            .push(InboundSubstreamState::WaitingMessage {
-                first: true,
-                connection_id: connec_unique_id,
-                substream: protocol,
-            });
     }
 
     /// Takes the given [`KadRequestMsg`] and composes it into an outbound request-response protocol
@@ -617,15 +623,15 @@ impl ConnectionHandler for Handler {
             HandlerIn::Reset(request_id) => {
                 if let Some(state) = self
                     .inbound_substreams
-                    .iter_mut()
-                    .find(|state| match state {
+                    .iter_mut_of_type::<InboundSubstreamState>()
+                    .find(|state| match **state {
                         InboundSubstreamState::WaitingBehaviour(conn_id, _, _) => {
-                            conn_id == &request_id.connec_unique_id
+                            conn_id == request_id.connec_unique_id
                         }
                         _ => false,
                     })
                 {
-                    state.close();
+                    state.get_mut().close();
                 }
             }
             HandlerIn::FindNodeReq { key, query_id } => {
@@ -764,8 +770,16 @@ impl ConnectionHandler for Handler {
                 Poll::Pending => {}
             }
 
-            if let Poll::Ready(Some(event)) = self.inbound_substreams.poll_next_unpin(cx) {
-                return Poll::Ready(event);
+            if let Poll::Ready(Some(event_result)) = self.inbound_substreams.poll_next_unpin(cx) {
+                match event_result {
+                    Ok(event) => return Poll::Ready(event),
+                    Err(_stream_set_timeout) => {
+                        tracing::trace!(
+                            "Inbound substream timed out waiting for peer, send, or close"
+                        );
+                        continue;
+                    }
+                }
             }
 
             if self.outbound_substreams.len() < MAX_NUM_STREAMS {
@@ -849,8 +863,11 @@ fn compute_new_protocol_status(
 
 impl Handler {
     fn answer_pending_request(&mut self, request_id: RequestId, mut msg: KadResponseMsg) {
-        for state in self.inbound_substreams.iter_mut() {
-            match state.try_answer_with(request_id, msg) {
+        for state in self
+            .inbound_substreams
+            .iter_mut_of_type::<InboundSubstreamState>()
+        {
+            match state.get_mut().try_answer_with(request_id, msg) {
                 Ok(()) => return,
                 Err(m) => {
                     msg = m;
@@ -1007,7 +1024,6 @@ impl futures::Stream for InboundSubstreamState {
                     }
                 },
                 InboundSubstreamState::Poisoned { .. } => unreachable!(),
-                InboundSubstreamState::Cancelled => return Poll::Ready(None),
             }
         }
     }