feat(stream): add Control::open_stream_on_connection

protocols/stream/CHANGELOG.md

@@ -1,5 +1,8 @@
 ## 0.5.0-alpha
 
+- Add `Control::open_stream_on_connection` to open an outbound stream on a
+  specific connection instead of a randomly chosen one.
+  See [PR 6487](https://github.com/libp2p/rust-libp2p/pull/6487).
 - Raise MSRV to 1.88.0.
   See [PR 6273](https://github.com/libp2p/rust-libp2p/pull/6273).
 

protocols/stream/src/control.rs

@@ -11,7 +11,7 @@ use futures::{
     channel::{mpsc, oneshot},
 };
 use libp2p_identity::PeerId;
-use libp2p_swarm::{Stream, StreamProtocol};
+use libp2p_swarm::{ConnectionId, Stream, StreamProtocol};
 
 use crate::{AlreadyRegistered, handler::NewStream, shared::Shared};
 
@@ -41,27 +41,41 @@ impl Control {
     /// time which is enforced by requiring `&mut self`.
     ///
     /// This backpressure mechanism breaks if you clone [`Control`]s excessively.
+    ///
+    /// > Alternatively, [`Control::open_stream_on_connection`] opens on a
+    /// > specific connection instead of a randomly chosen one.
     pub async fn open_stream(
         &mut self,
         peer: PeerId,
         protocol: StreamProtocol,
     ) -> Result<Stream, OpenStreamError> {
         tracing::debug!(%peer, "Requesting new stream");
 
-        let mut new_stream_sender = Shared::lock(&self.shared).sender(peer);
-
-        let (sender, receiver) = oneshot::channel();
-
-        new_stream_sender
-            .send(NewStream { protocol, sender })
-            .await
-            .map_err(|e| io::Error::new(io::ErrorKind::ConnectionReset, e))?;
-
-        let stream = receiver
-            .await
-            .map_err(|e| io::Error::new(io::ErrorKind::ConnectionReset, e))??;
+        let new_stream_sender = Shared::lock(&self.shared).sender(peer);
+        request_stream(new_stream_sender, protocol).await
+    }
 
-        Ok(stream)
+    /// Like [`Control::open_stream`] but opens on a specific `connection`
+    /// instead of picking one of the peer's connections. The `connection` is
+    /// the [`ConnectionId`] from a [`libp2p_swarm::SwarmEvent::ConnectionEstablished`].
+    /// Errors with [`io::ErrorKind::NotConnected`] if that connection is not, or
+    /// no longer, established.
+    pub async fn open_stream_on_connection(
+        &mut self,
+        peer: PeerId,
+        connection: ConnectionId,
+        protocol: StreamProtocol,
+    ) -> Result<Stream, OpenStreamError> {
+        tracing::debug!(%peer, %connection, "Requesting new stream on connection");
+
+        let new_stream_sender = Shared::lock(&self.shared)
+            .senders
+            .get(&connection)
+            .cloned()
+            .ok_or_else(|| {
+                io::Error::new(io::ErrorKind::NotConnected, "connection not established")
+            })?;
+        request_stream(new_stream_sender, protocol).await
     }
 
     /// Accept inbound streams for the provided protocol.
@@ -75,6 +89,23 @@ impl Control {
     }
 }
 
+/// Send a [`NewStream`] request on `new_stream_sender` and await the opened stream.
+async fn request_stream(
+    mut new_stream_sender: mpsc::Sender<NewStream>,
+    protocol: StreamProtocol,
+) -> Result<Stream, OpenStreamError> {
+    let (sender, receiver) = oneshot::channel();
+
+    new_stream_sender
+        .send(NewStream { protocol, sender })
+        .await
+        .map_err(|e| io::Error::new(io::ErrorKind::ConnectionReset, e))?;
+
+    receiver
+        .await
+        .map_err(|e| io::Error::new(io::ErrorKind::ConnectionReset, e))?
+}
+
 /// Errors while opening a new stream.
 #[derive(Debug)]
 #[non_exhaustive]

protocols/stream/src/shared.rs

@@ -20,7 +20,7 @@ pub(crate) struct Shared {
     supported_inbound_protocols: HashMap<StreamProtocol, mpsc::Sender<(PeerId, Stream)>>,
 
     connections: HashMap<ConnectionId, PeerId>,
-    senders: HashMap<ConnectionId, mpsc::Sender<NewStream>>,
+    pub(crate) senders: HashMap<ConnectionId, mpsc::Sender<NewStream>>,
 
     /// Tracks channel pairs for a peer whilst we are dialing them.
     pending_channels: HashMap<PeerId, (mpsc::Sender<NewStream>, mpsc::Receiver<NewStream>)>,

protocols/stream/tests/lib.rs

@@ -3,7 +3,7 @@ use std::io;
 use futures::{AsyncReadExt as _, AsyncWriteExt as _, StreamExt as _};
 use libp2p_identity::PeerId;
 use libp2p_stream as stream;
-use libp2p_swarm::{StreamProtocol, Swarm};
+use libp2p_swarm::{ConnectionId, StreamProtocol, Swarm, SwarmEvent};
 use libp2p_swarm_test::SwarmExt as _;
 use stream::OpenStreamError;
 use tracing::level_filters::LevelFilter;
@@ -78,3 +78,66 @@ async fn dial_errors_are_propagated() {
     assert_eq!(e.kind(), io::ErrorKind::NotConnected);
     assert_eq!("Dial error: no addresses for peer.", e.to_string());
 }
+
+#[tokio::test]
+async fn open_stream_on_connection_targets_an_existing_connection() {
+    let mut swarm1 = Swarm::new_ephemeral_tokio(|_| stream::Behaviour::new());
+    let mut swarm2 = Swarm::new_ephemeral_tokio(|_| stream::Behaviour::new());
+
+    let mut control = swarm1.behaviour().new_control();
+    let mut incoming = swarm2.behaviour().new_control().accept(PROTOCOL).unwrap();
+
+    let (addr, _) = swarm2.listen().with_memory_addr_external().await;
+    let swarm2_peer_id = *swarm2.local_peer_id();
+
+    tokio::spawn(async move {
+        while let Some((_, mut stream)) = incoming.next().await {
+            stream.write_all(&[42]).await.unwrap();
+            stream.close().await.unwrap();
+        }
+    });
+    tokio::spawn(swarm2.loop_on_next());
+
+    // Dial and capture the established connection id.
+    swarm1.dial(addr).unwrap();
+    let conn_id = loop {
+        if let SwarmEvent::ConnectionEstablished { connection_id, .. } =
+            swarm1.next_swarm_event().await
+        {
+            break connection_id;
+        }
+    };
+    tokio::spawn(swarm1.loop_on_next());
+
+    let mut stream = control
+        .open_stream_on_connection(swarm2_peer_id, conn_id, PROTOCOL)
+        .await
+        .unwrap();
+    let mut buf = [0u8; 1];
+    stream.read_exact(&mut buf).await.unwrap();
+    assert_eq!([42], buf);
+}
+
+#[tokio::test]
+async fn open_stream_on_connection_errors_on_unknown_connection() {
+    let mut swarm1 = Swarm::new_ephemeral_tokio(|_| stream::Behaviour::new());
+    let mut swarm2 = Swarm::new_ephemeral_tokio(|_| stream::Behaviour::new());
+
+    let mut control = swarm1.behaviour().new_control();
+    swarm2.listen().with_memory_addr_external().await;
+    swarm1.connect(&mut swarm2).await;
+    let swarm2_peer_id = *swarm2.local_peer_id();
+
+    tokio::spawn(swarm1.loop_on_next());
+    tokio::spawn(swarm2.loop_on_next());
+
+    let bogus = ConnectionId::new_unchecked(99999);
+    let error = control
+        .open_stream_on_connection(swarm2_peer_id, bogus, PROTOCOL)
+        .await
+        .unwrap_err();
+    let OpenStreamError::Io(e) = error else {
+        panic!("unexpected error: {error}")
+    };
+    assert_eq!(e.kind(), io::ErrorKind::NotConnected);
+}