#![cfg(not(target_arch = "wasm32"))]
use anyhow::Result;
use futures::StreamExt;
use iroh::{
endpoint::{Connection, RecvStream, SendStream},
protocol::{AcceptError, ProtocolHandler, Router},
Endpoint, EndpointAddr, EndpointId, Watcher,
};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::{broadcast, mpsc, oneshot, RwLock};
use tokio_stream::wrappers::BroadcastStream;
use crate::heartbeat::{
classify_incoming_uni, parse_incoming_pong, respond_to_ping, HealthTransition, HeartbeatConfig,
IncomingUniClassification, IrohConnectionProbe, IrohHeartbeatManager, IrohProbeRegistry,
PrefixedRecvStream,
};
use crate::iroh_connection_policy::{
decide_inbound_install, decide_outbound_install, should_start_outbound_dial,
should_wait_for_canonical_inbound, ExistingConnectionState, IrohConnectionDirection,
IrohConnectionInstallDecision, NONCANONICAL_OUTBOUND_DIAL_GRACE_MS,
};
pub fn native_iroh_transport_config() -> iroh::endpoint::QuicTransportConfig {
use iroh::endpoint::VarInt;
let idle_timeout = std::time::Duration::from_secs(120)
.try_into()
.expect("120 seconds is a valid Iroh idle timeout");
let keep_alive_secs = if cfg!(any(target_os = "ios", target_os = "android")) {
90
} else {
25
};
iroh::endpoint::QuicTransportConfig::builder()
.max_idle_timeout(Some(idle_timeout))
.keep_alive_interval(std::time::Duration::from_secs(keep_alive_secs))
.stream_receive_window(VarInt::from_u32(8 * 1024 * 1024))
.receive_window(VarInt::from_u32(16 * 1024 * 1024))
.send_window(16 * 1024 * 1024)
.build()
}
async fn should_break_accept_loop(connection: &Connection) -> bool {
if connection.close_reason().is_some() {
return true;
}
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
connection.close_reason().is_some()
}
fn local_prefers_outbound(local_endpoint_id: EndpointId, remote_endpoint_id: EndpointId) -> bool {
local_endpoint_id.to_string() > remote_endpoint_id.to_string()
}
async fn outbound_dial_precheck(
endpoint: &Endpoint,
endpoint_id: EndpointId,
connections: &Arc<RwLock<HashMap<EndpointId, Connection>>>,
connection_inserted_at: &Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
) -> bool {
let prefers_outbound = local_prefers_outbound(endpoint.id(), endpoint_id);
let existing = {
let conns = connections.read().await;
if let Some(existing) = conns.get(&endpoint_id) {
let direction = connection_inserted_at
.read()
.await
.get(&endpoint_id)
.map(|metadata| metadata.direction)
.unwrap_or(IrohConnectionDirection::Outbound);
Some((existing.close_reason().is_none(), direction))
} else {
None
}
};
if !should_start_outbound_dial(existing, prefers_outbound) {
return true;
}
if !should_wait_for_canonical_inbound(existing, prefers_outbound) {
return false;
}
tokio::time::sleep(std::time::Duration::from_millis(
NONCANONICAL_OUTBOUND_DIAL_GRACE_MS,
))
.await;
connections
.read()
.await
.get(&endpoint_id)
.is_some_and(|connection| connection.close_reason().is_none())
}
fn manual_disconnect_notice_key(endpoint_id: EndpointId, transport_stable_id: u64) -> String {
format!("{endpoint_id}:{transport_stable_id}")
}
fn heartbeat_connection_key(endpoint_id: EndpointId, transport_stable_id: u64) -> String {
format!("{endpoint_id}:{transport_stable_id}")
}
#[derive(Debug, Clone, Copy)]
struct IrohConnectionMetadata {
inserted_at: Instant,
direction: IrohConnectionDirection,
}
#[derive(Debug)]
pub enum IncomingStreamType {
Bi(SendStream, RecvStream),
Uni(PrefixedRecvStream),
}
#[derive(Debug)]
pub struct IncomingStream {
pub endpoint_id: EndpointId,
pub transport_stable_id: u64,
pub stream: IncomingStreamType,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum ConnectEvent {
Connected,
Closed { error: Option<String> },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum ReplacementConnectEvent {
Connected { transport_stable_id: u64 },
Closed { error: Option<String> },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum AcceptEvent {
Accepted {
endpoint_id: EndpointId,
transport_stable_id: u64,
},
Closed {
endpoint_id: EndpointId,
transport_stable_id: u64,
error: Option<String>,
was_locally_closed: bool,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ExternalConnectionInstallOutcome {
Installed {
transport_stable_id: u64,
},
KeptExisting {
fresh_transport_stable_id: u64,
kept_transport_stable_id: u64,
},
}
#[derive(Debug, Clone, Copy)]
struct InboundReplacementAuthorization {
transport_id: u64,
expires_at: Instant,
}
type InboundReplacementAuthorizations =
Arc<RwLock<HashMap<EndpointId, InboundReplacementAuthorization>>>;
fn selected_custom_transport_id(connection: &Connection) -> Option<u64> {
connection.paths().iter().find_map(|path| {
if !path.is_selected() {
return None;
}
match path.remote_addr() {
iroh::TransportAddr::Custom(addr) => Some(addr.id()),
_ => None,
}
})
}
fn custom_transport_arbitration_override(
existing_alive: bool,
existing_custom_transport_id: Option<u64>,
fresh_custom_transport_id: Option<u64>,
authorized_replacement_transport_id: Option<u64>,
) -> Option<IrohConnectionInstallDecision> {
if authorized_replacement_transport_id
.is_some_and(|expected| fresh_custom_transport_id == Some(expected))
{
return Some(IrohConnectionInstallDecision::ReplaceExisting {
close_existing_reason: "native-custom-transport-upgrade",
});
}
if existing_alive
&& existing_custom_transport_id.is_some()
&& fresh_custom_transport_id != existing_custom_transport_id
{
return Some(IrohConnectionInstallDecision::KeepExisting {
close_fresh_reason: "existing-custom-transport-preferred",
});
}
None
}
async fn consume_authorized_inbound_replacement(
connection: &Connection,
authorizations: &InboundReplacementAuthorizations,
) -> Option<u64> {
let endpoint_id = connection.remote_id();
let deadline = Instant::now() + std::time::Duration::from_secs(1);
loop {
let authorization = authorizations.read().await.get(&endpoint_id).copied();
let Some(authorization) = authorization else {
return None;
};
if Instant::now() >= authorization.expires_at {
let mut authorizations = authorizations.write().await;
if authorizations.get(&endpoint_id).is_some_and(|current| {
current.transport_id == authorization.transport_id
&& current.expires_at == authorization.expires_at
}) {
authorizations.remove(&endpoint_id);
}
return None;
}
if selected_custom_transport_id(connection) == Some(authorization.transport_id) {
let mut authorizations = authorizations.write().await;
if authorizations.get(&endpoint_id).is_some_and(|current| {
current.transport_id == authorization.transport_id
&& current.expires_at == authorization.expires_at
}) {
authorizations.remove(&endpoint_id);
return Some(authorization.transport_id);
}
return None;
}
if Instant::now() >= deadline || connection.close_reason().is_some() {
return None;
}
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
}
}
#[derive(Debug, Clone)]
pub struct PlutoniumProtocol {
event_sender: broadcast::Sender<AcceptEvent>,
stream_sender: async_channel::Sender<IncomingStream>,
connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
local_endpoint_id: EndpointId,
manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
probe_registry: IrohProbeRegistry,
heartbeat_manager: Option<IrohHeartbeatManager>,
heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
heartbeat_config: HeartbeatConfig,
inbound_replacement_authorizations: InboundReplacementAuthorizations,
}
async fn run_connection_loop(
connection: Connection,
event_sender: broadcast::Sender<AcceptEvent>,
stream_sender: async_channel::Sender<IncomingStream>,
connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
local_endpoint_id: EndpointId,
connection_direction: IrohConnectionDirection,
manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
probe_registry: IrohProbeRegistry,
heartbeat_manager: Option<IrohHeartbeatManager>,
heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
heartbeat_config: HeartbeatConfig,
replacement_transport_id: Option<u64>,
mut install_outcome_sender: Option<oneshot::Sender<ExternalConnectionInstallOutcome>>,
) -> std::result::Result<(), AcceptError> {
let endpoint_id = connection.remote_id();
let stable_id = connection.stable_id();
let endpoint_key = endpoint_id.to_string();
let heartbeat_key = heartbeat_connection_key(endpoint_id, stable_id as u64);
let mut kept_existing = None;
{
let mut conns = connections.write().await;
let mut insert_times = connection_inserted_at.write().await;
if let Some(previous) = conns.get(&endpoint_id).cloned() {
let previous_direction = insert_times
.get(&endpoint_id)
.map(|metadata| metadata.direction)
.unwrap_or(connection_direction);
let previous_age_ms = insert_times
.get(&endpoint_id)
.map(|metadata| metadata.inserted_at.elapsed().as_millis() as u64)
.unwrap_or(0);
let existing = Some(ExistingConnectionState {
same_stable_id: previous.stable_id() == stable_id,
alive: previous.close_reason().is_none(),
direction: previous_direction,
age_ms: previous_age_ms,
});
let prefers_outbound = local_prefers_outbound(local_endpoint_id, endpoint_id);
let fresh_custom_transport_id = selected_custom_transport_id(&connection);
let existing_custom_transport_id = selected_custom_transport_id(&previous);
let forced_replacement = replacement_transport_id
.is_some_and(|expected| fresh_custom_transport_id == Some(expected));
let decision = if forced_replacement && previous.stable_id() == stable_id {
IrohConnectionInstallDecision::Install
} else if let Some(decision) = custom_transport_arbitration_override(
existing.as_ref().is_some_and(|state| state.alive),
existing_custom_transport_id,
fresh_custom_transport_id,
replacement_transport_id,
) {
decision
} else {
match connection_direction {
IrohConnectionDirection::Inbound => {
decide_inbound_install(existing, prefers_outbound)
}
IrohConnectionDirection::Outbound => {
decide_outbound_install(existing, prefers_outbound)
}
}
};
match decision {
IrohConnectionInstallDecision::Install => {}
IrohConnectionInstallDecision::ReplaceExisting {
close_existing_reason,
} => {
previous.close(0u8.into(), close_existing_reason.as_bytes());
}
IrohConnectionInstallDecision::KeepExisting { close_fresh_reason } => {
kept_existing = Some(ExternalConnectionInstallOutcome::KeptExisting {
fresh_transport_stable_id: stable_id as u64,
kept_transport_stable_id: previous.stable_id() as u64,
});
connection.close(0u8.into(), close_fresh_reason.as_bytes());
}
}
}
if kept_existing.is_none() {
conns.insert(endpoint_id, connection.clone());
insert_times.insert(
endpoint_id,
IrohConnectionMetadata {
inserted_at: Instant::now(),
direction: connection_direction,
},
);
}
}
if let Some(outcome) = kept_existing {
if let Some(sender) = install_outcome_sender.take() {
let _ = sender.send(outcome);
}
return Ok(());
}
if let Some(sender) = install_outcome_sender.take() {
let _ = sender.send(ExternalConnectionInstallOutcome::Installed {
transport_stable_id: stable_id as u64,
});
}
if let (Some(mgr), Some(tx)) = (&heartbeat_manager, &heartbeat_health_tx) {
mgr.start_connection(
heartbeat_key.clone(),
connection.clone(),
heartbeat_config.clone(),
tx.clone(),
)
.await;
}
event_sender
.send(AcceptEvent::Accepted {
endpoint_id,
transport_stable_id: stable_id as u64,
})
.ok();
loop {
tokio::select! {
biased;
res = connection.accept_bi() => {
match res {
Ok((send, recv)) => {
let _ = stream_sender.send(IncomingStream {
endpoint_id,
transport_stable_id: stable_id as u64,
stream: IncomingStreamType::Bi(send, recv),
}).await;
}
Err(_) => {
if should_break_accept_loop(&connection).await {
break;
}
},
}
}
res = connection.accept_uni() => {
match res {
Ok(recv) => {
let conn_clone = connection.clone();
let mgr_clone = heartbeat_manager.clone();
let probe_registry = probe_registry.clone();
let endpoint_key = endpoint_key.clone();
let heartbeat_key = heartbeat_key.clone();
let transport_stable_id = stable_id as u64;
let notice_key = manual_disconnect_notice_key(
endpoint_id,
transport_stable_id,
);
let notices = manual_disconnect_notices.clone();
let application_streams = stream_sender.clone();
tokio::spawn(async move {
match classify_incoming_uni(recv).await {
IncomingUniClassification::Control { type_id, payload } => {
use crate::heartbeat::codec;
if type_id == codec::TYPE_PING {
if !respond_to_ping(&conn_clone, &payload).await {
eprintln!(
"[OpenRTC][iroh-probe] failed to send pong endpoint_id={} transport_stable_id={}",
endpoint_key,
transport_stable_id,
);
}
} else if type_id == codec::TYPE_PONG {
if let Some(pong) = parse_incoming_pong(&payload) {
probe_registry
.deliver_pong(
&endpoint_key,
transport_stable_id,
&pong,
)
.await;
if let Some(manager) = mgr_clone {
manager.deliver_pong(&heartbeat_key, pong).await;
}
}
} else if type_id == codec::TYPE_MANUAL_DISCONNECT {
notices.write().await.insert(notice_key);
conn_clone.close(
0u8.into(),
crate::lifecycle_reason::REASON_MANUAL_DISCONNECT
.as_bytes(),
);
}
}
IncomingUniClassification::Application(recv) => {
let _ = application_streams
.send(IncomingStream {
endpoint_id,
transport_stable_id,
stream: IncomingStreamType::Uni(recv),
})
.await;
}
IncomingUniClassification::MalformedControl => {}
}
});
}
Err(_) => {
if should_break_accept_loop(&connection).await {
break;
}
},
}
}
_ = connection.closed() => break,
}
}
if let Some(ref mgr) = heartbeat_manager {
mgr.stop_connection(&heartbeat_key).await;
}
let close_reason = connection.close_reason();
let close_reason_debug = close_reason.as_ref().map(|reason| format!("{:?}", reason));
let was_locally_closed = matches!(
close_reason,
Some(iroh::endpoint::ConnectionError::LocallyClosed)
);
event_sender
.send(AcceptEvent::Closed {
endpoint_id,
transport_stable_id: stable_id as u64,
error: close_reason_debug,
was_locally_closed,
})
.ok();
{
let mut conns = connections.write().await;
let should_remove = conns
.get(&endpoint_id)
.map(|current| current.stable_id() == stable_id)
.unwrap_or(false);
if should_remove {
conns.remove(&endpoint_id);
connection_inserted_at.write().await.remove(&endpoint_id);
}
}
Ok(())
}
impl PlutoniumProtocol {
pub const ALPN: &[u8] = b"plutonium/p2p/1";
fn new(
event_sender: broadcast::Sender<AcceptEvent>,
stream_sender: async_channel::Sender<IncomingStream>,
connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
local_endpoint_id: EndpointId,
manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
probe_registry: IrohProbeRegistry,
inbound_replacement_authorizations: InboundReplacementAuthorizations,
) -> Self {
Self {
event_sender,
stream_sender,
connections,
connection_inserted_at,
local_endpoint_id,
manual_disconnect_notices,
probe_registry,
heartbeat_manager: None,
heartbeat_health_tx: None,
heartbeat_config: HeartbeatConfig::default(),
inbound_replacement_authorizations,
}
}
pub fn with_heartbeat(
mut self,
manager: IrohHeartbeatManager,
health_tx: mpsc::Sender<HealthTransition>,
config: HeartbeatConfig,
) -> Self {
self.heartbeat_manager = Some(manager);
self.heartbeat_health_tx = Some(health_tx);
self.heartbeat_config = config;
self
}
async fn handle_connection(
self,
connection: Connection,
) -> std::result::Result<(), AcceptError> {
let replacement_transport_id = consume_authorized_inbound_replacement(
&connection,
&self.inbound_replacement_authorizations,
)
.await;
run_connection_loop(
connection,
self.event_sender.clone(),
self.stream_sender.clone(),
self.connections.clone(),
self.connection_inserted_at.clone(),
self.local_endpoint_id,
IrohConnectionDirection::Inbound,
self.manual_disconnect_notices.clone(),
self.probe_registry.clone(),
self.heartbeat_manager.clone(),
self.heartbeat_health_tx.clone(),
self.heartbeat_config.clone(),
replacement_transport_id,
None,
)
.await
}
}
impl ProtocolHandler for PlutoniumProtocol {
#[allow(refining_impl_trait)]
fn accept(
&self,
connection: Connection,
) -> impl n0_future::Future<Output = std::result::Result<(), AcceptError>> + std::marker::Send
{
let proto = self.clone();
async move { proto.handle_connection(connection).await }
}
}
async fn connect(
endpoint: &Endpoint,
endpoint_id: EndpointId,
event_sender: async_channel::Sender<ConnectEvent>,
connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
stream_sender: async_channel::Sender<IncomingStream>,
manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
probe_registry: IrohProbeRegistry,
heartbeat_manager: Option<IrohHeartbeatManager>,
heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
heartbeat_config: HeartbeatConfig,
) -> Result<()> {
if outbound_dial_precheck(endpoint, endpoint_id, &connections, &connection_inserted_at).await {
event_sender.send(ConnectEvent::Connected).await?;
return Ok(());
}
let connection = endpoint
.connect(endpoint_id, PlutoniumProtocol::ALPN)
.await?;
let stable_id = connection.stable_id();
{
let mut conns = connections.write().await;
let mut insert_times = connection_inserted_at.write().await;
if let Some(previous) = conns.get(&endpoint_id).cloned() {
let previous_direction = insert_times
.get(&endpoint_id)
.map(|metadata| metadata.direction)
.unwrap_or(IrohConnectionDirection::Outbound);
let previous_age_ms = insert_times
.get(&endpoint_id)
.map(|metadata| metadata.inserted_at.elapsed().as_millis() as u64)
.unwrap_or(0);
match decide_outbound_install(
Some(ExistingConnectionState {
same_stable_id: previous.stable_id() == stable_id,
alive: previous.close_reason().is_none(),
direction: previous_direction,
age_ms: previous_age_ms,
}),
local_prefers_outbound(endpoint.id(), endpoint_id),
) {
IrohConnectionInstallDecision::Install => {}
IrohConnectionInstallDecision::ReplaceExisting {
close_existing_reason,
} => {
previous.close(0u8.into(), close_existing_reason.as_bytes());
}
IrohConnectionInstallDecision::KeepExisting { close_fresh_reason } => {
connection.close(0u8.into(), close_fresh_reason.as_bytes());
event_sender.send(ConnectEvent::Connected).await?;
let _ = event_sender
.send(ConnectEvent::Closed { error: None })
.await;
return Ok(());
}
}
}
conns.insert(endpoint_id, connection.clone());
insert_times.insert(
endpoint_id,
IrohConnectionMetadata {
inserted_at: Instant::now(),
direction: IrohConnectionDirection::Outbound,
},
);
}
event_sender.send(ConnectEvent::Connected).await?;
let (accept_tx, _) = broadcast::channel(1);
run_connection_loop(
connection,
accept_tx,
stream_sender,
connections.clone(),
connection_inserted_at,
endpoint.id(),
IrohConnectionDirection::Outbound,
manual_disconnect_notices,
probe_registry,
heartbeat_manager,
heartbeat_health_tx,
heartbeat_config,
None,
None,
)
.await
.ok();
event_sender
.send(ConnectEvent::Closed { error: None })
.await?;
Ok(())
}
async fn connect_addr(
endpoint: &Endpoint,
endpoint_id: EndpointId,
endpoint_addr: EndpointAddr,
event_sender: async_channel::Sender<ConnectEvent>,
connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
stream_sender: async_channel::Sender<IncomingStream>,
manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
probe_registry: IrohProbeRegistry,
heartbeat_manager: Option<IrohHeartbeatManager>,
heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
heartbeat_config: HeartbeatConfig,
) -> Result<()> {
if outbound_dial_precheck(endpoint, endpoint_id, &connections, &connection_inserted_at).await {
event_sender.send(ConnectEvent::Connected).await?;
return Ok(());
}
let connection = endpoint
.connect(endpoint_addr, PlutoniumProtocol::ALPN)
.await?;
let stable_id = connection.stable_id();
{
let mut conns = connections.write().await;
let mut insert_times = connection_inserted_at.write().await;
if let Some(previous) = conns.get(&endpoint_id).cloned() {
let previous_direction = insert_times
.get(&endpoint_id)
.map(|metadata| metadata.direction)
.unwrap_or(IrohConnectionDirection::Outbound);
let previous_age_ms = insert_times
.get(&endpoint_id)
.map(|metadata| metadata.inserted_at.elapsed().as_millis() as u64)
.unwrap_or(0);
match decide_outbound_install(
Some(ExistingConnectionState {
same_stable_id: previous.stable_id() == stable_id,
alive: previous.close_reason().is_none(),
direction: previous_direction,
age_ms: previous_age_ms,
}),
local_prefers_outbound(endpoint.id(), endpoint_id),
) {
IrohConnectionInstallDecision::Install => {}
IrohConnectionInstallDecision::ReplaceExisting {
close_existing_reason,
} => {
previous.close(0u8.into(), close_existing_reason.as_bytes());
}
IrohConnectionInstallDecision::KeepExisting { close_fresh_reason } => {
connection.close(0u8.into(), close_fresh_reason.as_bytes());
event_sender.send(ConnectEvent::Connected).await?;
let _ = event_sender
.send(ConnectEvent::Closed { error: None })
.await;
return Ok(());
}
}
}
conns.insert(endpoint_id, connection.clone());
insert_times.insert(
endpoint_id,
IrohConnectionMetadata {
inserted_at: Instant::now(),
direction: IrohConnectionDirection::Outbound,
},
);
}
event_sender.send(ConnectEvent::Connected).await?;
let (accept_tx, _) = broadcast::channel(1);
run_connection_loop(
connection,
accept_tx,
stream_sender,
connections.clone(),
connection_inserted_at,
endpoint.id(),
IrohConnectionDirection::Outbound,
manual_disconnect_notices,
probe_registry,
heartbeat_manager,
heartbeat_health_tx,
heartbeat_config,
None,
None,
)
.await
.ok();
event_sender
.send(ConnectEvent::Closed { error: None })
.await?;
Ok(())
}
async fn install_replacement_connection(
endpoint: &Endpoint,
connection: Connection,
replacement_transport_id: u64,
event_sender: async_channel::Sender<ReplacementConnectEvent>,
connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
stream_sender: async_channel::Sender<IncomingStream>,
manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
probe_registry: IrohProbeRegistry,
heartbeat_manager: Option<IrohHeartbeatManager>,
heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
heartbeat_config: HeartbeatConfig,
) -> Result<()> {
fn replacement_path_summary(connection: &Connection) -> Vec<String> {
connection
.paths()
.iter()
.map(|path| {
format!(
"selected={} remote_addr={:?}",
path.is_selected(),
path.remote_addr()
)
})
.collect()
}
eprintln!(
"[OpenRTC][custom-transport] replacement admission waiting stable_id={} remote_endpoint_id={} expected_transport_id={} paths={:?}",
connection.stable_id(),
connection.remote_id(),
replacement_transport_id,
replacement_path_summary(&connection)
);
let selected_expected_transport = || {
connection.paths().iter().any(|path| {
path.is_selected()
&& matches!(
path.remote_addr(),
iroh::TransportAddr::Custom(addr) if addr.id() == replacement_transport_id
)
})
};
let selected = tokio::time::timeout(std::time::Duration::from_secs(5), async {
loop {
if selected_expected_transport() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
}
})
.await
.is_ok();
if !selected {
eprintln!(
"[OpenRTC][custom-transport] replacement admission rejected stable_id={} remote_endpoint_id={} expected_transport_id={} paths={:?}",
connection.stable_id(),
connection.remote_id(),
replacement_transport_id,
replacement_path_summary(&connection)
);
connection.close(0u8.into(), b"custom-transport-not-selected");
anyhow::bail!(
"replacement connection did not select custom transport id {replacement_transport_id}"
);
}
eprintln!(
"[OpenRTC][custom-transport] replacement admission accepted stable_id={} remote_endpoint_id={} transport_id={} paths={:?}",
connection.stable_id(),
connection.remote_id(),
replacement_transport_id,
replacement_path_summary(&connection)
);
let (accept_tx, _) = broadcast::channel(1);
let (install_outcome_sender, install_outcome_receiver) = oneshot::channel();
let connection_loop = tokio::spawn(run_connection_loop(
connection,
accept_tx,
stream_sender,
connections,
connection_inserted_at,
endpoint.id(),
IrohConnectionDirection::Outbound,
manual_disconnect_notices,
probe_registry,
heartbeat_manager,
heartbeat_health_tx,
heartbeat_config,
Some(replacement_transport_id),
Some(install_outcome_sender),
));
let install_outcome = install_outcome_receiver
.await
.map_err(|_| anyhow::anyhow!("replacement connection install task ended"))?;
let transport_stable_id = match install_outcome {
ExternalConnectionInstallOutcome::Installed {
transport_stable_id,
} => transport_stable_id,
ExternalConnectionInstallOutcome::KeptExisting {
fresh_transport_stable_id,
kept_transport_stable_id,
} => {
anyhow::bail!(
"replacement connection lost arbitration fresh_stable_id={} kept_stable_id={}",
fresh_transport_stable_id,
kept_transport_stable_id,
);
}
};
event_sender
.send(ReplacementConnectEvent::Connected {
transport_stable_id,
})
.await?;
let _ = connection_loop.await;
let _ = event_sender
.send(ReplacementConnectEvent::Closed { error: None })
.await;
Ok(())
}
#[allow(unused_variables)]
async fn connect_replacement_addr(
endpoint: &Endpoint,
_endpoint_id: EndpointId,
endpoint_addr: EndpointAddr,
replacement_transport_id: u64,
event_sender: async_channel::Sender<ReplacementConnectEvent>,
connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
stream_sender: async_channel::Sender<IncomingStream>,
manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
probe_registry: IrohProbeRegistry,
heartbeat_manager: Option<IrohHeartbeatManager>,
heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
heartbeat_config: HeartbeatConfig,
) -> Result<()> {
#[cfg(not(openrtc_iroh_preferred_transport_api))]
anyhow::bail!(
"custom transport replacement requires a host Iroh build with the preferred-transport API"
);
#[cfg(openrtc_iroh_preferred_transport_api)]
{
let preferred_transport_addr = endpoint_addr
.addrs
.iter()
.find(|addr| {
matches!(
addr,
iroh::TransportAddr::Custom(custom)
if custom.id() == replacement_transport_id
)
})
.cloned()
.ok_or_else(|| {
anyhow::anyhow!(
"replacement address does not contain custom transport id {replacement_transport_id}"
)
})?;
eprintln!(
"[OpenRTC][custom-transport] replacement dial pinned remote_endpoint_id={} transport_id={} addr={:?}",
endpoint_addr.id, replacement_transport_id, preferred_transport_addr
);
let connection = endpoint
.connect_with_opts(
endpoint_addr,
PlutoniumProtocol::ALPN,
iroh::endpoint::ConnectOptions::new()
.with_preferred_transport_addr(preferred_transport_addr),
)
.await?
.await?;
install_replacement_connection(
endpoint,
connection,
replacement_transport_id,
event_sender,
connections,
connection_inserted_at,
stream_sender,
manual_disconnect_notices,
probe_registry,
heartbeat_manager,
heartbeat_health_tx,
heartbeat_config,
)
.await
}
}
#[derive(Debug, Clone)]
pub struct IrohNativeNode {
endpoint: Endpoint,
#[allow(dead_code)]
router: Option<Router>,
accept_events: broadcast::Sender<AcceptEvent>,
connections: Arc<RwLock<HashMap<EndpointId, Connection>>>,
connection_inserted_at: Arc<RwLock<HashMap<EndpointId, IrohConnectionMetadata>>>,
incoming_streams: async_channel::Sender<IncomingStream>,
incoming_streams_receiver: async_channel::Receiver<IncomingStream>,
manual_disconnect_notices: Arc<RwLock<HashSet<String>>>,
probe_registry: IrohProbeRegistry,
heartbeat_manager: Option<IrohHeartbeatManager>,
heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
heartbeat_config: HeartbeatConfig,
inbound_replacement_authorizations: InboundReplacementAuthorizations,
}
impl IrohNativeNode {
pub async fn spawn_with_endpoint(endpoint: Endpoint) -> Result<Self> {
Self::spawn_with_endpoint_config(endpoint, true).await
}
pub async fn spawn_with_endpoint_no_router(endpoint: Endpoint) -> Result<Self> {
Self::spawn_with_endpoint_config(endpoint, false).await
}
async fn spawn_with_endpoint_config(endpoint: Endpoint, spawn_router: bool) -> Result<Self> {
Self::spawn_with_endpoint_config_and_heartbeat(
endpoint,
spawn_router,
None,
None,
HeartbeatConfig::default(),
)
.await
}
pub async fn spawn_with_heartbeat(
endpoint: Endpoint,
heartbeat_manager: IrohHeartbeatManager,
heartbeat_health_tx: mpsc::Sender<HealthTransition>,
heartbeat_config: HeartbeatConfig,
) -> Result<Self> {
Self::spawn_with_endpoint_config_and_heartbeat(
endpoint,
true,
Some(heartbeat_manager),
Some(heartbeat_health_tx),
heartbeat_config,
)
.await
}
async fn spawn_with_endpoint_config_and_heartbeat(
endpoint: Endpoint,
spawn_router: bool,
heartbeat_manager: Option<IrohHeartbeatManager>,
heartbeat_health_tx: Option<mpsc::Sender<HealthTransition>>,
heartbeat_config: HeartbeatConfig,
) -> Result<Self> {
let (event_sender, _) = broadcast::channel(128);
let (stream_sender, stream_receiver) = async_channel::bounded(64);
let connections = Arc::new(RwLock::new(HashMap::new()));
let connection_inserted_at = Arc::new(RwLock::new(HashMap::new()));
let manual_disconnect_notices = Arc::new(RwLock::new(HashSet::new()));
let probe_registry = IrohProbeRegistry::new();
let inbound_replacement_authorizations = Arc::new(RwLock::new(HashMap::new()));
let router = if spawn_router {
let proto = PlutoniumProtocol::new(
event_sender.clone(),
stream_sender.clone(),
connections.clone(),
connection_inserted_at.clone(),
endpoint.id(),
manual_disconnect_notices.clone(),
probe_registry.clone(),
inbound_replacement_authorizations.clone(),
);
let proto = if let (Some(mgr), Some(tx)) =
(heartbeat_manager.clone(), heartbeat_health_tx.clone())
{
proto.with_heartbeat(mgr, tx, heartbeat_config.clone())
} else {
proto
};
Some(
Router::builder(endpoint.clone())
.accept(PlutoniumProtocol::ALPN, proto)
.spawn(),
)
} else {
None
};
Ok(Self {
endpoint,
router,
accept_events: event_sender,
connections,
connection_inserted_at,
incoming_streams: stream_sender,
incoming_streams_receiver: stream_receiver,
manual_disconnect_notices,
probe_registry,
heartbeat_manager,
heartbeat_health_tx,
heartbeat_config,
inbound_replacement_authorizations,
})
}
pub fn endpoint(&self) -> &Endpoint {
&self.endpoint
}
pub async fn node_addr(&self) -> Result<EndpointAddr> {
Ok(self.endpoint.watch_addr().get())
}
pub async fn is_connected(&self, endpoint_id: EndpointId) -> bool {
let conns = self.connections.read().await;
if let Some(conn) = conns.get(&endpoint_id) {
conn.close_reason().is_none()
} else {
false
}
}
pub fn accept_events(&self) -> futures::stream::BoxStream<'static, AcceptEvent> {
let receiver = self.accept_events.subscribe();
Box::pin(
BroadcastStream::new(receiver).filter_map(|event| futures::future::ready(event.ok())),
)
}
pub fn connect(
&self,
endpoint_id: EndpointId,
) -> futures::stream::BoxStream<'static, ConnectEvent> {
let (event_sender, event_receiver) = async_channel::bounded(16);
let endpoint = self.endpoint.clone();
let connections = self.connections.clone();
let connection_inserted_at = self.connection_inserted_at.clone();
let stream_sender = self.incoming_streams.clone();
let hb_mgr = self.heartbeat_manager.clone();
let hb_tx = self.heartbeat_health_tx.clone();
let hb_cfg = self.heartbeat_config.clone();
let manual_notices = self.manual_disconnect_notices.clone();
let probe_registry = self.probe_registry.clone();
tokio::spawn(async move {
let result = connect(
&endpoint,
endpoint_id,
event_sender.clone(),
connections,
connection_inserted_at,
stream_sender,
manual_notices,
probe_registry,
hb_mgr,
hb_tx,
hb_cfg,
)
.await;
if let Err(error) = result {
let _ = event_sender
.send(ConnectEvent::Closed {
error: Some(error.to_string()),
})
.await;
}
});
Box::pin(event_receiver)
}
pub fn connect_addr(
&self,
endpoint_id: EndpointId,
endpoint_addr: EndpointAddr,
) -> futures::stream::BoxStream<'static, ConnectEvent> {
let (event_sender, event_receiver) = async_channel::bounded(16);
let endpoint = self.endpoint.clone();
let connections = self.connections.clone();
let connection_inserted_at = self.connection_inserted_at.clone();
let stream_sender = self.incoming_streams.clone();
let hb_mgr = self.heartbeat_manager.clone();
let hb_tx = self.heartbeat_health_tx.clone();
let hb_cfg = self.heartbeat_config.clone();
let manual_notices = self.manual_disconnect_notices.clone();
let probe_registry = self.probe_registry.clone();
tokio::spawn(async move {
let result = connect_addr(
&endpoint,
endpoint_id,
endpoint_addr,
event_sender.clone(),
connections,
connection_inserted_at,
stream_sender,
manual_notices,
probe_registry,
hb_mgr,
hb_tx,
hb_cfg,
)
.await;
if let Err(error) = result {
let _ = event_sender
.send(ConnectEvent::Closed {
error: Some(error.to_string()),
})
.await;
}
});
Box::pin(event_receiver)
}
pub fn connect_replacement_addr(
&self,
endpoint_id: EndpointId,
endpoint_addr: EndpointAddr,
replacement_transport_id: u64,
) -> futures::stream::BoxStream<'static, ReplacementConnectEvent> {
let (event_sender, event_receiver) = async_channel::bounded(16);
let endpoint = self.endpoint.clone();
let connections = self.connections.clone();
let connection_inserted_at = self.connection_inserted_at.clone();
let stream_sender = self.incoming_streams.clone();
let hb_mgr = self.heartbeat_manager.clone();
let hb_tx = self.heartbeat_health_tx.clone();
let hb_cfg = self.heartbeat_config.clone();
let manual_notices = self.manual_disconnect_notices.clone();
let probe_registry = self.probe_registry.clone();
tokio::spawn(async move {
if let Err(error) = connect_replacement_addr(
&endpoint,
endpoint_id,
endpoint_addr,
replacement_transport_id,
event_sender.clone(),
connections,
connection_inserted_at,
stream_sender,
manual_notices,
probe_registry,
hb_mgr,
hb_tx,
hb_cfg,
)
.await
{
let _ = event_sender
.send(ReplacementConnectEvent::Closed {
error: Some(error.to_string()),
})
.await;
}
});
Box::pin(event_receiver)
}
pub async fn disconnect(&self, endpoint_id: EndpointId) -> Result<()> {
self.disconnect_with_reason(
endpoint_id,
crate::lifecycle_reason::REASON_DISCONNECTED_BY_USER,
)
.await
}
pub async fn disconnect_with_reason(
&self,
endpoint_id: EndpointId,
reason: &str,
) -> Result<()> {
let connection = {
let mut conns = self.connections.write().await;
conns.remove(&endpoint_id)
};
if let Some(conn) = connection {
if std::env::var("PLUTO_RTC_TEARDOWN_TRACE").is_ok() {
eprintln!(
"[PlutoRTC][teardown-trace] NativeNode::disconnect endpoint_id={}",
endpoint_id
);
}
conn.close(1u8.into(), reason.as_bytes());
}
Ok(())
}
pub async fn disconnect_with_reason_if_current(
&self,
endpoint_id: EndpointId,
expected_transport_stable_id: u64,
reason: &str,
) -> Result<bool> {
let connection = {
let mut connections = self.connections.write().await;
let is_current = connections.get(&endpoint_id).is_some_and(|connection| {
connection.stable_id() as u64 == expected_transport_stable_id
});
if !is_current {
return Ok(false);
}
connections.remove(&endpoint_id)
};
if let Some(connection) = connection {
connection.close(1u8.into(), reason.as_bytes());
return Ok(true);
}
Ok(false)
}
pub async fn open_bi(&self, endpoint_id: EndpointId) -> Result<(SendStream, RecvStream)> {
let (_, send, recv) = self.open_bi_with_transport_stable_id(endpoint_id).await?;
Ok((send, recv))
}
pub async fn open_bi_with_transport_stable_id(
&self,
endpoint_id: EndpointId,
) -> Result<(u64, SendStream, RecvStream)> {
let connection = {
let conns = self.connections.read().await;
conns.get(&endpoint_id).cloned()
};
if let Some(conn) = connection {
let transport_stable_id = conn.stable_id() as u64;
let (send, recv) = conn.open_bi().await?;
Ok((transport_stable_id, send, recv))
} else {
Err(anyhow::anyhow!("No active connection to {}", endpoint_id))
}
}
pub async fn open_uni(&self, endpoint_id: EndpointId) -> Result<SendStream> {
let connection = {
let conns = self.connections.read().await;
conns.get(&endpoint_id).cloned()
};
if let Some(conn) = connection {
let send = conn.open_uni().await?;
Ok(send)
} else {
Err(anyhow::anyhow!("No active connection to {}", endpoint_id))
}
}
pub async fn probe_connection(
&self,
endpoint_id: EndpointId,
timeout: std::time::Duration,
) -> Option<IrohConnectionProbe> {
let connection = self.connections.read().await.get(&endpoint_id).cloned()?;
let probed_stable_id = connection.stable_id() as u64;
let responsive = self
.probe_registry
.probe(
&endpoint_id.to_string(),
probed_stable_id,
&connection,
timeout,
)
.await;
let current_stable_id = self
.connections
.read()
.await
.get(&endpoint_id)
.map(|current| current.stable_id() as u64)?;
if current_stable_id != probed_stable_id {
return Some(IrohConnectionProbe {
transport_stable_id: probed_stable_id,
responsive: false,
});
}
Some(IrohConnectionProbe {
transport_stable_id: probed_stable_id,
responsive,
})
}
pub fn incoming_streams_stream(&self) -> async_channel::Receiver<IncomingStream> {
self.incoming_streams_receiver.clone()
}
pub(crate) async fn authorize_inbound_replacement(
&self,
endpoint_id: EndpointId,
transport_id: u64,
ttl: std::time::Duration,
) -> Instant {
let expires_at = Instant::now() + ttl;
self.inbound_replacement_authorizations
.write()
.await
.insert(
endpoint_id,
InboundReplacementAuthorization {
transport_id,
expires_at,
},
);
expires_at
}
pub(crate) async fn revoke_inbound_replacement_if_current(
&self,
endpoint_id: EndpointId,
transport_id: u64,
expires_at: Instant,
) -> bool {
let mut authorizations = self.inbound_replacement_authorizations.write().await;
if authorizations.get(&endpoint_id).is_some_and(|current| {
current.transport_id == transport_id && current.expires_at == expires_at
}) {
authorizations.remove(&endpoint_id);
true
} else {
false
}
}
pub async fn accept_external_connection(&self, connection: Connection) -> Result<()> {
let replacement_transport_id = consume_authorized_inbound_replacement(
&connection,
&self.inbound_replacement_authorizations,
)
.await;
run_connection_loop(
connection,
self.accept_events.clone(),
self.incoming_streams.clone(),
self.connections.clone(),
self.connection_inserted_at.clone(),
self.endpoint.id(),
IrohConnectionDirection::Inbound,
self.manual_disconnect_notices.clone(),
self.probe_registry.clone(),
self.heartbeat_manager.clone(),
self.heartbeat_health_tx.clone(),
self.heartbeat_config.clone(),
replacement_transport_id,
None,
)
.await
.map_err(|e| anyhow::anyhow!(e.to_string()))
}
pub(crate) async fn accept_external_connection_with_install_notifier(
&self,
connection: Connection,
install_outcome_sender: oneshot::Sender<ExternalConnectionInstallOutcome>,
) -> Result<()> {
let replacement_transport_id = consume_authorized_inbound_replacement(
&connection,
&self.inbound_replacement_authorizations,
)
.await;
run_connection_loop(
connection,
self.accept_events.clone(),
self.incoming_streams.clone(),
self.connections.clone(),
self.connection_inserted_at.clone(),
self.endpoint.id(),
IrohConnectionDirection::Inbound,
self.manual_disconnect_notices.clone(),
self.probe_registry.clone(),
self.heartbeat_manager.clone(),
self.heartbeat_health_tx.clone(),
self.heartbeat_config.clone(),
replacement_transport_id,
Some(install_outcome_sender),
)
.await
.map_err(|e| anyhow::anyhow!(e.to_string()))
}
pub async fn get_connection(&self, endpoint_id: EndpointId) -> Option<Connection> {
let conns = self.connections.read().await;
conns.get(&endpoint_id).cloned()
}
pub async fn active_endpoint_ids(&self) -> Vec<EndpointId> {
let conns = self.connections.read().await;
conns.keys().cloned().collect()
}
pub async fn take_manual_disconnect_notice(
&self,
endpoint_id: EndpointId,
transport_stable_id: u64,
) -> bool {
self.manual_disconnect_notices
.write()
.await
.remove(&manual_disconnect_notice_key(
endpoint_id,
transport_stable_id,
))
}
}
#[cfg(test)]
mod tests {
use super::*;
use iroh::endpoint::Endpoint;
use tokio::io::AsyncWriteExt;
use tokio::time::{sleep, timeout, Duration};
async fn setup_protocol_endpoint() -> (
Router,
PlutoniumProtocol,
broadcast::Receiver<AcceptEvent>,
async_channel::Receiver<IncomingStream>,
Arc<RwLock<HashMap<EndpointId, Connection>>>,
) {
let (event_tx, event_rx) = broadcast::channel(16);
let (stream_tx, stream_rx) = async_channel::unbounded();
let connections = Arc::new(RwLock::new(HashMap::new()));
let connection_inserted_at = Arc::new(RwLock::new(HashMap::new()));
let endpoint = Endpoint::builder(iroh::endpoint::presets::N0)
.alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
.bind()
.await
.unwrap();
let protocol = PlutoniumProtocol::new(
event_tx.clone(),
stream_tx.clone(),
connections.clone(),
connection_inserted_at,
endpoint.id(),
Arc::new(RwLock::new(HashSet::new())),
IrohProbeRegistry::new(),
Arc::new(RwLock::new(HashMap::new())),
);
let router = Router::builder(endpoint)
.accept(PlutoniumProtocol::ALPN, Arc::new(protocol.clone()))
.spawn();
(router, protocol, event_rx, stream_rx, connections)
}
#[test]
fn manual_disconnect_notice_is_scoped_to_transport_generation() {
let endpoint_id = iroh::SecretKey::generate().public();
let old_transport = manual_disconnect_notice_key(endpoint_id, 41);
let replacement_transport = manual_disconnect_notice_key(endpoint_id, 42);
let mut notices = HashSet::from([old_transport.clone()]);
assert_ne!(old_transport, replacement_transport);
assert!(!notices.remove(&replacement_transport));
assert!(notices.remove(&old_transport));
}
#[test]
fn authorized_custom_transport_replaces_live_base_connection() {
assert_eq!(
custom_transport_arbitration_override(true, None, Some(4344901), Some(4344901)),
Some(IrohConnectionInstallDecision::ReplaceExisting {
close_existing_reason: "native-custom-transport-upgrade",
})
);
}
#[test]
fn live_custom_transport_wins_against_late_non_custom_duplicate() {
assert_eq!(
custom_transport_arbitration_override(true, Some(4344901), None, None),
Some(IrohConnectionInstallDecision::KeepExisting {
close_fresh_reason: "existing-custom-transport-preferred",
})
);
}
#[test]
fn unauthorized_custom_transport_uses_normal_direction_arbitration() {
assert_eq!(
custom_transport_arbitration_override(true, None, Some(4344901), None),
None
);
}
#[tokio::test]
async fn delayed_replacement_cleanup_cannot_revoke_a_newer_grant() {
let endpoint = Endpoint::builder(iroh::endpoint::presets::N0)
.relay_mode(iroh::RelayMode::Disabled)
.alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
.bind()
.await
.expect("endpoint");
let node = IrohNativeNode::spawn_with_endpoint_no_router(endpoint)
.await
.expect("native node");
let remote = iroh::SecretKey::generate().public();
let old_expiry = node
.authorize_inbound_replacement(remote, 4344901, Duration::from_secs(1))
.await;
let new_expiry = node
.authorize_inbound_replacement(remote, 4344901, Duration::from_secs(2))
.await;
assert!(
!node
.revoke_inbound_replacement_if_current(remote, 4344901, old_expiry)
.await
);
assert!(
node.revoke_inbound_replacement_if_current(remote, 4344901, new_expiry)
.await
);
}
#[tokio::test]
async fn node_addr_does_not_wait_for_an_unavailable_relay() {
let endpoint = Endpoint::builder(iroh::endpoint::presets::N0)
.relay_mode(iroh::RelayMode::Disabled)
.alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
.bind()
.await
.expect("bind relay-disabled endpoint");
let expected_id = endpoint.id();
let node = IrohNativeNode::spawn_with_endpoint(endpoint)
.await
.expect("spawn relay-disabled node");
let address = timeout(Duration::from_secs(1), node.node_addr())
.await
.expect("node_addr must not wait for relay readiness")
.expect("read current endpoint address");
assert_eq!(address.id, expected_id);
}
#[tokio::test]
async fn test_two_nodes_connect_and_exchange_streams() {
let (r1, _proto1, _events1, _streams1, _) = setup_protocol_endpoint().await;
let (r2, _proto2, mut events2, streams2, _) = setup_protocol_endpoint().await;
let ep1 = r1.endpoint();
let ep2 = r2.endpoint();
let addr2 = ep2.addr();
let conn_res = ep1.connect(addr2, PlutoniumProtocol::ALPN).await.unwrap();
let event = timeout(Duration::from_secs(5), events2.recv())
.await
.unwrap()
.unwrap();
match event {
AcceptEvent::Accepted { endpoint_id, .. } => assert_eq!(endpoint_id, ep1.id()),
_ => panic!("Expected AcceptEvent::Accepted"),
}
let (mut send1, _recv1) = conn_res.open_bi().await.unwrap();
send1.write_all(b"hello node2").await.unwrap();
let incoming = timeout(Duration::from_secs(5), streams2.recv())
.await
.unwrap()
.unwrap();
assert_eq!(incoming.endpoint_id, ep1.id());
let mut recv2 = match incoming.stream {
IncomingStreamType::Bi(_, r) => r,
_ => panic!("Expected Bi stream"),
};
let mut buf = [0u8; 11];
recv2.read_exact(&mut buf).await.unwrap();
assert_eq!(&buf, b"hello node2");
}
#[tokio::test]
async fn retiring_losing_control_candidate_preserves_buffered_verdict() {
let (r1, _proto1, _events1, _streams1, _) = setup_protocol_endpoint().await;
let (r2, _proto2, mut events2, streams2, _) = setup_protocol_endpoint().await;
let ep1 = r1.endpoint();
let ep2 = r2.endpoint();
let connection = ep1
.connect(ep2.addr(), PlutoniumProtocol::ALPN)
.await
.expect("connect protocol endpoints");
timeout(Duration::from_secs(5), events2.recv())
.await
.expect("host accept timeout")
.expect("host accept event");
let (mut dialer_send, mut dialer_recv) =
connection.open_bi().await.expect("open candidate stream");
dialer_send
.write_all(b"candidate")
.await
.expect("activate candidate stream");
dialer_send.flush().await.expect("flush candidate stream");
let incoming = timeout(Duration::from_secs(5), streams2.recv())
.await
.expect("host stream timeout")
.expect("host stream");
let IncomingStreamType::Bi(mut host_send, host_recv) = incoming.stream else {
panic!("expected bidirectional candidate stream");
};
let verdict = b"session-token-approved";
host_send
.write_all(verdict)
.await
.expect("write approval verdict");
host_send.flush().await.expect("flush approval verdict");
let (retired, received) = tokio::join!(
crate::client::retire_losing_native_control_candidate(host_send, host_recv),
dialer_recv.read_to_end(256),
);
retired.expect("retire candidate after peer acknowledges verdict");
assert_eq!(
received.expect("dialer reads losing-stream verdict"),
verdict,
);
}
#[tokio::test]
async fn retiring_replaced_control_send_preserves_buffered_verdict() {
let (r1, _proto1, _events1, _streams1, _) = setup_protocol_endpoint().await;
let (r2, _proto2, mut events2, streams2, _) = setup_protocol_endpoint().await;
let connection = r1
.endpoint()
.connect(r2.endpoint().addr(), PlutoniumProtocol::ALPN)
.await
.expect("connect protocol endpoints");
timeout(Duration::from_secs(5), events2.recv())
.await
.expect("host accept timeout")
.expect("host accept event");
let (mut dialer_send, mut dialer_recv) =
connection.open_bi().await.expect("open displaced stream");
dialer_send
.write_all(b"candidate")
.await
.expect("activate displaced stream");
dialer_send.flush().await.expect("flush displaced stream");
let incoming = timeout(Duration::from_secs(5), streams2.recv())
.await
.expect("host stream timeout")
.expect("host stream");
let IncomingStreamType::Bi(mut host_send, _host_recv) = incoming.stream else {
panic!("expected bidirectional displaced stream");
};
let verdict = b"session-token-approved";
host_send
.write_all(verdict)
.await
.expect("write approval verdict");
host_send.flush().await.expect("host should flush approval");
let displaced_send = Arc::new(tokio::sync::Mutex::new(host_send));
let (retired, received) = tokio::join!(
crate::client::retire_replaced_native_control_send(displaced_send),
dialer_recv.read_to_end(256),
);
retired.expect("retire replaced stream after peer acknowledges verdict");
assert_eq!(
received.expect("dialer reads replaced-stream verdict"),
verdict,
);
}
#[tokio::test]
async fn test_healthy_connection_not_replaced() {
let (r1, _proto1, _events1, _streams1, _) = setup_protocol_endpoint().await;
let (r2, _proto2, mut events2, _streams2, conns2) = setup_protocol_endpoint().await;
let ep1 = r1.endpoint();
let ep2 = r2.endpoint();
let addr2 = ep2.addr();
let addr1 = ep1.addr();
let _conn1 = ep1.connect(addr2, PlutoniumProtocol::ALPN).await.unwrap();
let _ = timeout(Duration::from_secs(5), events2.recv())
.await
.unwrap()
.unwrap();
let active_count = conns2.read().await.len();
assert_eq!(active_count, 1);
let original_stable_id = conns2.read().await.get(&ep1.id()).unwrap().stable_id();
let _conn2 = ep2.connect(addr1, PlutoniumProtocol::ALPN).await.unwrap();
sleep(Duration::from_millis(100)).await;
let current_conn = conns2.read().await.get(&ep1.id()).unwrap().clone();
assert_eq!(current_conn.stable_id(), original_stable_id);
}
#[tokio::test]
async fn idle_symmetric_iroh_heartbeat_send_uni_open_rate_bounded() {
use crate::heartbeat::idle_symmetric_heartbeat_max_send_uni_opens_upper_bound;
use crate::heartbeat::iroh_heartbeat::test_counters;
use std::sync::atomic::Ordering;
test_counters::reset_heartbeat_send_uni_count();
let tick = Duration::from_millis(200);
let heartbeat_config = HeartbeatConfig {
tick_interval: tick,
suspect_after: Duration::from_secs(10),
stale_after: Duration::from_secs(30),
send_timeout: Duration::from_secs(2),
};
let (tx1, _rx1) = mpsc::channel::<HealthTransition>(32);
let (tx2, _rx2) = mpsc::channel::<HealthTransition>(32);
let mgr1 = IrohHeartbeatManager::new();
let mgr2 = IrohHeartbeatManager::new();
let ep1 = Endpoint::builder(iroh::endpoint::presets::N0)
.alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
.bind()
.await
.unwrap();
let ep2 = Endpoint::builder(iroh::endpoint::presets::N0)
.alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
.bind()
.await
.unwrap();
let node1 = IrohNativeNode::spawn_with_heartbeat(ep1, mgr1, tx1, heartbeat_config.clone())
.await
.unwrap();
let node2 = IrohNativeNode::spawn_with_heartbeat(ep2, mgr2, tx2, heartbeat_config)
.await
.unwrap();
let remote_id = node2.endpoint().id();
let remote_addr = node2.node_addr().await.unwrap();
let mut conn_stream = node1.connect_addr(remote_id, remote_addr);
let connected = timeout(Duration::from_secs(5), async {
while let Some(ev) = conn_stream.next().await {
match ev {
ConnectEvent::Connected => return true,
ConnectEvent::Closed { .. } => return false,
}
}
false
})
.await
.unwrap();
assert!(connected, "expected outbound connect to reach Connected");
let observe = Duration::from_millis(900);
sleep(observe).await;
let observed = test_counters::HEARTBEAT_SEND_UNI_COUNT.load(Ordering::SeqCst);
let bound = idle_symmetric_heartbeat_max_send_uni_opens_upper_bound(observe, tick);
assert!(
observed <= bound,
"heartbeat send_uni opens should stay within idle symmetric model (observed={} bound={})",
observed,
bound
);
assert!(
observed >= 4,
"expected some heartbeat uni traffic after idle window (observed={})",
observed
);
}
#[tokio::test]
async fn active_probe_requires_remote_control_loop_round_trip() {
let ep1 = Endpoint::builder(iroh::endpoint::presets::N0)
.alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
.bind()
.await
.expect("bind first endpoint");
let ep2 = Endpoint::builder(iroh::endpoint::presets::N0)
.alpns(vec![PlutoniumProtocol::ALPN.to_vec()])
.bind()
.await
.expect("bind second endpoint");
let node1 = IrohNativeNode::spawn_with_endpoint(ep1)
.await
.expect("spawn first node");
let node2 = IrohNativeNode::spawn_with_endpoint(ep2)
.await
.expect("spawn second node");
let remote_id = node2.endpoint().id();
let mut events = node1.connect_addr(
remote_id,
node2.node_addr().await.expect("second node address"),
);
let connected = timeout(Duration::from_secs(5), async {
while let Some(event) = events.next().await {
if matches!(event, ConnectEvent::Connected) {
return true;
}
}
false
})
.await
.expect("connect timeout");
assert!(connected, "expected physical connection");
let probe = node1
.probe_connection(remote_id, Duration::from_secs(2))
.await
.expect("current physical generation");
assert!(
probe.responsive,
"remote runtime must return the typed pong"
);
assert_ne!(probe.transport_stable_id, 0);
}
#[tokio::test]
async fn application_uni_stream_prefix_is_replayed_after_control_classification() {
let (r1, _proto1, _events1, _streams1, _) = setup_protocol_endpoint().await;
let (r2, _proto2, mut events2, streams2, _) = setup_protocol_endpoint().await;
let connection = r1
.endpoint()
.connect(r2.endpoint().addr(), PlutoniumProtocol::ALPN)
.await
.expect("connect endpoints");
timeout(Duration::from_secs(5), events2.recv())
.await
.expect("accept event timeout")
.expect("accept event");
let payload = b"application-uni-payload-that-is-not-control";
let mut send = connection.open_uni().await.expect("open app uni stream");
send.write_all(payload)
.await
.expect("write app uni payload");
send.finish().expect("finish app uni payload");
let incoming = timeout(Duration::from_secs(5), streams2.recv())
.await
.expect("application stream timeout")
.expect("application stream");
let IncomingStreamType::Uni(mut recv) = incoming.stream else {
panic!("expected application uni stream");
};
let mut received = Vec::new();
tokio::io::AsyncReadExt::read_to_end(&mut recv, &mut received)
.await
.expect("read replayed application payload");
assert_eq!(received, payload);
}
}