From 0830b241a58fd119c1d121512a3b4d3dea04e734 Mon Sep 17 00:00:00 2001 From: Mick van Dijke <12992260+mickvandijke@users.noreply.github.com> Date: Tue, 28 Jul 2026 22:56:02 +0200 Subject: [PATCH 1/8] fix(relay): preserve verified tunnel lifecycle --- src/link_transport_impl.rs | 19 +- src/masque/mod.rs | 2 +- src/masque/relay_server.rs | 79 +++++- src/masque/relay_socket.rs | 162 +++++++++-- src/nat_traversal_api.rs | 565 ++++++++++++++++++++++++------------- src/p2p_endpoint.rs | 85 +++++- 6 files changed, 678 insertions(+), 234 deletions(-) diff --git a/src/link_transport_impl.rs b/src/link_transport_impl.rs index e772e176..483924f8 100644 --- a/src/link_transport_impl.rs +++ b/src/link_transport_impl.rs @@ -660,11 +660,13 @@ impl LinkTransport for P2pLinkTransport { // Get the underlying QUIC connection by address match endpoint.get_quic_connection(&socket_addr).await { Ok(Some(conn)) => { - // Extract peer public key from TLS identity - let public_key = conn - .peer_identity() - .and_then(|id| id.downcast::>().ok()) - .map(|boxed| *boxed); + // rustls exposes the authenticated identity as a + // certificate vector whose first entry is the RFC + // 7250 ML-DSA SPKI, not as a bare `Vec`. + let public_key = + crate::p2p_endpoint::extract_public_key_bytes_from_connection( + &conn, + ); let link_conn = P2pLinkConn::new(conn, public_key, socket_addr); Some((Ok(link_conn), endpoint)) } @@ -725,11 +727,8 @@ impl LinkTransport for P2pLinkTransport { .map_err(|e| LinkError::ConnectionFailed(e.to_string()))? .ok_or_else(|| LinkError::ConnectionFailed("Connection not found".to_string()))?; - // Extract peer public key from TLS identity - let public_key = conn - .peer_identity() - .and_then(|id| id.downcast::>().ok()) - .map(|boxed| *boxed); + // Preserve the identity authenticated by this exact connection. + let public_key = crate::p2p_endpoint::extract_public_key_bytes_from_connection(&conn); Ok(P2pLinkConn::new(conn, public_key, connected_addr)) }) diff --git a/src/masque/mod.rs b/src/masque/mod.rs index 1354d798..b8a4cced 100644 --- a/src/masque/mod.rs +++ b/src/masque/mod.rs @@ -127,4 +127,4 @@ pub use relay_server::{ pub use relay_session::{ RelayPeerId, RelaySession, RelaySessionConfig, RelaySessionState, RelaySessionStats, }; -pub use relay_socket::{MasqueRelaySocket, RawRelayStreams}; +pub use relay_socket::{MasqueRelaySocket, RawRelayStreams, RelayTunnelControl}; diff --git a/src/masque/relay_server.rs b/src/masque/relay_server.rs index 3790dcdc..b2c8299f 100644 --- a/src/masque/relay_server.rs +++ b/src/masque/relay_server.rs @@ -32,7 +32,7 @@ //! ``` use bytes::Bytes; -use parking_lot::RwLock as ParkingRwLock; +use parking_lot::{Mutex as ParkingMutex, RwLock as ParkingRwLock}; use std::collections::HashMap; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; use std::sync::Arc; @@ -474,6 +474,20 @@ struct Reservation { released_at: Instant, } +/// RAII lease keeping a relay's parent QUIC connection out of ordinary +/// peer-table pruning for the lifetime of one CONNECT-UDP forwarding stream. +pub(crate) struct RelayControlConnectionGuard { + server: Arc, + stable_id: usize, +} + +impl Drop for RelayControlConnectionGuard { + fn drop(&mut self) { + self.server + .unprotect_relay_control_connection(self.stable_id); + } +} + /// MASQUE Relay Server /// /// Manages multiple relay sessions and coordinates datagram forwarding @@ -526,6 +540,14 @@ pub struct MasqueRelayServer { /// aborting its reader/writer tasks, so the socket is exclusively owned by the /// time it can be leased. forwarding: RwLock>, + /// QUIC connections currently carrying one or more CONNECT-UDP forwarding + /// streams, keyed by Quinn's stable connection id. + /// + /// Ordinary peer-table maintenance must not force-close these connections: + /// doing so destroys an otherwise healthy, canary-verified relay. Values are + /// reference counts because one authenticated connection may carry multiple + /// relay streams. + relay_control_connections: ParkingMutex>, /// Mutex stripes serializing concurrent CONNECTs from the same authenticated /// identity (ADR-011), indexed by the first fingerprint byte. Length is /// `RELAY_PEER_LOCK_STRIPES`. @@ -589,6 +611,7 @@ impl MasqueRelayServer { upnp_mappings: RwLock::new(HashMap::new()), reservations: RwLock::new(HashMap::new()), forwarding: RwLock::new(HashMap::new()), + relay_control_connections: ParkingMutex::new(HashMap::new()), peer_locks: (0..RELAY_PEER_LOCK_STRIPES) .map(|_| Mutex::new(())) .collect(), @@ -661,6 +684,7 @@ impl MasqueRelayServer { upnp_mappings: RwLock::new(HashMap::new()), reservations: RwLock::new(HashMap::new()), forwarding: RwLock::new(HashMap::new()), + relay_control_connections: ParkingMutex::new(HashMap::new()), peer_locks: (0..RELAY_PEER_LOCK_STRIPES) .map(|_| Mutex::new(())) .collect(), @@ -682,6 +706,40 @@ impl MasqueRelayServer { } } + /// Protect a QUIC connection from ordinary peer-table pruning while it + /// carries a live CONNECT-UDP forwarding stream. + pub(crate) fn protect_relay_control_connection( + self: &Arc, + stable_id: usize, + ) -> RelayControlConnectionGuard { + let mut protected = self.relay_control_connections.lock(); + *protected.entry(stable_id).or_insert(0) += 1; + drop(protected); + RelayControlConnectionGuard { + server: Arc::clone(self), + stable_id, + } + } + + /// Whether `stable_id` currently carries a live CONNECT-UDP stream. + pub(crate) fn is_relay_control_connection(&self, stable_id: usize) -> bool { + self.relay_control_connections + .lock() + .get(&stable_id) + .is_some_and(|count| *count > 0) + } + + fn unprotect_relay_control_connection(&self, stable_id: usize) { + let mut protected = self.relay_control_connections.lock(); + let Some(count) = protected.get_mut(&stable_id) else { + return; + }; + *count -= 1; + if *count == 0 { + protected.remove(&stable_id); + } + } + /// Whether this node is willing to serve as a relay. pub fn is_relay_serving_enabled(&self) -> bool { self.relay_serving_enabled.load(Ordering::Acquire) @@ -2225,6 +2283,25 @@ pub struct SessionInfo { #[allow(clippy::unwrap_used, clippy::expect_used)] mod tests { use super::*; + + #[test] + fn relay_control_connection_protection_is_reference_counted() { + let server = Arc::new(MasqueRelayServer::new( + MasqueRelayConfig::default(), + test_addr(9000), + )); + let stable_id = 42; + + assert!(!server.is_relay_control_connection(stable_id)); + let first = server.protect_relay_control_connection(stable_id); + let second = server.protect_relay_control_connection(stable_id); + assert!(server.is_relay_control_connection(stable_id)); + + drop(first); + assert!(server.is_relay_control_connection(stable_id)); + drop(second); + assert!(!server.is_relay_control_connection(stable_id)); + } use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; fn test_addr(port: u16) -> SocketAddr { diff --git a/src/masque/relay_socket.rs b/src/masque/relay_socket.rs index 4c2440be..fac28e7b 100644 --- a/src/masque/relay_socket.rs +++ b/src/masque/relay_socket.rs @@ -45,7 +45,8 @@ use std::future::Future; use std::io::{self, IoSliceMut}; use std::net::SocketAddr; use std::pin::Pin; -use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Weak}; use std::task::{Context, Poll}; use std::time::Duration; use tokio::sync::{Notify, mpsc}; @@ -100,6 +101,83 @@ pub struct RawRelayStreams { pub recv_stream: crate::high_level::RecvStream, } +/// Owns the background tasks and CONNECT-UDP streams backing a relay socket. +/// +/// Calling [`shutdown`](Self::shutdown) aborts and joins every tunnel task. +/// Aborting the reader and writer tasks drops their QUIC stream halves, which +/// promptly tells the relay server to close the associated MASQUE session and +/// release its capacity slot. Shutdown is idempotent. +#[derive(Debug)] +pub struct RelayTunnelControl { + tasks: PlMutex>>, + closed: Notify, + is_closed: AtomicBool, +} + +impl RelayTunnelControl { + fn new() -> Arc { + Arc::new(Self { + tasks: PlMutex::new(Vec::new()), + closed: Notify::new(), + is_closed: AtomicBool::new(false), + }) + } + + fn register(&self, handle: tokio::task::JoinHandle<()>) { + if self.is_closed() { + handle.abort(); + return; + } + + let mut tasks = self.tasks.lock(); + if self.is_closed() { + handle.abort(); + } else { + tasks.push(handle); + } + } + + fn mark_closed(&self) { + if !self.is_closed.swap(true, Ordering::AcqRel) { + self.closed.notify_waiters(); + } + } + + /// Returns whether the tunnel has failed or has been explicitly shut down. + pub fn is_closed(&self) -> bool { + self.is_closed.load(Ordering::Acquire) + } + + /// Wait until the tunnel reader exits or shutdown is requested. + pub async fn closed(&self) { + loop { + if self.is_closed() { + return; + } + let notified = self.closed.notified(); + if self.is_closed() { + return; + } + notified.await; + } + } + + /// Stop the tunnel and wait for all task-owned QUIC streams to be dropped. + pub async fn shutdown(&self) { + self.mark_closed(); + let handles = { + let mut tasks = self.tasks.lock(); + std::mem::take(&mut *tasks) + }; + for handle in &handles { + handle.abort(); + } + for handle in handles { + let _ = handle.await; + } + } +} + /// A virtual UDP socket backed entirely by a MASQUE relay tunnel. /// /// All traffic — both outgoing and incoming — flows through the relay @@ -166,9 +244,9 @@ impl MasqueRelaySocket { /// - A keepalive ticker that injects zero-length frames so the /// NAT conntrack entry stays alive on idle connections. /// - /// Returns the socket alongside a [`Notify`] that fires exactly once - /// when the reader task exits (tunnel failure). Callers that own the - /// backing endpoint should use this to trigger a **graceful** close + /// Returns the socket alongside a [`RelayTunnelControl`] that reports + /// tunnel failure and provides explicit teardown. Callers that own the + /// backing endpoint should use it to trigger a **graceful** close /// — `Endpoint::close(code, reason)` sends CONNECTION_CLOSE frames /// to every connection before the endpoint driver future is dropped. /// Without this, the driver's `Drop` impl fires last and cascades @@ -180,10 +258,10 @@ impl MasqueRelaySocket { relay_public_addr: SocketAddr, _relay_server_addr: SocketAddr, original_socket: Arc, - ) -> (Arc, Arc) { + ) -> (Arc, Arc) { let (send_tx, mut send_rx) = mpsc::channel::(SEND_QUEUE_CAPACITY); let (recv_tx, recv_rx) = mpsc::channel::<(Bytes, SocketAddr)>(RECV_QUEUE_CAPACITY); - let closed = Arc::new(Notify::new()); + let control = RelayTunnelControl::new(); let send_capacity_freed = Arc::new(Notify::new()); let target_mtu: Arc> = Arc::new(DashMap::new()); @@ -201,8 +279,8 @@ impl MasqueRelaySocket { // Background task: read length-prefixed frames from relay stream // and forward decoded (payload, source) pairs to `poll_recv`. // Holds the payload as `Bytes` throughout — no Vec round-trip. - let closed_reader = Arc::clone(&closed); - tokio::spawn(async move { + let weak_control: Weak = Arc::downgrade(&control); + let reader_handle = tokio::spawn(async move { loop { let mut len_buf = [0u8; 4]; if let Err(e) = recv_stream.read_exact(&mut len_buf).await { @@ -305,17 +383,17 @@ impl MasqueRelaySocket { // Dropping `recv_tx` here wakes any pending `poll_recv` // with Poll::Ready(None), signalling end-of-stream. // - // Signal any watcher waiting on the close notification so it - // can initiate a graceful endpoint close BEFORE the driver's - // `Drop` fires. `notify_waiters` wakes every current waiter - // exactly once; subsequent waits see `Pending`, which is - // fine — the shutdown only needs to run once. - closed_reader.notify_waiters(); + // Signal the owner before the endpoint driver is dropped so it + // can gracefully close connections accepted through the tunnel. + if let Some(control) = weak_control.upgrade() { + control.mark_closed(); + } }); + control.register(reader_handle); // Background task: write queued outbound packets to relay stream. let writer_capacity = Arc::clone(&send_capacity_freed); - tokio::spawn(async move { + let writer_handle = tokio::spawn(async move { while let Some(encoded) = send_rx.recv().await { // `recv` completing means the channel just freed a // slot. Wake any poller parked on full-queue @@ -358,6 +436,7 @@ impl MasqueRelaySocket { drop(send_rx); writer_capacity.notify_waiters(); }); + control.register(writer_handle); // Background task: periodic keepalive pings. // Sends a zero-length frame through the writer channel to keep @@ -365,7 +444,7 @@ impl MasqueRelaySocket { // connection. The writer encodes it as a 4-byte `[0,0,0,0]` // length prefix with no payload; the relay server skips it. let keepalive_tx = send_tx; - tokio::spawn(async move { + let keepalive_handle = tokio::spawn(async move { let mut tick = tokio::time::interval(RELAY_KEEPALIVE_INTERVAL); tick.tick().await; // skip immediate first tick loop { @@ -377,8 +456,9 @@ impl MasqueRelaySocket { } } }); + control.register(keepalive_handle); - (socket, closed) + (socket, control) } /// Remaining capacity in the outbound send channel. Exposed for @@ -647,3 +727,51 @@ impl UdpPoller for TunnelPoller { } } } + +#[cfg(test)] +mod relay_tunnel_control_tests { + use super::RelayTunnelControl; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + + struct DropMarker(Arc); + + impl Drop for DropMarker { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } + } + + #[tokio::test] + async fn shutdown_aborts_registered_tasks_and_wakes_waiters() { + let control = RelayTunnelControl::new(); + let dropped = Arc::new(AtomicBool::new(false)); + let marker = DropMarker(Arc::clone(&dropped)); + control.register(tokio::spawn(async move { + let _marker = marker; + std::future::pending::<()>().await; + })); + + let waiter_control = Arc::clone(&control); + let waiter = tokio::spawn(async move { + waiter_control.closed().await; + }); + + control.shutdown().await; + let waiter_result = waiter.await; + + assert!(waiter_result.is_ok()); + assert!(control.is_closed()); + assert!(dropped.load(Ordering::Acquire)); + } + + #[tokio::test] + async fn shutdown_is_idempotent() { + let control = RelayTunnelControl::new(); + + control.shutdown().await; + control.shutdown().await; + + assert!(control.is_closed()); + } +} diff --git a/src/nat_traversal_api.rs b/src/nat_traversal_api.rs index f9267632..5c394ae1 100644 --- a/src/nat_traversal_api.rs +++ b/src/nat_traversal_api.rs @@ -361,6 +361,11 @@ pub struct RelaySession { pub established_at: std::time::Instant, /// Relay server address pub relay_addr: SocketAddr, + /// Whether this session created a dedicated control connection. + /// + /// Sessions established on an existing peer connection must not close + /// that shared connection when their CONNECT-UDP stream is torn down. + pub owns_connection: bool, } impl RelaySession { @@ -376,6 +381,19 @@ impl RelaySession { } } +/// Resources owned by the current proactive relay allocation. +/// +/// The relay endpoint is usable while the allocation is provisional, but its +/// address is not exposed through `relay_public_addr` until publication. +struct ProactiveRelay { + relay_server_addr: SocketAddr, + public_addr: SocketAddr, + generation: u64, + endpoint: Arc, + tunnel: Arc, + published: bool, +} + /// Event from the constrained engine with transport address context /// /// This wrapper adds the transport address to engine events so that P2pEndpoint @@ -482,16 +500,19 @@ pub struct NatTraversalEndpoint { /// bootstrap window is the only time we care about growing the local /// candidate set from OBSERVED_ADDRESS frames. bootstrap_complete: Arc, - /// Relay address to re-advertise to new peers (set after proactive relay setup) + /// Relay address exposed after a proactive relay passes its canary gate. relay_public_addr: Arc>>, /// Monotonic generation for proactive relay state. /// - /// Incremented whenever relay state is established or reset. Any - /// advertise/re-advertise path that snapshots a relay address also carries - /// this generation and drops work if a newer relay state superseded it. + /// Incremented whenever relay state is established or reset. A publish + /// verdict must match the current generation so it cannot expose a + /// superseded relay. relay_generation: Arc, - /// Peers already advertised the relay address to - relay_advertised_peers: Arc>>, + /// Current proactive relay allocation, including provisional allocations. + proactive_relay: Arc>>, + /// Serializes prepare, publish, and abort so stale verdicts cannot race a + /// newer proactive relay generation. + proactive_relay_lifecycle: Arc>, /// Task handles for transport listener tasks /// Used for cleanup on shutdown transport_listener_handles: Arc>>>, @@ -1733,9 +1754,8 @@ impl NatTraversalEndpoint { bootstrap_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), relay_public_addr: Arc::new(std::sync::Mutex::new(None)), relay_generation: Arc::new(std::sync::atomic::AtomicU64::new(0)), - relay_advertised_peers: Arc::new(std::sync::Mutex::new( - std::collections::HashSet::new(), - )), + proactive_relay: Arc::new(std::sync::Mutex::new(None)), + proactive_relay_lifecycle: Arc::new(TokioMutex::new(())), transport_listener_handles: Arc::new(ParkingMutex::new(Vec::new())), constrained_engine, constrained_event_tx: constrained_event_tx.clone(), @@ -2175,9 +2195,8 @@ impl NatTraversalEndpoint { bootstrap_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), relay_public_addr: Arc::new(std::sync::Mutex::new(None)), relay_generation: Arc::new(std::sync::atomic::AtomicU64::new(0)), - relay_advertised_peers: Arc::new(std::sync::Mutex::new( - std::collections::HashSet::new(), - )), + proactive_relay: Arc::new(std::sync::Mutex::new(None)), + proactive_relay_lifecycle: Arc::new(TokioMutex::new(())), transport_listener_handles: Arc::new(ParkingMutex::new(Vec::new())), constrained_engine, constrained_event_tx: constrained_event_tx.clone(), @@ -3478,6 +3497,13 @@ impl NatTraversalEndpoint { addr, request ); + // This connection now carries relay traffic and + // must not be force-closed by ordinary peer-table + // pruning. The RAII guard releases protection on + // every exit path, including allocation failures. + let _relay_control_guard = + server.protect_relay_control_connection(stable_id); + // Handle the request via relay server match server .handle_connect_request(&request, addr, peer_id) @@ -3957,11 +3983,6 @@ impl NatTraversalEndpoint { .wrapping_add(1) } - fn relay_generation_matches(&self, generation: u64, relay_addr: SocketAddr) -> bool { - self.current_relay_generation() == generation - && self.relay_public_addr.lock().ok().and_then(|g| *g) == Some(relay_addr) - } - /// Check if the proactive relay session is still alive. Returns true if /// no relay was established (nothing to monitor) or the relay is healthy. /// Returns false if a relay was established but the underlying QUIC @@ -3972,6 +3993,14 @@ impl NatTraversalEndpoint { None => return true, // No relay — nothing to monitor }; + if let Ok(state) = self.proactive_relay.lock() + && let Some(state) = state.as_ref() + && state.public_addr == relay_addr + && state.tunnel.is_closed() + { + return false; + } + // Check the specific session for the advertised relay address. // Other relay sessions may exist but are irrelevant — peers are // using relay_addr, so that's the one that must be healthy. @@ -3998,9 +4027,6 @@ impl NatTraversalEndpoint { if let Ok(mut addr) = self.relay_public_addr.lock() { *addr = None; } - if let Ok(mut peers) = self.relay_advertised_peers.lock() { - peers.clear(); - } // Remove dead sessions self.relay_sessions.retain(|_, session| session.is_active()); info!( @@ -4052,36 +4078,18 @@ impl NatTraversalEndpoint { "relay session: establishing CONNECT-UDP Bind session" ); - // Prefer reusing an existing peer connection to the relay. - // The relay server's handle_relay_requests is spawned for each ACCEPTED - // connection, so using the existing connection ensures a handler is - // already listening for bidi streams. - let connection = if let Some(existing) = self.connections.get(&relay_addr) { - if existing.close_reason().is_none() { - info!( - relay = %relay_addr, - stable_id = existing.stable_id(), - "relay session: reusing existing peer connection to relay" - ); - existing.clone() - } else { - // Existing connection is dead — fall back to creating a new one - debug!( - relay = %relay_addr, - close_reason = ?existing.close_reason(), - "relay session: existing peer connection is closed; cold-dialing relay" - ); - drop(existing); - self.connect_new_to_relay(relay_addr).await? - } - } else { - // No existing connection — create one - debug!( - relay = %relay_addr, - "relay session: no existing peer connection; cold-dialing relay" - ); - self.connect_new_to_relay(relay_addr).await? - }; + // A relay allocation owns its parent QUIC connection. Borrowing an + // ordinary DHT connection made relay lifetime depend on routing-table + // pruning: removing that peer also destroyed an otherwise healthy, + // canary-verified CONNECT-UDP stream. The relay server protects this + // dedicated connection from its own peer-table pruning while the + // forwarding stream is live. + debug!( + relay = %relay_addr, + "relay session: creating dedicated control connection" + ); + let connection = self.connect_new_to_relay(relay_addr).await?; + let owns_connection = true; // Cap on the end-to-end CONNECT-UDP handshake (open_bi → write // request → read_exact response). Without a cap, reusing a peer @@ -4199,6 +4207,7 @@ impl NatTraversalEndpoint { public_address, established_at: std::time::Instant::now(), relay_addr, + owns_connection, }; // DashMap provides lock-free .insert() @@ -4741,6 +4750,20 @@ impl NatTraversalEndpoint { Ok(None) } + /// Whether the currently tracked connection at `addr` carries a live + /// CONNECT-UDP forwarding stream. + pub(crate) fn is_relay_control_connection(&self, addr: &SocketAddr) -> bool { + let Some(server) = self.relay_server.as_ref() else { + return false; + }; + socket_addr_variants(*addr).into_iter().any(|candidate| { + self.connections.get(&candidate).is_some_and(|connection| { + connection.close_reason().is_none() + && server.is_relay_control_connection(connection.stable_id()) + }) + }) + } + /// Get the receiver for accepted connection addresses. /// The P2pEndpoint's incoming_connection_forwarder uses this to register /// accepted connections in connected_peers. @@ -4911,6 +4934,20 @@ impl NatTraversalEndpoint { continue; } } + if expected_stable_id.is_none() + && self.relay_server.as_ref().is_some_and(|server| { + server.is_relay_control_connection(entry.value().stable_id()) + }) + && entry.value().close_reason().is_none() + { + info!( + stable_id = entry.value().stable_id(), + peer = %candidate, + "remove_connection: keeping live relay control connection" + ); + drop(entry); + continue; + } if entry.value().close_reason().is_none() { // Force the connection shut so its Quinn resources are @@ -5040,7 +5077,7 @@ impl NatTraversalEndpoint { is_symmetric } - /// Set up a proactive relay as a supplementary inbound path. + /// Prepare a proactive relay as a supplementary inbound path. /// /// Establishes a MASQUE relay session with `bootstrap_addr`, creates /// a **second** Quinn endpoint backed by the relay tunnel socket, @@ -5048,34 +5085,36 @@ impl NatTraversalEndpoint { /// original UDP socket are never touched — direct connections /// continue to work exactly as before. /// - /// Peers that connect to the relay-allocated address land on the - /// relay endpoint. Accepted connections are inserted into the - /// shared `connections` DashMap so the rest of the stack sees them - /// identically to direct connections. - pub async fn setup_proactive_relay( + /// The returned allocation is provisional: it can accept canary + /// connections but is not exposed by [`relay_public_addr`](Self::relay_public_addr) + /// Relay addresses are never sent through connection-level `ADD_ADDRESS`; + /// the authenticated upper layer publishes them after activation. + pub async fn prepare_proactive_relay( &self, bootstrap_addr: SocketAddr, ) -> Result { - info!( - "Setting up proactive relay via bootstrap {} for symmetric NAT", - bootstrap_addr - ); + let _lifecycle_guard = self.proactive_relay_lifecycle.lock().await; - // Step 1: Establish relay session with bootstrap - let (public_addr, raw_streams) = self.establish_relay_session(bootstrap_addr).await?; - let relay_public_addr = public_addr.ok_or_else(|| { - NatTraversalError::ConnectionFailed("Relay did not provide public address".to_string()) - })?; - let raw_streams = raw_streams.ok_or_else(|| { - NatTraversalError::ConnectionFailed("Relay did not provide socket".to_string()) - })?; + // A new acquisition supersedes any prior allocation. Teardown happens + // before the new CONNECT-UDP request so rejected canaries cannot + // accumulate relay-server capacity. + let previous = self + .proactive_relay + .lock() + .map_err(|_| NatTraversalError::ConfigError("mutex poisoned".to_string()))? + .take(); + if let Some(previous) = previous { + self.teardown_proactive_relay(previous, b"relay superseded") + .await; + } info!( - "Relay session established, public address: {}", - relay_public_addr + "Preparing proactive relay via bootstrap {} for symmetric NAT", + bootstrap_addr ); - // Step 2: Get the main endpoint's original socket (kept alive so the + // Validate every local prerequisite before acquiring scarce capacity + // from the relay server. // relay connection's own QUIC traffic — which flows over the real UDP // socket — continues to work). let main_endpoint = self.inner_endpoint.as_ref().ok_or_else(|| { @@ -5086,22 +5125,6 @@ impl NatTraversalEndpoint { NatTraversalError::ConnectionFailed(format!("Failed to get original socket: {e}",)) })?; - // Step 3: Build a tunnel-only relay socket and a second Quinn - // endpoint. The main endpoint is never touched. - // - // `closed_notify` fires when the relay reader task exits because - // its QUIC stream failed. A watcher spawned below uses this to - // call `Endpoint::close(...)` on the relay endpoint BEFORE the - // driver's Drop fires — connections get a clean CONNECTION_CLOSE - // instead of the "endpoint driver future was dropped" cascade. - let (relay_socket, closed_notify) = crate::masque::MasqueRelaySocket::new( - raw_streams.send_stream, - raw_streams.recv_stream, - relay_public_addr, - bootstrap_addr, - original_socket, - ); - let server_config = self .relay_server_config .lock() @@ -5126,39 +5149,77 @@ impl NatTraversalEndpoint { NatTraversalError::ConfigError("No async runtime available".to_string()) })?; - let relay_endpoint = InnerEndpoint::new_with_abstract_socket( + // Acquire the relay only after local validation is complete. + let (public_addr, raw_streams) = self.establish_relay_session(bootstrap_addr).await?; + let Some(relay_public_addr) = public_addr else { + self.remove_matching_relay_session(bootstrap_addr, None, b"invalid relay allocation"); + return Err(NatTraversalError::ConnectionFailed( + "Relay did not provide public address".to_string(), + )); + }; + let Some(raw_streams) = raw_streams else { + self.remove_matching_relay_session( + bootstrap_addr, + Some(relay_public_addr), + b"relay allocation missing streams", + ); + return Err(NatTraversalError::ConnectionFailed( + "Relay did not provide socket".to_string(), + )); + }; + + info!( + "Relay session established, public address: {}", + relay_public_addr + ); + + // Build a tunnel-only relay socket and a second Quinn endpoint. The + // main endpoint remains untouched. + let (relay_socket, tunnel) = crate::masque::MasqueRelaySocket::new( + raw_streams.send_stream, + raw_streams.recv_stream, + relay_public_addr, + bootstrap_addr, + original_socket, + ); + + let relay_endpoint = match InnerEndpoint::new_with_abstract_socket( EndpointConfig::default(), server_config, relay_socket, runtime, - ) - .map_err(|e| { - NatTraversalError::ConnectionFailed(format!("Failed to create relay endpoint: {e}")) - })?; + ) { + Ok(endpoint) => endpoint, + Err(error) => { + tunnel.shutdown().await; + self.remove_matching_relay_session( + bootstrap_addr, + Some(relay_public_addr), + b"relay endpoint creation failed", + ); + return Err(NatTraversalError::ConnectionFailed(format!( + "Failed to create relay endpoint: {error}" + ))); + } + }; info!("Relay endpoint created (relay addr: {})", relay_public_addr); let relay_generation = self.next_relay_generation(); self.relay_setup_attempted .store(true, std::sync::atomic::Ordering::Relaxed); - if let Ok(mut addr) = self.relay_public_addr.lock() { - *addr = Some(relay_public_addr); - } - if let Ok(mut peers) = self.relay_advertised_peers.lock() { - peers.clear(); - } info!( relay_generation, relay_addr = %relay_public_addr, - "Proactive relay generation established" + "Proactive relay generation prepared" ); // Share the relay endpoint between the accept loop and the // graceful-close watcher. let relay_endpoint = Arc::new(relay_endpoint); - // Spawn the tunnel-death watcher: when the MASQUE reader task - // signals the `Notify`, explicitly close the relay endpoint with a + // Spawn the tunnel-death watcher: when the MASQUE reader task exits, + // explicitly close the relay endpoint with a // meaningful application-level error code. `Endpoint::close` // dispatches a `ConnectionEvent::Close` to every active connection // on the endpoint — they emit CONNECTION_CLOSE and shut down @@ -5167,17 +5228,20 @@ impl NatTraversalEndpoint { // cascading errors. Dropping `relay_endpoint` after `close()` is // safe — `EndpointRef::drop` only decrements a refcount. { - let relay_endpoint = Arc::clone(&relay_endpoint); + let relay_endpoint = Arc::downgrade(&relay_endpoint); + let tunnel = Arc::clone(&tunnel); tokio::spawn(async move { - closed_notify.notified().await; + tunnel.closed().await; info!( "MASQUE tunnel for relay {} died — closing relay endpoint gracefully", relay_public_addr ); - relay_endpoint.close( - crate::VarInt::from_u32(RELAY_TUNNEL_LOST_CODE), - b"relay tunnel lost", - ); + if let Some(relay_endpoint) = relay_endpoint.upgrade() { + relay_endpoint.close( + crate::VarInt::from_u32(RELAY_TUNNEL_LOST_CODE), + b"relay tunnel lost", + ); + } }); } @@ -5185,41 +5249,204 @@ impl NatTraversalEndpoint { // accepted connections into the shared connections DashMap. self.spawn_relay_endpoint_accept_loop(Arc::clone(&relay_endpoint), relay_public_addr); - // Step 5: Advertise the relay address to all connected peers - let mut advertised = 0; - for entry in self.connections.iter() { - if !self.relay_generation_matches(relay_generation, relay_public_addr) { - debug!( - relay_generation, - relay_addr = %relay_public_addr, - "Skipping stale relay advertisement after relay generation changed" - ); - break; + let state = ProactiveRelay { + relay_server_addr: bootstrap_addr, + public_addr: relay_public_addr, + generation: relay_generation, + endpoint: relay_endpoint, + tunnel, + published: false, + }; + if let Err(state) = self.install_proactive_relay(state) { + self.teardown_proactive_relay(state, b"relay state mutex poisoned") + .await; + return Err(NatTraversalError::ConfigError("mutex poisoned".to_string())); + } + + Ok(relay_public_addr) + } + + /// Activate a prepared proactive relay after its external canary succeeds. + /// + /// The address is checked against the current generation so a late canary + /// verdict cannot publish a superseded allocation. + pub async fn publish_proactive_relay( + &self, + relay_public_addr: SocketAddr, + ) -> Result<(), NatTraversalError> { + let _lifecycle_guard = self.proactive_relay_lifecycle.lock().await; + + let relay_generation = { + let mut state = self + .proactive_relay + .lock() + .map_err(|_| NatTraversalError::ConfigError("mutex poisoned".to_string()))?; + let state = state.as_mut().ok_or_else(|| { + NatTraversalError::ConnectionFailed( + "No proactive relay allocation is prepared".to_string(), + ) + })?; + if state.public_addr != relay_public_addr { + return Err(NatTraversalError::ConnectionFailed(format!( + "Prepared relay address {} does not match {}", + state.public_addr, relay_public_addr + ))); } - let conn = entry.value().clone(); - match conn.send_nat_address_advertisement(relay_public_addr, 100) { - Ok(_) => { - advertised += 1; - if let Ok(mut peers) = self.relay_advertised_peers.lock() { - peers.insert(*entry.key()); - } - } - Err(e) => { - debug!( - "Failed to advertise relay address to {}: {}", - entry.key(), - e - ); - } + if state.tunnel.is_closed() { + return Err(NatTraversalError::ConnectionFailed(format!( + "Prepared relay tunnel for {relay_public_addr} is closed" + ))); + } + if state.published { + return Ok(()); + } + + state.published = true; + state.generation + }; + + if self.current_relay_generation() != relay_generation { + return Err(NatTraversalError::ConnectionFailed( + "Prepared relay generation was superseded".to_string(), + )); + } + *self + .relay_public_addr + .lock() + .map_err(|_| NatTraversalError::ConfigError("mutex poisoned".to_string()))? = + Some(relay_public_addr); + + // Relay addresses are intentionally not broadcast through the + // connection-level ADD_ADDRESS extension. A relay allocation belongs + // to the upper-layer peer identity, not to every existing transport + // path. Broadcasting it here causes all connected peers to validate + // the fresh relay endpoint at once, competing with real and canary + // handshakes. Saorsa-core publishes the canary-verified address through + // its authenticated, sequenced PublishAddressSet path instead. + info!( + relay_generation, + relay_addr = %relay_public_addr, + "Activated canary-verified proactive relay" + ); + Ok(()) + } + + /// Abort a prepared or published proactive relay. + /// + /// A mismatched or already-removed address is treated as a stale, + /// idempotent abort and leaves the current generation untouched. + pub async fn abort_proactive_relay( + &self, + relay_public_addr: SocketAddr, + ) -> Result<(), NatTraversalError> { + let _lifecycle_guard = self.proactive_relay_lifecycle.lock().await; + let state = { + let mut current = self + .proactive_relay + .lock() + .map_err(|_| NatTraversalError::ConfigError("mutex poisoned".to_string()))?; + match current.as_ref() { + Some(state) if state.public_addr == relay_public_addr => current.take(), + _ => None, + } + }; + + if let Some(state) = state { + self.teardown_proactive_relay(state, b"relay allocation aborted") + .await; + } + Ok(()) + } + + /// Legacy eager setup: prepare and immediately publish. + /// + /// Canary-gated callers should use [`prepare_proactive_relay`](Self::prepare_proactive_relay) + /// followed by either [`publish_proactive_relay`](Self::publish_proactive_relay) + /// or [`abort_proactive_relay`](Self::abort_proactive_relay). + pub async fn setup_proactive_relay( + &self, + relay_addr: SocketAddr, + ) -> Result { + let allocated = self.prepare_proactive_relay(relay_addr).await?; + if let Err(error) = self.publish_proactive_relay(allocated).await { + let _ = self.abort_proactive_relay(allocated).await; + return Err(error); + } + Ok(allocated) + } + + async fn teardown_proactive_relay(&self, state: ProactiveRelay, reason: &'static [u8]) { + state + .endpoint + .close(crate::VarInt::from_u32(RELAY_TUNNEL_LOST_CODE), reason); + state.tunnel.shutdown().await; + + if tokio::time::timeout(Duration::from_secs(1), state.endpoint.wait_idle()) + .await + .is_err() + { + debug!( + relay_addr = %state.public_addr, + "Timed out draining proactive relay endpoint during teardown" + ); + } + + self.remove_matching_relay_session( + state.relay_server_addr, + Some(state.public_addr), + reason, + ); + + if self.current_relay_generation() == state.generation { + self.next_relay_generation(); + if let Ok(mut addr) = self.relay_public_addr.lock() + && *addr == Some(state.public_addr) + { + *addr = None; } + self.relay_setup_attempted + .store(false, std::sync::atomic::Ordering::Release); } info!( - "Advertised relay address {} to {} peers", - relay_public_addr, advertised + relay_addr = %state.public_addr, + relay_server = %state.relay_server_addr, + published = state.published, + "Proactive relay torn down" ); + } - Ok(relay_public_addr) + fn install_proactive_relay(&self, state: ProactiveRelay) -> Result<(), ProactiveRelay> { + match self.proactive_relay.lock() { + Ok(mut current) => { + *current = Some(state); + Ok(()) + } + Err(_) => Err(state), + } + } + + fn remove_matching_relay_session( + &self, + relay_server_addr: SocketAddr, + public_addr: Option, + reason: &'static [u8], + ) { + let matching_session = self + .relay_sessions + .get(&relay_server_addr) + .is_some_and(|session| session.public_address == public_addr); + if !matching_session { + return; + } + + if let Some((_, session)) = self.relay_sessions.remove(&relay_server_addr) + && session.owns_connection + { + session + .connection + .close(crate::VarInt::from_u32(RELAY_TUNNEL_LOST_CODE), reason); + } } /// Spawn an accept loop for the relay endpoint. @@ -5781,6 +6008,15 @@ impl NatTraversalEndpoint { self.incoming_notify.notify_waiters(); self.shutdown_notify.notify_waiters(); + let proactive_addr = self + .proactive_relay + .lock() + .ok() + .and_then(|state| state.as_ref().map(|relay| relay.public_addr)); + if let Some(proactive_addr) = proactive_addr { + let _ = self.abort_proactive_relay(proactive_addr).await; + } + // Best-effort UPnP teardown. The endpoint is the sole owner of // the service (the discovery manager only holds a read-only // `UpnpStateRx`), so we can move it out and call its async @@ -6587,70 +6823,11 @@ impl NatTraversalEndpoint { // `reachability::driver`, which calls `setup_proactive_relay` // unconditionally after bootstrap via the transport handle. // - // This function is retained for the re-advertise pass below, - // which pushes a previously-acquired relay address to peers - // that connected after initial setup. - // // The deleted auto-setup block walked bootstrap nodes and called // `setup_proactive_relay` directly from this event loop; the new // acquisition driver in saorsa-core performs the same walk // through the XOR-closest set of the DHT instead. - // Re-advertise relay address to peers that connected after initial setup - { - let relay_addr = self.relay_public_addr.lock().ok().and_then(|g| *g); - if let Some(relay_addr) = relay_addr { - let relay_generation = self.current_relay_generation(); - let unadvertised: Vec = { - let advertised = self - .relay_advertised_peers - .lock() - .unwrap_or_else(|e| e.into_inner()); - self.connections - .iter() - .filter(|e| { - !advertised.contains(e.key()) && e.value().close_reason().is_none() - }) - .map(|e| *e.key()) - .collect() - }; - if !unadvertised.is_empty() { - info!( - "Relay re-advertise: {} new peers to notify about {}", - unadvertised.len(), - relay_addr - ); - } - for peer_addr in unadvertised { - if !self.relay_generation_matches(relay_generation, relay_addr) { - debug!( - relay_generation, - relay_addr = %relay_addr, - "Relay re-advertise: dropping stale generation" - ); - break; - } - if let Some(mut entry) = self.connections.get_mut(&peer_addr) { - match entry - .value_mut() - .send_nat_address_advertisement(relay_addr, 100) - { - Ok(_) => { - info!( - "Re-advertised relay {} to new peer {}", - relay_addr, peer_addr - ); - if let Ok(mut a) = self.relay_advertised_peers.lock() { - a.insert(peer_addr); - } - } - Err(_) => {} - } - } - } - } - } - Ok(()) } diff --git a/src/p2p_endpoint.rs b/src/p2p_endpoint.rs index 4a0afd32..99b09982 100644 --- a/src/p2p_endpoint.rs +++ b/src/p2p_endpoint.rs @@ -259,7 +259,7 @@ fn pending_dial_remaining_wait(wait_timeout: Duration, age: Duration) -> Option< /// peer identity, if TLS-based authentication was used. /// /// Returns `None` for unauthenticated or constrained connections. -fn extract_public_key_bytes_from_connection( +pub(crate) fn extract_public_key_bytes_from_connection( connection: &crate::high_level::Connection, ) -> Option> { let identity = connection.peer_identity()?; @@ -1203,6 +1203,23 @@ async fn do_cleanup_connection( ) -> bool { let variants = socket_addr_variants(*addr); + // A live CONNECT-UDP stream makes this connection part of an established + // relay's data plane. Ordinary disconnect APIs and the stale-peer reaper + // must not force-close it; the relay forwarding guard releases this pin + // when the stream actually ends. Reader-exit cleanup still proceeds for a + // connection that has already failed. + if expected_stable_id.is_none() + && variants + .iter() + .any(|candidate| inner.is_relay_control_connection(candidate)) + { + info!( + peer = %addr, + "do_cleanup_connection: keeping live relay control connection" + ); + return false; + } + // Step 1: Try to remove from the NAT traversal layer first (lock-free // DashMap). When the caller is a reader-exit firing for an older // connection that has since been replaced via simultaneous-open, the @@ -2746,7 +2763,7 @@ impl P2pEndpoint { // single connection and is discarded after the dial completes, // so we don't wire up a graceful-close watcher — the single // connection is the natural unit of recovery at this layer. - let (relay_socket, _closed) = crate::masque::MasqueRelaySocket::new( + let (relay_socket, _tunnel) = crate::masque::MasqueRelaySocket::new( raw_streams.send_stream, raw_streams.recv_stream, relay_public_addr, @@ -3413,15 +3430,14 @@ impl P2pEndpoint { self.inner.set_relay_serving_enabled(enabled); } - /// Establish a proactive MASQUE relay session with `relay_addr` as a - /// supplementary inbound path. + /// Establish and immediately publish a proactive MASQUE relay session with + /// `relay_addr` as a supplementary inbound path. /// - /// This is the caller-driven entry point for ADR-014-style relay acquisition - /// in saorsa-core. It delegates to [`NatTraversalEndpoint::setup_proactive_relay`], - /// which establishes the MASQUE `CONNECT-UDP` session, creates a **second** - /// Quinn endpoint backed by the relay tunnel, and advertises the - /// relay-allocated address to all currently connected peers. The main - /// endpoint and its original UDP socket are never touched. + /// This legacy convenience entry point prepares and publishes in one call. + /// Canary-gated callers should use + /// [`prepare_proactive_relay`](Self::prepare_proactive_relay), followed by + /// either [`publish_proactive_relay`](Self::publish_proactive_relay) or + /// [`abort_proactive_relay`](Self::abort_proactive_relay). /// /// On success, the returned `SocketAddr` is the relay-allocated public /// address the caller should publish as its contact address in the DHT @@ -3440,6 +3456,42 @@ impl P2pEndpoint { Ok(allocated) } + /// Prepare a proactive MASQUE relay without publishing its address. + /// + /// The relay endpoint is live so callers can run inbound canary probes, + /// but the address is not exposed through `relay_public_addr` or emitted as + /// `RelayEstablished` until + /// [`publish_proactive_relay`](Self::publish_proactive_relay) is called. + /// Relay addresses are not sent through connection-level `ADD_ADDRESS`; + /// authenticated upper layers own their publication. + pub async fn prepare_proactive_relay( + &self, + relay_addr: SocketAddr, + ) -> Result { + Ok(self.inner.prepare_proactive_relay(relay_addr).await?) + } + + /// Activate the current provisional proactive relay. + pub async fn publish_proactive_relay( + &self, + relay_public_addr: SocketAddr, + ) -> Result<(), EndpointError> { + self.inner + .publish_proactive_relay(relay_public_addr) + .await?; + Ok(()) + } + + /// Abort a provisional or published proactive relay and release its + /// MASQUE session and relay-server capacity. + pub async fn abort_proactive_relay( + &self, + relay_public_addr: SocketAddr, + ) -> Result<(), EndpointError> { + self.inner.abort_proactive_relay(relay_public_addr).await?; + Ok(()) + } + /// Get list of connected peers pub async fn connected_peers(&self) -> Vec { self.connected_peers @@ -3970,7 +4022,18 @@ impl P2pEndpoint { // can tell upper layers exactly which relay to stop // advertising. let dead = inner.relay_public_addr(); - inner.reset_relay_state(); + if let Some(relay_addr) = dead { + if let Err(error) = inner.abort_proactive_relay(relay_addr).await { + warn!( + relay_addr = %relay_addr, + %error, + "Failed to tear down unhealthy proactive relay" + ); + inner.reset_relay_state(); + } + } else { + inner.reset_relay_state(); + } relay_event_sent = false; if let Some(relay_addr) = dead { info!( From 4a2ceb96c424f12d76422b06674a45c2472cda71 Mon Sep 17 00:00:00 2001 From: Mick van Dijke <12992260+mickvandijke@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:30:26 +0200 Subject: [PATCH 2/8] fix(relay): finalize provisional relay lifecycle Carry the authoritative accepted connection into LinkTransport, keep provisional and published relay allocations behind one identity-scoped lifecycle, and remove packet-by-packet relay tracing from the hot path.\n\nCompletes WithAutonomi/saorsa-core#138 and supports WithAutonomi/saorsa-core#136.\n\nSemVer: breaking --- src/link_transport_impl.rs | 81 +++++--- src/masque/mod.rs | 3 +- src/masque/relay_server.rs | 40 +--- src/masque/relay_socket.rs | 32 +-- src/nat_traversal_api.rs | 400 +++++++++++++++++-------------------- src/p2p_endpoint.rs | 88 ++++---- 6 files changed, 288 insertions(+), 356 deletions(-) diff --git a/src/link_transport_impl.rs b/src/link_transport_impl.rs index 483924f8..abc88dca 100644 --- a/src/link_transport_impl.rs +++ b/src/link_transport_impl.rs @@ -650,37 +650,20 @@ impl LinkTransport for P2pLinkTransport { endpoint, |endpoint| async move { // Wait for an incoming connection - if let Some(peer_conn) = endpoint.accept().await { + if let Some((peer_conn, conn)) = endpoint.accept_with_connection().await { // Extract SocketAddr from TransportAddr let socket_addr = peer_conn .remote_addr .as_socket_addr() .unwrap_or_else(|| SocketAddr::from(([0, 0, 0, 0], 0))); - // Get the underlying QUIC connection by address - match endpoint.get_quic_connection(&socket_addr).await { - Ok(Some(conn)) => { - // rustls exposes the authenticated identity as a - // certificate vector whose first entry is the RFC - // 7250 ML-DSA SPKI, not as a bare `Vec`. - let public_key = - crate::p2p_endpoint::extract_public_key_bytes_from_connection( - &conn, - ); - let link_conn = P2pLinkConn::new(conn, public_key, socket_addr); - Some((Ok(link_conn), endpoint)) - } - Ok(None) => { - // Connection not found, try again - Some(( - Err(LinkError::ConnectionFailed( - "Connection not found".to_string(), - )), - endpoint, - )) - } - Err(e) => Some((Err(LinkError::ConnectionFailed(e.to_string())), endpoint)), - } + // rustls exposes the authenticated identity as a + // certificate vector whose first entry is the RFC 7250 + // ML-DSA SPKI, not as a bare `Vec`. + let public_key = + crate::p2p_endpoint::extract_public_key_bytes_from_connection(&conn); + let link_conn = P2pLinkConn::new(conn, public_key, socket_addr); + Some((Ok(link_conn), endpoint)) } else { // Endpoint is shutting down None @@ -1440,6 +1423,54 @@ mod tests { assert!(state.capabilities.is_empty()); } + #[tokio::test] + async fn accept_uses_the_authoritative_connection_handle() { + let bind_addr: SocketAddr = "127.0.0.1:0".parse().expect("valid bind address"); + let server_endpoint = Arc::new( + P2pEndpoint::new( + P2pConfig::builder() + .bind_addr(bind_addr) + .build() + .expect("valid server config"), + ) + .await + .expect("server endpoint"), + ); + let server_addr = server_endpoint.local_addr().expect("server address"); + let server = P2pLinkTransport::from_endpoint(Arc::clone(&server_endpoint)); + let client = P2pEndpoint::new( + P2pConfig::builder() + .bind_addr(bind_addr) + .build() + .expect("valid client config"), + ) + .await + .expect("client endpoint"); + + let mut incoming = server.accept(ProtocolId::DEFAULT); + client + .connect(server_addr) + .await + .expect("client connection"); + let accepted = tokio::time::timeout(std::time::Duration::from_secs(10), incoming.next()) + .await + .expect("accept timed out") + .expect("accept stream ended") + .expect("accepted connection"); + + assert_eq!( + accepted.remote_addr().ip(), + client.local_addr().unwrap().ip() + ); + assert!( + accepted.peer_public_key().is_some(), + "accepted handle must retain its authenticated identity" + ); + + client.shutdown().await; + server_endpoint.shutdown().await; + } + // ========================================================================= // Phase 3: SharedTransport Tests // ========================================================================= diff --git a/src/masque/mod.rs b/src/masque/mod.rs index b8a4cced..e72597d7 100644 --- a/src/masque/mod.rs +++ b/src/masque/mod.rs @@ -127,4 +127,5 @@ pub use relay_server::{ pub use relay_session::{ RelayPeerId, RelaySession, RelaySessionConfig, RelaySessionState, RelaySessionStats, }; -pub use relay_socket::{MasqueRelaySocket, RawRelayStreams, RelayTunnelControl}; +pub(crate) use relay_socket::RelayTunnelControl; +pub use relay_socket::{MasqueRelaySocket, RawRelayStreams}; diff --git a/src/masque/relay_server.rs b/src/masque/relay_server.rs index b2c8299f..efc37ee4 100644 --- a/src/masque/relay_server.rs +++ b/src/masque/relay_server.rs @@ -1379,12 +1379,6 @@ impl MasqueRelayServer { match socket.recv_from(&mut buf).await { Ok((len, source)) => { let payload = Bytes::copy_from_slice(&buf[..len]); - tracing::trace!( - session_id, - source = %source, - len, - "RELAY_TUNNEL[srv]: dgram-loop dir1 recv UDP → forwarding to relay-client" - ); // Encode as uncompressed datagram (includes source address // so client can decode without context registration) @@ -1452,24 +1446,10 @@ impl MasqueRelayServer { }; match resolved { Some((target, payload)) => { - tracing::trace!( - session_id, - target = %target, - len = payload.len(), - "RELAY_TUNNEL[srv]: dgram-loop dir2 recv from relay-client → sendto target" - ); server2.stats.record_bytes(payload.len() as u64); server2.stats.record_datagram(); match socket2.send_to(&payload, target).await { - Ok(n) => { - tracing::trace!( - session_id, - target = %target, - len = payload.len(), - sent = n, - "RELAY_TUNNEL[srv]: dgram-loop dir2 sendto OK" - ); - } + Ok(_) => {} Err(e) => { tracing::warn!( session_id, @@ -1611,10 +1591,6 @@ impl MasqueRelayServer { match socket.recv_from(&mut buf).await { Ok((len, source)) => { let payload = Bytes::copy_from_slice(&buf[..len]); - tracing::trace!( - session_id, source = %source, len, - "RELAY_TUNNEL[srv]: stream-loop dir1 recv UDP → forwarding to relay-client" - ); let datagram = UncompressedDatagram::new(VarInt::from_u32(0), source, payload); let encoded = datagram.encode(); @@ -1775,26 +1751,14 @@ impl MasqueRelayServer { let mut cursor = Bytes::from(frame_buf); match UncompressedDatagram::decode(&mut cursor) { Ok(datagram) => { - tracing::trace!( - session_id, target = %datagram.target, - len = datagram.payload.len(), - "RELAY_TUNNEL[srv]: stream-loop dir2 recv from relay-client → sendto target" - ); stats2.record_bytes(datagram.payload.len() as u64); stats2.record_datagram(); let target = datagram.target; let payload_len = datagram.payload.len(); match socket2.send_to(&datagram.payload, target).await { - Ok(n) => { + Ok(_) => { // Confirmed forwarded to the third-party target. stats2.record_forwarded_to_target(payload_len as u64, 1); - tracing::trace!( - session_id, - target = %target, - len = payload_len, - sent = n, - "RELAY_TUNNEL[srv]: stream-loop dir2 sendto OK" - ); } Err(e) if is_message_too_large(&e) => { // Path-MTU exceeded. Emit a PmtuUpdate diff --git a/src/masque/relay_socket.rs b/src/masque/relay_socket.rs index fac28e7b..4fdcc960 100644 --- a/src/masque/relay_socket.rs +++ b/src/masque/relay_socket.rs @@ -108,7 +108,7 @@ pub struct RawRelayStreams { /// promptly tells the relay server to close the associated MASQUE session and /// release its capacity slot. Shutdown is idempotent. #[derive(Debug)] -pub struct RelayTunnelControl { +pub(crate) struct RelayTunnelControl { tasks: PlMutex>>, closed: Notify, is_closed: AtomicBool, @@ -144,12 +144,12 @@ impl RelayTunnelControl { } /// Returns whether the tunnel has failed or has been explicitly shut down. - pub fn is_closed(&self) -> bool { + pub(crate) fn is_closed(&self) -> bool { self.is_closed.load(Ordering::Acquire) } /// Wait until the tunnel reader exits or shutdown is requested. - pub async fn closed(&self) { + pub(crate) async fn closed(&self) { loop { if self.is_closed() { return; @@ -163,7 +163,7 @@ impl RelayTunnelControl { } /// Stop the tunnel and wait for all task-owned QUIC streams to be dropped. - pub async fn shutdown(&self) { + pub(crate) async fn shutdown(&self) { self.mark_closed(); let handles = { let mut tasks = self.tasks.lock(); @@ -252,7 +252,7 @@ impl MasqueRelaySocket { /// Without this, the driver's `Drop` impl fires last and cascades /// a cryptic `"endpoint driver future was dropped"` into every /// connection accepted through this tunnel. - pub fn new( + pub(crate) fn new( mut send_stream: crate::high_level::SendStream, mut recv_stream: crate::high_level::RecvStream, relay_public_addr: SocketAddr, @@ -358,14 +358,6 @@ impl MasqueRelaySocket { Ok(datagram) => { // `datagram.payload` is a zero-copy slice of // the original frame buffer — no clone needed. - let inbound_source = datagram.target; - let inbound_len = datagram.payload.len(); - tracing::trace!( - relay = %relay_public_addr, - source = %inbound_source, - len = inbound_len, - "RELAY_TUNNEL[clt]: decoded inbound frame → enqueue for Quinn poll_recv" - ); if recv_tx .send((datagram.payload, datagram.target)) .await @@ -507,14 +499,6 @@ impl AsyncUdpSocket for MasqueRelaySocket { // of `segment_size` bytes. Each segment must be sent as its // own tunnel frame — the relay server has a per-frame size // limit and cannot handle the entire batch as one. - tracing::trace!( - relay = %self.relay_public_addr, - destination = %transmit.destination, - len = transmit.contents.len(), - segment_size = ?transmit.segment_size, - "RELAY_TUNNEL[clt]: try_send → enqueue outbound for relay-server" - ); - // Per-target MTU enforcement: if a previous PmtuUpdate control // frame told us the egress path to this destination caps at // `mtu` bytes, drop oversized packets here so they never reach @@ -599,12 +583,6 @@ impl AsyncUdpSocket for MasqueRelaySocket { recv_meta.dst_ip = None; meta[filled] = recv_meta; - tracing::trace!( - source = %source, - len, - "RELAY_TUNNEL: recv from tunnel queue" - ); - filled += 1; } Poll::Ready(None) => { diff --git a/src/nat_traversal_api.rs b/src/nat_traversal_api.rs index 5c394ae1..fec18d80 100644 --- a/src/nat_traversal_api.rs +++ b/src/nat_traversal_api.rs @@ -361,11 +361,6 @@ pub struct RelaySession { pub established_at: std::time::Instant, /// Relay server address pub relay_addr: SocketAddr, - /// Whether this session created a dedicated control connection. - /// - /// Sessions established on an existing peer connection must not close - /// that shared connection when their CONNECT-UDP stream is torn down. - pub owns_connection: bool, } impl RelaySession { @@ -381,17 +376,76 @@ impl RelaySession { } } -/// Resources owned by the current proactive relay allocation. +/// Opaque identity of a prepared proactive relay allocation. /// -/// The relay endpoint is usable while the allocation is provisional, but its -/// address is not exposed through `relay_public_addr` until publication. +/// Callers must pass this exact handle to +/// [`NatTraversalEndpoint::publish_proactive_relay`] or +/// [`NatTraversalEndpoint::abort_proactive_relay`]. The allocation id prevents +/// a late canary verdict from acting on a newer allocation that reused the same +/// stable relay address. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PreparedRelay { + public_addr: SocketAddr, + allocation_id: u64, +} + +static NEXT_RELAY_ALLOCATION_ID: AtomicU64 = AtomicU64::new(0); + +impl PreparedRelay { + fn new(public_addr: SocketAddr) -> Self { + Self { + public_addr, + allocation_id: NEXT_RELAY_ALLOCATION_ID + .fetch_add(1, Ordering::AcqRel) + .wrapping_add(1), + } + } + + /// Public address allocated for this provisional relay. + pub fn public_addr(self) -> SocketAddr { + self.public_addr + } + + /// Create an independent handle for a mock relay allocation. + /// + /// This supports implementations of relay-establishment abstractions in + /// downstream tests. It cannot publish or abort an endpoint allocation + /// because its opaque identity does not match that endpoint's current + /// handle. + #[doc(hidden)] + pub fn detached(public_addr: SocketAddr) -> Self { + Self::new(public_addr) + } +} + +/// Resources owned by one proactive relay allocation. struct ProactiveRelay { relay_server_addr: SocketAddr, - public_addr: SocketAddr, - generation: u64, + handle: PreparedRelay, endpoint: Arc, tunnel: Arc, - published: bool, +} + +/// Complete proactive-relay lifecycle state. +/// +/// Keeping visibility and owned resources in one enum makes publication a +/// single, infallible state transition rather than a sequence of updates across +/// independent mutexes and atomics. +#[derive(Default)] +enum RelayLifecycleState { + #[default] + None, + Provisional(ProactiveRelay), + Published(ProactiveRelay), +} + +impl RelayLifecycleState { + fn published_addr(&self) -> Option { + match self { + Self::Published(relay) => Some(relay.handle.public_addr()), + Self::None | Self::Provisional(_) => None, + } + } } /// Event from the constrained engine with transport address context @@ -490,8 +544,6 @@ pub struct NatTraversalEndpoint { /// one handler running even when both the accept-side and dial-side /// spawn paths fire for it. Shared across all spawn sites. relay_handler_connections: Arc>, - /// Whether symmetric NAT relay setup has been attempted (one-shot) - relay_setup_attempted: Arc, /// Flipped once the external bootstrap phase is over (set by /// `P2pEndpoint::connect_known_peers` right before it broadcasts /// `P2pEvent::BootstrapStatus`). While false, the discovery polling @@ -500,19 +552,8 @@ pub struct NatTraversalEndpoint { /// bootstrap window is the only time we care about growing the local /// candidate set from OBSERVED_ADDRESS frames. bootstrap_complete: Arc, - /// Relay address exposed after a proactive relay passes its canary gate. - relay_public_addr: Arc>>, - /// Monotonic generation for proactive relay state. - /// - /// Incremented whenever relay state is established or reset. A publish - /// verdict must match the current generation so it cannot expose a - /// superseded relay. - relay_generation: Arc, - /// Current proactive relay allocation, including provisional allocations. - proactive_relay: Arc>>, - /// Serializes prepare, publish, and abort so stale verdicts cannot race a - /// newer proactive relay generation. - proactive_relay_lifecycle: Arc>, + /// Current proactive relay allocation and its publication state. + relay_lifecycle: Arc>, /// Task handles for transport listener tasks /// Used for cleanup on shutdown transport_listener_handles: Arc>>>, @@ -1750,12 +1791,8 @@ impl NatTraversalEndpoint { peer_address_update_rx: TokioMutex::new(peer_addr_rx), relay_server_config: Arc::new(std::sync::Mutex::new(relay_server_config)), relay_handler_connections: Arc::new(dashmap::DashSet::new()), - relay_setup_attempted: Arc::new(std::sync::atomic::AtomicBool::new(false)), bootstrap_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), - relay_public_addr: Arc::new(std::sync::Mutex::new(None)), - relay_generation: Arc::new(std::sync::atomic::AtomicU64::new(0)), - proactive_relay: Arc::new(std::sync::Mutex::new(None)), - proactive_relay_lifecycle: Arc::new(TokioMutex::new(())), + relay_lifecycle: Arc::new(TokioMutex::new(RelayLifecycleState::None)), transport_listener_handles: Arc::new(ParkingMutex::new(Vec::new())), constrained_engine, constrained_event_tx: constrained_event_tx.clone(), @@ -1937,7 +1974,6 @@ impl NatTraversalEndpoint { let connections_clone = endpoint.connections.clone(); let local_session_id = DiscoverySessionId::Local; - let relay_setup_attempted_clone = endpoint.relay_setup_attempted.clone(); let relay_server_clone = endpoint.relay_server.clone(); let advertise = endpoint.advertise_external_addresses; let bootstrap_complete_clone = endpoint.bootstrap_complete.clone(); @@ -1949,7 +1985,6 @@ impl NatTraversalEndpoint { connections_clone, event_callback_for_poll, local_session_id, - relay_setup_attempted_clone, relay_server_clone, advertise, bootstrap_complete_clone, @@ -2191,12 +2226,8 @@ impl NatTraversalEndpoint { peer_address_update_rx: TokioMutex::new(peer_addr_rx), relay_server_config: Arc::new(std::sync::Mutex::new(relay_server_config)), relay_handler_connections: Arc::new(dashmap::DashSet::new()), - relay_setup_attempted: Arc::new(std::sync::atomic::AtomicBool::new(false)), bootstrap_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), - relay_public_addr: Arc::new(std::sync::Mutex::new(None)), - relay_generation: Arc::new(std::sync::atomic::AtomicU64::new(0)), - proactive_relay: Arc::new(std::sync::Mutex::new(None)), - proactive_relay_lifecycle: Arc::new(TokioMutex::new(())), + relay_lifecycle: Arc::new(TokioMutex::new(RelayLifecycleState::None)), transport_listener_handles: Arc::new(ParkingMutex::new(Vec::new())), constrained_engine, constrained_event_tx: constrained_event_tx.clone(), @@ -2378,7 +2409,6 @@ impl NatTraversalEndpoint { let connections_clone = endpoint.connections.clone(); let local_session_id = DiscoverySessionId::Local; - let relay_setup_attempted_clone = endpoint.relay_setup_attempted.clone(); let relay_server_clone = endpoint.relay_server.clone(); let advertise = endpoint.advertise_external_addresses; let bootstrap_complete_clone = endpoint.bootstrap_complete.clone(); @@ -2390,7 +2420,6 @@ impl NatTraversalEndpoint { connections_clone, event_callback_for_poll, local_session_id, - relay_setup_attempted_clone, relay_server_clone, advertise, bootstrap_complete_clone, @@ -3610,7 +3639,6 @@ impl NatTraversalEndpoint { connections: Arc>, event_callback: Option>, local_session_id: DiscoverySessionId, - relay_setup_attempted: Arc, relay_server: Option>, advertise_external_addresses: bool, bootstrap_complete: Arc, @@ -3769,11 +3797,8 @@ impl NatTraversalEndpoint { // 2. Send ADD_ADDRESS to all peers for newly discovered addresses // (Critical for CGNAT - peers need to know our external address to hole-punch back) - // Skip if relay is active — only the relay address should be advertised. // Skip entirely for outbound-only clients that don't need to be reachable. - if advertise_external_addresses - && !relay_setup_attempted.load(std::sync::atomic::Ordering::Relaxed) - { + if advertise_external_addresses { for addr in &new_addresses { broadcast_address_to_peers(&connections, *addr, 100); } @@ -3967,37 +3992,26 @@ impl NatTraversalEndpoint { } } - /// Get the relay public address, if a proactive relay has been established. - pub fn relay_public_addr(&self) -> Option { - self.relay_public_addr.lock().ok().and_then(|g| *g) - } - - fn current_relay_generation(&self) -> u64 { - self.relay_generation - .load(std::sync::atomic::Ordering::Acquire) - } - - fn next_relay_generation(&self) -> u64 { - self.relay_generation - .fetch_add(1, std::sync::atomic::Ordering::AcqRel) - .wrapping_add(1) + /// Get the relay public address after its allocation has been published. + pub async fn relay_public_addr(&self) -> Option { + self.relay_lifecycle.lock().await.published_addr() } /// Check if the proactive relay session is still alive. Returns true if /// no relay was established (nothing to monitor) or the relay is healthy. /// Returns false if a relay was established but the underlying QUIC /// connection has closed. - pub fn is_relay_healthy(&self) -> bool { - let relay_addr = match self.relay_public_addr.lock().ok().and_then(|g| *g) { - Some(addr) => addr, - None => return true, // No relay — nothing to monitor + pub async fn is_relay_healthy(&self) -> bool { + let lifecycle = self.relay_lifecycle.lock().await; + let relay = match &*lifecycle { + RelayLifecycleState::Published(relay) => relay, + RelayLifecycleState::None | RelayLifecycleState::Provisional(_) => { + return true; + } }; + let relay_addr = relay.handle.public_addr(); - if let Ok(state) = self.proactive_relay.lock() - && let Some(state) = state.as_ref() - && state.public_addr == relay_addr - && state.tunnel.is_closed() - { + if relay.tunnel.is_closed() { return false; } @@ -4010,29 +4024,31 @@ impl NatTraversalEndpoint { } } - // No matching session found warn!( - "Relay session for {} is dead — resetting for re-establishment", + "Relay session for {} is dead — re-establishment required", relay_addr ); false } - /// Reset relay state so the next poll cycle can re-establish. Called when - /// the relay session is detected as dead. - pub fn reset_relay_state(&self) { - let generation = self.next_relay_generation(); - self.relay_setup_attempted - .store(false, std::sync::atomic::Ordering::Relaxed); - if let Ok(mut addr) = self.relay_public_addr.lock() { - *addr = None; + pub(crate) async fn published_relay_handle(&self) -> Option { + let lifecycle = self.relay_lifecycle.lock().await; + match &*lifecycle { + RelayLifecycleState::Published(relay) => Some(relay.handle), + RelayLifecycleState::None | RelayLifecycleState::Provisional(_) => None, + } + } + + pub(crate) async fn abort_current_proactive_relay(&self) { + let mut lifecycle = self.relay_lifecycle.lock().await; + let state = std::mem::take(&mut *lifecycle); + match state { + RelayLifecycleState::Provisional(relay) | RelayLifecycleState::Published(relay) => { + self.teardown_proactive_relay(relay, b"relay allocation aborted") + .await; + } + RelayLifecycleState::None => {} } - // Remove dead sessions - self.relay_sessions.retain(|_, session| session.is_active()); - info!( - generation, - "Relay state reset — will re-establish on next poll cycle" - ); } /// Check if relay fallback is available @@ -4089,7 +4105,6 @@ impl NatTraversalEndpoint { "relay session: creating dedicated control connection" ); let connection = self.connect_new_to_relay(relay_addr).await?; - let owns_connection = true; // Cap on the end-to-end CONNECT-UDP handshake (open_bi → write // request → read_exact response). Without a cap, reusing a peer @@ -4207,7 +4222,6 @@ impl NatTraversalEndpoint { public_address, established_at: std::time::Instant::now(), relay_addr, - owns_connection, }; // DashMap provides lock-free .insert() @@ -5092,20 +5106,20 @@ impl NatTraversalEndpoint { pub async fn prepare_proactive_relay( &self, bootstrap_addr: SocketAddr, - ) -> Result { - let _lifecycle_guard = self.proactive_relay_lifecycle.lock().await; + ) -> Result { + let mut lifecycle = self.relay_lifecycle.lock().await; // A new acquisition supersedes any prior allocation. Teardown happens // before the new CONNECT-UDP request so rejected canaries cannot // accumulate relay-server capacity. - let previous = self - .proactive_relay - .lock() - .map_err(|_| NatTraversalError::ConfigError("mutex poisoned".to_string()))? - .take(); - if let Some(previous) = previous { - self.teardown_proactive_relay(previous, b"relay superseded") - .await; + let previous = std::mem::take(&mut *lifecycle); + match previous { + RelayLifecycleState::Provisional(previous) + | RelayLifecycleState::Published(previous) => { + self.teardown_proactive_relay(previous, b"relay superseded") + .await; + } + RelayLifecycleState::None => {} } info!( @@ -5205,13 +5219,11 @@ impl NatTraversalEndpoint { info!("Relay endpoint created (relay addr: {})", relay_public_addr); - let relay_generation = self.next_relay_generation(); - self.relay_setup_attempted - .store(true, std::sync::atomic::Ordering::Relaxed); + let handle = PreparedRelay::new(relay_public_addr); info!( - relay_generation, + allocation_id = handle.allocation_id, relay_addr = %relay_public_addr, - "Proactive relay generation prepared" + "Proactive relay allocation prepared" ); // Share the relay endpoint between the accept loop and the @@ -5251,70 +5263,52 @@ impl NatTraversalEndpoint { let state = ProactiveRelay { relay_server_addr: bootstrap_addr, - public_addr: relay_public_addr, - generation: relay_generation, + handle, endpoint: relay_endpoint, tunnel, - published: false, }; - if let Err(state) = self.install_proactive_relay(state) { - self.teardown_proactive_relay(state, b"relay state mutex poisoned") - .await; - return Err(NatTraversalError::ConfigError("mutex poisoned".to_string())); - } + *lifecycle = RelayLifecycleState::Provisional(state); - Ok(relay_public_addr) + Ok(handle) } /// Activate a prepared proactive relay after its external canary succeeds. /// - /// The address is checked against the current generation so a late canary - /// verdict cannot publish a superseded allocation. + /// The opaque allocation identity is checked against the current state so + /// a late canary verdict cannot publish a superseded allocation. pub async fn publish_proactive_relay( &self, - relay_public_addr: SocketAddr, + prepared: PreparedRelay, ) -> Result<(), NatTraversalError> { - let _lifecycle_guard = self.proactive_relay_lifecycle.lock().await; - - let relay_generation = { - let mut state = self - .proactive_relay - .lock() - .map_err(|_| NatTraversalError::ConfigError("mutex poisoned".to_string()))?; - let state = state.as_mut().ok_or_else(|| { - NatTraversalError::ConnectionFailed( - "No proactive relay allocation is prepared".to_string(), - ) - })?; - if state.public_addr != relay_public_addr { - return Err(NatTraversalError::ConnectionFailed(format!( - "Prepared relay address {} does not match {}", - state.public_addr, relay_public_addr - ))); + let mut lifecycle = self.relay_lifecycle.lock().await; + let current = std::mem::take(&mut *lifecycle); + let relay = match current { + RelayLifecycleState::Provisional(relay) if relay.handle == prepared => relay, + RelayLifecycleState::Published(relay) if relay.handle == prepared => { + *lifecycle = RelayLifecycleState::Published(relay); + return Ok(()); } - if state.tunnel.is_closed() { + other => { + *lifecycle = other; return Err(NatTraversalError::ConnectionFailed(format!( - "Prepared relay tunnel for {relay_public_addr} is closed" + "Prepared relay allocation {} at {} is stale or no longer active", + prepared.allocation_id, + prepared.public_addr() ))); } - if state.published { - return Ok(()); - } - - state.published = true; - state.generation }; - - if self.current_relay_generation() != relay_generation { - return Err(NatTraversalError::ConnectionFailed( - "Prepared relay generation was superseded".to_string(), - )); + if relay.tunnel.is_closed() { + *lifecycle = RelayLifecycleState::Provisional(relay); + return Err(NatTraversalError::ConnectionFailed(format!( + "Prepared relay tunnel for {} is closed", + prepared.public_addr() + ))); } - *self - .relay_public_addr - .lock() - .map_err(|_| NatTraversalError::ConfigError("mutex poisoned".to_string()))? = - Some(relay_public_addr); + + // No fallible work follows this assignment. Publication is therefore a + // single transactional transition: retries can never observe + // `Published` without the address also being visible through the state. + *lifecycle = RelayLifecycleState::Published(relay); // Relay addresses are intentionally not broadcast through the // connection-level ADD_ADDRESS extension. A relay allocation belongs @@ -5324,8 +5318,8 @@ impl NatTraversalEndpoint { // handshakes. Saorsa-core publishes the canary-verified address through // its authenticated, sequenced PublishAddressSet path instead. info!( - relay_generation, - relay_addr = %relay_public_addr, + allocation_id = prepared.allocation_id, + relay_addr = %prepared.public_addr(), "Activated canary-verified proactive relay" ); Ok(()) @@ -5333,27 +5327,27 @@ impl NatTraversalEndpoint { /// Abort a prepared or published proactive relay. /// - /// A mismatched or already-removed address is treated as a stale, - /// idempotent abort and leaves the current generation untouched. + /// A mismatched or already-removed handle is treated as a stale, + /// idempotent abort and leaves the current allocation untouched. pub async fn abort_proactive_relay( &self, - relay_public_addr: SocketAddr, + prepared: PreparedRelay, ) -> Result<(), NatTraversalError> { - let _lifecycle_guard = self.proactive_relay_lifecycle.lock().await; - let state = { - let mut current = self - .proactive_relay - .lock() - .map_err(|_| NatTraversalError::ConfigError("mutex poisoned".to_string()))?; - match current.as_ref() { - Some(state) if state.public_addr == relay_public_addr => current.take(), - _ => None, + let mut lifecycle = self.relay_lifecycle.lock().await; + let current = std::mem::take(&mut *lifecycle); + match current { + RelayLifecycleState::Provisional(relay) | RelayLifecycleState::Published(relay) + if relay.handle == prepared => + { + self.teardown_proactive_relay(relay, b"relay allocation aborted") + .await; + } + other => { + // A stale abort is intentionally idempotent, but unlike the + // former address-only check it cannot tear down a newer relay + // that reclaimed the same public socket. + *lifecycle = other; } - }; - - if let Some(state) = state { - self.teardown_proactive_relay(state, b"relay allocation aborted") - .await; } Ok(()) } @@ -5367,12 +5361,12 @@ impl NatTraversalEndpoint { &self, relay_addr: SocketAddr, ) -> Result { - let allocated = self.prepare_proactive_relay(relay_addr).await?; - if let Err(error) = self.publish_proactive_relay(allocated).await { - let _ = self.abort_proactive_relay(allocated).await; + let prepared = self.prepare_proactive_relay(relay_addr).await?; + if let Err(error) = self.publish_proactive_relay(prepared).await { + let _ = self.abort_proactive_relay(prepared).await; return Err(error); } - Ok(allocated) + Ok(prepared.public_addr()) } async fn teardown_proactive_relay(&self, state: ProactiveRelay, reason: &'static [u8]) { @@ -5386,46 +5380,25 @@ impl NatTraversalEndpoint { .is_err() { debug!( - relay_addr = %state.public_addr, + relay_addr = %state.handle.public_addr(), "Timed out draining proactive relay endpoint during teardown" ); } self.remove_matching_relay_session( state.relay_server_addr, - Some(state.public_addr), + Some(state.handle.public_addr()), reason, ); - if self.current_relay_generation() == state.generation { - self.next_relay_generation(); - if let Ok(mut addr) = self.relay_public_addr.lock() - && *addr == Some(state.public_addr) - { - *addr = None; - } - self.relay_setup_attempted - .store(false, std::sync::atomic::Ordering::Release); - } - info!( - relay_addr = %state.public_addr, + relay_addr = %state.handle.public_addr(), relay_server = %state.relay_server_addr, - published = state.published, + allocation_id = state.handle.allocation_id, "Proactive relay torn down" ); } - fn install_proactive_relay(&self, state: ProactiveRelay) -> Result<(), ProactiveRelay> { - match self.proactive_relay.lock() { - Ok(mut current) => { - *current = Some(state); - Ok(()) - } - Err(_) => Err(state), - } - } - fn remove_matching_relay_session( &self, relay_server_addr: SocketAddr, @@ -5440,9 +5413,7 @@ impl NatTraversalEndpoint { return; } - if let Some((_, session)) = self.relay_sessions.remove(&relay_server_addr) - && session.owns_connection - { + if let Some((_, session)) = self.relay_sessions.remove(&relay_server_addr) { session .connection .close(crate::VarInt::from_u32(RELAY_TUNNEL_LOST_CODE), reason); @@ -6008,14 +5979,7 @@ impl NatTraversalEndpoint { self.incoming_notify.notify_waiters(); self.shutdown_notify.notify_waiters(); - let proactive_addr = self - .proactive_relay - .lock() - .ok() - .and_then(|state| state.as_ref().map(|relay| relay.public_addr)); - if let Some(proactive_addr) = proactive_addr { - let _ = self.abort_proactive_relay(proactive_addr).await; - } + self.abort_current_proactive_relay().await; // Best-effort UPnP teardown. The endpoint is the sole owner of // the service (the discovery manager only holds a read-only @@ -7828,16 +7792,6 @@ impl NatTraversalEndpoint { addr: SocketAddr, candidate: &CandidateAddress, ) -> Result<(), NatTraversalError> { - // After relay setup, suppress automatic candidate advertisements. - // The relay address is the only reachable address for this node; - // advertising NATted addresses would overwrite it in peers' DHTs. - if self - .relay_setup_attempted - .load(std::sync::atomic::Ordering::Relaxed) - { - return Ok(()); - } - debug!( "Sending candidate advertisement to {}: {}", addr, candidate.address @@ -8043,6 +7997,18 @@ impl crate::TokenStore for DefaultTokenStore { mod tests { use super::*; + #[test] + fn prepared_relay_identity_distinguishes_reused_public_address() { + let public_addr = "203.0.113.7:9000".parse().expect("test address"); + let first = PreparedRelay::detached(public_addr); + let first_copy = first; + let replacement = PreparedRelay::detached(public_addr); + + assert_eq!(first, first_copy); + assert_ne!(first, replacement); + assert_eq!(first.public_addr(), replacement.public_addr()); + } + #[test] fn test_nat_traversal_config_default() { let config = NatTraversalConfig::default(); diff --git a/src/p2p_endpoint.rs b/src/p2p_endpoint.rs index 99b09982..bdbc8089 100644 --- a/src/p2p_endpoint.rs +++ b/src/p2p_endpoint.rs @@ -74,7 +74,7 @@ use crate::happy_eyeballs::{self, HappyEyeballsConfig}; pub use crate::nat_traversal_api::TraversalPhase; use crate::nat_traversal_api::{ NatTraversalEndpoint, NatTraversalError, NatTraversalEvent, NatTraversalStatistics, - RELAY_TUNNEL_INITIAL_MTU, + PreparedRelay, RELAY_TUNNEL_INITIAL_MTU, }; use crate::shared::{normalize_socket_addr, socket_addr_variants}; use crate::transport::{ProtocolEngine, TransportAddr, TransportRegistry}; @@ -2886,6 +2886,19 @@ impl P2pEndpoint { /// This method races the inner accept against the shutdown token, so it /// will return promptly when `shutdown()` is called. pub async fn accept(&self) -> Option { + self.accept_with_connection() + .await + .map(|(peer, _connection)| peer) + } + + /// Accept an incoming peer while retaining its authoritative QUIC handle. + /// + /// The link-transport adapter needs the exact accepted handle. Looking it + /// up again by address races connection cleanup when a short-lived peer + /// closes immediately after the handshake. + pub(crate) async fn accept_with_connection( + &self, + ) -> Option<(PeerConnection, crate::high_level::Connection)> { if self.shutdown.is_cancelled() { return None; } @@ -2899,6 +2912,8 @@ impl P2pEndpoint { Ok((remote_addr, connection)) => { // Extract public key from TLS handshake let remote_public_key = extract_public_key_bytes_from_connection(&connection); + let reader_connection = connection.clone(); + let accepted_connection = connection.clone(); // They initiated the connection to us = Server side if let Err(e) = @@ -2918,26 +2933,11 @@ impl P2pEndpoint { last_activity: Instant::now(), }; - // Spawn background reader task BEFORE storing in connected_peers - // to prevent race where recv() misses early data - match self.inner.get_connection(&remote_addr) { - Ok(Some(conn)) => { - info!("accept: spawning reader task for {}", remote_addr); - self.spawn_reader_task(remote_addr, conn).await; - } - Ok(None) => { - error!( - "accept: get_connection({}) returned None — NO reader task spawned!", - remote_addr - ); - } - Err(e) => { - error!( - "accept: get_connection({}) failed: {} — NO reader task spawned!", - remote_addr, e - ); - } - } + // The accepted connection is authoritative. Using it directly + // avoids racing a peer's immediate close against a redundant + // lookup in the shared connection map. + info!("accept: spawning reader task for {}", remote_addr); + self.spawn_reader_task(remote_addr, reader_connection).await; self.connected_peers .write() @@ -2959,7 +2959,7 @@ impl P2pEndpoint { Side::Server, ); - Some(peer_conn) + Some((peer_conn, accepted_connection)) } Err(e) => { debug!("Accept failed: {}", e); @@ -3416,8 +3416,8 @@ impl P2pEndpoint { /// relay's underlying QUIC connection is still open. Returns `false` if a /// relay was established but the session has died — the caller should /// rebind. - pub fn is_relay_healthy(&self) -> bool { - self.inner.is_relay_healthy() + pub async fn is_relay_healthy(&self) -> bool { + self.inner.is_relay_healthy().await } /// Enable or disable relay serving on this node's MASQUE relay server. @@ -3467,18 +3467,16 @@ impl P2pEndpoint { pub async fn prepare_proactive_relay( &self, relay_addr: SocketAddr, - ) -> Result { + ) -> Result { Ok(self.inner.prepare_proactive_relay(relay_addr).await?) } /// Activate the current provisional proactive relay. pub async fn publish_proactive_relay( &self, - relay_public_addr: SocketAddr, + prepared: PreparedRelay, ) -> Result<(), EndpointError> { - self.inner - .publish_proactive_relay(relay_public_addr) - .await?; + self.inner.publish_proactive_relay(prepared).await?; Ok(()) } @@ -3486,9 +3484,9 @@ impl P2pEndpoint { /// MASQUE session and relay-server capacity. pub async fn abort_proactive_relay( &self, - relay_public_addr: SocketAddr, + prepared: PreparedRelay, ) -> Result<(), EndpointError> { - self.inner.abort_proactive_relay(relay_public_addr).await?; + self.inner.abort_proactive_relay(prepared).await?; Ok(()) } @@ -4002,7 +4000,8 @@ impl P2pEndpoint { // Upper layers use this to trigger a DHT self-lookup for // relay address propagation. if !relay_event_sent { - if let Some(relay_addr) = inner.relay_public_addr() { + if let Some(prepared) = inner.published_relay_handle().await { + let relay_addr = prepared.public_addr(); info!( "Relay established at {} — emitting RelayEstablished event", relay_addr @@ -4013,29 +4012,22 @@ impl P2pEndpoint { } // Monitor relay health. If the relay session died (connection - // closed, server restarted, etc.), reset state so the next - // poll cycle re-establishes through a (potentially different) - // relay candidate. The RelayEstablished flag is also reset so - // upper layers re-publish the new address. - if relay_event_sent && !inner.is_relay_healthy() { - // Snapshot the dead address BEFORE reset clears it so we - // can tell upper layers exactly which relay to stop - // advertising. - let dead = inner.relay_public_addr(); - if let Some(relay_addr) = dead { - if let Err(error) = inner.abort_proactive_relay(relay_addr).await { + // closed, server restarted, etc.), tear down its one lifecycle + // state so the upper layer can acquire a replacement. + if relay_event_sent && !inner.is_relay_healthy().await { + let dead = inner.published_relay_handle().await; + if let Some(prepared) = dead { + if let Err(error) = inner.abort_proactive_relay(prepared).await { warn!( - relay_addr = %relay_addr, + relay_addr = %prepared.public_addr(), %error, "Failed to tear down unhealthy proactive relay" ); - inner.reset_relay_state(); } - } else { - inner.reset_relay_state(); } relay_event_sent = false; - if let Some(relay_addr) = dead { + if let Some(prepared) = dead { + let relay_addr = prepared.public_addr(); info!( "Relay tunnel at {} is unhealthy — emitting RelayLost event", relay_addr From 478ee3d36d8016156a8f33bd45844497fd3e4c9a Mon Sep 17 00:00:00 2001 From: Mick van Dijke <12992260+mickvandijke@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:04:03 +0200 Subject: [PATCH 3/8] fix(relay): harden provisional relay lifecycle --- src/lib.rs | 2 + src/masque/connect.rs | 58 +++- src/masque/relay_server.rs | 41 +++ src/masque/relay_socket.rs | 19 +- src/nat_traversal_api.rs | 535 ++++++++++++++++++++++++++++++------- src/p2p_endpoint.rs | 24 +- src/relay_allocation.rs | 393 +++++++++++++++++++++++++++ 7 files changed, 972 insertions(+), 100 deletions(-) create mode 100644 src/relay_allocation.rs diff --git a/src/lib.rs b/src/lib.rs index 0c010d8f..c3c27aa5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -258,6 +258,8 @@ pub mod metrics; /// TURN-style relay protocol for NAT traversal fallback pub mod relay; +mod relay_allocation; +pub use relay_allocation::{RelayAllocationReceipt, RelayAllocationReceiptError}; /// MASQUE CONNECT-UDP Bind protocol for fully connectable P2P nodes pub mod masque; diff --git a/src/masque/connect.rs b/src/masque/connect.rs index cd6a17ed..86a4e1c6 100644 --- a/src/masque/connect.rs +++ b/src/masque/connect.rs @@ -275,6 +275,9 @@ pub struct ConnectUdpResponse { pub proxy_public_address: Option, /// Human-readable reason phrase pub reason: Option, + /// Relay-signed proof that this allocation belongs to the authenticated + /// client. Present for successful authenticated bind allocations. + pub allocation_receipt: Option, /// Relay-internal: the session id created for a successful CONNECT. NOT part /// of the wire format (it is not encoded/decoded); the relay sets it so the /// connection handler can start forwarding the *exact* session it created, @@ -300,6 +303,7 @@ impl ConnectUdpResponse { status: Self::STATUS_OK, proxy_public_address: public_addr, reason: None, + allocation_receipt: None, session_id: None, } } @@ -310,6 +314,7 @@ impl ConnectUdpResponse { status, proxy_public_address: None, reason: Some(reason.into()), + allocation_receipt: None, session_id: None, } } @@ -353,14 +358,14 @@ impl ConnectUdpResponse { /// Encode the response as wire format /// - /// Format: [status (2)] [flags (1)] [addr if present] + /// Format: [status (2)] [flags (1)] [addr] [reason] [allocation receipt] pub fn encode(&self) -> Bytes { let mut buf = BytesMut::new(); // Status code buf.put_u16(self.status); - // Flags: bit 0 = has address, bit 1 = has reason + // Flags: bit 0 = has address, bit 1 = has reason, bit 2 = has receipt let mut flags: u8 = 0; if self.proxy_public_address.is_some() { flags |= 0x01; @@ -368,6 +373,9 @@ impl ConnectUdpResponse { if self.reason.is_some() { flags |= 0x02; } + if self.allocation_receipt.is_some() { + flags |= 0x04; + } buf.put_u8(flags); // Public address if present @@ -394,6 +402,14 @@ impl ConnectUdpResponse { buf.put_slice(reason_bytes); } + if let Some(receipt) = &self.allocation_receipt { + let receipt = receipt.encode(); + if let Ok(length) = VarInt::from_u64(receipt.len() as u64) { + length.encode(&mut buf); + } + buf.put_slice(&receipt); + } + buf.freeze() } @@ -407,6 +423,7 @@ impl ConnectUdpResponse { let flags = buf.get_u8(); let has_addr = (flags & 0x01) != 0; let has_reason = (flags & 0x02) != 0; + let has_receipt = (flags & 0x04) != 0; let proxy_public_address = if has_addr { if buf.remaining() < 1 { @@ -457,10 +474,31 @@ impl ConnectUdpResponse { None }; + let allocation_receipt = if has_receipt { + let receipt_len = VarInt::decode(buf) + .map_err(|_| ConnectError::InvalidResponse("invalid receipt length".into()))? + .into_inner() as usize; + if buf.remaining() < receipt_len { + return Err(ConnectError::InvalidResponse( + "missing allocation receipt".into(), + )); + } + let mut receipt = vec![0; receipt_len]; + buf.copy_to_slice(&mut receipt); + Some( + crate::RelayAllocationReceipt::decode(&receipt).map_err(|_| { + ConnectError::InvalidResponse("invalid allocation receipt".into()) + })?, + ) + } else { + None + }; + Ok(Self { status, proxy_public_address, reason, + allocation_receipt, // Not part of the wire format; only set by the relay on the response // object it returns locally. session_id: None, @@ -569,6 +607,22 @@ mod tests { assert_eq!(original, decoded); } + #[test] + fn test_response_roundtrip_with_allocation_receipt() { + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(203, 0, 113, 50)), 9000); + let (public_key, secret_key) = crate::generate_ml_dsa_keypair().expect("test identity"); + let mut original = ConnectUdpResponse::success(Some(addr)); + original.allocation_receipt = Some( + crate::RelayAllocationReceipt::issue(&public_key, &secret_key, [7; 32], addr, 42) + .expect("allocation receipt"), + ); + + let encoded = original.encode(); + let decoded = ConnectUdpResponse::decode(&mut encoded.clone()).unwrap(); + + assert_eq!(original, decoded); + } + #[test] fn test_response_roundtrip_success_no_addr() { let original = ConnectUdpResponse::success(None); diff --git a/src/masque/relay_server.rs b/src/masque/relay_server.rs index efc37ee4..59a6543c 100644 --- a/src/masque/relay_server.rs +++ b/src/masque/relay_server.rs @@ -55,6 +55,7 @@ use crate::masque::{ }; use crate::relay::error::{RelayError, RelayResult, SessionErrorKind}; use crate::upnp::{UpnpConfig, UpnpMappingService}; +use crate::{MlDsaPublicKey, MlDsaSecretKey, RelayAllocationReceipt}; /// Interval at which both sides of a relay stream send a zero-length /// keepalive frame. Keeps the NAT conntrack entry alive (default @@ -576,6 +577,9 @@ pub struct MasqueRelayServer { /// When set, the server refuses inbound clients whose source IP matches /// one of our current upstream relays (see [`IpPolicy`]). ip_policy: Option>, + /// Identity used to sign allocations for third-party reachability + /// witnesses. It is the same identity presented by this endpoint in TLS. + allocation_signer: Option<(MlDsaPublicKey, MlDsaSecretKey)>, } impl std::fmt::Debug for MasqueRelayServer { @@ -616,6 +620,7 @@ impl MasqueRelayServer { .map(|_| Mutex::new(())) .collect(), ip_policy: None, + allocation_signer: None, } } @@ -689,9 +694,19 @@ impl MasqueRelayServer { .map(|_| Mutex::new(())) .collect(), ip_policy: None, + allocation_signer: None, } } + /// Configure the endpoint identity used to sign relay allocations. + pub fn set_allocation_signer( + &mut self, + public_key: MlDsaPublicKey, + secret_key: MlDsaSecretKey, + ) { + self.allocation_signer = Some((public_key, secret_key)); + } + /// Enable or disable relay serving. /// /// Called by the ADR-014 reachability classifier: public nodes leave this @@ -1075,6 +1090,31 @@ impl MasqueRelayServer { // Create new session with the bound socket let session_id = self.next_session_id.fetch_add(1, Ordering::SeqCst); + let allocation_receipt = match (&self.allocation_signer, peer_id) { + (Some((public_key, secret_key)), Some(target_peer_id)) => { + match RelayAllocationReceipt::issue( + public_key, + secret_key, + target_peer_id, + advertised_address, + session_id, + ) { + Ok(receipt) => Some(receipt), + Err(error) => { + tracing::error!( + %error, + client = %client_addr, + "Failed to sign relay allocation" + ); + return Ok(ConnectUdpResponse::error( + 500, + "Failed to sign relay allocation", + )); + } + } + } + _ => None, + }; let mut session = RelaySession::new( session_id, self.config.session_config.clone(), @@ -1127,6 +1167,7 @@ impl MasqueRelayServer { // session rather than re-looking-up by client address (which races with a // same-address reconnect). Not part of the wire format. let mut response = ConnectUdpResponse::success(Some(advertised_address)); + response.allocation_receipt = allocation_receipt; response.session_id = Some(session_id); Ok(response) } diff --git a/src/masque/relay_socket.rs b/src/masque/relay_socket.rs index 4fdcc960..243ce6e1 100644 --- a/src/masque/relay_socket.rs +++ b/src/masque/relay_socket.rs @@ -178,6 +178,12 @@ impl RelayTunnelControl { } } +fn mark_writer_exit(control: &Weak) { + if let Some(control) = control.upgrade() { + control.mark_closed(); + } +} + /// A virtual UDP socket backed entirely by a MASQUE relay tunnel. /// /// All traffic — both outgoing and incoming — flows through the relay @@ -385,6 +391,7 @@ impl MasqueRelaySocket { // Background task: write queued outbound packets to relay stream. let writer_capacity = Arc::clone(&send_capacity_freed); + let writer_control = Arc::downgrade(&control); let writer_handle = tokio::spawn(async move { while let Some(encoded) = send_rx.recv().await { // `recv` completing means the channel just freed a @@ -427,6 +434,7 @@ impl MasqueRelaySocket { // of waiting forever. drop(send_rx); writer_capacity.notify_waiters(); + mark_writer_exit(&writer_control); }); control.register(writer_handle); @@ -708,7 +716,7 @@ impl UdpPoller for TunnelPoller { #[cfg(test)] mod relay_tunnel_control_tests { - use super::RelayTunnelControl; + use super::{RelayTunnelControl, mark_writer_exit}; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; @@ -752,4 +760,13 @@ mod relay_tunnel_control_tests { assert!(control.is_closed()); } + + #[test] + fn writer_exit_marks_tunnel_closed() { + let control = RelayTunnelControl::new(); + + mark_writer_exit(&Arc::downgrade(&control)); + + assert!(control.is_closed()); + } } diff --git a/src/nat_traversal_api.rs b/src/nat_traversal_api.rs index fec18d80..aa54418f 100644 --- a/src/nat_traversal_api.rs +++ b/src/nat_traversal_api.rs @@ -275,7 +275,7 @@ use parking_lot::{Mutex as ParkingMutex, RwLock as ParkingRwLock}; use tokio::{ net::UdpSocket, - sync::{Mutex as TokioMutex, mpsc}, + sync::{Mutex as TokioMutex, mpsc, oneshot}, time::{sleep, timeout}, }; @@ -361,6 +361,8 @@ pub struct RelaySession { pub established_at: std::time::Instant, /// Relay server address pub relay_addr: SocketAddr, + /// Relay-signed proof for this exact allocation, when available. + pub allocation_receipt: Option, } impl RelaySession { @@ -424,6 +426,7 @@ struct ProactiveRelay { handle: PreparedRelay, endpoint: Arc, tunnel: Arc, + allocation_receipt: Option, } /// Complete proactive-relay lifecycle state. @@ -448,6 +451,299 @@ impl RelayLifecycleState { } } +struct RelayHealthSnapshot { + relay_addr: SocketAddr, + tunnel: Arc, +} + +enum RelayLifecycleCommand { + BeginPrepare { + reply: oneshot::Sender<(u64, Option)>, + }, + CompletePrepare { + generation: u64, + relay: ProactiveRelay, + reply: oneshot::Sender>, + }, + Publish { + prepared: PreparedRelay, + reply: oneshot::Sender>, + }, + TakeMatching { + prepared: PreparedRelay, + reply: oneshot::Sender>, + }, + TakeCurrent { + reply: oneshot::Sender>, + }, + PublishedAddr { + reply: oneshot::Sender>, + }, + PublishedHandle { + reply: oneshot::Sender>, + }, + Receipt { + prepared: PreparedRelay, + reply: oneshot::Sender>, + }, + Health { + reply: oneshot::Sender>, + }, +} + +/// Serializes proactive-relay state transitions without holding a mutex across +/// network acquisition or teardown. Slow I/O stays with the caller; this actor +/// only owns short, deterministic state changes. +#[derive(Clone)] +struct RelayLifecycleHandle { + commands: mpsc::Sender, + published: Arc, +} + +impl RelayLifecycleHandle { + fn new() -> Self { + let (commands, mut receiver) = mpsc::channel(16); + let published = Arc::new(AtomicBool::new(false)); + let actor_published = Arc::clone(&published); + tokio::spawn(async move { + let mut generation = 0u64; + let mut state = RelayLifecycleState::None; + while let Some(command) = receiver.recv().await { + match command { + RelayLifecycleCommand::BeginPrepare { reply } => { + generation = generation.wrapping_add(1); + actor_published.store(false, Ordering::Release); + let previous = match std::mem::take(&mut state) { + RelayLifecycleState::Provisional(relay) + | RelayLifecycleState::Published(relay) => Some(relay), + RelayLifecycleState::None => None, + }; + let _ = reply.send((generation, previous)); + } + RelayLifecycleCommand::CompletePrepare { + generation: completed_generation, + relay, + reply, + } => { + if generation == completed_generation + && matches!(state, RelayLifecycleState::None) + { + state = RelayLifecycleState::Provisional(relay); + let _ = reply.send(Ok(())); + } else { + let _ = reply.send(Err(relay)); + } + } + RelayLifecycleCommand::Publish { prepared, reply } => { + let current = std::mem::take(&mut state); + match current { + RelayLifecycleState::Provisional(relay) if relay.handle == prepared => { + if relay.tunnel.is_closed() { + state = RelayLifecycleState::Provisional(relay); + let _ = reply.send(Err(format!( + "Prepared relay tunnel for {} is closed", + prepared.public_addr() + ))); + } else { + state = RelayLifecycleState::Published(relay); + actor_published.store(true, Ordering::Release); + let _ = reply.send(Ok(())); + } + } + RelayLifecycleState::Published(relay) if relay.handle == prepared => { + state = RelayLifecycleState::Published(relay); + actor_published.store(true, Ordering::Release); + let _ = reply.send(Ok(())); + } + other => { + state = other; + let _ = reply.send(Err(format!( + "Prepared relay allocation {} at {} is stale or no longer active", + prepared.allocation_id, + prepared.public_addr() + ))); + } + } + } + RelayLifecycleCommand::TakeMatching { prepared, reply } => { + let current = std::mem::take(&mut state); + match current { + RelayLifecycleState::Provisional(relay) + | RelayLifecycleState::Published(relay) + if relay.handle == prepared => + { + generation = generation.wrapping_add(1); + actor_published.store(false, Ordering::Release); + let _ = reply.send(Some(relay)); + } + other => { + state = other; + let _ = reply.send(None); + } + } + } + RelayLifecycleCommand::TakeCurrent { reply } => { + generation = generation.wrapping_add(1); + actor_published.store(false, Ordering::Release); + let relay = match std::mem::take(&mut state) { + RelayLifecycleState::Provisional(relay) + | RelayLifecycleState::Published(relay) => Some(relay), + RelayLifecycleState::None => None, + }; + let _ = reply.send(relay); + } + RelayLifecycleCommand::PublishedAddr { reply } => { + let _ = reply.send(state.published_addr()); + } + RelayLifecycleCommand::PublishedHandle { reply } => { + let handle = match &state { + RelayLifecycleState::Published(relay) => Some(relay.handle), + RelayLifecycleState::None | RelayLifecycleState::Provisional(_) => None, + }; + let _ = reply.send(handle); + } + RelayLifecycleCommand::Receipt { prepared, reply } => { + let receipt = match &state { + RelayLifecycleState::Provisional(relay) + | RelayLifecycleState::Published(relay) + if relay.handle == prepared => + { + relay.allocation_receipt.clone() + } + RelayLifecycleState::None + | RelayLifecycleState::Provisional(_) + | RelayLifecycleState::Published(_) => None, + }; + let _ = reply.send(receipt); + } + RelayLifecycleCommand::Health { reply } => { + let health = match &state { + RelayLifecycleState::Published(relay) => Some(RelayHealthSnapshot { + relay_addr: relay.handle.public_addr(), + tunnel: Arc::clone(&relay.tunnel), + }), + RelayLifecycleState::None | RelayLifecycleState::Provisional(_) => None, + }; + let _ = reply.send(health); + } + } + } + }); + Self { + commands, + published, + } + } + + fn is_published(&self) -> bool { + self.published.load(Ordering::Acquire) + } + + async fn begin_prepare(&self) -> Result<(u64, Option), String> { + let (reply, response) = oneshot::channel(); + self.commands + .send(RelayLifecycleCommand::BeginPrepare { reply }) + .await + .map_err(|_| "relay lifecycle actor stopped".to_string())?; + response + .await + .map_err(|_| "relay lifecycle actor stopped".to_string()) + } + + async fn complete_prepare( + &self, + generation: u64, + relay: ProactiveRelay, + ) -> Result, String> { + let (reply, response) = oneshot::channel(); + self.commands + .send(RelayLifecycleCommand::CompletePrepare { + generation, + relay, + reply, + }) + .await + .map_err(|_| "relay lifecycle actor stopped".to_string())?; + response + .await + .map_err(|_| "relay lifecycle actor stopped".to_string()) + } + + async fn publish(&self, prepared: PreparedRelay) -> Result<(), String> { + let (reply, response) = oneshot::channel(); + self.commands + .send(RelayLifecycleCommand::Publish { prepared, reply }) + .await + .map_err(|_| "relay lifecycle actor stopped".to_string())?; + response + .await + .map_err(|_| "relay lifecycle actor stopped".to_string())? + } + + async fn take_matching( + &self, + prepared: PreparedRelay, + ) -> Result, String> { + let (reply, response) = oneshot::channel(); + self.commands + .send(RelayLifecycleCommand::TakeMatching { prepared, reply }) + .await + .map_err(|_| "relay lifecycle actor stopped".to_string())?; + response + .await + .map_err(|_| "relay lifecycle actor stopped".to_string()) + } + + async fn take_current(&self) -> Option { + let (reply, response) = oneshot::channel(); + if self + .commands + .send(RelayLifecycleCommand::TakeCurrent { reply }) + .await + .is_err() + { + return None; + } + response.await.ok().flatten() + } + + async fn published_addr(&self) -> Option { + let (reply, response) = oneshot::channel(); + self.commands + .send(RelayLifecycleCommand::PublishedAddr { reply }) + .await + .ok()?; + response.await.ok().flatten() + } + + async fn published_handle(&self) -> Option { + let (reply, response) = oneshot::channel(); + self.commands + .send(RelayLifecycleCommand::PublishedHandle { reply }) + .await + .ok()?; + response.await.ok().flatten() + } + + async fn receipt(&self, prepared: PreparedRelay) -> Option { + let (reply, response) = oneshot::channel(); + self.commands + .send(RelayLifecycleCommand::Receipt { prepared, reply }) + .await + .ok()?; + response.await.ok().flatten() + } + + async fn health(&self) -> Option { + let (reply, response) = oneshot::channel(); + self.commands + .send(RelayLifecycleCommand::Health { reply }) + .await + .ok()?; + response.await.ok().flatten() + } +} + /// Event from the constrained engine with transport address context /// /// This wrapper adds the transport address to engine events so that P2pEndpoint @@ -553,7 +849,7 @@ pub struct NatTraversalEndpoint { /// candidate set from OBSERVED_ADDRESS frames. bootstrap_complete: Arc, /// Current proactive relay allocation and its publication state. - relay_lifecycle: Arc>, + relay_lifecycle: RelayLifecycleHandle, /// Task handles for transport listener tasks /// Used for cleanup on shutdown transport_listener_handles: Arc>>>, @@ -1724,7 +2020,10 @@ impl NatTraversalEndpoint { require_authentication: true, ..MasqueRelayConfig::default() }; - let server = MasqueRelayServer::new(relay_config, local_addr); + let mut server = MasqueRelayServer::new(relay_config, local_addr); + if let Some((public_key, secret_key)) = config.identity_key.clone() { + server.set_allocation_signer(public_key, secret_key); + } info!( "Created MASQUE relay server on {} (symmetric P2P node)", local_addr @@ -1792,7 +2091,7 @@ impl NatTraversalEndpoint { relay_server_config: Arc::new(std::sync::Mutex::new(relay_server_config)), relay_handler_connections: Arc::new(dashmap::DashSet::new()), bootstrap_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), - relay_lifecycle: Arc::new(TokioMutex::new(RelayLifecycleState::None)), + relay_lifecycle: RelayLifecycleHandle::new(), transport_listener_handles: Arc::new(ParkingMutex::new(Vec::new())), constrained_engine, constrained_event_tx: constrained_event_tx.clone(), @@ -1977,6 +2276,7 @@ impl NatTraversalEndpoint { let relay_server_clone = endpoint.relay_server.clone(); let advertise = endpoint.advertise_external_addresses; let bootstrap_complete_clone = endpoint.bootstrap_complete.clone(); + let relay_lifecycle = endpoint.relay_lifecycle.clone(); tokio::spawn(async move { Self::poll_discovery( discovery_manager_clone, @@ -1988,6 +2288,7 @@ impl NatTraversalEndpoint { relay_server_clone, advertise, bootstrap_complete_clone, + relay_lifecycle, ) .await; }); @@ -2159,7 +2460,10 @@ impl NatTraversalEndpoint { require_authentication: true, ..MasqueRelayConfig::default() }; - let server = MasqueRelayServer::new(relay_config, local_addr); + let mut server = MasqueRelayServer::new(relay_config, local_addr); + if let Some((public_key, secret_key)) = config.identity_key.clone() { + server.set_allocation_signer(public_key, secret_key); + } info!( "Created MASQUE relay server on {} (symmetric P2P node)", local_addr @@ -2227,7 +2531,7 @@ impl NatTraversalEndpoint { relay_server_config: Arc::new(std::sync::Mutex::new(relay_server_config)), relay_handler_connections: Arc::new(dashmap::DashSet::new()), bootstrap_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), - relay_lifecycle: Arc::new(TokioMutex::new(RelayLifecycleState::None)), + relay_lifecycle: RelayLifecycleHandle::new(), transport_listener_handles: Arc::new(ParkingMutex::new(Vec::new())), constrained_engine, constrained_event_tx: constrained_event_tx.clone(), @@ -2412,6 +2716,7 @@ impl NatTraversalEndpoint { let relay_server_clone = endpoint.relay_server.clone(); let advertise = endpoint.advertise_external_addresses; let bootstrap_complete_clone = endpoint.bootstrap_complete.clone(); + let relay_lifecycle = endpoint.relay_lifecycle.clone(); tokio::spawn(async move { Self::poll_discovery( discovery_manager_clone, @@ -2423,6 +2728,7 @@ impl NatTraversalEndpoint { relay_server_clone, advertise, bootstrap_complete_clone, + relay_lifecycle, ) .await; }); @@ -3642,6 +3948,7 @@ impl NatTraversalEndpoint { relay_server: Option>, advertise_external_addresses: bool, bootstrap_complete: Arc, + relay_lifecycle: RelayLifecycleHandle, ) { use tokio::time::{Duration, interval}; @@ -3798,7 +4105,7 @@ impl NatTraversalEndpoint { // 2. Send ADD_ADDRESS to all peers for newly discovered addresses // (Critical for CGNAT - peers need to know our external address to hole-punch back) // Skip entirely for outbound-only clients that don't need to be reachable. - if advertise_external_addresses { + if advertise_external_addresses && !relay_lifecycle.is_published() { for addr in &new_addresses { broadcast_address_to_peers(&connections, *addr, 100); } @@ -3852,7 +4159,10 @@ impl NatTraversalEndpoint { // first observer so hole-punch coordination can // start before quorum. See the OBSERVED_ADDRESS // path above for the rationale. - if check.new_observer && advertise_external_addresses { + if check.new_observer + && advertise_external_addresses + && !relay_lifecycle.is_published() + { broadcast_address_to_peers( &connections, candidate.address, @@ -3994,7 +4304,7 @@ impl NatTraversalEndpoint { /// Get the relay public address after its allocation has been published. pub async fn relay_public_addr(&self) -> Option { - self.relay_lifecycle.lock().await.published_addr() + self.relay_lifecycle.published_addr().await } /// Check if the proactive relay session is still alive. Returns true if @@ -4002,15 +4312,9 @@ impl NatTraversalEndpoint { /// Returns false if a relay was established but the underlying QUIC /// connection has closed. pub async fn is_relay_healthy(&self) -> bool { - let lifecycle = self.relay_lifecycle.lock().await; - let relay = match &*lifecycle { - RelayLifecycleState::Published(relay) => relay, - RelayLifecycleState::None | RelayLifecycleState::Provisional(_) => { - return true; - } + let Some(relay) = self.relay_lifecycle.health().await else { + return true; }; - let relay_addr = relay.handle.public_addr(); - if relay.tunnel.is_closed() { return false; } @@ -4019,35 +4323,37 @@ impl NatTraversalEndpoint { // Other relay sessions may exist but are irrelevant — peers are // using relay_addr, so that's the one that must be healthy. for entry in self.relay_sessions.iter() { - if entry.value().public_address == Some(relay_addr) { + if entry.value().public_address == Some(relay.relay_addr) { return entry.value().is_active(); } } warn!( "Relay session for {} is dead — re-establishment required", - relay_addr + relay.relay_addr ); false } pub(crate) async fn published_relay_handle(&self) -> Option { - let lifecycle = self.relay_lifecycle.lock().await; - match &*lifecycle { - RelayLifecycleState::Published(relay) => Some(relay.handle), - RelayLifecycleState::None | RelayLifecycleState::Provisional(_) => None, - } + self.relay_lifecycle.published_handle().await + } + + /// Return the relay-signed receipt for a live allocation. + /// + /// Matching the opaque handle prevents a stale acquisition from obtaining + /// a receipt for a newer allocation that reused the same socket address. + pub async fn proactive_relay_receipt( + &self, + prepared: PreparedRelay, + ) -> Option { + self.relay_lifecycle.receipt(prepared).await } pub(crate) async fn abort_current_proactive_relay(&self) { - let mut lifecycle = self.relay_lifecycle.lock().await; - let state = std::mem::take(&mut *lifecycle); - match state { - RelayLifecycleState::Provisional(relay) | RelayLifecycleState::Published(relay) => { - self.teardown_proactive_relay(relay, b"relay allocation aborted") - .await; - } - RelayLifecycleState::None => {} + if let Some(relay) = self.relay_lifecycle.take_current().await { + self.teardown_proactive_relay(relay, b"relay allocation aborted") + .await; } } @@ -4072,8 +4378,14 @@ impl NatTraversalEndpoint { pub async fn establish_relay_session( &self, relay_addr: SocketAddr, - ) -> Result<(Option, Option), NatTraversalError> - { + ) -> Result< + ( + Option, + Option, + Option, + ), + NatTraversalError, + > { // Check if we already have an active session to this relay // DashMap provides lock-free .get() that returns Option> if let Some(session) = self.relay_sessions.get(&relay_addr) { @@ -4083,7 +4395,11 @@ impl NatTraversalEndpoint { public_address = ?session.public_address, "relay session: reusing active CONNECT-UDP session" ); - return Ok((session.public_address, None)); + return Ok(( + session.public_address, + None, + session.allocation_receipt.clone(), + )); } } @@ -4191,7 +4507,7 @@ impl NatTraversalEndpoint { if !response.is_success() { let reason = response.reason.unwrap_or_else(|| "unknown".to_string()); // Distinguish "at capacity" (client should walk to next candidate) from - // other failure modes. See ADR-014 in saorsa-core for the 2-client cap + // other failure modes. See ADR-016 in saorsa-core for the four-client cap // design. if response.status == MASQUE_RELAY_FULL_STATUS { return Err(NatTraversalError::RelayAtCapacity { reason }); @@ -4203,6 +4519,7 @@ impl NatTraversalEndpoint { } let public_address = response.proxy_public_address; + let allocation_receipt = response.allocation_receipt.clone(); info!( "Relay session established with public address: {:?}", @@ -4222,6 +4539,7 @@ impl NatTraversalEndpoint { public_address, established_at: std::time::Instant::now(), relay_addr, + allocation_receipt: allocation_receipt.clone(), }; // DashMap provides lock-free .insert() @@ -4236,7 +4554,7 @@ impl NatTraversalEndpoint { } } - Ok((public_address, raw_streams)) + Ok((public_address, raw_streams, allocation_receipt)) } /// Create a fresh QUIC connection to a relay server. @@ -4274,6 +4592,44 @@ impl NatTraversalEndpoint { Ok(connection) } + /// Open one fresh authenticated QUIC connection without registering it in + /// any peer, address, or connection-deduplication map. + /// + /// The returned identity is extracted before this method closes exactly + /// the connection it created. This is intentionally separate from normal + /// application dialing so a reachability probe can never borrow or tear + /// down a live DHT/application connection. + pub async fn probe_fresh_authenticated( + &self, + target: SocketAddr, + ) -> Result, NatTraversalError> { + let endpoint = self.inner_endpoint.as_ref().ok_or_else(|| { + NatTraversalError::ConfigError("QUIC endpoint not initialized".to_string()) + })?; + let server_name = target.ip().to_string(); + let connecting = endpoint.connect(target, &server_name).map_err(|error| { + NatTraversalError::ConnectionFailed(format!( + "Failed to initiate fresh authenticated probe: {error}" + )) + })?; + let connection = timeout(self.config.coordination_timeout, connecting) + .await + .map_err(|_| NatTraversalError::Timeout)? + .map_err(|error| { + NatTraversalError::ConnectionFailed(format!( + "Fresh authenticated probe failed: {error}" + )) + })?; + + let public_key = Self::extract_public_key_from_connection(&connection); + connection.close(crate::VarInt::from_u32(0), b"reachability probe complete"); + public_key.ok_or_else(|| { + NatTraversalError::ConnectionFailed( + "Fresh authenticated probe did not expose a peer identity".to_string(), + ) + }) + } + /// Get active relay sessions pub fn relay_sessions(&self) -> Arc> { self.relay_sessions.clone() @@ -5107,19 +5463,18 @@ impl NatTraversalEndpoint { &self, bootstrap_addr: SocketAddr, ) -> Result { - let mut lifecycle = self.relay_lifecycle.lock().await; - - // A new acquisition supersedes any prior allocation. Teardown happens - // before the new CONNECT-UDP request so rejected canaries cannot - // accumulate relay-server capacity. - let previous = std::mem::take(&mut *lifecycle); - match previous { - RelayLifecycleState::Provisional(previous) - | RelayLifecycleState::Published(previous) => { - self.teardown_proactive_relay(previous, b"relay superseded") - .await; - } - RelayLifecycleState::None => {} + let (generation, previous) = self + .relay_lifecycle + .begin_prepare() + .await + .map_err(NatTraversalError::ConnectionFailed)?; + + // A new acquisition supersedes any prior allocation. The actor releases + // ownership before teardown, so lifecycle queries and shutdown are not + // blocked behind network I/O. + if let Some(previous) = previous { + self.teardown_proactive_relay(previous, b"relay superseded") + .await; } info!( @@ -5164,7 +5519,8 @@ impl NatTraversalEndpoint { })?; // Acquire the relay only after local validation is complete. - let (public_addr, raw_streams) = self.establish_relay_session(bootstrap_addr).await?; + let (public_addr, raw_streams, allocation_receipt) = + self.establish_relay_session(bootstrap_addr).await?; let Some(relay_public_addr) = public_addr else { self.remove_matching_relay_session(bootstrap_addr, None, b"invalid relay allocation"); return Err(NatTraversalError::ConnectionFailed( @@ -5266,8 +5622,23 @@ impl NatTraversalEndpoint { handle, endpoint: relay_endpoint, tunnel, + allocation_receipt, }; - *lifecycle = RelayLifecycleState::Provisional(state); + match self + .relay_lifecycle + .complete_prepare(generation, state) + .await + .map_err(NatTraversalError::ConnectionFailed)? + { + Ok(()) => {} + Err(stale) => { + self.teardown_proactive_relay(stale, b"relay acquisition superseded") + .await; + return Err(NatTraversalError::ConnectionFailed( + "Relay acquisition was superseded before completion".to_string(), + )); + } + } Ok(handle) } @@ -5280,35 +5651,10 @@ impl NatTraversalEndpoint { &self, prepared: PreparedRelay, ) -> Result<(), NatTraversalError> { - let mut lifecycle = self.relay_lifecycle.lock().await; - let current = std::mem::take(&mut *lifecycle); - let relay = match current { - RelayLifecycleState::Provisional(relay) if relay.handle == prepared => relay, - RelayLifecycleState::Published(relay) if relay.handle == prepared => { - *lifecycle = RelayLifecycleState::Published(relay); - return Ok(()); - } - other => { - *lifecycle = other; - return Err(NatTraversalError::ConnectionFailed(format!( - "Prepared relay allocation {} at {} is stale or no longer active", - prepared.allocation_id, - prepared.public_addr() - ))); - } - }; - if relay.tunnel.is_closed() { - *lifecycle = RelayLifecycleState::Provisional(relay); - return Err(NatTraversalError::ConnectionFailed(format!( - "Prepared relay tunnel for {} is closed", - prepared.public_addr() - ))); - } - - // No fallible work follows this assignment. Publication is therefore a - // single transactional transition: retries can never observe - // `Published` without the address also being visible through the state. - *lifecycle = RelayLifecycleState::Published(relay); + self.relay_lifecycle + .publish(prepared) + .await + .map_err(NatTraversalError::ConnectionFailed)?; // Relay addresses are intentionally not broadcast through the // connection-level ADD_ADDRESS extension. A relay allocation belongs @@ -5333,21 +5679,14 @@ impl NatTraversalEndpoint { &self, prepared: PreparedRelay, ) -> Result<(), NatTraversalError> { - let mut lifecycle = self.relay_lifecycle.lock().await; - let current = std::mem::take(&mut *lifecycle); - match current { - RelayLifecycleState::Provisional(relay) | RelayLifecycleState::Published(relay) - if relay.handle == prepared => - { - self.teardown_proactive_relay(relay, b"relay allocation aborted") - .await; - } - other => { - // A stale abort is intentionally idempotent, but unlike the - // former address-only check it cannot tear down a newer relay - // that reclaimed the same public socket. - *lifecycle = other; - } + if let Some(relay) = self + .relay_lifecycle + .take_matching(prepared) + .await + .map_err(NatTraversalError::ConnectionFailed)? + { + self.teardown_proactive_relay(relay, b"relay allocation aborted") + .await; } Ok(()) } @@ -7792,6 +8131,10 @@ impl NatTraversalEndpoint { addr: SocketAddr, candidate: &CandidateAddress, ) -> Result<(), NatTraversalError> { + if self.relay_lifecycle.is_published() { + debug!("Suppressing ADD_ADDRESS candidate advertisement while a relay is published"); + return Ok(()); + } debug!( "Sending candidate advertisement to {}: {}", addr, candidate.address diff --git a/src/p2p_endpoint.rs b/src/p2p_endpoint.rs index bdbc8089..1a9a17b0 100644 --- a/src/p2p_endpoint.rs +++ b/src/p2p_endpoint.rs @@ -2731,7 +2731,7 @@ impl P2pEndpoint { ); // Step 1: Establish relay session (control plane handshake) - let (public_addr, raw_streams) = self + let (public_addr, raw_streams, _allocation_receipt) = self .inner .establish_relay_session(relay_addr) .await @@ -3490,6 +3490,28 @@ impl P2pEndpoint { Ok(()) } + /// Return the relay-signed receipt for a live proactive allocation. + pub async fn proactive_relay_receipt( + &self, + prepared: PreparedRelay, + ) -> Option { + self.inner.proactive_relay_receipt(prepared).await + } + + /// Perform a fresh, isolated authenticated reachability probe. + /// + /// This connection never enters the ordinary peer or address maps and is + /// closed by the same call that created it. + pub async fn probe_fresh_authenticated( + &self, + target: SocketAddr, + ) -> Result, EndpointError> { + self.inner + .probe_fresh_authenticated(target) + .await + .map_err(EndpointError::NatTraversal) + } + /// Get list of connected peers pub async fn connected_peers(&self) -> Vec { self.connected_peers diff --git a/src/relay_allocation.rs b/src/relay_allocation.rs new file mode 100644 index 00000000..65c1f32b --- /dev/null +++ b/src/relay_allocation.rs @@ -0,0 +1,393 @@ +// Copyright 2024 Saorsa Labs Ltd. +// +// This Saorsa Network Software is licensed under the General Public License (GPL), version 3. +// Please see the file LICENSE-GPL, or visit for the full text. +// +// Full details available at https://saorsalabs.com/licenses + +//! Cryptographic proof that a relay issued a specific allocation. + +use serde::{Deserialize, Serialize}; +use std::net::{IpAddr, SocketAddr}; +use std::time::{SystemTime, UNIX_EPOCH}; +use thiserror::Error; + +use crate::crypto::pqc::MlDsaOperations; +use crate::crypto::pqc::ml_dsa::MlDsa65; +use crate::crypto::pqc::types::{MlDsaPublicKey, MlDsaSecretKey, MlDsaSignature}; +use crate::crypto::raw_public_keys::pqc::fingerprint_public_key; + +const RECEIPT_VERSION: u8 = 1; +const RECEIPT_DOMAIN: &[u8] = b"SAORSA_RELAY_ALLOCATION_V1"; +const RECEIPT_LIFETIME_SECS: u64 = 24 * 60 * 60; + +/// A relay-signed binding between an authenticated client and one allocation. +/// +/// Witnesses must validate this receipt before attempting a canary dial. This +/// prevents a requester from turning the canary service into an arbitrary +/// reflected dial primitive. +#[derive(Clone, Serialize, Deserialize)] +pub struct RelayAllocationReceipt { + version: u8, + target_peer_id: [u8; 32], + relayer_peer_id: [u8; 32], + relay_addr: SocketAddr, + allocation_id: u64, + expires_at_unix_secs: u64, + relayer_public_key: Vec, + signature: Vec, +} + +impl std::fmt::Debug for RelayAllocationReceipt { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("RelayAllocationReceipt") + .field("version", &self.version) + .field("target_peer_id", &hex::encode(self.target_peer_id)) + .field("relayer_peer_id", &hex::encode(self.relayer_peer_id)) + .field("relay_addr", &self.relay_addr) + .field("allocation_id", &self.allocation_id) + .field("expires_at_unix_secs", &self.expires_at_unix_secs) + .finish_non_exhaustive() + } +} + +impl PartialEq for RelayAllocationReceipt { + fn eq(&self, other: &Self) -> bool { + self.version == other.version + && self.target_peer_id == other.target_peer_id + && self.relayer_peer_id == other.relayer_peer_id + && self.relay_addr == other.relay_addr + && self.allocation_id == other.allocation_id + && self.expires_at_unix_secs == other.expires_at_unix_secs + && self.relayer_public_key == other.relayer_public_key + && self.signature == other.signature + } +} + +impl Eq for RelayAllocationReceipt {} + +/// Why a relay-allocation receipt could not be issued or verified. +#[derive(Debug, Error, Clone, PartialEq, Eq)] +pub enum RelayAllocationReceiptError { + /// The receipt uses an unsupported wire version. + #[error("unsupported relay allocation receipt version {0}")] + UnsupportedVersion(u8), + /// The receipt is not bound to the requesting peer. + #[error("relay allocation receipt target does not match requester")] + TargetMismatch, + /// The receipt is not signed by the claimed relayer. + #[error("relay allocation receipt relayer does not match claim")] + RelayerMismatch, + /// The receipt is for a different relay allocation address. + #[error("relay allocation receipt address does not match request")] + AddressMismatch, + /// The receipt has expired. + #[error("relay allocation receipt expired")] + Expired, + /// The local clock could not be represented as Unix time. + #[error("system clock is before the Unix epoch")] + InvalidClock, + /// The embedded ML-DSA material is malformed. + #[error("relay allocation receipt contains invalid ML-DSA material")] + InvalidCryptoMaterial, + /// The receipt wire representation is malformed. + #[error("relay allocation receipt encoding is invalid")] + InvalidEncoding, + /// The ML-DSA signature is invalid. + #[error("relay allocation receipt signature is invalid")] + InvalidSignature, + /// Signing failed. + #[error("failed to sign relay allocation receipt")] + SigningFailed, +} + +impl RelayAllocationReceipt { + /// Issue a receipt for an allocation made to an authenticated client. + pub fn issue( + relayer_public_key: &MlDsaPublicKey, + relayer_secret_key: &MlDsaSecretKey, + target_peer_id: [u8; 32], + relay_addr: SocketAddr, + allocation_id: u64, + ) -> Result { + let now = unix_time_secs(SystemTime::now())?; + let relayer_peer_id = fingerprint_public_key(relayer_public_key); + let mut receipt = Self { + version: RECEIPT_VERSION, + target_peer_id, + relayer_peer_id, + relay_addr, + allocation_id, + expires_at_unix_secs: now.saturating_add(RECEIPT_LIFETIME_SECS), + relayer_public_key: relayer_public_key.as_bytes().to_vec(), + signature: Vec::new(), + }; + let signature = MlDsa65::new() + .sign(relayer_secret_key, &receipt.signing_message()) + .map_err(|_| RelayAllocationReceiptError::SigningFailed)?; + receipt.signature = signature.as_bytes().to_vec(); + Ok(receipt) + } + + /// Validate the receipt and all bindings supplied by a canary requester. + pub fn verify( + &self, + target_peer_id: [u8; 32], + relayer_peer_id: [u8; 32], + relay_addr: SocketAddr, + now: SystemTime, + ) -> Result<(), RelayAllocationReceiptError> { + if self.version != RECEIPT_VERSION { + return Err(RelayAllocationReceiptError::UnsupportedVersion( + self.version, + )); + } + if self.target_peer_id != target_peer_id { + return Err(RelayAllocationReceiptError::TargetMismatch); + } + if self.relayer_peer_id != relayer_peer_id { + return Err(RelayAllocationReceiptError::RelayerMismatch); + } + if self.relay_addr != relay_addr { + return Err(RelayAllocationReceiptError::AddressMismatch); + } + if unix_time_secs(now)? >= self.expires_at_unix_secs { + return Err(RelayAllocationReceiptError::Expired); + } + + let public_key = MlDsaPublicKey::from_bytes(&self.relayer_public_key) + .map_err(|_| RelayAllocationReceiptError::InvalidCryptoMaterial)?; + if fingerprint_public_key(&public_key) != self.relayer_peer_id { + return Err(RelayAllocationReceiptError::RelayerMismatch); + } + let signature = MlDsaSignature::from_bytes(&self.signature) + .map_err(|_| RelayAllocationReceiptError::InvalidCryptoMaterial)?; + match MlDsa65::new().verify(&public_key, &self.signing_message(), &signature) { + Ok(true) => Ok(()), + Ok(false) => Err(RelayAllocationReceiptError::InvalidSignature), + Err(_) => Err(RelayAllocationReceiptError::InvalidCryptoMaterial), + } + } + + /// Authenticated client fingerprint bound into this receipt. + pub fn target_peer_id(&self) -> [u8; 32] { + self.target_peer_id + } + + /// Authenticated relay fingerprint bound into this receipt. + pub fn relayer_peer_id(&self) -> [u8; 32] { + self.relayer_peer_id + } + + /// Public allocation address bound into this receipt. + pub fn relay_addr(&self) -> SocketAddr { + self.relay_addr + } + + pub(crate) fn encode(&self) -> Vec { + let mut encoded = Vec::with_capacity( + 1 + 32 + + 32 + + 1 + + 16 + + 2 + + 8 + + 8 + + 2 + + self.relayer_public_key.len() + + 2 + + self.signature.len(), + ); + encoded.push(self.version); + encoded.extend_from_slice(&self.target_peer_id); + encoded.extend_from_slice(&self.relayer_peer_id); + match self.relay_addr.ip() { + IpAddr::V4(ip) => { + encoded.push(4); + encoded.extend_from_slice(&ip.octets()); + } + IpAddr::V6(ip) => { + encoded.push(6); + encoded.extend_from_slice(&ip.octets()); + } + } + encoded.extend_from_slice(&self.relay_addr.port().to_be_bytes()); + encoded.extend_from_slice(&self.allocation_id.to_be_bytes()); + encoded.extend_from_slice(&self.expires_at_unix_secs.to_be_bytes()); + encoded.extend_from_slice(&(self.relayer_public_key.len() as u16).to_be_bytes()); + encoded.extend_from_slice(&self.relayer_public_key); + encoded.extend_from_slice(&(self.signature.len() as u16).to_be_bytes()); + encoded.extend_from_slice(&self.signature); + encoded + } + + pub(crate) fn decode(mut encoded: &[u8]) -> Result { + fn take<'a>( + encoded: &mut &'a [u8], + length: usize, + ) -> Result<&'a [u8], RelayAllocationReceiptError> { + if encoded.len() < length { + return Err(RelayAllocationReceiptError::InvalidEncoding); + } + let (value, remainder) = encoded.split_at(length); + *encoded = remainder; + Ok(value) + } + + let version = take(&mut encoded, 1)?[0]; + let target_peer_id = take(&mut encoded, 32)? + .try_into() + .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?; + let relayer_peer_id = take(&mut encoded, 32)? + .try_into() + .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?; + let ip = match take(&mut encoded, 1)?[0] { + 4 => { + let octets: [u8; 4] = take(&mut encoded, 4)? + .try_into() + .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?; + IpAddr::V4(std::net::Ipv4Addr::from(octets)) + } + 6 => { + let octets: [u8; 16] = take(&mut encoded, 16)? + .try_into() + .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?; + IpAddr::V6(std::net::Ipv6Addr::from(octets)) + } + _ => return Err(RelayAllocationReceiptError::InvalidEncoding), + }; + let port = u16::from_be_bytes( + take(&mut encoded, 2)? + .try_into() + .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?, + ); + let allocation_id = u64::from_be_bytes( + take(&mut encoded, 8)? + .try_into() + .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?, + ); + let expires_at_unix_secs = u64::from_be_bytes( + take(&mut encoded, 8)? + .try_into() + .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?, + ); + let public_key_len = u16::from_be_bytes( + take(&mut encoded, 2)? + .try_into() + .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?, + ) as usize; + let relayer_public_key = take(&mut encoded, public_key_len)?.to_vec(); + let signature_len = u16::from_be_bytes( + take(&mut encoded, 2)? + .try_into() + .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?, + ) as usize; + let signature = take(&mut encoded, signature_len)?.to_vec(); + if !encoded.is_empty() { + return Err(RelayAllocationReceiptError::InvalidEncoding); + } + + Ok(Self { + version, + target_peer_id, + relayer_peer_id, + relay_addr: SocketAddr::new(ip, port), + allocation_id, + expires_at_unix_secs, + relayer_public_key, + signature, + }) + } + + fn signing_message(&self) -> Vec { + let mut message = Vec::with_capacity(RECEIPT_DOMAIN.len() + 100); + message.extend_from_slice(RECEIPT_DOMAIN); + message.push(self.version); + message.extend_from_slice(&self.target_peer_id); + message.extend_from_slice(&self.relayer_peer_id); + match self.relay_addr.ip() { + IpAddr::V4(ip) => { + message.push(4); + message.extend_from_slice(&ip.octets()); + } + IpAddr::V6(ip) => { + message.push(6); + message.extend_from_slice(&ip.octets()); + } + } + message.extend_from_slice(&self.relay_addr.port().to_be_bytes()); + message.extend_from_slice(&self.allocation_id.to_be_bytes()); + message.extend_from_slice(&self.expires_at_unix_secs.to_be_bytes()); + message + } +} + +fn unix_time_secs(time: SystemTime) -> Result { + time.duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs()) + .map_err(|_| RelayAllocationReceiptError::InvalidClock) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::crypto::raw_public_keys::pqc::generate_ml_dsa_keypair; + use std::net::{Ipv4Addr, SocketAddrV4}; + + #[test] + fn receipt_verifies_only_for_exact_binding() { + let (public_key, secret_key) = generate_ml_dsa_keypair().expect("keypair"); + let relayer = fingerprint_public_key(&public_key); + let target = [7; 32]; + let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 12_345)); + let receipt = RelayAllocationReceipt::issue(&public_key, &secret_key, target, addr, 42) + .expect("receipt"); + + assert!( + receipt + .verify(target, relayer, addr, SystemTime::now()) + .is_ok() + ); + assert_eq!( + receipt.verify([8; 32], relayer, addr, SystemTime::now()), + Err(RelayAllocationReceiptError::TargetMismatch) + ); + assert_eq!( + receipt.verify(target, [9; 32], addr, SystemTime::now()), + Err(RelayAllocationReceiptError::RelayerMismatch) + ); + } + + #[test] + fn tampering_invalidates_signature() { + let (public_key, secret_key) = generate_ml_dsa_keypair().expect("keypair"); + let relayer = fingerprint_public_key(&public_key); + let target = [7; 32]; + let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 12_345)); + let mut receipt = RelayAllocationReceipt::issue(&public_key, &secret_key, target, addr, 42) + .expect("receipt"); + receipt.allocation_id += 1; + + assert_eq!( + receipt.verify(target, relayer, addr, SystemTime::now()), + Err(RelayAllocationReceiptError::InvalidSignature) + ); + } + + #[test] + fn receipt_expires_at_its_signed_deadline() { + let (public_key, secret_key) = generate_ml_dsa_keypair().expect("keypair"); + let relayer = fingerprint_public_key(&public_key); + let target = [7; 32]; + let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 12_345)); + let receipt = RelayAllocationReceipt::issue(&public_key, &secret_key, target, addr, 42) + .expect("receipt"); + let deadline = UNIX_EPOCH + std::time::Duration::from_secs(receipt.expires_at_unix_secs); + + assert_eq!( + receipt.verify(target, relayer, addr, deadline), + Err(RelayAllocationReceiptError::Expired) + ); + } +} From 81b809d0c9d6cb5d6eb25c7d024843585e01b68b Mon Sep 17 00:00:00 2001 From: Mick van Dijke <12992260+mickvandijke@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:29:49 +0200 Subject: [PATCH 4/8] fix(relay): make ownership cancellation-safe --- src/masque/relay_socket.rs | 43 +++++++++++++- src/nat_traversal_api.rs | 112 +++++++++++++++++++++++++++---------- 2 files changed, 124 insertions(+), 31 deletions(-) diff --git a/src/masque/relay_socket.rs b/src/masque/relay_socket.rs index 243ce6e1..22ff7366 100644 --- a/src/masque/relay_socket.rs +++ b/src/masque/relay_socket.rs @@ -164,6 +164,22 @@ impl RelayTunnelControl { /// Stop the tunnel and wait for all task-owned QUIC streams to be dropped. pub(crate) async fn shutdown(&self) { + let handles = self.abort_tasks(); + for handle in handles { + let _ = handle.await; + } + } + + /// Stop the tunnel without waiting for task cancellation to complete. + /// + /// This is the cancellation-safe fallback used by relay ownership guards: + /// `Drop` cannot await, but it must still ensure the task-owned QUIC stream + /// halves are scheduled for prompt release. + pub(crate) fn shutdown_now(&self) { + drop(self.abort_tasks()); + } + + fn abort_tasks(&self) -> Vec> { self.mark_closed(); let handles = { let mut tasks = self.tasks.lock(); @@ -172,9 +188,7 @@ impl RelayTunnelControl { for handle in &handles { handle.abort(); } - for handle in handles { - let _ = handle.await; - } + handles } } @@ -761,6 +775,29 @@ mod relay_tunnel_control_tests { assert!(control.is_closed()); } + #[tokio::test] + async fn shutdown_now_aborts_registered_tasks_without_an_await() { + let control = RelayTunnelControl::new(); + let dropped = Arc::new(AtomicBool::new(false)); + let marker = DropMarker(Arc::clone(&dropped)); + control.register(tokio::spawn(async move { + let _marker = marker; + std::future::pending::<()>().await; + })); + + tokio::task::yield_now().await; + control.shutdown_now(); + + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while !dropped.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("aborted tunnel task should release its owned stream state"); + assert!(control.is_closed()); + } + #[test] fn writer_exit_marks_tunnel_closed() { let control = RelayTunnelControl::new(); diff --git a/src/nat_traversal_api.rs b/src/nat_traversal_api.rs index aa54418f..6058de99 100644 --- a/src/nat_traversal_api.rs +++ b/src/nat_traversal_api.rs @@ -427,6 +427,75 @@ struct ProactiveRelay { endpoint: Arc, tunnel: Arc, allocation_receipt: Option, + relay_sessions: Arc>, + relay_session_stable_id: usize, + cleanup_armed: bool, +} + +impl ProactiveRelay { + async fn teardown(mut self, reason: &'static [u8]) { + self.endpoint + .close(crate::VarInt::from_u32(RELAY_TUNNEL_LOST_CODE), reason); + self.tunnel.shutdown().await; + + if tokio::time::timeout(Duration::from_secs(1), self.endpoint.wait_idle()) + .await + .is_err() + { + debug!( + relay_addr = %self.handle.public_addr(), + "Timed out draining proactive relay endpoint during teardown" + ); + } + + self.remove_owned_session(reason); + + info!( + relay_addr = %self.handle.public_addr(), + relay_server = %self.relay_server_addr, + allocation_id = self.handle.allocation_id, + "Proactive relay torn down" + ); + self.cleanup_armed = false; + } + + fn remove_owned_session(&self, reason: &'static [u8]) { + let matching_session = self + .relay_sessions + .get(&self.relay_server_addr) + .is_some_and(|session| { + session.public_address == Some(self.handle.public_addr()) + && session.connection.stable_id() == self.relay_session_stable_id + }); + if !matching_session { + return; + } + + if let Some((_, session)) = self.relay_sessions.remove(&self.relay_server_addr) { + session + .connection + .close(crate::VarInt::from_u32(RELAY_TUNNEL_LOST_CODE), reason); + } + } +} + +impl Drop for ProactiveRelay { + fn drop(&mut self) { + if !self.cleanup_armed { + return; + } + + warn!( + relay_addr = %self.handle.public_addr(), + allocation_id = self.handle.allocation_id, + "Relay owner dropped before graceful teardown completed; forcing cleanup" + ); + let reason = b"relay owner dropped"; + self.endpoint + .close(crate::VarInt::from_u32(RELAY_TUNNEL_LOST_CODE), reason); + self.tunnel.shutdown_now(); + self.remove_owned_session(reason); + } } /// Complete proactive-relay lifecycle state. @@ -5471,7 +5540,8 @@ impl NatTraversalEndpoint { // A new acquisition supersedes any prior allocation. The actor releases // ownership before teardown, so lifecycle queries and shutdown are not - // blocked behind network I/O. + // blocked behind network I/O. The ownership guard forces cleanup if + // this future is cancelled before graceful teardown completes. if let Some(previous) = previous { self.teardown_proactive_relay(previous, b"relay superseded") .await; @@ -5537,6 +5607,15 @@ impl NatTraversalEndpoint { "Relay did not provide socket".to_string(), )); }; + let relay_session_stable_id = self + .relay_sessions + .get(&bootstrap_addr) + .map(|session| session.connection.stable_id()) + .ok_or_else(|| { + NatTraversalError::ConnectionFailed( + "Relay allocation lost its owning control session".to_string(), + ) + })?; info!( "Relay session established, public address: {}", @@ -5623,6 +5702,9 @@ impl NatTraversalEndpoint { endpoint: relay_endpoint, tunnel, allocation_receipt, + relay_sessions: Arc::clone(&self.relay_sessions), + relay_session_stable_id, + cleanup_armed: true, }; match self .relay_lifecycle @@ -5709,33 +5791,7 @@ impl NatTraversalEndpoint { } async fn teardown_proactive_relay(&self, state: ProactiveRelay, reason: &'static [u8]) { - state - .endpoint - .close(crate::VarInt::from_u32(RELAY_TUNNEL_LOST_CODE), reason); - state.tunnel.shutdown().await; - - if tokio::time::timeout(Duration::from_secs(1), state.endpoint.wait_idle()) - .await - .is_err() - { - debug!( - relay_addr = %state.handle.public_addr(), - "Timed out draining proactive relay endpoint during teardown" - ); - } - - self.remove_matching_relay_session( - state.relay_server_addr, - Some(state.handle.public_addr()), - reason, - ); - - info!( - relay_addr = %state.handle.public_addr(), - relay_server = %state.relay_server_addr, - allocation_id = state.handle.allocation_id, - "Proactive relay torn down" - ); + state.teardown(reason).await; } fn remove_matching_relay_session( From fdb4cb25d1bf1e75b5d13c94d3e2fb3b5a45cc6d Mon Sep 17 00:00:00 2001 From: Mick van Dijke <12992260+mickvandijke@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:31:50 +0200 Subject: [PATCH 5/8] fix(relay): close remaining ownership races --- src/nat_traversal_api.rs | 293 ++++++++++++++++++++++++++++----------- 1 file changed, 214 insertions(+), 79 deletions(-) diff --git a/src/nat_traversal_api.rs b/src/nat_traversal_api.rs index 6058de99..eed42b9b 100644 --- a/src/nat_traversal_api.rs +++ b/src/nat_traversal_api.rs @@ -378,6 +378,86 @@ impl RelaySession { } } +/// Cancellation-safe ownership of one exact entry in `relay_sessions`. +/// +/// The guard is armed before the entry is inserted, so every subsequent await +/// is protected. Ownership can be transferred into [`ProactiveRelay`] or +/// explicitly disarmed when the session is intentionally left managed by the +/// endpoint-wide session table. +struct RelaySessionOwner { + relay_server_addr: SocketAddr, + public_address: Option, + relay_sessions: Arc>, + stable_id: usize, + cleanup_armed: bool, +} + +impl RelaySessionOwner { + fn new( + relay_server_addr: SocketAddr, + public_address: Option, + relay_sessions: Arc>, + stable_id: usize, + ) -> Self { + Self { + relay_server_addr, + public_address, + relay_sessions, + stable_id, + cleanup_armed: true, + } + } + + fn disarm(mut self) { + self.cleanup_armed = false; + } + + fn remove_owned_session(&mut self, reason: &'static [u8]) { + if !self.cleanup_armed { + return; + } + + let public_address = self.public_address; + let stable_id = self.stable_id; + let removed = self + .relay_sessions + .remove_if(&self.relay_server_addr, |_, session| { + session.public_address == public_address + && session.connection.stable_id() == stable_id + }); + + if let Some((_, session)) = removed { + session + .connection + .close(crate::VarInt::from_u32(RELAY_TUNNEL_LOST_CODE), reason); + } + self.cleanup_armed = false; + } +} + +impl Drop for RelaySessionOwner { + fn drop(&mut self) { + if !self.cleanup_armed { + return; + } + + warn!( + relay_server = %self.relay_server_addr, + relay_addr = ?self.public_address, + stable_id = self.stable_id, + "Relay session owner dropped before ownership transfer; forcing cleanup" + ); + self.remove_owned_session(b"relay session owner dropped"); + } +} + +struct EstablishedRelaySession { + public_address: Option, + raw_streams: Option, + allocation_receipt: Option, + owner: RelaySessionOwner, +} + /// Opaque identity of a prepared proactive relay allocation. /// /// Callers must pass this exact handle to @@ -422,13 +502,11 @@ impl PreparedRelay { /// Resources owned by one proactive relay allocation. struct ProactiveRelay { - relay_server_addr: SocketAddr, handle: PreparedRelay, endpoint: Arc, tunnel: Arc, allocation_receipt: Option, - relay_sessions: Arc>, - relay_session_stable_id: usize, + relay_session_owner: RelaySessionOwner, cleanup_armed: bool, } @@ -448,35 +526,16 @@ impl ProactiveRelay { ); } - self.remove_owned_session(reason); + self.relay_session_owner.remove_owned_session(reason); info!( relay_addr = %self.handle.public_addr(), - relay_server = %self.relay_server_addr, + relay_server = %self.relay_session_owner.relay_server_addr, allocation_id = self.handle.allocation_id, "Proactive relay torn down" ); self.cleanup_armed = false; } - - fn remove_owned_session(&self, reason: &'static [u8]) { - let matching_session = self - .relay_sessions - .get(&self.relay_server_addr) - .is_some_and(|session| { - session.public_address == Some(self.handle.public_addr()) - && session.connection.stable_id() == self.relay_session_stable_id - }); - if !matching_session { - return; - } - - if let Some((_, session)) = self.relay_sessions.remove(&self.relay_server_addr) { - session - .connection - .close(crate::VarInt::from_u32(RELAY_TUNNEL_LOST_CODE), reason); - } - } } impl Drop for ProactiveRelay { @@ -494,7 +553,7 @@ impl Drop for ProactiveRelay { self.endpoint .close(crate::VarInt::from_u32(RELAY_TUNNEL_LOST_CODE), reason); self.tunnel.shutdown_now(); - self.remove_owned_session(reason); + self.relay_session_owner.remove_owned_session(reason); } } @@ -4455,20 +4514,44 @@ impl NatTraversalEndpoint { ), NatTraversalError, > { + let EstablishedRelaySession { + public_address, + raw_streams, + allocation_receipt, + owner, + } = self.establish_owned_relay_session(relay_addr).await?; + owner.disarm(); + Ok((public_address, raw_streams, allocation_receipt)) + } + + async fn establish_owned_relay_session( + &self, + relay_addr: SocketAddr, + ) -> Result { // Check if we already have an active session to this relay // DashMap provides lock-free .get() that returns Option> if let Some(session) = self.relay_sessions.get(&relay_addr) { if session.is_active() { + let public_address = session.public_address; + let allocation_receipt = session.allocation_receipt.clone(); + let stable_id = session.connection.stable_id(); debug!( relay = %relay_addr, - public_address = ?session.public_address, + public_address = ?public_address, "relay session: reusing active CONNECT-UDP session" ); - return Ok(( - session.public_address, - None, - session.allocation_receipt.clone(), - )); + drop(session); + return Ok(EstablishedRelaySession { + public_address, + raw_streams: None, + allocation_receipt, + owner: RelaySessionOwner::new( + relay_addr, + public_address, + Arc::clone(&self.relay_sessions), + stable_id, + ), + }); } } @@ -4602,7 +4685,16 @@ impl NatTraversalEndpoint { recv_stream, }); - // Store the session + // Arm ownership before publishing the session into the map. There are + // no suspension points between these operations, and every later await + // is protected by the guard. + let stable_id = connection.stable_id(); + let owner = RelaySessionOwner::new( + relay_addr, + public_address, + Arc::clone(&self.relay_sessions), + stable_id, + ); let session = RelaySession { connection, public_address, @@ -4623,7 +4715,12 @@ impl NatTraversalEndpoint { } } - Ok((public_address, raw_streams, allocation_receipt)) + Ok(EstablishedRelaySession { + public_address, + raw_streams, + allocation_receipt, + owner, + }) } /// Create a fresh QUIC connection to a relay server. @@ -5589,33 +5686,22 @@ impl NatTraversalEndpoint { })?; // Acquire the relay only after local validation is complete. - let (public_addr, raw_streams, allocation_receipt) = - self.establish_relay_session(bootstrap_addr).await?; + let EstablishedRelaySession { + public_address: public_addr, + raw_streams, + allocation_receipt, + owner: relay_session_owner, + } = self.establish_owned_relay_session(bootstrap_addr).await?; let Some(relay_public_addr) = public_addr else { - self.remove_matching_relay_session(bootstrap_addr, None, b"invalid relay allocation"); return Err(NatTraversalError::ConnectionFailed( "Relay did not provide public address".to_string(), )); }; let Some(raw_streams) = raw_streams else { - self.remove_matching_relay_session( - bootstrap_addr, - Some(relay_public_addr), - b"relay allocation missing streams", - ); return Err(NatTraversalError::ConnectionFailed( "Relay did not provide socket".to_string(), )); }; - let relay_session_stable_id = self - .relay_sessions - .get(&bootstrap_addr) - .map(|session| session.connection.stable_id()) - .ok_or_else(|| { - NatTraversalError::ConnectionFailed( - "Relay allocation lost its owning control session".to_string(), - ) - })?; info!( "Relay session established, public address: {}", @@ -5641,11 +5727,6 @@ impl NatTraversalEndpoint { Ok(endpoint) => endpoint, Err(error) => { tunnel.shutdown().await; - self.remove_matching_relay_session( - bootstrap_addr, - Some(relay_public_addr), - b"relay endpoint creation failed", - ); return Err(NatTraversalError::ConnectionFailed(format!( "Failed to create relay endpoint: {error}" ))); @@ -5697,13 +5778,11 @@ impl NatTraversalEndpoint { self.spawn_relay_endpoint_accept_loop(Arc::clone(&relay_endpoint), relay_public_addr); let state = ProactiveRelay { - relay_server_addr: bootstrap_addr, handle, endpoint: relay_endpoint, tunnel, allocation_receipt, - relay_sessions: Arc::clone(&self.relay_sessions), - relay_session_stable_id, + relay_session_owner, cleanup_armed: true, }; match self @@ -5794,27 +5873,6 @@ impl NatTraversalEndpoint { state.teardown(reason).await; } - fn remove_matching_relay_session( - &self, - relay_server_addr: SocketAddr, - public_addr: Option, - reason: &'static [u8], - ) { - let matching_session = self - .relay_sessions - .get(&relay_server_addr) - .is_some_and(|session| session.public_address == public_addr); - if !matching_session { - return; - } - - if let Some((_, session)) = self.relay_sessions.remove(&relay_server_addr) { - session - .connection - .close(crate::VarInt::from_u32(RELAY_TUNNEL_LOST_CODE), reason); - } - } - /// Spawn an accept loop for the relay endpoint. /// /// Accepted connections are inserted into the shared `connections` @@ -8408,6 +8466,83 @@ mod tests { assert_eq!(first.public_addr(), replacement.public_addr()); } + #[tokio::test] + async fn relay_session_owner_only_removes_its_exact_session() { + let endpoint_config = || NatTraversalConfig { + bind_addr: Some("127.0.0.1:0".parse().expect("test bind address")), + ..Default::default() + }; + let server = NatTraversalEndpoint::new(endpoint_config(), None, None) + .await + .expect("server endpoint"); + let client = NatTraversalEndpoint::new(endpoint_config(), None, None) + .await + .expect("client endpoint"); + let server_addr = server + .get_endpoint() + .expect("server transport endpoint") + .local_addr() + .expect("server address"); + let server_name = server_addr.ip().to_string(); + + let original = client + .connect_to(&server_name, server_addr) + .await + .expect("original connection"); + let replacement = client + .connect_to(&server_name, server_addr) + .await + .expect("replacement connection"); + assert_ne!(original.stable_id(), replacement.stable_id()); + + let public_addr = Some("203.0.113.7:9000".parse().expect("public address")); + let sessions = Arc::new(dashmap::DashMap::new()); + sessions.insert( + server_addr, + RelaySession { + connection: original.clone(), + public_address: public_addr, + established_at: std::time::Instant::now(), + relay_addr: server_addr, + allocation_receipt: None, + }, + ); + let stale_owner = RelaySessionOwner::new( + server_addr, + public_addr, + Arc::clone(&sessions), + original.stable_id(), + ); + + sessions.insert( + server_addr, + RelaySession { + connection: replacement.clone(), + public_address: public_addr, + established_at: std::time::Instant::now(), + relay_addr: server_addr, + allocation_receipt: None, + }, + ); + drop(stale_owner); + + let current = sessions.get(&server_addr).expect("replacement retained"); + assert_eq!(current.connection.stable_id(), replacement.stable_id()); + drop(current); + + let replacement_owner = RelaySessionOwner::new( + server_addr, + public_addr, + Arc::clone(&sessions), + replacement.stable_id(), + ); + drop(replacement_owner); + assert!(sessions.is_empty(), "exact owner must remove its session"); + + client.shutdown().await.expect("client shutdown"); + server.shutdown().await.expect("server shutdown"); + } + #[test] fn test_nat_traversal_config_default() { let config = NatTraversalConfig::default(); From d960473f165107f15e9a7c53dca015bf2033ce17 Mon Sep 17 00:00:00 2001 From: Mick van Dijke <12992260+mickvandijke@users.noreply.github.com> Date: Fri, 31 Jul 2026 14:02:33 +0200 Subject: [PATCH 6/8] fix(relay): make shutdown terminal --- src/masque/relay_socket.rs | 5 + src/nat_traversal_api.rs | 197 ++++++++++++++++++++++++++++++++----- 2 files changed, 176 insertions(+), 26 deletions(-) diff --git a/src/masque/relay_socket.rs b/src/masque/relay_socket.rs index 22ff7366..721c6abb 100644 --- a/src/masque/relay_socket.rs +++ b/src/masque/relay_socket.rs @@ -123,6 +123,11 @@ impl RelayTunnelControl { }) } + #[cfg(test)] + pub(crate) fn detached() -> Arc { + Self::new() + } + fn register(&self, handle: tokio::task::JoinHandle<()>) { if self.is_closed() { handle.abort(); diff --git a/src/nat_traversal_api.rs b/src/nat_traversal_api.rs index eed42b9b..87b9fd2a 100644 --- a/src/nat_traversal_api.rs +++ b/src/nat_traversal_api.rs @@ -568,13 +568,14 @@ enum RelayLifecycleState { None, Provisional(ProactiveRelay), Published(ProactiveRelay), + Shutdown, } impl RelayLifecycleState { fn published_addr(&self) -> Option { match self { Self::Published(relay) => Some(relay.handle.public_addr()), - Self::None | Self::Provisional(_) => None, + Self::None | Self::Provisional(_) | Self::Shutdown => None, } } } @@ -586,7 +587,7 @@ struct RelayHealthSnapshot { enum RelayLifecycleCommand { BeginPrepare { - reply: oneshot::Sender<(u64, Option)>, + reply: oneshot::Sender), String>>, }, CompletePrepare { generation: u64, @@ -601,7 +602,7 @@ enum RelayLifecycleCommand { prepared: PreparedRelay, reply: oneshot::Sender>, }, - TakeCurrent { + Shutdown { reply: oneshot::Sender>, }, PublishedAddr { @@ -639,14 +640,19 @@ impl RelayLifecycleHandle { while let Some(command) = receiver.recv().await { match command { RelayLifecycleCommand::BeginPrepare { reply } => { - generation = generation.wrapping_add(1); - actor_published.store(false, Ordering::Release); let previous = match std::mem::take(&mut state) { RelayLifecycleState::Provisional(relay) | RelayLifecycleState::Published(relay) => Some(relay), RelayLifecycleState::None => None, + RelayLifecycleState::Shutdown => { + state = RelayLifecycleState::Shutdown; + let _ = reply.send(Err("relay lifecycle is shut down".to_string())); + continue; + } }; - let _ = reply.send((generation, previous)); + generation = generation.wrapping_add(1); + actor_published.store(false, Ordering::Release); + let _ = reply.send(Ok((generation, previous))); } RelayLifecycleCommand::CompletePrepare { generation: completed_generation, @@ -683,6 +689,10 @@ impl RelayLifecycleHandle { actor_published.store(true, Ordering::Release); let _ = reply.send(Ok(())); } + RelayLifecycleState::Shutdown => { + state = RelayLifecycleState::Shutdown; + let _ = reply.send(Err("relay lifecycle is shut down".to_string())); + } other => { state = other; let _ = reply.send(Err(format!( @@ -710,14 +720,15 @@ impl RelayLifecycleHandle { } } } - RelayLifecycleCommand::TakeCurrent { reply } => { + RelayLifecycleCommand::Shutdown { reply } => { generation = generation.wrapping_add(1); actor_published.store(false, Ordering::Release); - let relay = match std::mem::take(&mut state) { - RelayLifecycleState::Provisional(relay) - | RelayLifecycleState::Published(relay) => Some(relay), - RelayLifecycleState::None => None, - }; + let relay = + match std::mem::replace(&mut state, RelayLifecycleState::Shutdown) { + RelayLifecycleState::Provisional(relay) + | RelayLifecycleState::Published(relay) => Some(relay), + RelayLifecycleState::None | RelayLifecycleState::Shutdown => None, + }; let _ = reply.send(relay); } RelayLifecycleCommand::PublishedAddr { reply } => { @@ -726,7 +737,9 @@ impl RelayLifecycleHandle { RelayLifecycleCommand::PublishedHandle { reply } => { let handle = match &state { RelayLifecycleState::Published(relay) => Some(relay.handle), - RelayLifecycleState::None | RelayLifecycleState::Provisional(_) => None, + RelayLifecycleState::None + | RelayLifecycleState::Provisional(_) + | RelayLifecycleState::Shutdown => None, }; let _ = reply.send(handle); } @@ -740,7 +753,8 @@ impl RelayLifecycleHandle { } RelayLifecycleState::None | RelayLifecycleState::Provisional(_) - | RelayLifecycleState::Published(_) => None, + | RelayLifecycleState::Published(_) + | RelayLifecycleState::Shutdown => None, }; let _ = reply.send(receipt); } @@ -750,7 +764,9 @@ impl RelayLifecycleHandle { relay_addr: relay.handle.public_addr(), tunnel: Arc::clone(&relay.tunnel), }), - RelayLifecycleState::None | RelayLifecycleState::Provisional(_) => None, + RelayLifecycleState::None + | RelayLifecycleState::Provisional(_) + | RelayLifecycleState::Shutdown => None, }; let _ = reply.send(health); } @@ -775,7 +791,7 @@ impl RelayLifecycleHandle { .map_err(|_| "relay lifecycle actor stopped".to_string())?; response .await - .map_err(|_| "relay lifecycle actor stopped".to_string()) + .map_err(|_| "relay lifecycle actor stopped".to_string())? } async fn complete_prepare( @@ -822,11 +838,11 @@ impl RelayLifecycleHandle { .map_err(|_| "relay lifecycle actor stopped".to_string()) } - async fn take_current(&self) -> Option { + async fn shutdown(&self) -> Option { let (reply, response) = oneshot::channel(); if self .commands - .send(RelayLifecycleCommand::TakeCurrent { reply }) + .send(RelayLifecycleCommand::Shutdown { reply }) .await .is_err() { @@ -4478,13 +4494,6 @@ impl NatTraversalEndpoint { self.relay_lifecycle.receipt(prepared).await } - pub(crate) async fn abort_current_proactive_relay(&self) { - if let Some(relay) = self.relay_lifecycle.take_current().await { - self.teardown_proactive_relay(relay, b"relay allocation aborted") - .await; - } - } - /// Check if relay fallback is available pub async fn has_relay_fallback(&self) -> bool { match &self.relay_manager { @@ -6432,7 +6441,10 @@ impl NatTraversalEndpoint { self.incoming_notify.notify_waiters(); self.shutdown_notify.notify_waiters(); - self.abort_current_proactive_relay().await; + if let Some(relay) = self.relay_lifecycle.shutdown().await { + self.teardown_proactive_relay(relay, b"endpoint shutdown") + .await; + } // Best-effort UPnP teardown. The endpoint is the sole owner of // the service (the discovery manager only holds a read-only @@ -8454,6 +8466,36 @@ impl crate::TokenStore for DefaultTokenStore { mod tests { use super::*; + async fn detached_proactive_relay(public_addr: SocketAddr) -> ProactiveRelay { + let endpoint = NatTraversalEndpoint::new( + NatTraversalConfig { + bind_addr: Some("127.0.0.1:0".parse().expect("test bind address")), + ..Default::default() + }, + None, + None, + ) + .await + .expect("test endpoint"); + let relay_endpoint = Arc::new(endpoint.get_endpoint().expect("transport endpoint").clone()); + endpoint.shutdown().await.expect("test endpoint shutdown"); + + ProactiveRelay { + handle: PreparedRelay::new(public_addr), + endpoint: relay_endpoint, + tunnel: crate::masque::RelayTunnelControl::detached(), + allocation_receipt: None, + relay_session_owner: RelaySessionOwner { + relay_server_addr: "127.0.0.1:1".parse().expect("relay server address"), + public_address: Some(public_addr), + relay_sessions: Arc::new(dashmap::DashMap::new()), + stable_id: usize::MAX, + cleanup_armed: false, + }, + cleanup_armed: false, + } + } + #[test] fn prepared_relay_identity_distinguishes_reused_public_address() { let public_addr = "203.0.113.7:9000".parse().expect("test address"); @@ -8466,6 +8508,109 @@ mod tests { assert_eq!(first.public_addr(), replacement.public_addr()); } + #[tokio::test] + async fn relay_lifecycle_shutdown_rejects_inflight_and_later_commands() { + let lifecycle = RelayLifecycleHandle::new(); + let (generation, previous) = lifecycle.begin_prepare().await.expect("begin prepare"); + assert!(previous.is_none()); + + assert!( + lifecycle.shutdown().await.is_none(), + "shutdown before completion has no installed relay to drain" + ); + + let relay = + detached_proactive_relay("203.0.113.10:10000".parse().expect("relay public address")) + .await; + let prepared = relay.handle; + let completion = lifecycle + .complete_prepare(generation, relay) + .await + .expect("lifecycle actor"); + assert!( + completion.is_err(), + "an acquisition begun before shutdown must not install afterward" + ); + drop(completion); + + let later_begin = lifecycle.begin_prepare().await; + assert!( + matches!(later_begin, Err(ref error) if error.contains("shut down")), + "prepare after shutdown must be rejected" + ); + let later_publish = lifecycle.publish(prepared).await; + assert!( + matches!(later_publish, Err(ref error) if error.contains("shut down")), + "publish after shutdown must be rejected" + ); + } + + #[tokio::test] + async fn relay_lifecycle_shutdown_drains_installed_relay_once() { + let lifecycle = RelayLifecycleHandle::new(); + let relay = + detached_proactive_relay("203.0.113.11:10001".parse().expect("relay public address")) + .await; + let prepared = relay.handle; + let (generation, previous) = lifecycle.begin_prepare().await.expect("begin prepare"); + assert!(previous.is_none()); + let completion = lifecycle + .complete_prepare(generation, relay) + .await + .expect("lifecycle actor"); + assert!(completion.is_ok(), "relay must install before shutdown"); + lifecycle + .publish(prepared) + .await + .expect("publish before shutdown"); + assert!(lifecycle.is_published()); + + let drained = lifecycle + .shutdown() + .await + .expect("shutdown must return installed relay"); + assert_eq!(drained.handle, prepared); + drop(drained); + assert!(!lifecycle.is_published()); + assert!( + lifecycle.shutdown().await.is_none(), + "terminal shutdown must drain at most once" + ); + } + + #[tokio::test] + async fn endpoint_shutdown_rejects_relay_prepare_and_publish() { + let endpoint = NatTraversalEndpoint::new( + NatTraversalConfig { + bind_addr: Some("127.0.0.1:0".parse().expect("test bind address")), + ..Default::default() + }, + None, + None, + ) + .await + .expect("test endpoint"); + endpoint.shutdown().await.expect("endpoint shutdown"); + + let prepare = endpoint + .prepare_proactive_relay("127.0.0.1:9".parse().expect("relay address")) + .await; + assert!( + matches!(prepare, Err(NatTraversalError::ConnectionFailed(ref error)) if error.contains("shut down")), + "public prepare path must reject after shutdown without dialing" + ); + + let publish = endpoint + .publish_proactive_relay(PreparedRelay::detached( + "203.0.113.12:10002".parse().expect("relay public address"), + )) + .await; + assert!( + matches!(publish, Err(NatTraversalError::ConnectionFailed(ref error)) if error.contains("shut down")), + "public publish path must reject after shutdown" + ); + } + #[tokio::test] async fn relay_session_owner_only_removes_its_exact_session() { let endpoint_config = || NatTraversalConfig { From e24989f21a0bc90cd42bf45ffbdb0bc1700048eb Mon Sep 17 00:00:00 2001 From: Mick van Dijke <12992260+mickvandijke@users.noreply.github.com> Date: Mon, 3 Aug 2026 11:07:50 +0200 Subject: [PATCH 7/8] fix(relay): align allocation receipt peer identities --- src/lib.rs | 4 +++- src/nat_traversal_api.rs | 10 ++++------ src/relay_allocation.rs | 31 +++++++++++++++++++++++++------ 3 files changed, 32 insertions(+), 13 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index c3c27aa5..d4c56343 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -259,7 +259,9 @@ pub mod metrics; /// TURN-style relay protocol for NAT traversal fallback pub mod relay; mod relay_allocation; -pub use relay_allocation::{RelayAllocationReceipt, RelayAllocationReceiptError}; +pub use relay_allocation::{ + RelayAllocationReceipt, RelayAllocationReceiptError, relay_receipt_peer_id, +}; /// MASQUE CONNECT-UDP Bind protocol for fully connectable P2P nodes pub mod masque; diff --git a/src/nat_traversal_api.rs b/src/nat_traversal_api.rs index 87b9fd2a..c08d2cb4 100644 --- a/src/nat_traversal_api.rs +++ b/src/nat_traversal_api.rs @@ -6410,9 +6410,9 @@ impl NatTraversalEndpoint { None } - /// Derive the canonical 32-byte peer fingerprint (`AUTONOMI_PEER_ID_V2`) - /// from a connection's authenticated ML-DSA-65 identity, or `None` if the - /// peer is not PQC-authenticated. + /// Derive the overlay-compatible 32-byte peer ID used by relay allocation + /// receipts from a connection's authenticated ML-DSA-65 identity, or + /// `None` if the peer is not PQC-authenticated. /// /// This keys relay-port reservations to a cryptographic identity rather /// than an ephemeral socket address. The identity comes solely from the @@ -6421,9 +6421,7 @@ impl NatTraversalEndpoint { let spki = Self::extract_public_key_from_connection(connection)?; let public_key = crate::crypto::raw_public_keys::pqc::extract_public_key_from_spki(&spki).ok()?; - Some(crate::crypto::raw_public_keys::pqc::fingerprint_public_key( - &public_key, - )) + Some(crate::relay_receipt_peer_id(&public_key)) } /// Extract the raw SPKI bytes from a connection's TLS identity. diff --git a/src/relay_allocation.rs b/src/relay_allocation.rs index 65c1f32b..8616865c 100644 --- a/src/relay_allocation.rs +++ b/src/relay_allocation.rs @@ -15,12 +15,21 @@ use thiserror::Error; use crate::crypto::pqc::MlDsaOperations; use crate::crypto::pqc::ml_dsa::MlDsa65; use crate::crypto::pqc::types::{MlDsaPublicKey, MlDsaSecretKey, MlDsaSignature}; -use crate::crypto::raw_public_keys::pqc::fingerprint_public_key; const RECEIPT_VERSION: u8 = 1; const RECEIPT_DOMAIN: &[u8] = b"SAORSA_RELAY_ALLOCATION_V1"; const RECEIPT_LIFETIME_SECS: u64 = 24 * 60 * 60; +/// Derive the overlay peer identity bound into relay allocation receipts. +/// +/// The DHT identifies peers as `BLAKE3(raw ML-DSA public key)`. Transport's +/// internal `AUTONOMI_PEER_ID_V2` fingerprint uses a separate domain and must +/// not leak into this cross-layer receipt, otherwise a valid allocation can +/// never match the requester's DHT identity during canary verification. +pub fn relay_receipt_peer_id(public_key: &MlDsaPublicKey) -> [u8; 32] { + *blake3::hash(public_key.as_bytes()).as_bytes() +} + /// A relay-signed binding between an authenticated client and one allocation. /// /// Witnesses must validate this receipt before attempting a canary dial. This @@ -112,7 +121,7 @@ impl RelayAllocationReceipt { allocation_id: u64, ) -> Result { let now = unix_time_secs(SystemTime::now())?; - let relayer_peer_id = fingerprint_public_key(relayer_public_key); + let relayer_peer_id = relay_receipt_peer_id(relayer_public_key); let mut receipt = Self { version: RECEIPT_VERSION, target_peer_id, @@ -158,7 +167,7 @@ impl RelayAllocationReceipt { let public_key = MlDsaPublicKey::from_bytes(&self.relayer_public_key) .map_err(|_| RelayAllocationReceiptError::InvalidCryptoMaterial)?; - if fingerprint_public_key(&public_key) != self.relayer_peer_id { + if relay_receipt_peer_id(&public_key) != self.relayer_peer_id { return Err(RelayAllocationReceiptError::RelayerMismatch); } let signature = MlDsaSignature::from_bytes(&self.signature) @@ -335,10 +344,20 @@ mod tests { use crate::crypto::raw_public_keys::pqc::generate_ml_dsa_keypair; use std::net::{Ipv4Addr, SocketAddrV4}; + #[test] + fn receipt_peer_id_uses_overlay_identity_derivation() { + let (public_key, _) = generate_ml_dsa_keypair().expect("keypair"); + + assert_eq!( + relay_receipt_peer_id(&public_key), + *blake3::hash(public_key.as_bytes()).as_bytes() + ); + } + #[test] fn receipt_verifies_only_for_exact_binding() { let (public_key, secret_key) = generate_ml_dsa_keypair().expect("keypair"); - let relayer = fingerprint_public_key(&public_key); + let relayer = relay_receipt_peer_id(&public_key); let target = [7; 32]; let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 12_345)); let receipt = RelayAllocationReceipt::issue(&public_key, &secret_key, target, addr, 42) @@ -362,7 +381,7 @@ mod tests { #[test] fn tampering_invalidates_signature() { let (public_key, secret_key) = generate_ml_dsa_keypair().expect("keypair"); - let relayer = fingerprint_public_key(&public_key); + let relayer = relay_receipt_peer_id(&public_key); let target = [7; 32]; let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 12_345)); let mut receipt = RelayAllocationReceipt::issue(&public_key, &secret_key, target, addr, 42) @@ -378,7 +397,7 @@ mod tests { #[test] fn receipt_expires_at_its_signed_deadline() { let (public_key, secret_key) = generate_ml_dsa_keypair().expect("keypair"); - let relayer = fingerprint_public_key(&public_key); + let relayer = relay_receipt_peer_id(&public_key); let target = [7; 32]; let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 12_345)); let receipt = RelayAllocationReceipt::issue(&public_key, &secret_key, target, addr, 42) From 4b1ed67973763e4636694a9d9c0601ff41b9e933 Mon Sep 17 00:00:00 2001 From: Mick van Dijke <12992260+mickvandijke@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:26:34 +0200 Subject: [PATCH 8/8] fix(relay): remove untrusted allocation receipts --- src/lib.rs | 4 - src/masque/connect.rs | 63 +----- src/masque/relay_server.rs | 41 ---- src/nat_traversal_api.rs | 84 +------- src/p2p_endpoint.rs | 10 +- src/relay_allocation.rs | 412 ------------------------------------- 6 files changed, 17 insertions(+), 597 deletions(-) delete mode 100644 src/relay_allocation.rs diff --git a/src/lib.rs b/src/lib.rs index d4c56343..0c010d8f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -258,10 +258,6 @@ pub mod metrics; /// TURN-style relay protocol for NAT traversal fallback pub mod relay; -mod relay_allocation; -pub use relay_allocation::{ - RelayAllocationReceipt, RelayAllocationReceiptError, relay_receipt_peer_id, -}; /// MASQUE CONNECT-UDP Bind protocol for fully connectable P2P nodes pub mod masque; diff --git a/src/masque/connect.rs b/src/masque/connect.rs index 86a4e1c6..33b66e17 100644 --- a/src/masque/connect.rs +++ b/src/masque/connect.rs @@ -275,9 +275,6 @@ pub struct ConnectUdpResponse { pub proxy_public_address: Option, /// Human-readable reason phrase pub reason: Option, - /// Relay-signed proof that this allocation belongs to the authenticated - /// client. Present for successful authenticated bind allocations. - pub allocation_receipt: Option, /// Relay-internal: the session id created for a successful CONNECT. NOT part /// of the wire format (it is not encoded/decoded); the relay sets it so the /// connection handler can start forwarding the *exact* session it created, @@ -303,7 +300,6 @@ impl ConnectUdpResponse { status: Self::STATUS_OK, proxy_public_address: public_addr, reason: None, - allocation_receipt: None, session_id: None, } } @@ -314,7 +310,6 @@ impl ConnectUdpResponse { status, proxy_public_address: None, reason: Some(reason.into()), - allocation_receipt: None, session_id: None, } } @@ -358,14 +353,14 @@ impl ConnectUdpResponse { /// Encode the response as wire format /// - /// Format: [status (2)] [flags (1)] [addr] [reason] [allocation receipt] + /// Format: [status (2)] [flags (1)] [addr] [reason] pub fn encode(&self) -> Bytes { let mut buf = BytesMut::new(); // Status code buf.put_u16(self.status); - // Flags: bit 0 = has address, bit 1 = has reason, bit 2 = has receipt + // Flags: bit 0 = has address, bit 1 = has reason let mut flags: u8 = 0; if self.proxy_public_address.is_some() { flags |= 0x01; @@ -373,9 +368,6 @@ impl ConnectUdpResponse { if self.reason.is_some() { flags |= 0x02; } - if self.allocation_receipt.is_some() { - flags |= 0x04; - } buf.put_u8(flags); // Public address if present @@ -402,14 +394,6 @@ impl ConnectUdpResponse { buf.put_slice(reason_bytes); } - if let Some(receipt) = &self.allocation_receipt { - let receipt = receipt.encode(); - if let Ok(length) = VarInt::from_u64(receipt.len() as u64) { - length.encode(&mut buf); - } - buf.put_slice(&receipt); - } - buf.freeze() } @@ -423,7 +407,11 @@ impl ConnectUdpResponse { let flags = buf.get_u8(); let has_addr = (flags & 0x01) != 0; let has_reason = (flags & 0x02) != 0; - let has_receipt = (flags & 0x04) != 0; + if flags & !0x03 != 0 { + return Err(ConnectError::InvalidResponse( + "unsupported response flags".into(), + )); + } let proxy_public_address = if has_addr { if buf.remaining() < 1 { @@ -474,31 +462,10 @@ impl ConnectUdpResponse { None }; - let allocation_receipt = if has_receipt { - let receipt_len = VarInt::decode(buf) - .map_err(|_| ConnectError::InvalidResponse("invalid receipt length".into()))? - .into_inner() as usize; - if buf.remaining() < receipt_len { - return Err(ConnectError::InvalidResponse( - "missing allocation receipt".into(), - )); - } - let mut receipt = vec![0; receipt_len]; - buf.copy_to_slice(&mut receipt); - Some( - crate::RelayAllocationReceipt::decode(&receipt).map_err(|_| { - ConnectError::InvalidResponse("invalid allocation receipt".into()) - })?, - ) - } else { - None - }; - Ok(Self { status, proxy_public_address, reason, - allocation_receipt, // Not part of the wire format; only set by the relay on the response // object it returns locally. session_id: None, @@ -607,22 +574,6 @@ mod tests { assert_eq!(original, decoded); } - #[test] - fn test_response_roundtrip_with_allocation_receipt() { - let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(203, 0, 113, 50)), 9000); - let (public_key, secret_key) = crate::generate_ml_dsa_keypair().expect("test identity"); - let mut original = ConnectUdpResponse::success(Some(addr)); - original.allocation_receipt = Some( - crate::RelayAllocationReceipt::issue(&public_key, &secret_key, [7; 32], addr, 42) - .expect("allocation receipt"), - ); - - let encoded = original.encode(); - let decoded = ConnectUdpResponse::decode(&mut encoded.clone()).unwrap(); - - assert_eq!(original, decoded); - } - #[test] fn test_response_roundtrip_success_no_addr() { let original = ConnectUdpResponse::success(None); diff --git a/src/masque/relay_server.rs b/src/masque/relay_server.rs index 59a6543c..efc37ee4 100644 --- a/src/masque/relay_server.rs +++ b/src/masque/relay_server.rs @@ -55,7 +55,6 @@ use crate::masque::{ }; use crate::relay::error::{RelayError, RelayResult, SessionErrorKind}; use crate::upnp::{UpnpConfig, UpnpMappingService}; -use crate::{MlDsaPublicKey, MlDsaSecretKey, RelayAllocationReceipt}; /// Interval at which both sides of a relay stream send a zero-length /// keepalive frame. Keeps the NAT conntrack entry alive (default @@ -577,9 +576,6 @@ pub struct MasqueRelayServer { /// When set, the server refuses inbound clients whose source IP matches /// one of our current upstream relays (see [`IpPolicy`]). ip_policy: Option>, - /// Identity used to sign allocations for third-party reachability - /// witnesses. It is the same identity presented by this endpoint in TLS. - allocation_signer: Option<(MlDsaPublicKey, MlDsaSecretKey)>, } impl std::fmt::Debug for MasqueRelayServer { @@ -620,7 +616,6 @@ impl MasqueRelayServer { .map(|_| Mutex::new(())) .collect(), ip_policy: None, - allocation_signer: None, } } @@ -694,19 +689,9 @@ impl MasqueRelayServer { .map(|_| Mutex::new(())) .collect(), ip_policy: None, - allocation_signer: None, } } - /// Configure the endpoint identity used to sign relay allocations. - pub fn set_allocation_signer( - &mut self, - public_key: MlDsaPublicKey, - secret_key: MlDsaSecretKey, - ) { - self.allocation_signer = Some((public_key, secret_key)); - } - /// Enable or disable relay serving. /// /// Called by the ADR-014 reachability classifier: public nodes leave this @@ -1090,31 +1075,6 @@ impl MasqueRelayServer { // Create new session with the bound socket let session_id = self.next_session_id.fetch_add(1, Ordering::SeqCst); - let allocation_receipt = match (&self.allocation_signer, peer_id) { - (Some((public_key, secret_key)), Some(target_peer_id)) => { - match RelayAllocationReceipt::issue( - public_key, - secret_key, - target_peer_id, - advertised_address, - session_id, - ) { - Ok(receipt) => Some(receipt), - Err(error) => { - tracing::error!( - %error, - client = %client_addr, - "Failed to sign relay allocation" - ); - return Ok(ConnectUdpResponse::error( - 500, - "Failed to sign relay allocation", - )); - } - } - } - _ => None, - }; let mut session = RelaySession::new( session_id, self.config.session_config.clone(), @@ -1167,7 +1127,6 @@ impl MasqueRelayServer { // session rather than re-looking-up by client address (which races with a // same-address reconnect). Not part of the wire format. let mut response = ConnectUdpResponse::success(Some(advertised_address)); - response.allocation_receipt = allocation_receipt; response.session_id = Some(session_id); Ok(response) } diff --git a/src/nat_traversal_api.rs b/src/nat_traversal_api.rs index c08d2cb4..99345132 100644 --- a/src/nat_traversal_api.rs +++ b/src/nat_traversal_api.rs @@ -361,8 +361,6 @@ pub struct RelaySession { pub established_at: std::time::Instant, /// Relay server address pub relay_addr: SocketAddr, - /// Relay-signed proof for this exact allocation, when available. - pub allocation_receipt: Option, } impl RelaySession { @@ -454,7 +452,6 @@ impl Drop for RelaySessionOwner { struct EstablishedRelaySession { public_address: Option, raw_streams: Option, - allocation_receipt: Option, owner: RelaySessionOwner, } @@ -505,7 +502,6 @@ struct ProactiveRelay { handle: PreparedRelay, endpoint: Arc, tunnel: Arc, - allocation_receipt: Option, relay_session_owner: RelaySessionOwner, cleanup_armed: bool, } @@ -611,10 +607,6 @@ enum RelayLifecycleCommand { PublishedHandle { reply: oneshot::Sender>, }, - Receipt { - prepared: PreparedRelay, - reply: oneshot::Sender>, - }, Health { reply: oneshot::Sender>, }, @@ -743,21 +735,6 @@ impl RelayLifecycleHandle { }; let _ = reply.send(handle); } - RelayLifecycleCommand::Receipt { prepared, reply } => { - let receipt = match &state { - RelayLifecycleState::Provisional(relay) - | RelayLifecycleState::Published(relay) - if relay.handle == prepared => - { - relay.allocation_receipt.clone() - } - RelayLifecycleState::None - | RelayLifecycleState::Provisional(_) - | RelayLifecycleState::Published(_) - | RelayLifecycleState::Shutdown => None, - }; - let _ = reply.send(receipt); - } RelayLifecycleCommand::Health { reply } => { let health = match &state { RelayLifecycleState::Published(relay) => Some(RelayHealthSnapshot { @@ -869,15 +846,6 @@ impl RelayLifecycleHandle { response.await.ok().flatten() } - async fn receipt(&self, prepared: PreparedRelay) -> Option { - let (reply, response) = oneshot::channel(); - self.commands - .send(RelayLifecycleCommand::Receipt { prepared, reply }) - .await - .ok()?; - response.await.ok().flatten() - } - async fn health(&self) -> Option { let (reply, response) = oneshot::channel(); self.commands @@ -2164,10 +2132,7 @@ impl NatTraversalEndpoint { require_authentication: true, ..MasqueRelayConfig::default() }; - let mut server = MasqueRelayServer::new(relay_config, local_addr); - if let Some((public_key, secret_key)) = config.identity_key.clone() { - server.set_allocation_signer(public_key, secret_key); - } + let server = MasqueRelayServer::new(relay_config, local_addr); info!( "Created MASQUE relay server on {} (symmetric P2P node)", local_addr @@ -2604,10 +2569,7 @@ impl NatTraversalEndpoint { require_authentication: true, ..MasqueRelayConfig::default() }; - let mut server = MasqueRelayServer::new(relay_config, local_addr); - if let Some((public_key, secret_key)) = config.identity_key.clone() { - server.set_allocation_signer(public_key, secret_key); - } + let server = MasqueRelayServer::new(relay_config, local_addr); info!( "Created MASQUE relay server on {} (symmetric P2P node)", local_addr @@ -4483,17 +4445,6 @@ impl NatTraversalEndpoint { self.relay_lifecycle.published_handle().await } - /// Return the relay-signed receipt for a live allocation. - /// - /// Matching the opaque handle prevents a stale acquisition from obtaining - /// a receipt for a newer allocation that reused the same socket address. - pub async fn proactive_relay_receipt( - &self, - prepared: PreparedRelay, - ) -> Option { - self.relay_lifecycle.receipt(prepared).await - } - /// Check if relay fallback is available pub async fn has_relay_fallback(&self) -> bool { match &self.relay_manager { @@ -4515,22 +4466,15 @@ impl NatTraversalEndpoint { pub async fn establish_relay_session( &self, relay_addr: SocketAddr, - ) -> Result< - ( - Option, - Option, - Option, - ), - NatTraversalError, - > { + ) -> Result<(Option, Option), NatTraversalError> + { let EstablishedRelaySession { public_address, raw_streams, - allocation_receipt, owner, } = self.establish_owned_relay_session(relay_addr).await?; owner.disarm(); - Ok((public_address, raw_streams, allocation_receipt)) + Ok((public_address, raw_streams)) } async fn establish_owned_relay_session( @@ -4542,7 +4486,6 @@ impl NatTraversalEndpoint { if let Some(session) = self.relay_sessions.get(&relay_addr) { if session.is_active() { let public_address = session.public_address; - let allocation_receipt = session.allocation_receipt.clone(); let stable_id = session.connection.stable_id(); debug!( relay = %relay_addr, @@ -4553,7 +4496,6 @@ impl NatTraversalEndpoint { return Ok(EstablishedRelaySession { public_address, raw_streams: None, - allocation_receipt, owner: RelaySessionOwner::new( relay_addr, public_address, @@ -4680,7 +4622,6 @@ impl NatTraversalEndpoint { } let public_address = response.proxy_public_address; - let allocation_receipt = response.allocation_receipt.clone(); info!( "Relay session established with public address: {:?}", @@ -4709,7 +4650,6 @@ impl NatTraversalEndpoint { public_address, established_at: std::time::Instant::now(), relay_addr, - allocation_receipt: allocation_receipt.clone(), }; // DashMap provides lock-free .insert() @@ -4727,7 +4667,6 @@ impl NatTraversalEndpoint { Ok(EstablishedRelaySession { public_address, raw_streams, - allocation_receipt, owner, }) } @@ -5698,7 +5637,6 @@ impl NatTraversalEndpoint { let EstablishedRelaySession { public_address: public_addr, raw_streams, - allocation_receipt, owner: relay_session_owner, } = self.establish_owned_relay_session(bootstrap_addr).await?; let Some(relay_public_addr) = public_addr else { @@ -5790,7 +5728,6 @@ impl NatTraversalEndpoint { handle, endpoint: relay_endpoint, tunnel, - allocation_receipt, relay_session_owner, cleanup_armed: true, }; @@ -6410,9 +6347,9 @@ impl NatTraversalEndpoint { None } - /// Derive the overlay-compatible 32-byte peer ID used by relay allocation - /// receipts from a connection's authenticated ML-DSA-65 identity, or - /// `None` if the peer is not PQC-authenticated. + /// Derive the overlay-compatible 32-byte peer ID from a connection's + /// authenticated ML-DSA-65 identity, or `None` if the peer is not + /// PQC-authenticated. /// /// This keys relay-port reservations to a cryptographic identity rather /// than an ephemeral socket address. The identity comes solely from the @@ -6421,7 +6358,7 @@ impl NatTraversalEndpoint { let spki = Self::extract_public_key_from_connection(connection)?; let public_key = crate::crypto::raw_public_keys::pqc::extract_public_key_from_spki(&spki).ok()?; - Some(crate::relay_receipt_peer_id(&public_key)) + Some(*blake3::hash(public_key.as_bytes()).as_bytes()) } /// Extract the raw SPKI bytes from a connection's TLS identity. @@ -8482,7 +8419,6 @@ mod tests { handle: PreparedRelay::new(public_addr), endpoint: relay_endpoint, tunnel: crate::masque::RelayTunnelControl::detached(), - allocation_receipt: None, relay_session_owner: RelaySessionOwner { relay_server_addr: "127.0.0.1:1".parse().expect("relay server address"), public_address: Some(public_addr), @@ -8647,7 +8583,6 @@ mod tests { public_address: public_addr, established_at: std::time::Instant::now(), relay_addr: server_addr, - allocation_receipt: None, }, ); let stale_owner = RelaySessionOwner::new( @@ -8664,7 +8599,6 @@ mod tests { public_address: public_addr, established_at: std::time::Instant::now(), relay_addr: server_addr, - allocation_receipt: None, }, ); drop(stale_owner); diff --git a/src/p2p_endpoint.rs b/src/p2p_endpoint.rs index 1a9a17b0..07b7c18c 100644 --- a/src/p2p_endpoint.rs +++ b/src/p2p_endpoint.rs @@ -2731,7 +2731,7 @@ impl P2pEndpoint { ); // Step 1: Establish relay session (control plane handshake) - let (public_addr, raw_streams, _allocation_receipt) = self + let (public_addr, raw_streams) = self .inner .establish_relay_session(relay_addr) .await @@ -3490,14 +3490,6 @@ impl P2pEndpoint { Ok(()) } - /// Return the relay-signed receipt for a live proactive allocation. - pub async fn proactive_relay_receipt( - &self, - prepared: PreparedRelay, - ) -> Option { - self.inner.proactive_relay_receipt(prepared).await - } - /// Perform a fresh, isolated authenticated reachability probe. /// /// This connection never enters the ordinary peer or address maps and is diff --git a/src/relay_allocation.rs b/src/relay_allocation.rs deleted file mode 100644 index 8616865c..00000000 --- a/src/relay_allocation.rs +++ /dev/null @@ -1,412 +0,0 @@ -// Copyright 2024 Saorsa Labs Ltd. -// -// This Saorsa Network Software is licensed under the General Public License (GPL), version 3. -// Please see the file LICENSE-GPL, or visit for the full text. -// -// Full details available at https://saorsalabs.com/licenses - -//! Cryptographic proof that a relay issued a specific allocation. - -use serde::{Deserialize, Serialize}; -use std::net::{IpAddr, SocketAddr}; -use std::time::{SystemTime, UNIX_EPOCH}; -use thiserror::Error; - -use crate::crypto::pqc::MlDsaOperations; -use crate::crypto::pqc::ml_dsa::MlDsa65; -use crate::crypto::pqc::types::{MlDsaPublicKey, MlDsaSecretKey, MlDsaSignature}; - -const RECEIPT_VERSION: u8 = 1; -const RECEIPT_DOMAIN: &[u8] = b"SAORSA_RELAY_ALLOCATION_V1"; -const RECEIPT_LIFETIME_SECS: u64 = 24 * 60 * 60; - -/// Derive the overlay peer identity bound into relay allocation receipts. -/// -/// The DHT identifies peers as `BLAKE3(raw ML-DSA public key)`. Transport's -/// internal `AUTONOMI_PEER_ID_V2` fingerprint uses a separate domain and must -/// not leak into this cross-layer receipt, otherwise a valid allocation can -/// never match the requester's DHT identity during canary verification. -pub fn relay_receipt_peer_id(public_key: &MlDsaPublicKey) -> [u8; 32] { - *blake3::hash(public_key.as_bytes()).as_bytes() -} - -/// A relay-signed binding between an authenticated client and one allocation. -/// -/// Witnesses must validate this receipt before attempting a canary dial. This -/// prevents a requester from turning the canary service into an arbitrary -/// reflected dial primitive. -#[derive(Clone, Serialize, Deserialize)] -pub struct RelayAllocationReceipt { - version: u8, - target_peer_id: [u8; 32], - relayer_peer_id: [u8; 32], - relay_addr: SocketAddr, - allocation_id: u64, - expires_at_unix_secs: u64, - relayer_public_key: Vec, - signature: Vec, -} - -impl std::fmt::Debug for RelayAllocationReceipt { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter - .debug_struct("RelayAllocationReceipt") - .field("version", &self.version) - .field("target_peer_id", &hex::encode(self.target_peer_id)) - .field("relayer_peer_id", &hex::encode(self.relayer_peer_id)) - .field("relay_addr", &self.relay_addr) - .field("allocation_id", &self.allocation_id) - .field("expires_at_unix_secs", &self.expires_at_unix_secs) - .finish_non_exhaustive() - } -} - -impl PartialEq for RelayAllocationReceipt { - fn eq(&self, other: &Self) -> bool { - self.version == other.version - && self.target_peer_id == other.target_peer_id - && self.relayer_peer_id == other.relayer_peer_id - && self.relay_addr == other.relay_addr - && self.allocation_id == other.allocation_id - && self.expires_at_unix_secs == other.expires_at_unix_secs - && self.relayer_public_key == other.relayer_public_key - && self.signature == other.signature - } -} - -impl Eq for RelayAllocationReceipt {} - -/// Why a relay-allocation receipt could not be issued or verified. -#[derive(Debug, Error, Clone, PartialEq, Eq)] -pub enum RelayAllocationReceiptError { - /// The receipt uses an unsupported wire version. - #[error("unsupported relay allocation receipt version {0}")] - UnsupportedVersion(u8), - /// The receipt is not bound to the requesting peer. - #[error("relay allocation receipt target does not match requester")] - TargetMismatch, - /// The receipt is not signed by the claimed relayer. - #[error("relay allocation receipt relayer does not match claim")] - RelayerMismatch, - /// The receipt is for a different relay allocation address. - #[error("relay allocation receipt address does not match request")] - AddressMismatch, - /// The receipt has expired. - #[error("relay allocation receipt expired")] - Expired, - /// The local clock could not be represented as Unix time. - #[error("system clock is before the Unix epoch")] - InvalidClock, - /// The embedded ML-DSA material is malformed. - #[error("relay allocation receipt contains invalid ML-DSA material")] - InvalidCryptoMaterial, - /// The receipt wire representation is malformed. - #[error("relay allocation receipt encoding is invalid")] - InvalidEncoding, - /// The ML-DSA signature is invalid. - #[error("relay allocation receipt signature is invalid")] - InvalidSignature, - /// Signing failed. - #[error("failed to sign relay allocation receipt")] - SigningFailed, -} - -impl RelayAllocationReceipt { - /// Issue a receipt for an allocation made to an authenticated client. - pub fn issue( - relayer_public_key: &MlDsaPublicKey, - relayer_secret_key: &MlDsaSecretKey, - target_peer_id: [u8; 32], - relay_addr: SocketAddr, - allocation_id: u64, - ) -> Result { - let now = unix_time_secs(SystemTime::now())?; - let relayer_peer_id = relay_receipt_peer_id(relayer_public_key); - let mut receipt = Self { - version: RECEIPT_VERSION, - target_peer_id, - relayer_peer_id, - relay_addr, - allocation_id, - expires_at_unix_secs: now.saturating_add(RECEIPT_LIFETIME_SECS), - relayer_public_key: relayer_public_key.as_bytes().to_vec(), - signature: Vec::new(), - }; - let signature = MlDsa65::new() - .sign(relayer_secret_key, &receipt.signing_message()) - .map_err(|_| RelayAllocationReceiptError::SigningFailed)?; - receipt.signature = signature.as_bytes().to_vec(); - Ok(receipt) - } - - /// Validate the receipt and all bindings supplied by a canary requester. - pub fn verify( - &self, - target_peer_id: [u8; 32], - relayer_peer_id: [u8; 32], - relay_addr: SocketAddr, - now: SystemTime, - ) -> Result<(), RelayAllocationReceiptError> { - if self.version != RECEIPT_VERSION { - return Err(RelayAllocationReceiptError::UnsupportedVersion( - self.version, - )); - } - if self.target_peer_id != target_peer_id { - return Err(RelayAllocationReceiptError::TargetMismatch); - } - if self.relayer_peer_id != relayer_peer_id { - return Err(RelayAllocationReceiptError::RelayerMismatch); - } - if self.relay_addr != relay_addr { - return Err(RelayAllocationReceiptError::AddressMismatch); - } - if unix_time_secs(now)? >= self.expires_at_unix_secs { - return Err(RelayAllocationReceiptError::Expired); - } - - let public_key = MlDsaPublicKey::from_bytes(&self.relayer_public_key) - .map_err(|_| RelayAllocationReceiptError::InvalidCryptoMaterial)?; - if relay_receipt_peer_id(&public_key) != self.relayer_peer_id { - return Err(RelayAllocationReceiptError::RelayerMismatch); - } - let signature = MlDsaSignature::from_bytes(&self.signature) - .map_err(|_| RelayAllocationReceiptError::InvalidCryptoMaterial)?; - match MlDsa65::new().verify(&public_key, &self.signing_message(), &signature) { - Ok(true) => Ok(()), - Ok(false) => Err(RelayAllocationReceiptError::InvalidSignature), - Err(_) => Err(RelayAllocationReceiptError::InvalidCryptoMaterial), - } - } - - /// Authenticated client fingerprint bound into this receipt. - pub fn target_peer_id(&self) -> [u8; 32] { - self.target_peer_id - } - - /// Authenticated relay fingerprint bound into this receipt. - pub fn relayer_peer_id(&self) -> [u8; 32] { - self.relayer_peer_id - } - - /// Public allocation address bound into this receipt. - pub fn relay_addr(&self) -> SocketAddr { - self.relay_addr - } - - pub(crate) fn encode(&self) -> Vec { - let mut encoded = Vec::with_capacity( - 1 + 32 - + 32 - + 1 - + 16 - + 2 - + 8 - + 8 - + 2 - + self.relayer_public_key.len() - + 2 - + self.signature.len(), - ); - encoded.push(self.version); - encoded.extend_from_slice(&self.target_peer_id); - encoded.extend_from_slice(&self.relayer_peer_id); - match self.relay_addr.ip() { - IpAddr::V4(ip) => { - encoded.push(4); - encoded.extend_from_slice(&ip.octets()); - } - IpAddr::V6(ip) => { - encoded.push(6); - encoded.extend_from_slice(&ip.octets()); - } - } - encoded.extend_from_slice(&self.relay_addr.port().to_be_bytes()); - encoded.extend_from_slice(&self.allocation_id.to_be_bytes()); - encoded.extend_from_slice(&self.expires_at_unix_secs.to_be_bytes()); - encoded.extend_from_slice(&(self.relayer_public_key.len() as u16).to_be_bytes()); - encoded.extend_from_slice(&self.relayer_public_key); - encoded.extend_from_slice(&(self.signature.len() as u16).to_be_bytes()); - encoded.extend_from_slice(&self.signature); - encoded - } - - pub(crate) fn decode(mut encoded: &[u8]) -> Result { - fn take<'a>( - encoded: &mut &'a [u8], - length: usize, - ) -> Result<&'a [u8], RelayAllocationReceiptError> { - if encoded.len() < length { - return Err(RelayAllocationReceiptError::InvalidEncoding); - } - let (value, remainder) = encoded.split_at(length); - *encoded = remainder; - Ok(value) - } - - let version = take(&mut encoded, 1)?[0]; - let target_peer_id = take(&mut encoded, 32)? - .try_into() - .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?; - let relayer_peer_id = take(&mut encoded, 32)? - .try_into() - .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?; - let ip = match take(&mut encoded, 1)?[0] { - 4 => { - let octets: [u8; 4] = take(&mut encoded, 4)? - .try_into() - .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?; - IpAddr::V4(std::net::Ipv4Addr::from(octets)) - } - 6 => { - let octets: [u8; 16] = take(&mut encoded, 16)? - .try_into() - .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?; - IpAddr::V6(std::net::Ipv6Addr::from(octets)) - } - _ => return Err(RelayAllocationReceiptError::InvalidEncoding), - }; - let port = u16::from_be_bytes( - take(&mut encoded, 2)? - .try_into() - .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?, - ); - let allocation_id = u64::from_be_bytes( - take(&mut encoded, 8)? - .try_into() - .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?, - ); - let expires_at_unix_secs = u64::from_be_bytes( - take(&mut encoded, 8)? - .try_into() - .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?, - ); - let public_key_len = u16::from_be_bytes( - take(&mut encoded, 2)? - .try_into() - .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?, - ) as usize; - let relayer_public_key = take(&mut encoded, public_key_len)?.to_vec(); - let signature_len = u16::from_be_bytes( - take(&mut encoded, 2)? - .try_into() - .map_err(|_| RelayAllocationReceiptError::InvalidEncoding)?, - ) as usize; - let signature = take(&mut encoded, signature_len)?.to_vec(); - if !encoded.is_empty() { - return Err(RelayAllocationReceiptError::InvalidEncoding); - } - - Ok(Self { - version, - target_peer_id, - relayer_peer_id, - relay_addr: SocketAddr::new(ip, port), - allocation_id, - expires_at_unix_secs, - relayer_public_key, - signature, - }) - } - - fn signing_message(&self) -> Vec { - let mut message = Vec::with_capacity(RECEIPT_DOMAIN.len() + 100); - message.extend_from_slice(RECEIPT_DOMAIN); - message.push(self.version); - message.extend_from_slice(&self.target_peer_id); - message.extend_from_slice(&self.relayer_peer_id); - match self.relay_addr.ip() { - IpAddr::V4(ip) => { - message.push(4); - message.extend_from_slice(&ip.octets()); - } - IpAddr::V6(ip) => { - message.push(6); - message.extend_from_slice(&ip.octets()); - } - } - message.extend_from_slice(&self.relay_addr.port().to_be_bytes()); - message.extend_from_slice(&self.allocation_id.to_be_bytes()); - message.extend_from_slice(&self.expires_at_unix_secs.to_be_bytes()); - message - } -} - -fn unix_time_secs(time: SystemTime) -> Result { - time.duration_since(UNIX_EPOCH) - .map(|duration| duration.as_secs()) - .map_err(|_| RelayAllocationReceiptError::InvalidClock) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::crypto::raw_public_keys::pqc::generate_ml_dsa_keypair; - use std::net::{Ipv4Addr, SocketAddrV4}; - - #[test] - fn receipt_peer_id_uses_overlay_identity_derivation() { - let (public_key, _) = generate_ml_dsa_keypair().expect("keypair"); - - assert_eq!( - relay_receipt_peer_id(&public_key), - *blake3::hash(public_key.as_bytes()).as_bytes() - ); - } - - #[test] - fn receipt_verifies_only_for_exact_binding() { - let (public_key, secret_key) = generate_ml_dsa_keypair().expect("keypair"); - let relayer = relay_receipt_peer_id(&public_key); - let target = [7; 32]; - let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 12_345)); - let receipt = RelayAllocationReceipt::issue(&public_key, &secret_key, target, addr, 42) - .expect("receipt"); - - assert!( - receipt - .verify(target, relayer, addr, SystemTime::now()) - .is_ok() - ); - assert_eq!( - receipt.verify([8; 32], relayer, addr, SystemTime::now()), - Err(RelayAllocationReceiptError::TargetMismatch) - ); - assert_eq!( - receipt.verify(target, [9; 32], addr, SystemTime::now()), - Err(RelayAllocationReceiptError::RelayerMismatch) - ); - } - - #[test] - fn tampering_invalidates_signature() { - let (public_key, secret_key) = generate_ml_dsa_keypair().expect("keypair"); - let relayer = relay_receipt_peer_id(&public_key); - let target = [7; 32]; - let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 12_345)); - let mut receipt = RelayAllocationReceipt::issue(&public_key, &secret_key, target, addr, 42) - .expect("receipt"); - receipt.allocation_id += 1; - - assert_eq!( - receipt.verify(target, relayer, addr, SystemTime::now()), - Err(RelayAllocationReceiptError::InvalidSignature) - ); - } - - #[test] - fn receipt_expires_at_its_signed_deadline() { - let (public_key, secret_key) = generate_ml_dsa_keypair().expect("keypair"); - let relayer = relay_receipt_peer_id(&public_key); - let target = [7; 32]; - let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 12_345)); - let receipt = RelayAllocationReceipt::issue(&public_key, &secret_key, target, addr, 42) - .expect("receipt"); - let deadline = UNIX_EPOCH + std::time::Duration::from_secs(receipt.expires_at_unix_secs); - - assert_eq!( - receipt.verify(target, relayer, addr, deadline), - Err(RelayAllocationReceiptError::Expired) - ); - } -}