Added bytestream ALPN
This commit is contained in:
@@ -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()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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.
|
||||
///
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user