diff --git a/protocols/stream/CHANGELOG.md b/protocols/stream/CHANGELOG.md index 6b9bbc51325..298ec14c77f 100644 --- a/protocols/stream/CHANGELOG.md +++ b/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). diff --git a/protocols/stream/src/control.rs b/protocols/stream/src/control.rs index bea15778a18..36fb3713e56 100644 --- a/protocols/stream/src/control.rs +++ b/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,6 +41,9 @@ 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, @@ -48,20 +51,31 @@ impl Control { ) -> Result { 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 { + 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, + protocol: StreamProtocol, +) -> Result { + 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] diff --git a/protocols/stream/src/shared.rs b/protocols/stream/src/shared.rs index d75f0504689..5dee2e273d7 100644 --- a/protocols/stream/src/shared.rs +++ b/protocols/stream/src/shared.rs @@ -20,7 +20,7 @@ pub(crate) struct Shared { supported_inbound_protocols: HashMap>, connections: HashMap, - senders: HashMap>, + pub(crate) senders: HashMap>, /// Tracks channel pairs for a peer whilst we are dialing them. pending_channels: HashMap, mpsc::Receiver)>, diff --git a/protocols/stream/tests/lib.rs b/protocols/stream/tests/lib.rs index a00a00a8f62..76517b01481 100644 --- a/protocols/stream/tests/lib.rs +++ b/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); +}