From 555222d6d7527a1da8188814b0f308ee1562cb5d Mon Sep 17 00:00:00 2001 From: Ultradesu Date: Thu, 16 Jul 2026 18:35:21 +0300 Subject: [PATCH] Added bytestream ALPN --- crates/federation-net/src/config.rs | 57 +++- crates/federation-net/src/engine.rs | 328 ++++++++++++++++++++- crates/federation-net/src/error.rs | 9 + crates/federation-net/src/lib.rs | 3 +- crates/federation-net/tests/integration.rs | 109 +++++++ crates/music-dht/src/config.rs | 14 + crates/music-dht/src/lib.rs | 5 +- crates/music-dht/src/service.rs | 41 ++- crates/music-dht/tests/integration.rs | 69 +++++ 9 files changed, 630 insertions(+), 5 deletions(-) diff --git a/crates/federation-net/src/config.rs b/crates/federation-net/src/config.rs index 0b52c22..6c1436d 100644 --- a/crates/federation-net/src/config.rs +++ b/crates/federation-net/src/config.rs @@ -4,7 +4,7 @@ use std::path::PathBuf; use std::time::Duration; use crate::error::{NetworkError, Result}; -use crate::protocol::{NetworkId, SchemaId}; +use crate::protocol::{ALPN, NetworkId, SchemaId}; use crate::rendezvous::RendezvousConfig; /// Default maximum size of an encoded domain message (256 KiB). @@ -38,6 +38,9 @@ pub struct NetworkConfig { /// Automatic peer discovery over the mainline DHT; `None` disables it and /// peers are connected via tickets only. pub rendezvous: Option, + /// Auxiliary ALPN protocols on which this peer accepts raw byte streams + /// (see [`crate::NetworkEngine::stream_acceptor`]). + pub stream_protocols: Vec>, } impl NetworkConfig { @@ -61,6 +64,7 @@ pub struct NetworkConfigBuilder { event_channel_capacity: Option, max_concurrent_streams_per_peer: Option, rendezvous: Option, + stream_protocols: Vec>, } impl NetworkConfigBuilder { @@ -112,6 +116,14 @@ impl NetworkConfigBuilder { self } + /// Declares an auxiliary ALPN protocol on which this peer accepts raw + /// byte streams (see [`crate::NetworkEngine::stream_acceptor`]). May be + /// called multiple times, once per protocol. + pub fn stream_protocol(mut self, alpn: impl Into>) -> Self { + self.stream_protocols.push(alpn.into()); + self + } + /// Validates the configuration and builds a [`NetworkConfig`]. pub fn build(self) -> Result { let data_dir = self @@ -180,6 +192,25 @@ impl NetworkConfigBuilder { } } + for (index, alpn) in self.stream_protocols.iter().enumerate() { + if alpn.is_empty() { + return Err(NetworkError::InvalidConfig( + "stream protocol ALPN must not be empty".into(), + )); + } + if alpn.as_slice() == ALPN { + return Err(NetworkError::InvalidConfig( + "stream protocol ALPN must differ from the engine ALPN".into(), + )); + } + if self.stream_protocols[..index].contains(alpn) { + return Err(NetworkError::InvalidConfig(format!( + "duplicate stream protocol ALPN: {}", + String::from_utf8_lossy(alpn) + ))); + } + } + Ok(NetworkConfig { data_dir, network_id, @@ -189,6 +220,7 @@ impl NetworkConfigBuilder { event_channel_capacity, max_concurrent_streams_per_peer, rendezvous: self.rendezvous, + stream_protocols: self.stream_protocols, }) } } @@ -272,4 +304,27 @@ mod tests { }; assert!(base_builder().rendezvous(zero_ttl).build().is_err()); } + + #[test] + fn stream_protocols_are_validated() { + let config = base_builder().build().expect("valid config"); + assert!(config.stream_protocols.is_empty()); + + let config = base_builder() + .stream_protocol("my-app/blob/1") + .stream_protocol("my-app/other/1") + .build() + .expect("valid config"); + assert_eq!(config.stream_protocols.len(), 2); + + assert!(base_builder().stream_protocol("").build().is_err()); + assert!(base_builder().stream_protocol(ALPN).build().is_err()); + assert!( + base_builder() + .stream_protocol("dup/1") + .stream_protocol("dup/1") + .build() + .is_err() + ); + } } diff --git a/crates/federation-net/src/engine.rs b/crates/federation-net/src/engine.rs index 329d43c..299806b 100644 --- a/crates/federation-net/src/engine.rs +++ b/crates/federation-net/src/engine.rs @@ -50,6 +50,10 @@ const REJECT_LINGER: Duration = Duration::from_secs(3); const SHUTDOWN_TASK_GRACE: Duration = Duration::from_secs(5); /// Upper bound of the rendezvous round delay while no peer is connected yet. const RENDEZVOUS_LONELY_INTERVAL: Duration = Duration::from_secs(15); +/// Capacity of the queue of accepted-but-not-yet-consumed incoming byte +/// streams, per stream protocol. A full queue delays the handshake ack of +/// further incoming streams (natural backpressure). +const STREAM_ACCEPT_QUEUE: usize = 16; /// Bounds required of a domain message type. /// @@ -77,6 +81,76 @@ struct PeerState { generation: u64, } +/// A raw bidirectional byte stream to a peer, running over a dedicated +/// connection on an auxiliary ALPN. +/// +/// Returned by [`NetworkEngine::open_stream`] (outgoing) and by +/// [`StreamAcceptor::accept`] (incoming). Both sides passed the regular +/// network handshake, so the peer is authenticated and belongs to the same +/// network and schema. The connection behind the stream closes when both +/// halves are dropped. +pub struct ByteStream { + /// The authenticated peer on the other side of the stream. + pub peer_id: EndpointId, + /// The sending half of the stream. + pub send: SendStream, + /// The receiving half of the stream. + pub recv: RecvStream, + /// Keeps the dedicated connection alive for the lifetime of the stream. + _connection: Connection, +} + +impl fmt::Debug for ByteStream { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ByteStream") + .field("peer_id", &self.peer_id) + .finish_non_exhaustive() + } +} + +/// Receiver of incoming [`ByteStream`]s for one declared stream protocol. +/// +/// Obtained once per protocol from [`NetworkEngine::stream_acceptor`]. +/// Dropping the acceptor makes the engine refuse further incoming streams +/// on that protocol. +pub struct StreamAcceptor { + alpn: Vec, + rx: mpsc::Receiver, +} + +impl StreamAcceptor { + /// Waits for the next incoming stream. + /// + /// Returns `None` after the engine shut down. + pub async fn accept(&mut self) -> Option { + self.rx.recv().await + } + + /// The ALPN of the stream protocol this acceptor serves. + pub fn alpn(&self) -> &[u8] { + &self.alpn + } +} + +impl fmt::Debug for StreamAcceptor { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("StreamAcceptor") + .field("alpn", &String::from_utf8_lossy(&self.alpn)) + .finish_non_exhaustive() + } +} + +/// Producer/consumer pair behind one declared stream protocol. +struct StreamAcceptorSlot { + sender: mpsc::Sender, + /// Taken (once) by [`NetworkEngine::stream_acceptor`]. + receiver: Option>, +} + +fn alpn_display(alpn: &[u8]) -> String { + String::from_utf8_lossy(alpn).into_owned() +} + /// State shared between the engine handles, the protocol handler and all /// background tasks. struct Shared { @@ -84,6 +158,7 @@ struct Shared { endpoint: Endpoint, router: Mutex>, peers: Mutex>, + stream_acceptors: Mutex, StreamAcceptorSlot>>, events: Mutex>>>, tasks: Mutex>, next_generation: AtomicU64, @@ -402,6 +477,106 @@ impl Shared { Ok(()) } + /// Handles a freshly accepted connection on an auxiliary stream ALPN: + /// validates the stream handshake and hands the byte stream over to the + /// application through the protocol's acceptor. + /// + /// Unlike message connections, stream connections are not registered in + /// the peer registry: their lifecycle belongs entirely to the returned + /// [`ByteStream`]. + async fn handle_incoming_stream(&self, connection: &Connection, alpn: &[u8]) -> Result<()> { + let peer_id = connection.remote_id(); + if self.is_shutting_down() { + connection.close(CLOSE_CODE_SHUTDOWN, b"engine is shutting down"); + return Err(NetworkError::ShuttingDown); + } + debug!(peer = %peer_id, alpn = %alpn_display(alpn), "incoming byte stream"); + let request_timeout = self.config.request_timeout; + + let (mut send, mut recv) = timeout(request_timeout, connection.accept_bi()) + .await + .map_err(|_| NetworkError::Timeout) + .and_then(|res| { + res.map_err(|err| { + NetworkError::Transport(format!("failed to accept stream handshake: {err}")) + }) + }) + .inspect_err(|_| { + connection.close(CLOSE_CODE_HANDSHAKE_REJECTED, b"stream handshake failed"); + })?; + + let handshake: Result = timeout( + request_timeout, + wire::read_frame(&mut recv, MAX_HANDSHAKE_FRAME_SIZE), + ) + .await + .map_err(|_| NetworkError::Timeout) + .and_then(|res| res); + + let verdict = match &handshake { + Ok(handshake) => self.validate_handshake(handshake), + Err(_) => Some(HandshakeErrorCode::InvalidHandshake), + }; + if let Some(code) = verdict { + warn!(peer = %peer_id, code = ?code, "stream handshake rejected"); + let ack = HandshakeAck { + accepted: false, + error: Some(code), + }; + let sent = timeout( + request_timeout, + wire::write_frame(&mut send, &ack, MAX_HANDSHAKE_FRAME_SIZE), + ) + .await; + if matches!(sent, Ok(Ok(()))) { + let _ = send.finish(); + let _ = timeout(REJECT_LINGER, connection.closed()).await; + } + connection.close(CLOSE_CODE_HANDSHAKE_REJECTED, b"stream handshake rejected"); + return Err(handshake_code_to_error(code)); + } + + // Reserve a queue slot before acking so a stream the application + // will never consume is refused instead of silently dying after an + // accepted handshake. A full queue delays the ack (backpressure). + let sender = lock(&self.stream_acceptors) + .get(alpn) + .map(|slot| slot.sender.clone()) + .ok_or_else(|| NetworkError::UnknownStreamProtocol(alpn_display(alpn)))?; + let Ok(permit) = sender.reserve().await else { + connection.close(CLOSE_CODE_HANDSHAKE_REJECTED, b"stream acceptor closed"); + return Err(NetworkError::Transport( + "the stream acceptor was dropped by the application".to_string(), + )); + }; + + let ack = HandshakeAck { + accepted: true, + error: None, + }; + timeout( + request_timeout, + wire::write_frame(&mut send, &ack, MAX_HANDSHAKE_FRAME_SIZE), + ) + .await + .map_err(|_| NetworkError::Timeout) + .and_then(|res| res) + .inspect_err(|_| { + connection.close(CLOSE_CODE_HANDSHAKE_REJECTED, b"stream handshake failed"); + })?; + // The stream is deliberately not finished: the application payload + // follows on the same stream pair. + + info!(peer = %peer_id, alpn = %alpn_display(alpn), "byte stream accepted"); + permit.send(ByteStream { + peer_id, + send, + recv, + _connection: connection.clone(), + }); + Ok(()) + } + /// Dials `addr`, runs the client handshake and registers the connection. /// /// An existing healthy connection to the same peer is reused. The remote @@ -570,6 +745,43 @@ impl ProtocolHandler for FederationProtocol { } } +/// Protocol handler registered with the Iroh [`Router`] for one auxiliary +/// stream ALPN declared in the configuration. +struct StreamProtocol { + shared: Weak>, + alpn: Vec, +} + +impl fmt::Debug for StreamProtocol { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("StreamProtocol") + .field("alpn", &alpn_display(&self.alpn)) + .finish_non_exhaustive() + } +} + +impl ProtocolHandler for StreamProtocol { + async fn accept(&self, connection: Connection) -> Result<(), AcceptError> { + let Some(shared) = self.shared.upgrade() else { + connection.close(CLOSE_CODE_SHUTDOWN, b"engine is gone"); + return Ok(()); + }; + let peer_id = connection.remote_id(); + if let Err(err) = shared.handle_incoming_stream(&connection, &self.alpn).await { + shared + .report_protocol_error( + Some(peer_id), + format!( + "incoming byte stream on {} failed: {err}", + alpn_display(&self.alpn) + ), + ) + .await; + } + Ok(()) + } +} + /// A generic peer-to-peer network engine on top of Iroh. /// /// `M` is the application-defined domain message type; the engine treats it @@ -608,11 +820,26 @@ impl NetworkEngine { .map_err(|err| NetworkError::Transport(format!("failed to bind endpoint: {err}")))?; let endpoint_id = endpoint.id(); let (sender, receiver) = mpsc::channel(config.event_channel_capacity); + let stream_acceptors = config + .stream_protocols + .iter() + .map(|alpn| { + let (sender, receiver) = mpsc::channel(STREAM_ACCEPT_QUEUE); + ( + alpn.clone(), + StreamAcceptorSlot { + sender, + receiver: Some(receiver), + }, + ) + }) + .collect(); let shared = Arc::new(Shared { config, endpoint: endpoint.clone(), router: Mutex::new(None), peers: Mutex::new(HashMap::new()), + stream_acceptors: Mutex::new(stream_acceptors), events: Mutex::new(Some(sender)), tasks: Mutex::new(JoinSet::new()), next_generation: AtomicU64::new(0), @@ -622,7 +849,15 @@ impl NetworkEngine { let handler = FederationProtocol { shared: Arc::downgrade(&shared), }; - let router = Router::builder(endpoint).accept(ALPN, handler).spawn(); + let mut router_builder = Router::builder(endpoint).accept(ALPN, handler); + for alpn in &shared.config.stream_protocols { + let handler = StreamProtocol:: { + shared: Arc::downgrade(&shared), + alpn: alpn.clone(), + }; + router_builder = router_builder.accept(alpn.clone(), handler); + } + let router = router_builder.spawn(); *lock(&shared.router) = Some(router); if let Some(rendezvous) = shared.config.rendezvous.clone() { match RendezvousClient::new(shared.config.network_id, &rendezvous) { @@ -762,6 +997,97 @@ impl NetworkEngine { .map_err(|_| NetworkError::Timeout)? } + /// Opens a raw bidirectional byte stream to a peer over a dedicated + /// connection on the auxiliary ALPN `alpn`. + /// + /// The remote peer must have declared `alpn` in its configuration + /// ([`crate::NetworkConfigBuilder::stream_protocol`]) and hold on to the + /// matching [`StreamAcceptor`]. The stream begins with the regular + /// network handshake, so it can only be established between peers of the + /// same network and schema; afterwards the stream carries opaque + /// application bytes with no framing imposed by the engine. + /// + /// `target` can be a full [`EndpointAddr`] (e.g. from a ticket) or a + /// bare [`EndpointId`] when the transport already knows how to reach the + /// peer (an active or recent connection, or discovery). + pub async fn open_stream( + &self, + target: impl Into, + alpn: &[u8], + ) -> Result { + let shared = &self.shared; + shared.ensure_running()?; + let addr: EndpointAddr = target.into(); + let request_timeout = shared.config.request_timeout; + let connection = timeout(request_timeout, shared.endpoint.connect(addr, alpn)) + .await + .map_err(|_| NetworkError::Timeout)? + .map_err(|err| { + NetworkError::Transport(format!("failed to connect for a byte stream: {err}")) + })?; + let peer_id = connection.remote_id(); + let handshake = Handshake { + protocol_version: PROTOCOL_VERSION, + network_id: shared.config.network_id, + schema_id: shared.config.schema_id, + }; + let result = timeout(request_timeout, async { + let (mut send, mut recv) = connection.open_bi().await.map_err(|err| { + NetworkError::Transport(format!("failed to open the byte stream: {err}")) + })?; + wire::write_frame(&mut send, &handshake, MAX_HANDSHAKE_FRAME_SIZE).await?; + // The stream is deliberately not finished: application payload + // follows the handshake on the same stream. + let ack: HandshakeAck = wire::read_frame(&mut recv, MAX_HANDSHAKE_FRAME_SIZE).await?; + if ack.accepted { + Ok((send, recv)) + } else { + let code = ack.error.unwrap_or(HandshakeErrorCode::InvalidHandshake); + Err(handshake_code_to_error(code)) + } + }) + .await + .map_err(|_| NetworkError::Timeout) + .and_then(|res| res); + match result { + Ok((send, recv)) => { + debug!(peer = %peer_id, alpn = %alpn_display(alpn), "byte stream opened"); + Ok(ByteStream { + peer_id, + send, + recv, + _connection: connection, + }) + } + Err(err) => { + connection.close(CLOSE_CODE_HANDSHAKE_REJECTED, b"stream handshake failed"); + Err(err) + } + } + } + + /// Takes the acceptor of incoming byte streams for the declared stream + /// protocol `alpn`. + /// + /// The acceptor for a protocol can be taken exactly once. Streams whose + /// handshake succeeded are delivered in accept order; the queue is small + /// and bounded, so an unconsumed acceptor applies backpressure to + /// remote peers. + pub fn stream_acceptor(&self, alpn: &[u8]) -> Result { + let mut acceptors = lock(&self.shared.stream_acceptors); + let slot = acceptors + .get_mut(alpn) + .ok_or_else(|| NetworkError::UnknownStreamProtocol(alpn_display(alpn)))?; + let rx = slot + .receiver + .take() + .ok_or_else(|| NetworkError::StreamAcceptorTaken(alpn_display(alpn)))?; + Ok(StreamAcceptor { + alpn: alpn.to_vec(), + rx, + }) + } + /// Returns the ids of all currently connected peers. pub fn connected_peers(&self) -> Vec { lock(&self.shared.peers).keys().copied().collect() diff --git a/crates/federation-net/src/error.rs b/crates/federation-net/src/error.rs index 9d8ec76..a86754f 100644 --- a/crates/federation-net/src/error.rs +++ b/crates/federation-net/src/error.rs @@ -61,6 +61,15 @@ pub enum NetworkError { #[error("message rejected by peer: {0}")] MessageRejected(String), + /// The given ALPN was not declared as a stream protocol in the + /// configuration. + #[error("unknown stream protocol: {0}")] + UnknownStreamProtocol(String), + + /// The incoming-stream acceptor for this protocol was already taken. + #[error("stream acceptor already taken: {0}")] + StreamAcceptorTaken(String), + /// The engine is shutting down and no longer accepts operations. #[error("engine is shutting down")] ShuttingDown, diff --git a/crates/federation-net/src/lib.rs b/crates/federation-net/src/lib.rs index bf1c8a4..d31b038 100644 --- a/crates/federation-net/src/lib.rs +++ b/crates/federation-net/src/lib.rs @@ -65,7 +65,7 @@ pub use config::{ DEFAULT_EVENT_CHANNEL_CAPACITY, DEFAULT_MAX_CONCURRENT_STREAMS_PER_PEER, DEFAULT_MAX_MESSAGE_SIZE, DEFAULT_REQUEST_TIMEOUT, NetworkConfig, NetworkConfigBuilder, }; -pub use engine::{Message, NetworkEngine}; +pub use engine::{ByteStream, Message, NetworkEngine, StreamAcceptor}; pub use error::{NetworkError, Result}; pub use event::{ConnectionDirection, NetworkEvent, NetworkEventReceiver}; pub use protocol::{ALPN, NetworkId, PROTOCOL_VERSION, SchemaId}; @@ -73,6 +73,7 @@ pub use rendezvous::{DEFAULT_RENDEZVOUS_ENTRY_TTL, DEFAULT_RENDEZVOUS_INTERVAL, pub use ticket::{PeerTicket, TICKET_VERSION}; // Re-exported Iroh types that appear in the public API. +pub use iroh::endpoint::{RecvStream, SendStream}; pub use iroh::{EndpointAddr, EndpointId}; // Re-exported so applications can use the generic ticket helpers. pub use iroh_tickets::Ticket; diff --git a/crates/federation-net/tests/integration.rs b/crates/federation-net/tests/integration.rs index 159c893..e59ed10 100644 --- a/crates/federation-net/tests/integration.rs +++ b/crates/federation-net/tests/integration.rs @@ -359,3 +359,112 @@ async fn rendezvous_discovers_peers_without_tickets() { .await .expect("test timed out"); } + +#[tokio::test] +async fn byte_streams_between_peers() { + let _net = NET_LOCK.lock().await; + tokio::time::timeout(TEST_TIMEOUT, async { + const ECHO_ALPN: &[u8] = b"test/echo/1"; + + let dir_a = tempfile::tempdir().expect("tempdir"); + let dir_b = tempfile::tempdir().expect("tempdir"); + + // Peer A declares a stream protocol and echoes one message per stream. + let config_a = NetworkConfig::builder() + .data_dir(dir_a.path()) + .network_id(NetworkId::from_name("test-net")) + .schema_id(SchemaId::from_name("test-schema-v1")) + .request_timeout(Duration::from_secs(10)) + .stream_protocol(ECHO_ALPN) + .build() + .expect("valid config"); + let (engine_a, _events_a) = Engine::start(config_a).await.expect("engine a starts"); + let mut acceptor = engine_a.stream_acceptor(ECHO_ALPN).expect("acceptor"); + // The acceptor is exclusive and unknown protocols are rejected. + assert!(matches!( + engine_a.stream_acceptor(ECHO_ALPN), + Err(NetworkError::StreamAcceptorTaken(_)) + )); + assert!(matches!( + engine_a.stream_acceptor(b"test/unknown/1"), + Err(NetworkError::UnknownStreamProtocol(_)) + )); + let echo_task = tokio::spawn(async move { + let mut stream = acceptor.accept().await.expect("incoming stream"); + let mut buf = Vec::new(); + let mut chunk = [0u8; 1024]; + while let Some(n) = stream.recv.read(&mut chunk).await.expect("read") { + buf.extend_from_slice(&chunk[..n]); + } + stream.send.write_all(&buf).await.expect("write echo"); + stream.send.finish().expect("finish"); + // Keep the stream alive until the peer read everything. + let _ = stream.send.stopped().await; + (stream.peer_id, buf.len()) + }); + + let (engine_b, _events_b) = start(dir_b.path(), "test-net", "test-schema-v1").await; + let ticket = engine_a.ticket().await.expect("ticket"); + + // B opens a byte stream to A using the address from the ticket and + // sends a payload larger than one network frame. + let payload = vec![0xAB_u8; 512 * 1024]; + let mut stream = engine_b + .open_stream(ticket.endpoint_addr.clone(), ECHO_ALPN) + .await + .expect("open stream"); + assert_eq!(stream.peer_id, engine_a.endpoint_id()); + stream.send.write_all(&payload).await.expect("write"); + stream.send.finish().expect("finish"); + let mut echoed = Vec::new(); + let mut chunk = [0u8; 1024]; + while let Some(n) = stream.recv.read(&mut chunk).await.expect("read echo") { + echoed.extend_from_slice(&chunk[..n]); + } + assert_eq!(echoed, payload); + + let (peer_seen_by_a, len_seen_by_a) = echo_task.await.expect("echo task"); + assert_eq!(peer_seen_by_a, engine_b.endpoint_id()); + assert_eq!(len_seen_by_a, payload.len()); + + engine_a.shutdown().await.expect("shutdown a"); + engine_b.shutdown().await.expect("shutdown b"); + }) + .await + .expect("test timed out"); +} + +#[tokio::test] +async fn byte_stream_to_wrong_network_is_rejected() { + let _net = NET_LOCK.lock().await; + tokio::time::timeout(TEST_TIMEOUT, async { + const BLOB_ALPN: &[u8] = b"test/blob/1"; + + let dir_a = tempfile::tempdir().expect("tempdir"); + let dir_b = tempfile::tempdir().expect("tempdir"); + + let config_a = NetworkConfig::builder() + .data_dir(dir_a.path()) + .network_id(NetworkId::from_name("net-one")) + .schema_id(SchemaId::from_name("test-schema-v1")) + .request_timeout(Duration::from_secs(10)) + .stream_protocol(BLOB_ALPN) + .build() + .expect("valid config"); + let (engine_a, _events_a) = Engine::start(config_a).await.expect("engine a starts"); + let _acceptor = engine_a.stream_acceptor(BLOB_ALPN).expect("acceptor"); + + // B lives in a different network; the stream handshake must fail. + let (engine_b, _events_b) = start(dir_b.path(), "net-two", "test-schema-v1").await; + let ticket = engine_a.ticket().await.expect("ticket"); + let result = engine_b + .open_stream(ticket.endpoint_addr.clone(), BLOB_ALPN) + .await; + assert!(matches!(result, Err(NetworkError::NetworkMismatch))); + + engine_a.shutdown().await.expect("shutdown a"); + engine_b.shutdown().await.expect("shutdown b"); + }) + .await + .expect("test timed out"); +} diff --git a/crates/music-dht/src/config.rs b/crates/music-dht/src/config.rs index 3473c1c..12240b8 100644 --- a/crates/music-dht/src/config.rs +++ b/crates/music-dht/src/config.rs @@ -52,6 +52,9 @@ pub struct MusicDhtConfig { /// network find each other knowing nothing but the network id. `None` /// disables it; peers are then connected via tickets only. pub rendezvous: Option, + /// Auxiliary ALPN protocols on which this peer accepts raw byte streams + /// (see [`crate::MusicDhtService::stream_acceptor`]). + pub stream_protocols: Vec>, } impl MusicDhtConfig { @@ -76,6 +79,7 @@ pub struct MusicDhtConfigBuilder { transport_timeout: Option, dial_timeout: Option, rendezvous: Option, + stream_protocols: Vec>, } impl MusicDhtConfigBuilder { @@ -133,6 +137,15 @@ impl MusicDhtConfigBuilder { self } + /// Declares an auxiliary ALPN protocol on which this peer accepts raw + /// byte streams (see [`crate::MusicDhtService::stream_acceptor`]). May + /// be called multiple times, once per protocol. The ALPNs are validated + /// by the underlying network engine on service start. + pub fn stream_protocol(mut self, alpn: impl Into>) -> Self { + self.stream_protocols.push(alpn.into()); + self + } + /// Validates and builds the configuration. pub fn build(self) -> Result { let data_dir = self @@ -157,6 +170,7 @@ impl MusicDhtConfigBuilder { transport_timeout: self.transport_timeout.unwrap_or(DEFAULT_TRANSPORT_TIMEOUT), dial_timeout: self.dial_timeout.unwrap_or(DEFAULT_DIAL_TIMEOUT), rendezvous: self.rendezvous, + stream_protocols: self.stream_protocols, }; for (name, value) in [ ("republish_interval", config.republish_interval), diff --git a/crates/music-dht/src/lib.rs b/crates/music-dht/src/lib.rs index a3a8413..2ba9052 100644 --- a/crates/music-dht/src/lib.rs +++ b/crates/music-dht/src/lib.rs @@ -105,4 +105,7 @@ pub use service::{ }; // Re-exported types from the transport layer that appear in this API. -pub use federation_net::{EndpointId, NetworkId, PeerTicket, RendezvousConfig}; +pub use federation_net::{ + ByteStream, EndpointAddr, EndpointId, NetworkId, PeerTicket, RecvStream, RendezvousConfig, + SendStream, StreamAcceptor, +}; diff --git a/crates/music-dht/src/service.rs b/crates/music-dht/src/service.rs index aec69a7..b776174 100644 --- a/crates/music-dht/src/service.rs +++ b/crates/music-dht/src/service.rs @@ -4,7 +4,10 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::{Duration, Instant}; -use federation_net::{EndpointId, NetworkConfig, NetworkEngine, PeerTicket, SchemaId}; +use federation_net::{ + ByteStream, EndpointAddr, EndpointId, NetworkConfig, NetworkEngine, PeerTicket, SchemaId, + StreamAcceptor, +}; use tokio::sync::mpsc; use tokio::task::JoinHandle; use tracing::info; @@ -164,6 +167,9 @@ impl MusicDhtService { if let Some(rendezvous) = config.rendezvous.clone() { engine_builder = engine_builder.rendezvous(rendezvous); } + for alpn in &config.stream_protocols { + engine_builder = engine_builder.stream_protocol(alpn.clone()); + } let engine_config = engine_builder .build() .map_err(|err| MusicDhtError::Network(err.to_string()))?; @@ -241,6 +247,39 @@ impl MusicDhtService { self.node.engine.is_connected(peer) } + /// Takes the acceptor of incoming raw byte streams for the stream + /// protocol `alpn` declared in the configuration. + /// + /// The acceptor for a protocol can be taken exactly once; see + /// [`federation_net::NetworkEngine::stream_acceptor`]. + pub fn stream_acceptor(&self, alpn: &[u8]) -> Result { + self.node.ensure_running()?; + self.node.engine.stream_acceptor(alpn).map_err(Into::into) + } + + /// Opens a raw byte stream to `peer` on the auxiliary ALPN `alpn`. + /// + /// The peer is dialed using the address from its stored DHT contact + /// when one is known (which is the case for every peer that appears in + /// search results), falling back to whatever the transport itself knows + /// about the peer. See [`federation_net::NetworkEngine::open_stream`]. + pub async fn open_stream(&self, peer: EndpointId, alpn: &[u8]) -> Result { + self.node.ensure_running()?; + let contact_addr = self + .node + .known_contacts() + .into_iter() + .find(|contact| contact.peer_id == peer) + .and_then(|contact| contact.ticket.parse::().ok()) + .map(|ticket| ticket.endpoint_addr); + let addr: EndpointAddr = contact_addr.unwrap_or_else(|| peer.into()); + self.node + .engine + .open_stream(addr, alpn) + .await + .map_err(Into::into) + } + /// Synchronizes the published library with `specs`: the desired set of /// items this peer wants to share. /// diff --git a/crates/music-dht/tests/integration.rs b/crates/music-dht/tests/integration.rs index 6d1cc66..80d9ab0 100644 --- a/crates/music-dht/tests/integration.rs +++ b/crates/music-dht/tests/integration.rs @@ -187,3 +187,72 @@ async fn library_sync_and_distributed_search() { .await .expect("test timed out"); } + +/// Two nodes exchange raw bytes over an auxiliary stream protocol: the +/// requester finds an item through the DHT and then opens a byte stream to +/// its owner (the flow a file-transfer application follows). +#[tokio::test] +async fn byte_stream_to_item_owner() { + let _net = NET_LOCK.lock().await; + tokio::time::timeout(TEST_TIMEOUT, async { + const BLOB_ALPN: &[u8] = b"music-test/blob/1"; + + let dir_a = tempfile::tempdir().expect("tempdir"); + let dir_b = tempfile::tempdir().expect("tempdir"); + + // Node A owns the library and serves byte streams. + let config_a = MusicDhtConfig::builder() + .data_dir(dir_a.path()) + .network_id(NetworkId::from_name("stream-test-net")) + .request_timeout(Duration::from_secs(5)) + .lookup_timeout(Duration::from_secs(10)) + .stream_protocol(BLOB_ALPN) + .build() + .expect("valid config"); + let (node_a, _events_a) = MusicDhtService::start(config_a) + .await + .expect("service starts"); + let mut acceptor = node_a.stream_acceptor(BLOB_ALPN).expect("acceptor"); + let payload = b"pretend this is FLAC".to_vec(); + let served = payload.clone(); + let serve_task = tokio::spawn(async move { + let mut stream = acceptor.accept().await.expect("incoming stream"); + stream.send.write_all(&served).await.expect("write"); + stream.send.finish().expect("finish"); + let _ = stream.send.stopped().await; + stream.peer_id + }); + node_a + .sync_library(vec![spec("track:1", ItemKind::Track, "Teardrop", &["Massive Attack"])]) + .await + .expect("sync"); + + // Node B joins via ticket and finds the track and its owner. + let (node_b, _events_b) = start(dir_b.path(), "stream-test-net").await; + let ticket = node_a.ticket().await.expect("ticket"); + node_b.connect(ticket).await.expect("connect"); + let owner = search_until("the track by name", &node_b, "teardrop", |items| { + !items.is_empty() + }) + .await + .pop() + .expect("found item") + .owner; + assert_eq!(owner, node_a.endpoint_id()); + + // B streams the bytes from the owner. + let mut stream = node_b.open_stream(owner, BLOB_ALPN).await.expect("open stream"); + let mut received = Vec::new(); + let mut chunk = [0u8; 1024]; + while let Some(n) = stream.recv.read(&mut chunk).await.expect("read") { + received.extend_from_slice(&chunk[..n]); + } + assert_eq!(received, payload); + assert_eq!(serve_task.await.expect("serve task"), node_b.endpoint_id()); + + node_a.shutdown().await.expect("shutdown a"); + node_b.shutdown().await.expect("shutdown b"); + }) + .await + .expect("test timed out"); +}