Added bytestream ALPN

This commit is contained in:
Ultradesu
2026-07-16 18:35:21 +03:00
parent b9c61e2249
commit 555222d6d7
9 changed files with 630 additions and 5 deletions
+56 -1
View File
@@ -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<RendezvousConfig>,
/// Auxiliary ALPN protocols on which this peer accepts raw byte streams
/// (see [`crate::NetworkEngine::stream_acceptor`]).
pub stream_protocols: Vec<Vec<u8>>,
}
impl NetworkConfig {
@@ -61,6 +64,7 @@ pub struct NetworkConfigBuilder {
event_channel_capacity: Option<usize>,
max_concurrent_streams_per_peer: Option<usize>,
rendezvous: Option<RendezvousConfig>,
stream_protocols: Vec<Vec<u8>>,
}
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<Vec<u8>>) -> Self {
self.stream_protocols.push(alpn.into());
self
}
/// Validates the configuration and builds a [`NetworkConfig`].
pub fn build(self) -> Result<NetworkConfig> {
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()
);
}
}
+327 -1
View File
@@ -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<u8>,
rx: mpsc::Receiver<ByteStream>,
}
impl StreamAcceptor {
/// Waits for the next incoming stream.
///
/// Returns `None` after the engine shut down.
pub async fn accept(&mut self) -> Option<ByteStream> {
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<ByteStream>,
/// Taken (once) by [`NetworkEngine::stream_acceptor`].
receiver: Option<mpsc::Receiver<ByteStream>>,
}
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<M> {
@@ -84,6 +158,7 @@ struct Shared<M> {
endpoint: Endpoint,
router: Mutex<Option<Router>>,
peers: Mutex<HashMap<EndpointId, PeerState>>,
stream_acceptors: Mutex<HashMap<Vec<u8>, StreamAcceptorSlot>>,
events: Mutex<Option<mpsc::Sender<NetworkEvent<M>>>>,
tasks: Mutex<JoinSet<()>>,
next_generation: AtomicU64,
@@ -402,6 +477,106 @@ impl<M: Message> Shared<M> {
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<Handshake> = 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<M: Message> ProtocolHandler for FederationProtocol<M> {
}
}
/// Protocol handler registered with the Iroh [`Router`] for one auxiliary
/// stream ALPN declared in the configuration.
struct StreamProtocol<M> {
shared: Weak<Shared<M>>,
alpn: Vec<u8>,
}
impl<M> fmt::Debug for StreamProtocol<M> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("StreamProtocol")
.field("alpn", &alpn_display(&self.alpn))
.finish_non_exhaustive()
}
}
impl<M: Message> ProtocolHandler for StreamProtocol<M> {
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<M: Message> NetworkEngine<M> {
.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<M: Message> NetworkEngine<M> {
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::<M> {
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<M: Message> NetworkEngine<M> {
.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<EndpointAddr>,
alpn: &[u8],
) -> Result<ByteStream> {
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<StreamAcceptor> {
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<EndpointId> {
lock(&self.shared.peers).keys().copied().collect()
+9
View File
@@ -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,
+2 -1
View File
@@ -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;
+109
View File
@@ -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");
}
+14
View File
@@ -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<RendezvousConfig>,
/// Auxiliary ALPN protocols on which this peer accepts raw byte streams
/// (see [`crate::MusicDhtService::stream_acceptor`]).
pub stream_protocols: Vec<Vec<u8>>,
}
impl MusicDhtConfig {
@@ -76,6 +79,7 @@ pub struct MusicDhtConfigBuilder {
transport_timeout: Option<Duration>,
dial_timeout: Option<Duration>,
rendezvous: Option<RendezvousConfig>,
stream_protocols: Vec<Vec<u8>>,
}
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<Vec<u8>>) -> Self {
self.stream_protocols.push(alpn.into());
self
}
/// Validates and builds the configuration.
pub fn build(self) -> Result<MusicDhtConfig> {
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),
+4 -1
View File
@@ -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,
};
+40 -1
View File
@@ -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<StreamAcceptor> {
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<ByteStream> {
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::<PeerTicket>().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.
///
+69
View File
@@ -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");
}