From 691b705e7300d4163eb1c5dadfdb276d18faa4f6 Mon Sep 17 00:00:00 2001 From: Nidish Date: Wed, 2 Sep 2026 14:52:52 +0530 Subject: [PATCH 1/7] feat(truapi-server): share one localhost WS listener across native product executions --- .../kotlin/io/parity/truapi/TrUAPIHost.kt | 8 +- rust/crates/truapi-server/src/native.rs | 42 +- rust/crates/truapi-server/src/ws_bridge.rs | 1152 ++++++++++++++--- 3 files changed, 1027 insertions(+), 175 deletions(-) diff --git a/android/truapi-host/src/main/kotlin/io/parity/truapi/TrUAPIHost.kt b/android/truapi-host/src/main/kotlin/io/parity/truapi/TrUAPIHost.kt index 6385345aa..2d3516caa 100644 --- a/android/truapi-host/src/main/kotlin/io/parity/truapi/TrUAPIHost.kt +++ b/android/truapi-host/src/main/kotlin/io/parity/truapi/TrUAPIHost.kt @@ -938,11 +938,15 @@ class TrUAPIProductExecution internal constructor( ) : AutoCloseable { private val shutDown = AtomicBoolean(false) - /** Start this execution's independently authenticated localhost bridge. */ + /** + * Register this execution against the host runtime's shared localhost + * bridge, minting an independent authentication token. Every execution + * under the same host runtime connects through the same port. + */ @Throws(WsBridgeStartException::class) fun startWsBridge(bindPort: UShort = 0u): WsBridgeEndpoint = inner.startWsBridge(bindPort) - /** Stop the active bridge while leaving the execution reusable. */ + /** Revoke this execution's bridge registration while leaving it reusable. */ fun stopWsBridge() { inner.stopWsBridge() } diff --git a/rust/crates/truapi-server/src/native.rs b/rust/crates/truapi-server/src/native.rs index 193a5f6ce..1a15bdd45 100644 --- a/rust/crates/truapi-server/src/native.rs +++ b/rust/crates/truapi-server/src/native.rs @@ -43,7 +43,7 @@ use crate::native_renderer::{NativeRendererObserver, NativeRendererSubscription} use crate::runtime::sso_remote::sso_message_id; use crate::subscription::Spawner; #[cfg(feature = "ws-bridge")] -use crate::ws_bridge::{BridgeLogger, WsBridge, WsBridgeEndpoint, WsBridgeStartError}; +use crate::ws_bridge::{BridgeLogger, SharedWsBridge, WsBridgeEndpoint, WsBridgeStartError}; /// Host-thrown storage failure wrapping the canonical error payload, so its /// variants remain defined once in `truapi`. @@ -623,6 +623,11 @@ pub struct NativeTrUApiHostRuntime { events: Arc, #[cfg(feature = "ws-bridge")] spawner: Spawner, + /// One localhost WebSocket listener shared by every product execution + /// this host runtime opens. Starts lazily on the first execution's + /// `start_ws_bridge` call and lives for as long as this host runtime. + #[cfg(feature = "ws-bridge")] + ws_bridge: Arc, /// The one Worker execution per product; opening another replaces it. worker_executions: Mutex>>, } @@ -665,6 +670,8 @@ impl NativeTrUApiHostRuntime { events, #[cfg(feature = "ws-bridge")] spawner, + #[cfg(feature = "ws-bridge")] + ws_bridge: Arc::new(SharedWsBridge::new()), worker_executions: Mutex::new(HashMap::new()), })) } @@ -715,7 +722,9 @@ impl NativeTrUApiHostRuntime { chat_connection: Arc::new(crate::runtime::ActionChannel::chat()), renderer_connection: Arc::new(crate::runtime::ActionChannel::renderer()), #[cfg(feature = "ws-bridge")] - bridge: Mutex::new(None), + ws_bridge: self.ws_bridge.clone(), + #[cfg(feature = "ws-bridge")] + bridge_token: Mutex::new(None), #[cfg(feature = "ws-bridge")] product_control: Arc::new(Mutex::new(None)), }); @@ -1027,8 +1036,13 @@ pub struct NativeProductExecution { crate::runtime::ActionChannel, >, closed: AtomicBool, + /// The host runtime's shared listener; every execution under the same + /// host runtime clones this same `Arc`. + #[cfg(feature = "ws-bridge")] + ws_bridge: Arc, + /// This execution's own registered token, if its bridge is running. #[cfg(feature = "ws-bridge")] - bridge: Mutex>, + bridge_token: Mutex>, #[cfg(feature = "ws-bridge")] product_control: Arc>>, } @@ -1071,13 +1085,13 @@ impl NativeProductExecution { #[cfg(feature = "ws-bridge")] fn stop_bridge(&self) { - if let Some(mut bridge) = self - .bridge + if let Some(token) = self + .bridge_token .lock() .expect("native product bridge mutex poisoned") .take() { - bridge.stop(); + self.ws_bridge.revoke(&token); } *self .product_control @@ -1256,7 +1270,11 @@ impl NativeProductExecution { #[cfg(feature = "ws-bridge")] #[uniffi::export] impl NativeProductExecution { - /// Start this execution's independently authenticated localhost bridge. + /// Register this execution against the host runtime's shared localhost + /// bridge, minting an independent authentication token. Every execution + /// under the same host runtime connects through the same port; + /// `bind_port` only has an effect for the first execution to register + /// while the shared listener is still unstarted. pub fn start_ws_bridge(&self, bind_port: u16) -> Result { if self.closed.load(Ordering::Acquire) { return Err(WsBridgeStartError::Io( @@ -1264,7 +1282,7 @@ impl NativeProductExecution { )); } let mut guard = self - .bridge + .bridge_token .lock() .expect("native product bridge mutex poisoned"); if guard.is_some() { @@ -1288,12 +1306,14 @@ impl NativeProductExecution { .expect("native product control mutex poisoned") = Some(product_runtime.control()); product_runtime }); - let (bridge, endpoint) = WsBridge::start(bind_port, runtime_factory, logger)?; - *guard = Some(bridge); + let endpoint = self + .ws_bridge + .register(bind_port, runtime_factory, logger)?; + *guard = Some(endpoint.token.clone()); Ok(endpoint) } - /// Stop the active bridge while leaving the execution reusable. + /// Revoke this execution's bridge registration while leaving it reusable. pub fn stop_ws_bridge(&self) { self.stop_bridge(); } diff --git a/rust/crates/truapi-server/src/ws_bridge.rs b/rust/crates/truapi-server/src/ws_bridge.rs index ec47ca115..75a9ac15c 100644 --- a/rust/crates/truapi-server/src/ws_bridge.rs +++ b/rust/crates/truapi-server/src/ws_bridge.rs @@ -5,21 +5,40 @@ //! //! Feature-gated (`ws-bridge`) so wasm32 and no-tokio build paths stay lean. //! -//! Native bridges share one process-wide `tokio` runtime. Each [`WsBridge`] -//! owns only its accept loop and connection tasks; dropping or stopping one -//! bridge leaves the executor available to other products. +//! Native bridges share one process-wide `tokio` runtime, and every product +//! execution under one host runtime shares a single [`SharedWsBridge`] +//! listener: [`SharedWsBridge::register`] hands each execution its own +//! `{port, token}` endpoint on the one shared port, and +//! [`SharedWsBridge::revoke`] tears down only that execution's connections +//! when it closes, leaving the listener and every other execution's +//! connections untouched. //! //! Security model: the listener binds to `127.0.0.1` only, and every -//! connection must present the per-session 256-bit token (`?t=`, -//! drawn from the OS CSPRNG) before the WebSocket upgrade completes. The token -//! is the sole authentication gate and is compared in constant time. It is -//! handed only to the host's embedded WebView, so the bridge does not also pin -//! the `Origin` header (the WebView's origin is not known a priori). Inbound -//! messages are size-capped, and the per-connection outbound queue and the -//! total connection count are bounded to contain a misbehaving local peer. - +//! connection must present its registered per-execution 256-bit token +//! (`?t=`, drawn from the OS CSPRNG) before the WebSocket upgrade +//! completes. The handshake scans every currently registered token with a +//! constant-time comparison and does not exit early on a match or on a +//! duplicated `t=` parameter, so timing does not reveal which token (if any) +//! matched. Revoking a token removes it from the registry before existing +//! connections are aborted, so a handshake that has not yet matched a token +//! when it is revoked is rejected outright; a handshake that matched just +//! before the revocation may still receive an already-committed upgrade +//! response, but is aborted immediately afterward, before any wire traffic +//! flows. Tokens are handed only to the host's embedded WebView, so the +//! bridge does not also pin the `Origin` header (the WebView's origin is not +//! known a priori). Inbound messages are size-capped, and both the +//! per-execution and process-wide outbound queue and connection counts are +//! bounded to contain a misbehaving local peer. Each accepted connection's +//! handshake runs in its own task (bounded by [`HANDSHAKE_TIMEOUT`]), so a +//! peer that never completes one only ever stalls its own connection, never +//! the shared accept loop, another execution's connections, or the +//! listener's shutdown; the connection-count caps are reserved with a +//! compare-and-swap loop precisely because handshakes now run concurrently. + +use std::collections::HashMap; use std::io; use std::net::SocketAddr; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, OnceLock}; use futures::{SinkExt, StreamExt}; @@ -27,6 +46,7 @@ use rand::RngCore; use tokio::net::TcpListener; use tokio::runtime::{Handle, Runtime}; use tokio::sync::{mpsc, oneshot}; +use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::tungstenite::Message as WsMessage; use tokio_tungstenite::tungstenite::handshake::server::{ErrorResponse, Request, Response}; use tokio_tungstenite::tungstenite::http::{Response as HttpResponse, StatusCode}; @@ -36,10 +56,17 @@ use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; use crate::{FrameSink, ProductRuntime}; -/// Maximum simultaneous connections the bridge will service. The product uses -/// a single connection; the cap bounds resource use from a buggy or hostile -/// local peer opening many sockets. -const MAX_WS_BRIDGE_CONNECTIONS: usize = 32; +/// Maximum simultaneous connections a single registered execution may hold. +/// Each execution uses exactly one connection; the cap bounds resource use +/// from a buggy or hostile local peer opening many sockets against one +/// token. +const MAX_WS_CONNECTIONS_PER_EXECUTION: usize = 32; + +/// Maximum simultaneous connections across every execution sharing the +/// listener. Set well above any realistic concurrent-execution count (App, +/// Widget, Chat/Worker) while still bounding total resource use from a +/// misbehaving peer amplifying across many registered tokens. +const MAX_TOTAL_WS_CONNECTIONS: usize = 64; /// Bound on the per-connection outbound frame queue. A peer that stops reading /// cannot make the core buffer responses without limit; once the queue fills @@ -51,6 +78,15 @@ const OUTBOUND_QUEUE_CAP: usize = 4096; /// memory-amplification DoS well below tungstenite's 64 MiB default. const MAX_WS_MESSAGE_BYTES: usize = 8 << 20; +/// Ceiling on how long one connection's own setup task will wait for its +/// handshake (`authenticate_and_upgrade`) to resolve. Each connection's setup +/// runs in its own task, so a peer that never completes the handshake cannot +/// stall any other connection — this bound exists so a stalled setup task +/// (and the socket and file descriptor behind it) cannot linger forever, and +/// so the listener's shutdown never has to wait on one indefinitely. Generous +/// relative to a real localhost upgrade (sub-millisecond in practice). +const HANDSHAKE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); + /// Per-session descriptor returned to the host: product uses `port + token` /// to build its WebSocket URL (e.g. `ws://127.0.0.1:/?t=`). #[derive(Clone, Debug, uniffi::Record)] @@ -159,33 +195,205 @@ fn shared_native_executor() -> io::Result<(&'static SharedNativeExecutor, bool)> Ok((executor, initialized)) } -/// Running bridge handle. Drop or call [`WsBridge::stop`] to shut down. +/// One execution's registered runtime factory and live connection state. +struct RegistryEntry { + runtime_factory: Arc, + connection_count: Arc, + connections: Mutex, +} + +/// A connection can finish its handshake against an entry that is revoked a +/// moment later, before the accept loop gets to record its handle here. The +/// `revoked` flag and the handle list share one lock so revocation and +/// registration serialize against each other: whichever happens first is +/// what the other observes, so a connection admitted in that window is +/// always either recorded for the abort that already ran, or told to abort +/// itself immediately rather than being left running with no owner. +#[derive(Default)] +struct EntryConnections { + revoked: bool, + handles: Vec>, +} + +/// Token registry shared by the listener's accept loop and every +/// [`SharedWsBridge::register`]/[`SharedWsBridge::revoke`] call. +#[derive(Default)] +struct WsBridgeRegistry { + entries: Mutex>>, + total_connections: Arc, +} + +impl WsBridgeRegistry { + fn insert(&self, token: String, runtime_factory: Arc) { + self.entries + .lock() + .expect("ws bridge registry mutex poisoned") + .insert( + token, + Arc::new(RegistryEntry { + runtime_factory, + connection_count: Arc::new(AtomicUsize::new(0)), + connections: Mutex::new(EntryConnections::default()), + }), + ); + } + + /// Remove `token`'s entry, mark it revoked, and abort its live + /// connections. A handshake that matched this token just before the + /// removal, but has not yet registered its handle, observes the + /// `revoked` flag (set under the same lock as the abort loop below) and + /// aborts itself instead of registering. No-op for an unknown or + /// already-revoked token. + fn revoke(&self, token: &str) { + let Some(entry) = self + .entries + .lock() + .expect("ws bridge registry mutex poisoned") + .remove(token) + else { + return; + }; + let mut state = entry + .connections + .lock() + .expect("ws bridge registry entry mutex poisoned"); + state.revoked = true; + for handle in state.handles.iter() { + handle.abort(); + } + } + + /// Find the entry whose token matches `path_and_query`'s `?t=` value. + /// Scans every registered token without exiting early on a match, so + /// timing does not reveal which one (if any) matched. + fn find_matching(&self, path_and_query: Option<&str>) -> Option> { + let entries = self + .entries + .lock() + .expect("ws bridge registry mutex poisoned"); + let mut found = None; + for (token, entry) in entries.iter() { + if path_token_matches(path_and_query, token) { + found = Some(entry.clone()); + } + } + found + } + + /// Abort and drain every connection tracked under a still-registered + /// execution, returning the owned handles so the caller can await each + /// one to genuine completion rather than just requesting cancellation. + /// Safe to call more than once (or after `revoke` already drained some) + /// — later calls simply find nothing left to take for whatever was + /// already drained. A connection whose token was revoked (and thus whose + /// entry was already removed from the registry) moments before this + /// call is not covered here — `revoke` already aborted it directly, but + /// this method has no way to find and await it, so a caller cannot treat + /// its own completion as proof that connection has actually finished + /// unwinding too. + fn take_all_handles(&self) -> Vec> { + let entries = self + .entries + .lock() + .expect("ws bridge registry mutex poisoned"); + let mut all = Vec::new(); + for entry in entries.values() { + let mut state = entry + .connections + .lock() + .expect("ws bridge registry entry mutex poisoned"); + for handle in &state.handles { + handle.abort(); + } + all.append(&mut state.handles); + } + all + } +} + +/// Host-runtime-owned shared listener. Every product execution under the +/// same host runtime clones the same `Arc` and registers +/// against it; the underlying listener starts on the first registration and +/// lives for as long as the host runtime does. +#[derive(Default)] +pub struct SharedWsBridge { + inner: Mutex>, +} + +impl SharedWsBridge { + /// Construct an unstarted shared bridge. The listener binds lazily on + /// the first [`Self::register`] call. + pub fn new() -> Self { + Self::default() + } + + /// Ensure the shared listener is running and register a fresh + /// per-execution token against it. + /// + /// `bind_port` only takes effect for the first execution to register; + /// once the listener is up, every product connects through that same + /// port regardless of what a later caller requests. A later caller that + /// requested a specific (non-zero) port different from the one already + /// running gets a log line noting the request was ignored, rather than + /// silence. + pub fn register( + &self, + bind_port: u16, + runtime_factory: Arc, + logger: BridgeLogger, + ) -> Result { + let mut guard = self.inner.lock().expect("shared ws bridge mutex poisoned"); + if guard.is_none() { + *guard = Some(WsBridge::start(bind_port, logger.clone())?); + } else if bind_port != 0 { + let running_port = guard.as_ref().expect("just checked Some").port; + if bind_port != running_port { + logger( + "truapi.ws_bridge.bind_port_ignored", + &format!("requested={bind_port} running={running_port}"), + ); + } + } + Ok(guard + .as_ref() + .expect("shared bridge just inserted") + .register(runtime_factory)) + } + + /// Revoke one execution's token. No-op if the listener was never + /// started or the token is unknown. + pub fn revoke(&self, token: &str) { + if let Some(bridge) = self + .inner + .lock() + .expect("shared ws bridge mutex poisoned") + .as_ref() + { + bridge.revoke(token); + } + } +} + +/// Running listener handle. Drop or call [`WsBridge::stop`] to shut down. /// -/// The bridge's tasks run on the process-wide native executor. TrUAPI dispatch -/// futures are `Send`, so connections and independent frames from all products -/// can execute across the shared worker pool. -pub struct WsBridge { +/// The listener's tasks run on the process-wide native executor. TrUAPI +/// dispatch futures are `Send`, so connections and independent frames from +/// all registered executions can execute across the shared worker pool. +struct WsBridge { shutdown: Option>, stopped: Option>, accept_task: Option>, runtime_id: tokio::runtime::Id, + registry: Arc, + port: u16, } impl WsBridge { /// Bind a localhost listener and start the accept loop on the shared - /// native executor. Returns the [`WsBridgeEndpoint`] descriptor the host - /// hands to the product alongside the bridge handle. - pub fn start( - bind_port: u16, - runtime_factory: Arc, - logger: BridgeLogger, - ) -> io::Result<(Self, WsBridgeEndpoint)> { - let mut token_bytes = [0u8; 32]; - rand::thread_rng().fill_bytes(&mut token_bytes); - let token = hex::encode(token_bytes); - + /// native executor. + fn start(bind_port: u16, logger: BridgeLogger) -> io::Result { // Bind synchronously so we can surface bind errors and discover the - // actual port before returning the endpoint. + // actual port before returning. let std_listener = std::net::TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], bind_port)))?; std_listener.set_nonblocking(true)?; @@ -210,44 +418,57 @@ impl WsBridge { let _entered = handle.enter(); TcpListener::from_std(std_listener)? }; + let registry = Arc::new(WsBridgeRegistry::default()); let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); let (stopped_tx, stopped_rx) = std::sync::mpsc::channel::<()>(); - let accept_token = token.clone(); + let accept_registry = registry.clone(); let accept_logger = logger.clone(); let accept_task = handle.spawn(async move { - accept_loop( - listener, - runtime_factory, - accept_token, - accept_logger, - shutdown_rx, - ) - .await; + accept_loop(listener, accept_registry, accept_logger, shutdown_rx).await; let _ = stopped_tx.send(()); }); logger( "truapi.ws_bridge.started", - &format!( - "port={port} token_len={} runtime_id={runtime_id}", - token.len() - ), + &format!("port={port} runtime_id={runtime_id}"), ); - Ok(( - Self { - shutdown: Some(shutdown_tx), - stopped: Some(stopped_rx), - accept_task: Some(accept_task), - runtime_id, - }, - WsBridgeEndpoint { port, token }, - )) + Ok(Self { + shutdown: Some(shutdown_tx), + stopped: Some(stopped_rx), + accept_task: Some(accept_task), + runtime_id, + registry, + port, + }) + } + + /// Register one execution's runtime factory, minting a fresh token. + fn register(&self, runtime_factory: Arc) -> WsBridgeEndpoint { + let mut token_bytes = [0u8; 32]; + rand::thread_rng().fill_bytes(&mut token_bytes); + let token = hex::encode(token_bytes); + self.registry.insert(token.clone(), runtime_factory); + WsBridgeEndpoint { + port: self.port, + token, + } } - /// Signal this bridge's accept loop to exit without stopping the shared - /// native executor used by other products. - pub fn stop(&mut self) { + /// Revoke one execution's token and abort its live connections. + fn revoke(&self, token: &str) { + self.registry.revoke(token); + } + + /// Signal the accept loop to exit and abort every tracked connection + /// across every registered execution. + /// + /// Off the shared executor, this blocks until the accept loop has + /// actually drained and awaited every connection, so a caller there gets + /// a real "fully stopped" guarantee. From a task already running on the + /// shared executor, waiting is skipped instead (see below), so this + /// specific caller only gets a best-effort sweep, not that guarantee. + fn stop(&mut self) { if let Some(tx) = self.shutdown.take() { let _ = tx.send(()); } @@ -256,7 +477,11 @@ impl WsBridge { // where waiting preserves the existing "fully stopped on return" // behavior. Avoid blocking if a Rust caller drops the bridge from one // of the shared runtime's own workers, especially on a single-core - // runtime; the shutdown signal still lets the task clean itself up. + // runtime, where blocking here could deadlock against the very task + // this is waiting on. Sending on `shutdown` only schedules the accept + // loop to be re-polled, so on this path `stop` can return before the + // accept loop has even observed it, let alone drained anything — the + // fallback sweep below is what still cleans up whatever it can see. let called_from_shared_executor = Handle::try_current().is_ok_and(|handle| handle.id() == self.runtime_id); let stopped_cleanly = if called_from_shared_executor { @@ -274,6 +499,15 @@ impl WsBridge { { task.abort(); } + // Fallback sweep: off the shared executor, the accept loop's own + // shutdown branch already drained and awaited every connection + // before signaling `stopped`, so this finds nothing left. It only + // does real work on the shared-executor fast path above (which + // skips that wait) or if the accept task had to be force-aborted + // (e.g. a panic inside the loop before it reached its own shutdown + // branch) — in both cases this only aborts what it finds, it does + // not await it, since `stop` itself is not async. + drop(self.registry.take_all_handles()); } } @@ -285,21 +519,35 @@ impl Drop for WsBridge { async fn accept_loop( listener: TcpListener, - runtime_factory: Arc, - expected_token: String, + registry: Arc, logger: BridgeLogger, mut shutdown: oneshot::Receiver<()>, ) { - let mut handles: Vec> = Vec::new(); + // Each accepted connection's handshake runs in its own task (see + // `connection_setup`) so a stalled or hostile peer never blocks another + // connection's setup, only its own. This loop just tracks those setup + // tasks so shutdown can cancel and await whichever haven't resolved yet, + // in addition to draining every connection already registered. + let mut setup_tasks: Vec> = Vec::new(); loop { tokio::select! { _ = &mut shutdown => { logger("truapi.ws_bridge.shutdown", "accept loop exiting"); - for h in &handles { - h.abort(); + for task in &setup_tasks { + task.abort(); } - for h in handles { - let _ = h.await; + for task in setup_tasks { + let _ = task.await; + } + // `take_all_handles` both requests cancellation and hands back + // sole ownership of every handle still under a registered + // execution, so every one of them can genuinely be awaited + // here before this loop (and thus the listener's `stopped` + // signal) returns. It does not see a connection whose token + // was revoked (and thus removed from the registry) moments + // earlier — `revoke` already aborted that one directly. + for handle in registry.take_all_handles() { + let _ = handle.await; } break; } @@ -311,48 +559,218 @@ async fn accept_loop( continue; } }; - handles.retain(|h| !h.is_finished()); - if handles.len() >= MAX_WS_BRIDGE_CONNECTIONS { - logger("truapi.ws_bridge.connection_limit", &peer.to_string()); - drop(stream); - continue; - } - let runtime_factory = runtime_factory.clone(); + setup_tasks.retain(|task| !task.is_finished()); + let registry = registry.clone(); let logger = logger.clone(); - let expected = expected_token.clone(); - handles.push(tokio::spawn(async move { - handle_connection(stream, peer, runtime_factory, expected, logger).await; + setup_tasks.push(tokio::spawn(async move { + connection_setup(stream, peer, registry, logger).await; })); } } } } +/// Authenticate one accepted connection and, if the handshake succeeds, +/// register it and run its lifecycle. Spawned independently per connection +/// (from `accept_loop`) precisely so a stalled or hostile peer's handshake +/// blocks only this task, never another connection's setup, another +/// execution's connections, or the accept loop's ability to keep accepting. +/// +/// Because handshakes now run concurrently rather than one at a time, the +/// connection-count caps inside `authenticate_and_upgrade`'s handshake +/// callback are reserved with a compare-and-swap loop (`try_reserve`) rather +/// than a plain load-then-increment. The registration race this used to +/// avoid by serializing handshakes — a token revoked between a successful +/// handshake and its registration — is instead closed by the `revoked` flag +/// check below, which holds regardless of how many handshakes run at once. +async fn connection_setup( + stream: tokio::net::TcpStream, + peer: SocketAddr, + registry: Arc, + logger: BridgeLogger, +) { + let auth_result = tokio::time::timeout( + HANDSHAKE_TIMEOUT, + authenticate_and_upgrade(stream, peer, ®istry, logger.clone()), + ) + .await; + let Some((ws, entry, guard)) = (match auth_result { + Ok(resolved) => resolved, + Err(_) => { + logger("truapi.ws_bridge.handshake_timeout", &peer.to_string()); + return; + } + }) else { + return; + }; + let conn_logger = logger.clone(); + let conn_entry = entry.clone(); + let handle = tokio::spawn(async move { + let _guard = guard; + connection_lifecycle(ws, peer, conn_entry, conn_logger).await; + }); + let mut state = entry + .connections + .lock() + .expect("ws bridge registry entry mutex poisoned"); + state.handles.retain(|h| !h.is_finished()); + if state.revoked { + // The execution was revoked between this connection's successful + // handshake and this registration step; abort it now instead of + // leaving it running with no owner to revoke it later. + handle.abort(); + } else { + state.handles.push(handle); + } +} + +/// Decrements the global and per-execution connection counts on drop, which +/// runs whether the connection task ends normally or is aborted (e.g. by a +/// revoke), keeping the caps accurate under both. +struct ConnectionCountGuard { + total: Arc, + per_entry: Arc, +} + +impl Drop for ConnectionCountGuard { + fn drop(&mut self) { + self.total.fetch_sub(1, Ordering::AcqRel); + self.per_entry.fetch_sub(1, Ordering::AcqRel); + } +} + +/// Disposes a connection's `ProductRuntime` when this guard drops, however +/// that happens. `.abort()` — the only mechanism `revoke` and shutdown use to +/// end a connection — drops its task's future at whatever `.await` point it +/// was suspended at (almost always inside the read loop), which skips any +/// code written after that point, including an explicit `dispose()` call at +/// the end of the function. A value still on the stack at that point still +/// gets dropped, though, so holding the runtime here is what actually +/// disposes an aborted connection. `dispose()` is documented as idempotent, +/// so this is safe even alongside a path that also disposes explicitly. +struct DisposeGuard(Arc); + +impl Drop for DisposeGuard { + fn drop(&mut self) { + self.0.dispose(); + } +} + +/// A fully upgraded WebSocket connection, the execution entry its token +/// matched, and the connection-count guard reserved for it. +type AuthenticatedConnection = ( + WebSocketStream, + Arc, + ConnectionCountGuard, +); + +/// What the handshake callback found and reserved, shared between the +/// callback and the function's own return path. +type MatchedReservation = Arc, ConnectionCountGuard)>>>; + +/// Atomically reserve one slot by incrementing `counter` unless it is already +/// at `limit`, retrying under contention. Connection setup now runs +/// concurrently (one task per accepted connection), so this cannot be a +/// plain load-then-increment: two handshakes could otherwise both observe +/// room for the last slot and both take it. +fn try_reserve(counter: &AtomicUsize, limit: usize) -> bool { + let mut current = counter.load(Ordering::Acquire); + loop { + if current >= limit { + return false; + } + match counter.compare_exchange_weak( + current, + current + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return true, + Err(actual) => current = actual, + } + } +} + +/// Complete the WebSocket handshake, resolving it against the registry to +/// find the matching execution's entry and reserving its connection-count +/// slots for the life of the connection. +/// +/// Called from `connection_setup`, one independent task per accepted +/// connection, under [`HANDSHAKE_TIMEOUT`] — so a peer that never completes +/// its handshake only ever stalls that one task, not the accept loop or any +/// other connection's setup. +/// +/// The connection-count reservation happens inside the synchronous handshake +/// callback via [`try_reserve`], as soon as a token matches and both caps +/// have room, and a [`ConnectionCountGuard`] for it is constructed +/// immediately and stashed alongside the matched entry. Every way this +/// function can end — a successful upgrade, a post-callback handshake I/O +/// failure, or an outer timeout dropping this whole future — drops `matched` +/// and, with it, any reserved-but-unclaimed guard, so a failure after the +/// callback already committed a slot can never leak it. // `clippy::result_large_err` fires on the handshake callback because // tokio-tungstenite's `ErrorResponse` type carries the full HTTP response // (~136 bytes). The closure signature is dictated by tokio-tungstenite's // API, so the lint can only be silenced at the call site. #[allow(clippy::result_large_err)] -async fn handle_connection( +async fn authenticate_and_upgrade( stream: tokio::net::TcpStream, peer: SocketAddr, - runtime_factory: Arc, - expected_token: String, + registry: &Arc, logger: BridgeLogger, -) { +) -> Option { + let matched: MatchedReservation = Arc::new(Mutex::new(None)); + let auth_registry = registry.clone(); + let auth_matched = matched.clone(); let auth_logger = logger.clone(); - let callback = |req: &Request, resp: Response| -> Result { - if path_token_matches( - req.uri().path_and_query().map(|p| p.as_str()), - &expected_token, - ) { - Ok(resp) - } else { + let callback = move |req: &Request, resp: Response| -> Result { + let path_and_query = req.uri().path_and_query().map(|p| p.as_str()); + let Some(entry) = auth_registry.find_matching(path_and_query) else { auth_logger("truapi.ws_bridge.reject_unauthorized", &peer.to_string()); let mut err: ErrorResponse = HttpResponse::new(Some("invalid token".to_string())); *err.status_mut() = StatusCode::UNAUTHORIZED; - Err(err) + return Err(err); + }; + + // Reserve both caps' slots here, inside the synchronous handshake + // callback, so an over-cap peer is rejected at the HTTP upgrade + // rather than allowed to open a socket that gets dropped right + // after. Handshakes now run concurrently (one task per connection), + // so each reservation is a compare-and-swap, not a plain increment. + if !try_reserve(&entry.connection_count, MAX_WS_CONNECTIONS_PER_EXECUTION) { + auth_logger( + "truapi.ws_bridge.connection_limit_execution", + &peer.to_string(), + ); + let mut err: ErrorResponse = + HttpResponse::new(Some("execution connection limit reached".to_string())); + *err.status_mut() = StatusCode::SERVICE_UNAVAILABLE; + return Err(err); } + if !try_reserve(&auth_registry.total_connections, MAX_TOTAL_WS_CONNECTIONS) { + // Roll back the per-execution reservation above: it was never + // claimed, since the total cap is what ultimately rejected this + // attempt. + entry.connection_count.fetch_sub(1, Ordering::AcqRel); + auth_logger("truapi.ws_bridge.connection_limit_total", &peer.to_string()); + let mut err: ErrorResponse = + HttpResponse::new(Some("listener at capacity".to_string())); + *err.status_mut() = StatusCode::SERVICE_UNAVAILABLE; + return Err(err); + } + // Built right after both reservations above, so no return from this + // function — success, a later I/O failure, or an outer timeout — can + // drop `matched` without also dropping (and thus balancing) this + // guard. + let guard = ConnectionCountGuard { + total: auth_registry.total_connections.clone(), + per_entry: entry.connection_count.clone(), + }; + + *auth_matched + .lock() + .expect("ws bridge handshake mutex poisoned") = Some((entry, guard)); + Ok(resp) }; // Cap inbound message/frame size so a peer cannot force the runtime to @@ -368,15 +786,30 @@ async fn handle_connection( Ok(ws) => ws, Err(err) => { logger("truapi.ws_bridge.handshake_error", &err.to_string()); - return; + return None; } }; + let (entry, guard) = matched + .lock() + .expect("ws bridge handshake mutex poisoned") + .take() + .expect("a successful upgrade always resolved a matching registry entry"); logger("truapi.ws_bridge.connection_open", &peer.to_string()); + Some((ws, entry, guard)) +} + +async fn connection_lifecycle( + ws: WebSocketStream, + peer: SocketAddr, + entry: Arc, + logger: BridgeLogger, +) { let (mut sink, mut source) = ws.split(); let (out_tx, mut out_rx) = mpsc::channel::>(OUTBOUND_QUEUE_CAP); let frame_sink = Arc::new(WsFrameSink::new(out_tx)); - let product_runtime = Arc::new(runtime_factory.product_runtime(frame_sink)); + let product_runtime = Arc::new(entry.runtime_factory.product_runtime(frame_sink)); + let _dispose_guard = DisposeGuard(product_runtime.clone()); let pump_logger = logger.clone(); let pump = tokio::spawn(async move { @@ -429,16 +862,23 @@ async fn handle_connection( } // The connection is gone: cancel in-flight dispatches so long-pending - // handlers unwind instead of outliving the connection. + // handlers unwind instead of outliving the connection. `_dispose_guard` + // disposes `product_runtime` when it drops at the end of this function. for task in &in_flight { task.abort(); } - product_runtime.dispose(); let _ = pump.await; logger("truapi.ws_bridge.connection_closed", &peer.to_string()); } +/// Whether `path_and_query`'s `?t=` value (any occurrence) matches `expected`, +/// compared in constant time. Every `t=` pair present is checked — this does +/// not stop at the first match — so a peer padding the query with several +/// `t=` pairs cannot make a match resolve faster than a non-match and use +/// that timing to test a candidate token. The token length is fixed and +/// public, so a length mismatch may short-circuit; only the value comparison +/// must be constant time. fn path_token_matches(path_and_query: Option<&str>, expected: &str) -> bool { let Some(raw) = path_and_query else { return false; @@ -447,16 +887,17 @@ fn path_token_matches(path_and_query: Option<&str>, expected: &str) -> bool { Some(idx) => &raw[idx + 1..], None => return false, }; + let mut matched = false; for pair in query.split('&') { let (key, value) = match pair.split_once('=') { Some(kv) => kv, None => continue, }; if key == "t" && constant_time_eq(value.as_bytes(), expected.as_bytes()) { - return true; + matched = true; } } - false + matched } /// Constant-time byte-slice equality, used for the session-token check so a @@ -537,6 +978,20 @@ mod tests { Arc::new(move |sink| runtime.product_runtime(product.clone(), sink)) } + fn no_log() -> BridgeLogger { + Arc::new(|_, _| {}) + } + + fn connect(port: u16, token: &str) -> tokio::runtime::Runtime { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime"); + let url = format!("ws://127.0.0.1:{port}/?t={token}"); + rt.block_on(async { tokio_tungstenite::connect_async(&url).await.expect("dial") }); + rt + } + #[test] fn path_token_matches_exact() { assert!(path_token_matches(Some("/?t=abc"), "abc")); @@ -547,6 +1002,17 @@ mod tests { assert!(!path_token_matches(None, "abc")); } + /// A query with more than one `t=` pair is matched if ANY of them equals + /// the expected token, regardless of position — this is what closes the + /// timing differential a peer could otherwise create by padding the + /// query with non-matching `t=` pairs before or after the real one. + #[test] + fn path_token_matches_every_duplicated_t_pair_not_just_the_first() { + assert!(path_token_matches(Some("/?t=wrong&t=abc"), "abc")); + assert!(path_token_matches(Some("/?t=abc&t=wrong"), "abc")); + assert!(!path_token_matches(Some("/?t=wrong&t=alsowrong"), "abc")); + } + #[test] fn shared_executor_uses_multithread_scheduler() { let (executor, _) = shared_native_executor().expect("shared native executor"); @@ -599,8 +1065,7 @@ mod tests { #[test] fn drop_from_shared_executor_does_not_block_worker() { - let (bridge, _) = - WsBridge::start(0, test_runtime_factory(), Arc::new(|_, _| {})).expect("start bridge"); + let bridge = WsBridge::start(0, no_log()).expect("start bridge"); let (executor, _) = shared_native_executor().expect("shared native executor"); let (dropped_tx, dropped_rx) = std::sync::mpsc::channel(); @@ -614,15 +1079,14 @@ mod tests { .expect("dropping from an executor worker must not deadlock"); } - /// Spin the bridge up on `127.0.0.1:0`, dial it with a real - /// `tokio-tungstenite` client, send a known SCALE frame, and verify - /// the bridge echoes the SCALE-encoded `feature_supported` response. + /// Spin the shared listener up on `127.0.0.1:0`, register one execution, + /// dial it with a real `tokio-tungstenite` client, send a known SCALE + /// frame, and verify the bridge echoes the SCALE-encoded + /// `feature_supported` response. #[test] fn round_trip_feature_supported_through_bridge() { - let runtime_factory = test_runtime_factory(); - let logger: BridgeLogger = Arc::new(|_, _| {}); - let (mut bridge, endpoint) = - WsBridge::start(0, runtime_factory, logger).expect("start bridge"); + let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let endpoint = bridge.register(test_runtime_factory()); let url = format!("ws://127.0.0.1:{}/?t={}", endpoint.port, endpoint.token); // Use a fresh `tokio` runtime on the test thread so the client does @@ -682,97 +1146,461 @@ mod tests { )); assert_eq!(response.payload.value, expected.encode()); - bridge.stop(); + drop(bridge); } - /// Multiple product bridges use the same executor, and stopping one - /// product must not interrupt another product's bridge. + /// Two executions registered against the same shared listener land on + /// the same port with independent tokens, and a token only opens its own + /// execution's runtime, never the other's. #[test] - fn stopping_one_bridge_leaves_another_operational() { - let runtime_ids = Arc::new(Mutex::new(Vec::::new())); - let logger: BridgeLogger = { - let runtime_ids = runtime_ids.clone(); - Arc::new(move |marker, detail| { - if marker == "truapi.ws_bridge.started" - && let Some(runtime_id) = detail.split("runtime_id=").nth(1) - { - runtime_ids.lock().unwrap().push(runtime_id.to_string()); - } - }) - }; - let (mut first, _) = - WsBridge::start(0, test_runtime_factory(), logger.clone()).expect("first bridge"); - let (mut second, endpoint) = - WsBridge::start(0, test_runtime_factory(), logger).expect("second bridge"); + fn two_executions_share_one_port_with_isolated_tokens() { + let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let first = bridge.register(test_runtime_factory()); + let second = bridge.register(test_runtime_factory()); - let ids = runtime_ids.lock().unwrap().clone(); - assert_eq!(ids.len(), 2); - assert_eq!(ids[0], ids[1]); + assert_eq!(first.port, second.port); + assert_ne!(first.token, second.token); - first.stop(); + // Each token independently authenticates against the shared port. + connect(first.port, &first.token); + connect(second.port, &second.token); - let url = format!("ws://127.0.0.1:{}/?t={}", endpoint.port, endpoint.token); - let client = tokio::runtime::Builder::new_current_thread() + drop(bridge); + } + + /// A handshake presenting one execution's token must not be accepted as + /// belonging to a different execution's entry, and an unknown token is + /// rejected outright. + #[test] + fn wrong_or_unknown_token_is_rejected_at_handshake() { + let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let endpoint = bridge.register(test_runtime_factory()); + let _second = bridge.register(test_runtime_factory()); + + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime"); + + let url = format!("ws://127.0.0.1:{}/?t=bogus", endpoint.port); + let err = rt + .block_on(async { tokio_tungstenite::connect_async(&url).await }) + .expect_err("connection with an unknown token must be refused"); + let msg = format!("{err}"); + assert!( + msg.contains("401") || msg.to_lowercase().contains("unauthorized"), + "expected 401/unauthorized rejection, got: {msg}", + ); + + drop(bridge); + } + + /// Revoking one execution's token closes only its own connections and + /// rejects future handshakes against it, while a sibling execution on + /// the same shared listener keeps working. + #[test] + fn revoking_one_token_leaves_another_operational() { + let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let revoked = bridge.register(test_runtime_factory()); + let survives = bridge.register(test_runtime_factory()); + + let rt = tokio::runtime::Builder::new_current_thread() .enable_all() .build() .expect("test runtime"); - client.block_on(async { - let (mut ws, _) = tokio_tungstenite::connect_async(&url) + let revoked_url = format!("ws://127.0.0.1:{}/?t={}", revoked.port, revoked.token); + let mut revoked_ws = rt.block_on(async { + tokio_tungstenite::connect_async(&revoked_url) .await - .expect("second bridge remains reachable"); + .expect("dial revoked execution") + .0 + }); + + bridge.revoke(&revoked.token); + + // The already-open connection is torn down. Aborting the task drops + // the socket without a clean close handshake, so the client + // observes either end-of-stream or a read error, not necessarily a + // `Close` frame. + rt.block_on(async { + let deadline = tokio::time::sleep(std::time::Duration::from_secs(2)); + tokio::pin!(deadline); + loop { + tokio::select! { + _ = &mut deadline => panic!("revoked connection was not closed"), + frame = revoked_ws.next() => { + match frame { + None => break, + Some(Err(_)) => break, + Some(Ok(WsMessage::Close(_))) => continue, + Some(Ok(_)) => continue, + } + } + } + } + }); + + // ...and its token no longer authenticates. + let err = rt + .block_on(async { tokio_tungstenite::connect_async(&revoked_url).await }) + .expect_err("revoked token must be rejected"); + assert!(format!("{err}").to_lowercase().contains("unauthorized")); + + // The sibling execution is unaffected. + let survives_url = format!("ws://127.0.0.1:{}/?t={}", survives.port, survives.token); + rt.block_on(async { + let (mut ws, _) = tokio_tungstenite::connect_async(&survives_url) + .await + .expect("surviving execution remains reachable"); ws.close(None).await.expect("close client"); }); - second.stop(); + drop(bridge); } - /// A handshake with the wrong `?t=` token must be rejected at the HTTP - /// upgrade step with a 401, not silently dropped. + /// After a token is revoked, registering a fresh execution (simulating a + /// reconnect / restart) gets a brand new token that works independently + /// of the old one. #[test] - fn wrong_token_is_rejected_at_handshake() { - let runtime_factory = test_runtime_factory(); - let logger: BridgeLogger = Arc::new(|_, _| {}); - let (mut bridge, endpoint) = - WsBridge::start(0, runtime_factory, logger).expect("start bridge"); - let url = format!("ws://127.0.0.1:{}/?t=bogus", endpoint.port); + fn reconnecting_after_revoke_gets_a_fresh_token() { + let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let first = bridge.register(test_runtime_factory()); + bridge.revoke(&first.token); + + let second = bridge.register(test_runtime_factory()); + assert_ne!(first.token, second.token); + assert_eq!(first.port, second.port); + + connect(second.port, &second.token); + + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime"); + let first_url = format!("ws://127.0.0.1:{}/?t={}", first.port, first.token); + let err = rt + .block_on(async { tokio_tungstenite::connect_async(&first_url).await }) + .expect_err("the revoked token must stay rejected"); + assert!(format!("{err}").to_lowercase().contains("unauthorized")); + + drop(bridge); + } + + /// Dropping the shared listener (host runtime shutdown) tears down every + /// registered execution's connections, not just one. + #[test] + fn host_shutdown_closes_every_registered_execution() { + let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let first = bridge.register(test_runtime_factory()); + let second = bridge.register(test_runtime_factory()); let rt = tokio::runtime::Builder::new_current_thread() .enable_all() .build() .expect("test runtime"); + let (mut first_ws, mut second_ws) = rt.block_on(async { + let (first_ws, _) = tokio_tungstenite::connect_async(format!( + "ws://127.0.0.1:{}/?t={}", + first.port, first.token + )) + .await + .expect("dial first"); + let (second_ws, _) = tokio_tungstenite::connect_async(format!( + "ws://127.0.0.1:{}/?t={}", + second.port, second.token + )) + .await + .expect("dial second"); + (first_ws, second_ws) + }); + + drop(bridge); + + rt.block_on(async { + let deadline = tokio::time::sleep(std::time::Duration::from_secs(2)); + tokio::pin!(deadline); + tokio::select! { + _ = &mut deadline => panic!("connections were not closed on host shutdown"), + _ = async { + while first_ws.next().await.is_some() {} + while second_ws.next().await.is_some() {} + } => {} + } + }); + } + + /// `SharedWsBridge` starts its listener lazily on first registration and + /// hands every subsequent registration the same port. + #[test] + fn shared_ws_bridge_lazily_starts_and_reuses_its_port() { + let shared = SharedWsBridge::new(); + let first = shared + .register(0, test_runtime_factory(), no_log()) + .expect("first registration starts the listener"); + let second = shared + .register(0, test_runtime_factory(), no_log()) + .expect("second registration reuses it"); + + assert_eq!(first.port, second.port); + assert_ne!(first.token, second.token); + + connect(first.port, &first.token); + connect(second.port, &second.token); + + shared.revoke(&first.token); + // The second registration is untouched by revoking the first. + connect(second.port, &second.token); + } + + /// Once one execution has `MAX_WS_CONNECTIONS_PER_EXECUTION` live + /// connections, its next connection attempt is refused with a 503 even + /// though the shared listener's own total cap has plenty of room left. + #[test] + fn per_execution_cap_rejects_the_connection_past_the_limit() { + let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let endpoint = bridge.register(test_runtime_factory()); + let url = format!("ws://127.0.0.1:{}/?t={}", endpoint.port, endpoint.token); + + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime"); + // Keep every socket alive so the execution's connection count never + // drops back below the cap while the next attempt is made. + let _sockets = rt.block_on(async { + let mut sockets = Vec::new(); + for _ in 0..MAX_WS_CONNECTIONS_PER_EXECUTION { + let (ws, _) = tokio_tungstenite::connect_async(&url) + .await + .expect("dial under the per-execution cap"); + sockets.push(ws); + } + sockets + }); let err = rt .block_on(async { tokio_tungstenite::connect_async(&url).await }) - .expect_err("connection must be refused"); - let msg = format!("{err}"); + .expect_err("the connection past the per-execution cap must be refused"); + let msg = format!("{err}").to_lowercase(); assert!( - msg.contains("401") || msg.to_lowercase().contains("unauthorized"), - "expected 401/unauthorized rejection, got: {msg}", + msg.contains("503") || msg.contains("service unavailable"), + "expected a 503 rejection past the per-execution cap, got: {err}", ); - bridge.stop(); + drop(bridge); } - /// Dropping a `WsBridge` handle without an explicit `stop()` must still - /// shut its accept task down cleanly. `Drop::drop` calls `stop`, and a - /// second `stop` (from drop after the test's explicit one) is a no-op. + /// The shared listener's total connection cap is enforced independently + /// of any single execution's own cap: a brand-new execution (one that + /// has never opened a connection before, nowhere near its own + /// per-execution limit) is still refused once the shared budget is + /// gone. Simulates the budget already being exhausted by directly + /// setting the same atomic the real cap check reads, rather than + /// opening `MAX_TOTAL_WS_CONNECTIONS` real sockets just to reach it — + /// this exercises the exact same check with far less real I/O. #[test] - fn drop_calls_stop_idempotently() { - let runtime_factory = test_runtime_factory(); - let logger: BridgeLogger = Arc::new(|_, _| {}); - let (bridge, _endpoint) = - WsBridge::start(0, runtime_factory, logger).expect("start bridge"); - // Drop the bridge; the accept task must finish via Drop. + fn total_cap_rejects_the_connection_even_for_a_fresh_execution() { + let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let extra = bridge.register(test_runtime_factory()); + bridge + .registry + .total_connections + .store(MAX_TOTAL_WS_CONNECTIONS, Ordering::SeqCst); + + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime"); + + let extra_url = format!("ws://127.0.0.1:{}/?t={}", extra.port, extra.token); + let err = rt + .block_on(async { tokio_tungstenite::connect_async(&extra_url).await }) + .expect_err("a fresh execution must still be refused once the shared listener is full"); + let msg = format!("{err}").to_lowercase(); + assert!( + msg.contains("503") || msg.contains("service unavailable"), + "expected a 503 rejection past the total cap, got: {err}", + ); + drop(bridge); + } + + /// A panic while building one execution's product runtime (e.g. a bug in + /// its own logic) tears down only that connection, not a sibling's — this + /// bridge does not, say, hold a lock across the panicking call that a + /// sibling connection also needs. Tokio's per-task panic containment is + /// what provides the isolation this test observes, but the release + /// profile sets `panic = "abort"`, so that containment — and thus this + /// test's premise — is a property of test builds only; a panic here in + /// an actual release build aborts the whole host process regardless of + /// which task it originated in. + #[test] + fn a_panicking_execution_does_not_affect_a_sibling() { + let panicking_factory: Arc = + Arc::new(|_sink: Arc| -> ProductRuntime { + panic!("intentional test panic: simulating a failing product execution") + }); - // Build a second bridge and explicitly stop twice. The second - // call has no shutdown sender or accept task left to wait for, - // so it returns without panicking. - let runtime_factory = test_runtime_factory(); - let logger: BridgeLogger = Arc::new(|_, _| {}); - let (mut bridge, _endpoint) = - WsBridge::start(0, runtime_factory, logger).expect("start bridge"); - bridge.stop(); - bridge.stop(); + let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let failing = bridge.register(panicking_factory); + let healthy = bridge.register(test_runtime_factory()); + + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime"); + + // The handshake succeeds (the token is valid); the connection is then + // torn down once its factory panics while building the runtime. + rt.block_on(async { + let failing_url = format!("ws://127.0.0.1:{}/?t={}", failing.port, failing.token); + let (mut ws, _) = tokio_tungstenite::connect_async(&failing_url) + .await + .expect("handshake succeeds; the token itself is valid"); + let deadline = tokio::time::sleep(std::time::Duration::from_secs(2)); + tokio::pin!(deadline); + loop { + tokio::select! { + _ = &mut deadline => panic!("the panicking execution's connection was never closed"), + frame = ws.next() => match frame { + None | Some(Err(_)) => break, + Some(Ok(_)) => continue, + } + } + } + }); + + // The healthy sibling execution is unaffected. + let healthy_url = format!("ws://127.0.0.1:{}/?t={}", healthy.port, healthy.token); + rt.block_on(async { + let (mut ws, _) = tokio_tungstenite::connect_async(&healthy_url) + .await + .expect("sibling execution remains reachable"); + ws.close(None).await.expect("close client"); + }); + + drop(bridge); + } + + /// A peer that opens a TCP connection and never sends the HTTP upgrade + /// request — so its handshake never resolves — does not block a + /// sibling's connection attempt. Each accepted connection's handshake + /// runs in its own task; before that, everything shared the accept + /// loop's own inline handshake, so a stalled peer there would have + /// blocked every other execution's connections too. + #[test] + fn a_stalled_handshake_does_not_block_a_sibling_connection() { + let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let stalled = bridge.register(test_runtime_factory()); + let healthy = bridge.register(test_runtime_factory()); + + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime"); + + rt.block_on(async { + // A raw TCP connection that never sends any bytes: the accept + // loop sees it, but its handshake never resolves. + let _stalled_stream = tokio::net::TcpStream::connect(("127.0.0.1", stalled.port)) + .await + .expect("open a raw stream to the shared port"); + + let healthy_url = format!("ws://127.0.0.1:{}/?t={}", healthy.port, healthy.token); + let deadline = tokio::time::sleep(std::time::Duration::from_secs(2)); + tokio::pin!(deadline); + tokio::select! { + _ = &mut deadline => panic!( + "sibling connection was blocked by another connection's stalled handshake" + ), + result = tokio_tungstenite::connect_async(&healthy_url) => { + let (mut ws, _) = result.expect("sibling handshake must succeed promptly"); + ws.close(None).await.expect("close client"); + } + } + }); + + drop(bridge); + } + + /// Registers three tokens with distinguishable runtime factories and + /// confirms connecting through one token invokes only its own factory, + /// never a sibling's — ruling out a routing bug that only happens to + /// work by coincidence with exactly two registry entries. + #[test] + fn three_tokens_route_to_their_own_factory_only() { + fn tracked_factory(called: Arc) -> Arc { + let inner = test_runtime_factory(); + Arc::new(move |sink| { + called.fetch_add(1, Ordering::SeqCst); + inner.product_runtime(sink) + }) + } + + let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let calls: Vec> = (0..3).map(|_| Arc::new(AtomicUsize::new(0))).collect(); + let endpoints: Vec = calls + .iter() + .map(|called| bridge.register(tracked_factory(called.clone()))) + .collect(); + + // Connect through the middle token specifically, and wait for a real + // response: the server cannot answer without having already called + // `product_runtime()` on the connection task, which is what makes + // this a reliable synchronization point rather than racing the + // server's own handling of the just-completed handshake. + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime"); + let url = format!( + "ws://127.0.0.1:{}/?t={}", + endpoints[1].port, endpoints[1].token + ); + let ids = request_ids("system_feature_supported").expect("known request method"); + rt.block_on(async { + let (mut ws, _) = tokio_tungstenite::connect_async(&url).await.expect("dial"); + let request_frame = ProtocolMessage { + request_id: "p:1".into(), + payload: Payload { + id: ids.request_id, + value: HostFeatureSupportedRequest::V1( + v01::HostFeatureSupportedRequest::Chain { + genesis_hash: vec![0u8; 32], + }, + ) + .encode(), + }, + }; + ws.send(WsMessage::Binary(request_frame.encode())) + .await + .expect("send"); + loop { + match ws.next().await { + Some(Ok(WsMessage::Binary(_))) => break, + Some(Ok(_)) => continue, + Some(Err(err)) => panic!("ws error: {err}"), + None => panic!("connection closed before response"), + } + } + }); + + assert_eq!( + calls[0].load(Ordering::SeqCst), + 0, + "a sibling's factory must not be invoked" + ); + assert_eq!( + calls[1].load(Ordering::SeqCst), + 1, + "the matching token's own factory must be invoked exactly once" + ); + assert_eq!( + calls[2].load(Ordering::SeqCst), + 0, + "a sibling's factory must not be invoked" + ); + + drop(bridge); } } From 292e1c6761a51b8e60edcff280bd1ffa1919754f Mon Sep 17 00:00:00 2001 From: Nidish Date: Thu, 10 Sep 2026 18:25:44 +0530 Subject: [PATCH 2/7] fix(truapi-server): bound the handshake backlog and tighten bridge teardown --- rust/crates/truapi-server/src/host_core.rs | 42 ++- rust/crates/truapi-server/src/native.rs | 154 +++++++- rust/crates/truapi-server/src/ws_bridge.rs | 418 +++++++++++++++------ 3 files changed, 486 insertions(+), 128 deletions(-) diff --git a/rust/crates/truapi-server/src/host_core.rs b/rust/crates/truapi-server/src/host_core.rs index 7e8a9dbc5..709bced9b 100644 --- a/rust/crates/truapi-server/src/host_core.rs +++ b/rust/crates/truapi-server/src/host_core.rs @@ -1367,10 +1367,23 @@ impl ProductRuntime { // would poison this mutex and every later `receive_frame` would then panic // here, which is exactly the production-host-killing shape the debug tap // above was fixed for. - self.in_flight - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .insert(dispatch_id, abort_handle); + // + // The disposed check above is only a fast path: `dispose` can swap it true + // and drain `in_flight` at any point after that check returns and before + // this insert runs. Re-checking here, under the same lock `dispose` holds + // for its own swap-and-drain, closes that window - whichever runs first is + // what the other observes, so a dispatch that loses the race is turned away + // instead of running past a disposal that already happened. + { + let mut in_flight = self + .in_flight + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if self.disposed.load(Ordering::Acquire) { + return Ok(()); + } + in_flight.insert(dispatch_id, abort_handle); + } let transport: Arc = self.transport.clone(); let _ = Abortable::new(self.core.dispatch(message, transport), abort_registration).await; @@ -1461,17 +1474,28 @@ impl ProductRuntime { /// futures, and cancels active subscriptions. #[instrument(skip_all, fields(runtime.method = "product_runtime.dispose"))] pub fn dispose(&self) { - if self.disposed.swap(true, Ordering::AcqRel) { + // Already-disposed callers answer without taking the lock, which the + // authoritative swap below holds across an abort loop that can run + // arbitrary waker code. + if self.disposed.load(Ordering::Acquire) { return; } - for (_, handle) in self + // The swap and the drain share `in_flight`'s lock with `receive_frame`'s own + // disposed-check-then-insert, so whichever of the two critical sections runs + // first is what the other observes: a dispatch that inserted before this + // drain is caught by it, and one that hasn't inserted yet sees `disposed` + // already true and turns itself away instead of dispatching past disposal. + let mut in_flight = self .in_flight .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .drain() - { + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if self.disposed.swap(true, Ordering::AcqRel) { + return; + } + for (_, handle) in in_flight.drain() { handle.abort(); } + drop(in_flight); self.admin.product_runtime.detach_chat(); self.admin.product_runtime.detach_renderer(); self.host_subscriptions.close(); diff --git a/rust/crates/truapi-server/src/native.rs b/rust/crates/truapi-server/src/native.rs index 1a15bdd45..5dd42e763 100644 --- a/rust/crates/truapi-server/src/native.rs +++ b/rust/crates/truapi-server/src/native.rs @@ -1085,12 +1085,16 @@ impl NativeProductExecution { #[cfg(feature = "ws-bridge")] fn stop_bridge(&self) { - if let Some(token) = self + // Taken in its own statement so the guard is released before `revoke`, + // which blocks until this execution's connections have unwound. Held + // across that call, it would make a concurrent `start_ws_bridge` on + // this execution wait out the whole teardown. + let token = self .bridge_token .lock() .expect("native product bridge mutex poisoned") - .take() - { + .take(); + if let Some(token) = token { self.ws_bridge.revoke(&token); } *self @@ -4026,6 +4030,150 @@ mod tests { execution.stop_ws_bridge(); } + /// The PR's headline behavior, exercised through the real + /// `NativeProductExecution`/UniFFI-facing surface rather than only + /// `ws_bridge.rs`'s own lower-level unit tests: two executions under one + /// host runtime share the shared bridge's port with isolated tokens, and + /// stopping one leaves the other's live connection untouched. + #[cfg(feature = "ws-bridge")] + #[test] + fn two_executions_share_one_bridge_through_the_native_api() { + use futures::SinkExt; + use parity_scale_codec::Decode; + use tokio_tungstenite::tungstenite::Message as WsMessage; + use truapi::versioned::system::HostFeatureSupportedRequest; + + use crate::frame::{Payload, ProtocolMessage, request_ids}; + + let host = NativeTrUApiHostRuntime::with_runtime_config( + Arc::new(EventCallbacks::new()), + native_host_runtime_config(), + ) + .expect("host runtime config should be valid"); + let app = host + .open_product_execution( + Arc::new(EventCallbacks::new()), + None, + native_execution_config("shared.dot", ProductExecutionKind::App), + ) + .expect("App execution should open"); + let chat_host = Arc::new(EventCallbacks::new()); + let chat = host + .open_product_execution( + chat_host.clone(), + Some(chat_host), + native_execution_config("shared.dot", ProductExecutionKind::Worker), + ) + .expect("Chat execution should open"); + + let app_endpoint = app.start_ws_bridge(0).expect("start app bridge"); + let chat_endpoint = chat.start_ws_bridge(0).expect("start chat bridge"); + assert_eq!( + app_endpoint.port, chat_endpoint.port, + "both executions must share the one listener port" + ); + assert_ne!( + app_endpoint.token, chat_endpoint.token, + "each execution must get its own token" + ); + + let feature_ids = request_ids("system_feature_supported").expect("known request method"); + let round_trip = |request_id: &str| ProtocolMessage { + request_id: request_id.into(), + payload: Payload { + id: feature_ids.request_id, + value: HostFeatureSupportedRequest::V1(v01::HostFeatureSupportedRequest::Chain { + genesis_hash: vec![0u8; 32], + }) + .encode(), + }, + }; + async fn answer(ws: &mut S) -> ProtocolMessage + where + S: futures::Stream> + + Unpin, + { + tokio::time::timeout(std::time::Duration::from_secs(10), async { + loop { + match ws.next().await { + Some(Ok(WsMessage::Binary(bytes))) => { + break ProtocolMessage::decode(&mut &bytes[..]) + .expect("decode response"); + } + Some(Ok(_)) => continue, + Some(Err(err)) => panic!("ws error: {err}"), + None => panic!("connection closed before response"), + } + } + }) + .await + .expect("must answer") + } + + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime"); + + rt.block_on(async { + let app_url = format!( + "ws://127.0.0.1:{}/?t={}", + app_endpoint.port, app_endpoint.token + ); + let chat_url = format!( + "ws://127.0.0.1:{}/?t={}", + chat_endpoint.port, chat_endpoint.token + ); + + let (mut app_ws, _) = tokio_tungstenite::connect_async(&app_url) + .await + .expect("app dial"); + let (mut chat_ws, _) = tokio_tungstenite::connect_async(&chat_url) + .await + .expect("chat dial"); + + // Each execution's own token round-trips independently on the + // one shared port. + app_ws + .send(WsMessage::Binary(round_trip("app:1").encode())) + .await + .expect("send on app connection"); + assert_eq!(answer(&mut app_ws).await.request_id, "app:1"); + + chat_ws + .send(WsMessage::Binary(round_trip("chat:1").encode())) + .await + .expect("send on chat connection"); + assert_eq!(answer(&mut chat_ws).await.request_id, "chat:1"); + + // Stopping (revoking) App must not touch Chat's live connection, + // and must block until App's own connection has actually finished + // unwinding before returning. + app.stop_ws_bridge(); + + chat_ws + .send(WsMessage::Binary(round_trip("chat:2").encode())) + .await + .expect("send on chat connection after App stops"); + assert_eq!( + answer(&mut chat_ws).await.request_id, + "chat:2", + "Chat's connection must keep answering after a sibling execution stops" + ); + + // The token is removed from the registry synchronously inside + // `revoke`, so a fresh connect with it is refused with no retry + // loop. This says nothing about whether the aborted connection has + // finished unwinding. + assert!( + tokio_tungstenite::connect_async(&app_url).await.is_err(), + "a revoked token must not still be accepted" + ); + + chat.stop_ws_bridge(); + }); + } + fn native_host_runtime_no_session() -> Arc { let mut config = native_host_runtime_config(); config.local_session_secret = None; diff --git a/rust/crates/truapi-server/src/ws_bridge.rs b/rust/crates/truapi-server/src/ws_bridge.rs index 75a9ac15c..1c365b921 100644 --- a/rust/crates/truapi-server/src/ws_bridge.rs +++ b/rust/crates/truapi-server/src/ws_bridge.rs @@ -11,7 +11,16 @@ //! `{port, token}` endpoint on the one shared port, and //! [`SharedWsBridge::revoke`] tears down only that execution's connections //! when it closes, leaving the listener and every other execution's -//! connections untouched. +//! connections untouched. Off the shared executor, `revoke` blocks until +//! every connection it aborted has been joined, matching what +//! [`WsBridge::stop`] gives for the whole listener. Joining a task is not a +//! barrier on everything inside it: the destructor that disposes a +//! connection's `ProductRuntime` usually runs before the join resolves but +//! is not ordered against it, and tasks those connections detached — the +//! outbound pump, any dispatch already handed to the core — are cancelled by +//! that disposal rather than awaited. The wait is also unbounded, so a +//! connection wedged in non-yielding work blocks the caller until it +//! unwedges. //! //! Security model: the listener binds to `127.0.0.1` only, and every //! connection must present its registered per-execution 256-bit token @@ -21,21 +30,29 @@ //! duplicated `t=` parameter, so timing does not reveal which token (if any) //! matched. Revoking a token removes it from the registry before existing //! connections are aborted, so a handshake that has not yet matched a token -//! when it is revoked is rejected outright; a handshake that matched just +//! when it is revoked is rejected outright. A handshake that matched just //! before the revocation may still receive an already-committed upgrade -//! response, but is aborted immediately afterward, before any wire traffic -//! flows. Tokens are handed only to the host's embedded WebView, so the -//! bridge does not also pin the `Origin` header (the WebView's origin is not -//! known a priori). Inbound messages are size-capped, and both the -//! per-execution and process-wide outbound queue and connection counts are -//! bounded to contain a misbehaving local peer. Each accepted connection's -//! handshake runs in its own task (bounded by [`HANDSHAKE_TIMEOUT`]), so a -//! peer that never completes one only ever stalls its own connection, never -//! the shared accept loop, another execution's connections, or the -//! listener's shutdown; the connection-count caps are reserved with a -//! compare-and-swap loop precisely because handshakes now run concurrently. - -use std::collections::HashMap; +//! response, but its connection is never served: the revocation flag is read +//! under the same lock that would register the connection, and the task that +//! would read from the socket is spawned only on the branch that finds the +//! execution live, so a revoked token cannot carry a single frame. +//! +//! Tokens are handed only to the host's embedded WebView, so the bridge does +//! not also pin the `Origin` header (the WebView's origin is not known a +//! priori). Inbound messages are size-capped and each connection's outbound +//! queue is bounded, as are the per-execution and listener-wide counts of +//! admitted connections. A peer that never presents a token is invisible to +//! those counts, so the number of handshakes in flight is bounded separately +//! ([`MAX_PENDING_HANDSHAKES`]); worst-case sockets is the sum of the two +//! bounds. Each accepted connection's handshake runs in its own task +//! (bounded by [`HANDSHAKE_TIMEOUT`]) and a full handshake backlog evicts +//! its oldest entry rather than refusing the newcomer, so a peer that never +//! completes one only ever stalls its own connection, never the shared +//! accept loop, another execution's connections, or the listener's shutdown. +//! Because handshakes resolve concurrently, the connection-count caps are +//! reserved with a compare-and-swap loop rather than a read-then-increment. + +use std::collections::{HashMap, VecDeque}; use std::io; use std::net::SocketAddr; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -68,6 +85,16 @@ const MAX_WS_CONNECTIONS_PER_EXECUTION: usize = 32; /// misbehaving peer amplifying across many registered tokens. const MAX_TOTAL_WS_CONNECTIONS: usize = 64; +/// Maximum handshakes the listener will carry at once, counted separately +/// from [`MAX_TOTAL_WS_CONNECTIONS`]: a connection occupies a slot here only +/// until its handshake resolves, and one there only once a token has +/// matched. Worst-case sockets is therefore the sum of the two, not either +/// alone. The bound exists because a peer that never presents a token is +/// invisible to the connection caps, which are reserved inside the +/// handshake. Reaching it evicts the oldest handshake rather than turning +/// the newcomer away (see [`accept_loop`]). +const MAX_PENDING_HANDSHAKES: usize = 64; + /// Bound on the per-connection outbound frame queue. A peer that stops reading /// cannot make the core buffer responses without limit; once the queue fills /// the connection is treated as closed. @@ -203,12 +230,12 @@ struct RegistryEntry { } /// A connection can finish its handshake against an entry that is revoked a -/// moment later, before the accept loop gets to record its handle here. The -/// `revoked` flag and the handle list share one lock so revocation and -/// registration serialize against each other: whichever happens first is -/// what the other observes, so a connection admitted in that window is -/// always either recorded for the abort that already ran, or told to abort -/// itself immediately rather than being left running with no owner. +/// moment later. The `revoked` flag and the handle list share one lock, and +/// `connection_setup` both reads the flag and spawns the connection under a +/// single acquisition of it, so revocation and registration serialize: +/// whichever happens first is what the other observes. A revocation that +/// wins turns the connection away before any task reads from its socket; one +/// that loses finds the handle already recorded and aborts it. #[derive(Default)] struct EntryConnections { revoked: bool, @@ -239,19 +266,20 @@ impl WsBridgeRegistry { } /// Remove `token`'s entry, mark it revoked, and abort its live - /// connections. A handshake that matched this token just before the - /// removal, but has not yet registered its handle, observes the - /// `revoked` flag (set under the same lock as the abort loop below) and - /// aborts itself instead of registering. No-op for an unknown or - /// already-revoked token. - fn revoke(&self, token: &str) { + /// connections, handing their handles back so a caller can choose to + /// wait for them to actually finish unwinding. A handshake that matched + /// this token just before the removal, but has not yet registered its + /// handle, observes the `revoked` flag (set under the same lock as the + /// abort loop below) and aborts itself instead of registering. No-op + /// (empty result) for an unknown or already-revoked token. + fn revoke(&self, token: &str) -> Vec> { let Some(entry) = self .entries .lock() .expect("ws bridge registry mutex poisoned") .remove(token) else { - return; + return Vec::new(); }; let mut state = entry .connections @@ -261,6 +289,7 @@ impl WsBridgeRegistry { for handle in state.handles.iter() { handle.abort(); } + std::mem::take(&mut state.handles) } /// Find the entry whose token matches `path_and_query`'s `?t=` value. @@ -342,35 +371,59 @@ impl SharedWsBridge { runtime_factory: Arc, logger: BridgeLogger, ) -> Result { - let mut guard = self.inner.lock().expect("shared ws bridge mutex poisoned"); - if guard.is_none() { - *guard = Some(WsBridge::start(bind_port, logger.clone())?); - } else if bind_port != 0 { - let running_port = guard.as_ref().expect("just checked Some").port; - if bind_port != running_port { - logger( - "truapi.ws_bridge.bind_port_ignored", - &format!("requested={bind_port} running={running_port}"), - ); + // Every logger call below is deferred until after this lock is + // released. A host logger that blocks or calls back into + // `register`/`revoke` would otherwise stall or deadlock every other + // execution sharing this bridge, since this lock is the one thing + // every one of their `register`/`revoke` calls must take too. + let mut pending_logs: Vec<(&'static str, String)> = Vec::new(); + let endpoint = { + let mut guard = self.inner.lock().expect("shared ws bridge mutex poisoned"); + if guard.is_none() { + let (bridge, logs) = WsBridge::start(bind_port, logger.clone())?; + *guard = Some(bridge); + pending_logs = logs; + } else if bind_port != 0 { + let running_port = guard.as_ref().expect("just checked Some").port; + if bind_port != running_port { + pending_logs.push(( + "truapi.ws_bridge.bind_port_ignored", + format!("requested={bind_port} running={running_port}"), + )); + } } + guard + .as_ref() + .expect("shared bridge just inserted") + .register(runtime_factory) + }; + for (event, detail) in &pending_logs { + logger(event, detail); } - Ok(guard - .as_ref() - .expect("shared bridge just inserted") - .register(runtime_factory)) + Ok(endpoint) } /// Revoke one execution's token. No-op if the listener was never /// started or the token is unknown. + /// + /// Off the shared executor, this blocks until every one of that + /// execution's connection tasks has been joined — the same wait `stop` + /// performs for the whole bridge, scoped here to one execution. Joining + /// is not a barrier on the destructors inside those tasks, so a caller + /// cannot treat this returning as proof that every resource the + /// connection held is already released. From a task already running on + /// the shared executor, waiting is skipped to avoid deadlocking that + /// worker; the abort already happened, so those connections still unwind + /// on their own. pub fn revoke(&self, token: &str) { - if let Some(bridge) = self - .inner - .lock() - .expect("shared ws bridge mutex poisoned") - .as_ref() - { - bridge.revoke(token); - } + let wait = { + let guard = self.inner.lock().expect("shared ws bridge mutex poisoned"); + match guard.as_ref() { + Some(bridge) => bridge.revoke(token), + None => return, + } + }; + wait.block_until_finished(); } } @@ -391,7 +444,15 @@ struct WsBridge { impl WsBridge { /// Bind a localhost listener and start the accept loop on the shared /// native executor. - fn start(bind_port: u16, logger: BridgeLogger) -> io::Result { + /// + /// `logger` is still used for the accept loop's own ongoing (async, + /// lock-free) lifecycle events; the one-time startup events below are + /// instead returned for the caller to emit once it is no longer holding + /// [`SharedWsBridge::inner`]'s lock. + fn start( + bind_port: u16, + logger: BridgeLogger, + ) -> io::Result<(Self, Vec<(&'static str, String)>)> { // Bind synchronously so we can surface bind errors and discover the // actual port before returning. let std_listener = @@ -402,15 +463,6 @@ impl WsBridge { let (executor, initialized) = shared_native_executor()?; let handle = executor.handle(); let runtime_id = handle.id(); - if initialized { - logger( - "truapi.native.executor.started", - &format!( - "runtime_id={runtime_id} worker_threads={}", - executor.worker_threads() - ), - ); - } // Register the listener with the shared runtime's I/O driver before // returning so a successful start always yields a ready endpoint. @@ -418,6 +470,21 @@ impl WsBridge { let _entered = handle.enter(); TcpListener::from_std(std_listener)? }; + + // Queued only past the last fallible step. The executor is a process + // -wide `OnceLock`, so `initialized` is true for exactly one caller in + // the process; dropping its event on an error return would lose it for + // good rather than delay it. + let mut pending_logs: Vec<(&'static str, String)> = Vec::new(); + if initialized { + pending_logs.push(( + "truapi.native.executor.started", + format!( + "runtime_id={runtime_id} worker_threads={}", + executor.worker_threads() + ), + )); + } let registry = Arc::new(WsBridgeRegistry::default()); let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); let (stopped_tx, stopped_rx) = std::sync::mpsc::channel::<()>(); @@ -428,19 +495,22 @@ impl WsBridge { let _ = stopped_tx.send(()); }); - logger( + pending_logs.push(( "truapi.ws_bridge.started", - &format!("port={port} runtime_id={runtime_id}"), - ); + format!("port={port} runtime_id={runtime_id}"), + )); - Ok(Self { - shutdown: Some(shutdown_tx), - stopped: Some(stopped_rx), - accept_task: Some(accept_task), - runtime_id, - registry, - port, - }) + Ok(( + Self { + shutdown: Some(shutdown_tx), + stopped: Some(stopped_rx), + accept_task: Some(accept_task), + runtime_id, + registry, + port, + }, + pending_logs, + )) } /// Register one execution's runtime factory, minting a fresh token. @@ -455,9 +525,21 @@ impl WsBridge { } } - /// Revoke one execution's token and abort its live connections. - fn revoke(&self, token: &str) { - self.registry.revoke(token); + /// Revoke one execution's token and abort its live connections, without + /// blocking. The caller decides whether and how to wait for their tasks + /// to be joined, via the returned [`RevokeWait`]. + fn revoke(&self, token: &str) -> RevokeWait { + let connections = self.registry.revoke(token); + if connections.is_empty() { + return RevokeWait::Nothing; + } + // A task already running on this bridge's own executor must not block + // waiting on it: the connections still unwind on their own once + // aborted, so a caller here gets only a best-effort sweep. + if Handle::try_current().is_ok_and(|current| current.id() == self.runtime_id) { + return RevokeWait::Nothing; + } + RevokeWait::Handles(connections) } /// Signal the accept loop to exit and abort every tracked connection @@ -517,6 +599,42 @@ impl Drop for WsBridge { } } +/// What, if anything, a [`WsBridge::revoke`] caller should wait for. +enum RevokeWait { + /// Nothing was aborted, or waiting would block this bridge's own + /// executor. + Nothing, + /// These connections were just aborted; block until each of their tasks + /// has been joined. + Handles(Vec>), +} + +impl RevokeWait { + /// Block the calling thread until every aborted connection's task has + /// been joined, if any. Spawns a small joiner task on the shared executor + /// and blocks on a synchronous channel rather than awaiting directly, + /// since this is called from ordinary host threads with no executor of + /// their own. There is no deadline: a connection wedged in non-yielding + /// work holds the caller for as long as it stays wedged. + fn block_until_finished(self) { + let handles = match self { + RevokeWait::Nothing => return, + RevokeWait::Handles(handles) => handles, + }; + let Ok((executor, _)) = shared_native_executor() else { + return; + }; + let (done_tx, done_rx) = std::sync::mpsc::channel::<()>(); + executor.handle().spawn(async move { + for handle in handles { + let _ = handle.await; + } + let _ = done_tx.send(()); + }); + let _ = done_rx.recv(); + } +} + async fn accept_loop( listener: TcpListener, registry: Arc, @@ -525,10 +643,12 @@ async fn accept_loop( ) { // Each accepted connection's handshake runs in its own task (see // `connection_setup`) so a stalled or hostile peer never blocks another - // connection's setup, only its own. This loop just tracks those setup - // tasks so shutdown can cancel and await whichever haven't resolved yet, - // in addition to draining every connection already registered. - let mut setup_tasks: Vec> = Vec::new(); + // connection's setup, only its own. This loop tracks those setup tasks so + // shutdown can cancel and await whichever haven't resolved yet, in + // addition to draining every connection already registered, and so a full + // backlog can evict its oldest entry. Accept order is insertion order, so + // the front is the longest-running handshake. + let mut setup_tasks: VecDeque> = VecDeque::new(); loop { tokio::select! { _ = &mut shutdown => { @@ -560,9 +680,24 @@ async fn accept_loop( } }; setup_tasks.retain(|task| !task.is_finished()); + // A peer that never presents a token holds a task and a file + // descriptor for up to `HANDSHAKE_TIMEOUT` without ever + // reaching the connection caps, which are reserved inside the + // handshake. Bound that backlog by evicting its oldest entry + // rather than refusing this peer: a real localhost handshake + // resolves in well under a millisecond, so the oldest of a + // full backlog is one that has stalled, while refusing the + // newcomer would let a peer holding every slot lock the whole + // shared listener out for every execution. + if setup_tasks.len() >= MAX_PENDING_HANDSHAKES + && let Some(oldest) = setup_tasks.pop_front() + { + oldest.abort(); + logger("truapi.ws_bridge.handshake_backlog_evicted", &peer.to_string()); + } let registry = registry.clone(); let logger = logger.clone(); - setup_tasks.push(tokio::spawn(async move { + setup_tasks.push_back(tokio::spawn(async move { connection_setup(stream, peer, registry, logger).await; })); } @@ -576,13 +711,13 @@ async fn accept_loop( /// blocks only this task, never another connection's setup, another /// execution's connections, or the accept loop's ability to keep accepting. /// -/// Because handshakes now run concurrently rather than one at a time, the +/// Handshakes therefore resolve concurrently, which is why the /// connection-count caps inside `authenticate_and_upgrade`'s handshake /// callback are reserved with a compare-and-swap loop (`try_reserve`) rather -/// than a plain load-then-increment. The registration race this used to -/// avoid by serializing handshakes — a token revoked between a successful -/// handshake and its registration — is instead closed by the `revoked` flag -/// check below, which holds regardless of how many handshakes run at once. +/// than a plain load-then-increment, and why the `revoked` flag below is read +/// under the same lock that registers the connection: a token revoked while +/// this handshake was resolving must not end up with a task reading from its +/// socket, however many handshakes are in flight. async fn connection_setup( stream: tokio::net::TcpStream, peer: SocketAddr, @@ -603,24 +738,39 @@ async fn connection_setup( }) else { return; }; - let conn_logger = logger.clone(); - let conn_entry = entry.clone(); - let handle = tokio::spawn(async move { - let _guard = guard; - connection_lifecycle(ws, peer, conn_entry, conn_logger).await; - }); - let mut state = entry - .connections - .lock() - .expect("ws bridge registry entry mutex poisoned"); - state.handles.retain(|h| !h.is_finished()); - if state.revoked { - // The execution was revoked between this connection's successful - // handshake and this registration step; abort it now instead of - // leaving it running with no owner to revoke it later. - handle.abort(); - } else { - state.handles.push(handle); + // Read the revocation flag and spawn the connection under one acquisition + // of the lock `revoke` takes for its own flag-and-abort pass, so the two + // serialize: an execution revoked while this handshake was resolving never + // gets a task reading from the socket, and one still live cannot be + // revoked between the spawn and the handle landing in its list, where it + // would be left running with no owner. `tokio::spawn` only queues the + // task, so nothing host-supplied runs under the lock here. + let revoked = { + let mut state = entry + .connections + .lock() + .expect("ws bridge registry entry mutex poisoned"); + state.handles.retain(|h| !h.is_finished()); + if state.revoked { + true + } else { + let conn_logger = logger.clone(); + let conn_entry = entry.clone(); + state.handles.push(tokio::spawn(async move { + let _guard = guard; + connection_lifecycle(ws, peer, conn_entry, conn_logger).await; + })); + false + } + }; + // Logged after the lock is released: the logger is a host callback, and one + // that blocks or re-enters would otherwise stall every `revoke` and + // shutdown pass that needs this entry's lock. + if revoked { + logger( + "truapi.ws_bridge.connection_revoked_during_setup", + &peer.to_string(), + ); } } @@ -735,8 +885,8 @@ async fn authenticate_and_upgrade( // Reserve both caps' slots here, inside the synchronous handshake // callback, so an over-cap peer is rejected at the HTTP upgrade // rather than allowed to open a socket that gets dropped right - // after. Handshakes now run concurrently (one task per connection), - // so each reservation is a compare-and-swap, not a plain increment. + // after. Each reservation is a compare-and-swap because one task per + // connection means concurrent reservations against the same counter. if !try_reserve(&entry.connection_count, MAX_WS_CONNECTIONS_PER_EXECUTION) { auth_logger( "truapi.ws_bridge.connection_limit_execution", @@ -1065,7 +1215,7 @@ mod tests { #[test] fn drop_from_shared_executor_does_not_block_worker() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; let (executor, _) = shared_native_executor().expect("shared native executor"); let (dropped_tx, dropped_rx) = std::sync::mpsc::channel(); @@ -1085,7 +1235,7 @@ mod tests { /// `feature_supported` response. #[test] fn round_trip_feature_supported_through_bridge() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; let endpoint = bridge.register(test_runtime_factory()); let url = format!("ws://127.0.0.1:{}/?t={}", endpoint.port, endpoint.token); @@ -1154,7 +1304,7 @@ mod tests { /// execution's runtime, never the other's. #[test] fn two_executions_share_one_port_with_isolated_tokens() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; let first = bridge.register(test_runtime_factory()); let second = bridge.register(test_runtime_factory()); @@ -1173,7 +1323,7 @@ mod tests { /// rejected outright. #[test] fn wrong_or_unknown_token_is_rejected_at_handshake() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; let endpoint = bridge.register(test_runtime_factory()); let _second = bridge.register(test_runtime_factory()); @@ -1198,9 +1348,45 @@ mod tests { /// Revoking one execution's token closes only its own connections and /// rejects future handshakes against it, while a sibling execution on /// the same shared listener keeps working. + /// A peer holding every handshake slot open must not be able to lock a + /// legitimate connection out of the shared listener: a full backlog evicts + /// its oldest entry instead of refusing the newcomer. + #[test] + fn a_full_handshake_backlog_does_not_lock_out_a_new_connection() { + let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; + let endpoint = bridge.register(test_runtime_factory()); + + // Raw TCP connections that never send a byte, so each one occupies a + // handshake slot until it is evicted or times out. + let addr = format!("127.0.0.1:{}", endpoint.port); + let mut stalled = Vec::new(); + for _ in 0..MAX_PENDING_HANDSHAKES { + stalled.push(std::net::TcpStream::connect(&addr).expect("stall the backlog")); + } + + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime"); + let url = format!("ws://127.0.0.1:{}/?t={}", endpoint.port, endpoint.token); + rt.block_on(async { + let (mut ws, _) = tokio::time::timeout( + std::time::Duration::from_secs(10), + tokio_tungstenite::connect_async(&url), + ) + .await + .expect("a full backlog must not stall a new connection") + .expect("dial past a full backlog"); + ws.close(None).await.expect("close client"); + }); + + drop(stalled); + drop(bridge); + } + #[test] fn revoking_one_token_leaves_another_operational() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; let revoked = bridge.register(test_runtime_factory()); let survives = bridge.register(test_runtime_factory()); @@ -1263,7 +1449,7 @@ mod tests { /// of the old one. #[test] fn reconnecting_after_revoke_gets_a_fresh_token() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; let first = bridge.register(test_runtime_factory()); bridge.revoke(&first.token); @@ -1290,7 +1476,7 @@ mod tests { /// registered execution's connections, not just one. #[test] fn host_shutdown_closes_every_registered_execution() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; let first = bridge.register(test_runtime_factory()); let second = bridge.register(test_runtime_factory()); @@ -1357,7 +1543,7 @@ mod tests { /// though the shared listener's own total cap has plenty of room left. #[test] fn per_execution_cap_rejects_the_connection_past_the_limit() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; let endpoint = bridge.register(test_runtime_factory()); let url = format!("ws://127.0.0.1:{}/?t={}", endpoint.port, endpoint.token); @@ -1400,7 +1586,7 @@ mod tests { /// this exercises the exact same check with far less real I/O. #[test] fn total_cap_rejects_the_connection_even_for_a_fresh_execution() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; let extra = bridge.register(test_runtime_factory()); bridge .registry @@ -1441,7 +1627,7 @@ mod tests { panic!("intentional test panic: simulating a failing product execution") }); - let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; let failing = bridge.register(panicking_factory); let healthy = bridge.register(test_runtime_factory()); @@ -1490,7 +1676,7 @@ mod tests { /// blocked every other execution's connections too. #[test] fn a_stalled_handshake_does_not_block_a_sibling_connection() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; let stalled = bridge.register(test_runtime_factory()); let healthy = bridge.register(test_runtime_factory()); @@ -1537,7 +1723,7 @@ mod tests { }) } - let bridge = WsBridge::start(0, no_log()).expect("start bridge"); + let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; let calls: Vec> = (0..3).map(|_| Arc::new(AtomicUsize::new(0))).collect(); let endpoints: Vec = calls .iter() From 66faa0cfe749efb1c7d264e102307c893e8104a7 Mon Sep 17 00:00:00 2001 From: Nidish Date: Tue, 15 Sep 2026 15:26:15 +0530 Subject: [PATCH 3/7] fix(truapi-server): close the dispose race against an in-flight dispatch --- rust/crates/truapi-server/src/host_core.rs | 115 +++++++++++++++++++-- 1 file changed, 104 insertions(+), 11 deletions(-) diff --git a/rust/crates/truapi-server/src/host_core.rs b/rust/crates/truapi-server/src/host_core.rs index 709bced9b..e7e5d5dc7 100644 --- a/rust/crates/truapi-server/src/host_core.rs +++ b/rust/crates/truapi-server/src/host_core.rs @@ -1485,17 +1485,18 @@ impl ProductRuntime { // first is what the other observes: a dispatch that inserted before this // drain is caught by it, and one that hasn't inserted yet sees `disposed` // already true and turns itself away instead of dispatching past disposal. - let mut in_flight = self - .in_flight - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if self.disposed.swap(true, Ordering::AcqRel) { - return; - } - for (_, handle) in in_flight.drain() { - handle.abort(); + { + let mut in_flight = self + .in_flight + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if self.disposed.swap(true, Ordering::AcqRel) { + return; + } + for (_, handle) in in_flight.drain() { + handle.abort(); + } } - drop(in_flight); self.admin.product_runtime.detach_chat(); self.admin.product_runtime.detach_renderer(); self.host_subscriptions.close(); @@ -1611,7 +1612,7 @@ impl Transport for SinkTransport { #[cfg(test)] mod tests { use super::*; - use crate::frame::{Payload, ProtocolMessage, subscription_ids}; + use crate::frame::{Payload, ProtocolMessage, request_ids, subscription_ids}; use crate::host_logic::product_account::derive_identity_keypair; use crate::host_logic::sso::messages::{ RemoteMessage, RemoteMessageData, decode_incoming_sso_request, v1, @@ -2746,6 +2747,98 @@ mod tests { assert_eq!(response.payload.value, expected); } + /// A dispatch that passes `receive_frame`'s opening disposed check must not + /// run if `dispose` commits before it registers itself. The debug tap is the + /// only seam that can park a frame between those two points; the sink here + /// blocks on purpose, which its own contract forbids, so that the window is + /// reachable deterministically instead of by chance. + #[test] + fn a_dispatch_racing_dispose_does_not_reach_the_platform() { + struct ParkingDebugSink { + entered: std::sync::mpsc::SyncSender<()>, + release: Mutex>>, + } + + impl DebugSink for ParkingDebugSink { + fn emit(&self, _event: DebugEvent) { + let _ = self.entered.try_send(()); + if let Some(release) = self + .release + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take() + { + let _ = release.recv(); + } + } + } + + let navigations = Arc::new(Mutex::new(Vec::new())); + let platform = Arc::new(StubPlatform { + navigations: navigations.clone(), + ..Default::default() + }); + let (host_config, product) = runtime_config("myapp.dot"); + let runtime = Arc::new(ProductRuntime::from_platform_with_config( + platform, + host_config, + product, + test_spawner(), + Arc::new(RecordingSink::default()), + )); + + let (entered_tx, entered_rx) = std::sync::mpsc::sync_channel::<()>(1); + let (release_tx, release_rx) = std::sync::mpsc::channel::<()>(); + runtime.set_debug_sink( + ChannelId("race".to_string()), + Arc::new(ParkingDebugSink { + entered: entered_tx, + release: Mutex::new(Some(release_rx)), + }), + ); + + let ids = request_ids("system_navigate_to").expect("known request method"); + let frame = ProtocolMessage { + request_id: "nav:1".to_string(), + payload: Payload { + trait_id: ids.trait_id, + method_id: ids.method_id, + message_type: crate::frame::MESSAGE_TYPE_REQUEST, + value: truapi::versioned::system::HostNavigateToRequest::V1( + v01::HostNavigateToRequest { + url: "https://example.invalid/".to_string(), + }, + ) + .encode(), + }, + } + .encode(); + + let dispatching = { + let runtime = runtime.clone(); + std::thread::spawn(move || { + futures::executor::block_on(runtime.receive_frame(frame)).expect("receive frame"); + }) + }; + + // The frame is now parked inside the tap, past the opening disposed + // check and before it has registered itself. + entered_rx + .recv_timeout(std::time::Duration::from_secs(5)) + .expect("frame never reached the debug tap"); + runtime.dispose(); + let _ = release_tx.send(()); + dispatching.join().expect("dispatch thread panicked"); + + assert!( + navigations + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .is_empty(), + "a dispatch that lost the race with dispose still reached the platform" + ); + } + #[test] fn dispose_cancels_active_subscriptions() { let theme_stream_dropped = Arc::new(AtomicBool::new(false)); From 1bb9ba67f360e7e405ff4d51fe6ea0313f5a2355 Mon Sep 17 00:00:00 2001 From: Nidish Date: Tue, 15 Sep 2026 15:26:15 +0530 Subject: [PATCH 4/7] fix(truapi-server): dispose connections on close and bound each execution's share --- android/truapi-host/README.md | 2 +- rust/crates/truapi-server/src/native.rs | 4 +- rust/crates/truapi-server/src/ws_bridge.rs | 418 ++++++++++++--------- 3 files changed, 237 insertions(+), 187 deletions(-) diff --git a/android/truapi-host/README.md b/android/truapi-host/README.md index b98053a2d..627c42ff0 100644 --- a/android/truapi-host/README.md +++ b/android/truapi-host/README.md @@ -58,7 +58,7 @@ The public surface lives in [`src/main/kotlin/io/parity/truapi/TrUAPIHost.kt`](s - `HostStorage` - product-scoped read/write/clear interface the host backs with its own persistence. - `HostCoreStorage` - core-owned read/write/clear interface for auth session, pairing identity, and persisted permission decisions (`key` is a SCALE-encoded `CoreStorageKey`). - `LocalhostBridgeBootstrap` - JS snippet that publishes the WS bridge endpoint (`window.__truapi_localhost`) to the product page so it can dial back in. -- `TrUAPIHostRuntime` - process-owned runtime whose product executions share one authentication session. Open a connection per executable with `openProductExecution`, which returns a `TrUAPIProductExecution` carrying that connection's own WS bridge, permission authorization, theme/preimage/chain notifications, and the Chat controls below. +- `TrUAPIHostRuntime` - process-owned runtime whose product executions share one authentication session. Open a connection per executable with `openProductExecution`, which returns a `TrUAPIProductExecution` holding its own token on the runtime's shared WS bridge, permission authorization, theme/preimage/chain notifications, and the Chat controls below. - `ChatHostBridge` - native Chat storage and UI, implemented by hosts that serve the Chat modality and passed to `openProductExecution`. Hosts without it pass nothing and Chat calls answer unsupported. - `PocketHostBridge` - the host's Pocket card collection, implemented by hosts with a Pocket surface and passed as `pocket` to `openProductExecution`. The execution then offers `notifyPocketCardsChanged`. `removeCard` decides and removes together, returning `NativePocketRemoval.Removed`, `Absent` or `Privileged`, so a card cannot be pinned between the check and the removal. Like Chat, Pocket is reachable only from a Worker execution with an active session, so without `activateLocalSession` every Pocket call answers `Denied`. Hosts without the bridge pass nothing and Pocket calls answer unsupported. diff --git a/rust/crates/truapi-server/src/native.rs b/rust/crates/truapi-server/src/native.rs index 5dd42e763..6341ec6f3 100644 --- a/rust/crates/truapi-server/src/native.rs +++ b/rust/crates/truapi-server/src/native.rs @@ -4081,7 +4081,9 @@ mod tests { let round_trip = |request_id: &str| ProtocolMessage { request_id: request_id.into(), payload: Payload { - id: feature_ids.request_id, + trait_id: feature_ids.trait_id, + method_id: feature_ids.method_id, + message_type: crate::frame::MESSAGE_TYPE_REQUEST, value: HostFeatureSupportedRequest::V1(v01::HostFeatureSupportedRequest::Chain { genesis_hash: vec![0u8; 32], }) diff --git a/rust/crates/truapi-server/src/ws_bridge.rs b/rust/crates/truapi-server/src/ws_bridge.rs index 1c365b921..90da8ba14 100644 --- a/rust/crates/truapi-server/src/ws_bridge.rs +++ b/rust/crates/truapi-server/src/ws_bridge.rs @@ -5,52 +5,30 @@ //! //! Feature-gated (`ws-bridge`) so wasm32 and no-tokio build paths stay lean. //! -//! Native bridges share one process-wide `tokio` runtime, and every product -//! execution under one host runtime shares a single [`SharedWsBridge`] -//! listener: [`SharedWsBridge::register`] hands each execution its own -//! `{port, token}` endpoint on the one shared port, and -//! [`SharedWsBridge::revoke`] tears down only that execution's connections -//! when it closes, leaving the listener and every other execution's -//! connections untouched. Off the shared executor, `revoke` blocks until -//! every connection it aborted has been joined, matching what -//! [`WsBridge::stop`] gives for the whole listener. Joining a task is not a -//! barrier on everything inside it: the destructor that disposes a -//! connection's `ProductRuntime` usually runs before the join resolves but -//! is not ordered against it, and tasks those connections detached — the -//! outbound pump, any dispatch already handed to the core — are cancelled by -//! that disposal rather than awaited. The wait is also unbounded, so a -//! connection wedged in non-yielding work blocks the caller until it -//! unwedges. +//! Every product execution under one host runtime shares a single +//! [`SharedWsBridge`] listener on the process-wide `tokio` runtime. +//! [`SharedWsBridge::register`] hands each execution its own `{port, token}` +//! endpoint on that one port; [`SharedWsBridge::revoke`] tears down only that +//! execution's connections, leaving the listener and its siblings untouched. //! -//! Security model: the listener binds to `127.0.0.1` only, and every -//! connection must present its registered per-execution 256-bit token -//! (`?t=`, drawn from the OS CSPRNG) before the WebSocket upgrade -//! completes. The handshake scans every currently registered token with a -//! constant-time comparison and does not exit early on a match or on a -//! duplicated `t=` parameter, so timing does not reveal which token (if any) -//! matched. Revoking a token removes it from the registry before existing -//! connections are aborted, so a handshake that has not yet matched a token -//! when it is revoked is rejected outright. A handshake that matched just -//! before the revocation may still receive an already-committed upgrade -//! response, but its connection is never served: the revocation flag is read -//! under the same lock that would register the connection, and the task that -//! would read from the socket is spawned only on the branch that finds the -//! execution live, so a revoked token cannot carry a single frame. +//! Security model: the listener binds `127.0.0.1` only, and every connection +//! must present its execution's 256-bit token (`?t=`, from the OS +//! CSPRNG) before the upgrade completes. The handshake scans every registered +//! token with a constant-time comparison and exits early on neither a match +//! nor a duplicated `t=` pair, so timing reveals nothing about which token +//! matched. Revocation removes the token before aborting connections, and the +//! revoked flag is read under the same lock that registers a connection, so a +//! revoked token can never carry a frame. Tokens go only to the host's +//! embedded WebView, whose origin is not known a priori, so the `Origin` +//! header is not pinned. //! -//! Tokens are handed only to the host's embedded WebView, so the bridge does -//! not also pin the `Origin` header (the WebView's origin is not known a -//! priori). Inbound messages are size-capped and each connection's outbound -//! queue is bounded, as are the per-execution and listener-wide counts of -//! admitted connections. A peer that never presents a token is invisible to -//! those counts, so the number of handshakes in flight is bounded separately -//! ([`MAX_PENDING_HANDSHAKES`]); worst-case sockets is the sum of the two -//! bounds. Each accepted connection's handshake runs in its own task -//! (bounded by [`HANDSHAKE_TIMEOUT`]) and a full handshake backlog evicts -//! its oldest entry rather than refusing the newcomer, so a peer that never -//! completes one only ever stalls its own connection, never the shared -//! accept loop, another execution's connections, or the listener's shutdown. -//! Because handshakes resolve concurrently, the connection-count caps are -//! reserved with a compare-and-swap loop rather than a read-then-increment. +//! Bounds, all to contain a misbehaving local peer: inbound message size, each +//! connection's outbound queue, admitted connections per execution and +//! listener-wide, and — since a peer that never presents a token reaches none +//! of those — the number of handshakes in flight +//! ([`MAX_PENDING_HANDSHAKES`]). Worst-case sockets is the sum of the last +//! two. A full handshake backlog evicts its oldest entry rather than refusing +//! the newcomer, so one stalled peer cannot lock the listener. use std::collections::{HashMap, VecDeque}; use std::io; @@ -74,10 +52,16 @@ use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; use crate::{FrameSink, ProductRuntime}; /// Maximum simultaneous connections a single registered execution may hold. -/// Each execution uses exactly one connection; the cap bounds resource use -/// from a buggy or hostile local peer opening many sockets against one -/// token. -const MAX_WS_CONNECTIONS_PER_EXECUTION: usize = 32; +/// An execution uses one connection; the headroom absorbs reconnect churn, +/// and the cap bounds a buggy or hostile local peer opening many sockets +/// against one token. +/// +/// Kept low enough that +/// `MAX_TOTAL_WS_CONNECTIONS / MAX_WS_CONNECTIONS_PER_EXECUTION` executions +/// can each hold their full allowance at once. Were the per-execution cap +/// above that share, one execution could take enough of the shared budget to +/// starve its siblings, which a per-execution listener could not do. +const MAX_WS_CONNECTIONS_PER_EXECUTION: usize = 8; /// Maximum simultaneous connections across every execution sharing the /// listener. Set well above any realistic concurrent-execution count (App, @@ -129,7 +113,8 @@ pub struct WsBridgeEndpoint { #[derive(Debug, thiserror::Error, uniffi::Error)] #[uniffi(flat_error)] pub enum WsBridgeStartError { - /// A bridge is already running for this host. + /// This execution already registered a bridge token. The host's shared + /// listener being up is the normal case and not an error. #[error("ws bridge already running")] AlreadyRunning, /// Anything else (bind failure, runtime spin-up failure, ...). @@ -242,6 +227,17 @@ struct EntryConnections { handles: Vec>, } +impl EntryConnections { + /// Request cancellation of every tracked connection and hand back sole + /// ownership of their handles, so a caller can join what it aborted. + fn abort_and_take(&mut self) -> Vec> { + for handle in self.handles.iter() { + handle.abort(); + } + std::mem::take(&mut self.handles) + } +} + /// Token registry shared by the listener's accept loop and every /// [`SharedWsBridge::register`]/[`SharedWsBridge::revoke`] call. #[derive(Default)] @@ -286,22 +282,25 @@ impl WsBridgeRegistry { .lock() .expect("ws bridge registry entry mutex poisoned"); state.revoked = true; - for handle in state.handles.iter() { - handle.abort(); - } - std::mem::take(&mut state.handles) + state.abort_and_take() } /// Find the entry whose token matches `path_and_query`'s `?t=` value. /// Scans every registered token without exiting early on a match, so /// timing does not reveal which one (if any) matched. fn find_matching(&self, path_and_query: Option<&str>) -> Option> { - let entries = self + // Snapshot first: the scan below runs over peer-supplied request bytes + // once per registered token, and holding the registry lock across it + // would stall every concurrent handshake, registration and revocation. + let candidates: Vec<(String, Arc)> = self .entries .lock() - .expect("ws bridge registry mutex poisoned"); + .expect("ws bridge registry mutex poisoned") + .iter() + .map(|(token, entry)| (token.clone(), entry.clone())) + .collect(); let mut found = None; - for (token, entry) in entries.iter() { + for (token, entry) in candidates.iter() { if path_token_matches(path_and_query, token) { found = Some(entry.clone()); } @@ -309,17 +308,10 @@ impl WsBridgeRegistry { found } - /// Abort and drain every connection tracked under a still-registered - /// execution, returning the owned handles so the caller can await each - /// one to genuine completion rather than just requesting cancellation. - /// Safe to call more than once (or after `revoke` already drained some) - /// — later calls simply find nothing left to take for whatever was - /// already drained. A connection whose token was revoked (and thus whose - /// entry was already removed from the registry) moments before this - /// call is not covered here — `revoke` already aborted it directly, but - /// this method has no way to find and await it, so a caller cannot treat - /// its own completion as proof that connection has actually finished - /// unwinding too. + /// Abort and drain every connection under a still-registered execution, + /// handing back the owned handles so the caller can await them. Idempotent. + /// A connection revoked moments earlier is not covered: `revoke` removed + /// its entry and aborted it directly, so this cannot find it to await. fn take_all_handles(&self) -> Vec> { let entries = self .entries @@ -331,10 +323,7 @@ impl WsBridgeRegistry { .connections .lock() .expect("ws bridge registry entry mutex poisoned"); - for handle in &state.handles { - handle.abort(); - } - all.append(&mut state.handles); + all.append(&mut state.abort_and_take()); } all } @@ -406,33 +395,32 @@ impl SharedWsBridge { /// Revoke one execution's token. No-op if the listener was never /// started or the token is unknown. /// - /// Off the shared executor, this blocks until every one of that - /// execution's connection tasks has been joined — the same wait `stop` - /// performs for the whole bridge, scoped here to one execution. Joining - /// is not a barrier on the destructors inside those tasks, so a caller - /// cannot treat this returning as proof that every resource the - /// connection held is already released. From a task already running on - /// the shared executor, waiting is skipped to avoid deadlocking that - /// worker; the abort already happened, so those connections still unwind - /// on their own. + /// Off the shared executor this blocks until that execution's connection + /// tasks have been joined, the same wait `stop` performs listener-wide. + /// Joining is not a barrier on the destructors inside those tasks. On the + /// shared executor the wait is skipped to avoid deadlocking that worker. pub fn revoke(&self, token: &str) { - let wait = { + let aborted = { let guard = self.inner.lock().expect("shared ws bridge mutex poisoned"); match guard.as_ref() { Some(bridge) => bridge.revoke(token), None => return, } }; - wait.block_until_finished(); + join_aborted_connections(aborted); } } /// Running listener handle. Drop or call [`WsBridge::stop`] to shut down. /// +/// Crate-internal: hosts reach the bridge through [`SharedWsBridge`], which +/// owns the listener and hands out per-execution endpoints. Deliberately not +/// re-exported by the crate root. +/// /// The listener's tasks run on the process-wide native executor. TrUAPI /// dispatch futures are `Send`, so connections and independent frames from /// all registered executions can execute across the shared worker pool. -struct WsBridge { +pub(crate) struct WsBridge { shutdown: Option>, stopped: Option>, accept_task: Option>, @@ -525,21 +513,16 @@ impl WsBridge { } } - /// Revoke one execution's token and abort its live connections, without - /// blocking. The caller decides whether and how to wait for their tasks - /// to be joined, via the returned [`RevokeWait`]. - fn revoke(&self, token: &str) -> RevokeWait { + /// Revoke one execution's token and abort its live connections, returning + /// the tasks a caller may join. Empty when there was nothing to abort, or + /// when the caller is itself running on this bridge's executor and must + /// not block on it: those connections still unwind on their own. + fn revoke(&self, token: &str) -> Vec> { let connections = self.registry.revoke(token); - if connections.is_empty() { - return RevokeWait::Nothing; - } - // A task already running on this bridge's own executor must not block - // waiting on it: the connections still unwind on their own once - // aborted, so a caller here gets only a best-effort sweep. if Handle::try_current().is_ok_and(|current| current.id() == self.runtime_id) { - return RevokeWait::Nothing; + return Vec::new(); } - RevokeWait::Handles(connections) + connections } /// Signal the accept loop to exit and abort every tracked connection @@ -555,15 +538,10 @@ impl WsBridge { let _ = tx.send(()); } - // UniFFI hosts call stop synchronously from outside Rust's executor, - // where waiting preserves the existing "fully stopped on return" - // behavior. Avoid blocking if a Rust caller drops the bridge from one - // of the shared runtime's own workers, especially on a single-core - // runtime, where blocking here could deadlock against the very task - // this is waiting on. Sending on `shutdown` only schedules the accept - // loop to be re-polled, so on this path `stop` can return before the - // accept loop has even observed it, let alone drained anything — the - // fallback sweep below is what still cleans up whatever it can see. + // Blocking is safe only off the shared executor; on one of its own + // workers it could deadlock against the task being waited on. Skipping + // the wait means `stop` can return before the accept loop has even + // observed the signal, which is what the sweep below is for. let called_from_shared_executor = Handle::try_current().is_ok_and(|handle| handle.id() == self.runtime_id); let stopped_cleanly = if called_from_shared_executor { @@ -581,14 +559,9 @@ impl WsBridge { { task.abort(); } - // Fallback sweep: off the shared executor, the accept loop's own - // shutdown branch already drained and awaited every connection - // before signaling `stopped`, so this finds nothing left. It only - // does real work on the shared-executor fast path above (which - // skips that wait) or if the accept task had to be force-aborted - // (e.g. a panic inside the loop before it reached its own shutdown - // branch) — in both cases this only aborts what it finds, it does - // not await it, since `stop` itself is not async. + // Finds nothing when the accept loop drained on its own; does real work + // only on the non-blocking path above or after a force-abort. Aborts + // without awaiting, `stop` not being async. drop(self.registry.take_all_handles()); } } @@ -599,40 +572,26 @@ impl Drop for WsBridge { } } -/// What, if anything, a [`WsBridge::revoke`] caller should wait for. -enum RevokeWait { - /// Nothing was aborted, or waiting would block this bridge's own - /// executor. - Nothing, - /// These connections were just aborted; block until each of their tasks - /// has been joined. - Handles(Vec>), -} - -impl RevokeWait { - /// Block the calling thread until every aborted connection's task has - /// been joined, if any. Spawns a small joiner task on the shared executor - /// and blocks on a synchronous channel rather than awaiting directly, - /// since this is called from ordinary host threads with no executor of - /// their own. There is no deadline: a connection wedged in non-yielding - /// work holds the caller for as long as it stays wedged. - fn block_until_finished(self) { - let handles = match self { - RevokeWait::Nothing => return, - RevokeWait::Handles(handles) => handles, - }; - let Ok((executor, _)) = shared_native_executor() else { - return; - }; - let (done_tx, done_rx) = std::sync::mpsc::channel::<()>(); - executor.handle().spawn(async move { - for handle in handles { - let _ = handle.await; - } - let _ = done_tx.send(()); - }); - let _ = done_rx.recv(); +/// Block the calling thread until every aborted connection's task has been +/// joined. Spawns a small joiner task on the shared executor and blocks on a +/// synchronous channel rather than awaiting directly, since callers are +/// ordinary host threads with no executor of their own. There is no deadline: +/// a connection wedged in non-yielding work holds the caller until it unwedges. +fn join_aborted_connections(handles: Vec>) { + if handles.is_empty() { + return; } + let Ok((executor, _)) = shared_native_executor() else { + return; + }; + let (done_tx, done_rx) = std::sync::mpsc::channel::<()>(); + executor.handle().spawn(async move { + for handle in handles { + let _ = handle.await; + } + let _ = done_tx.send(()); + }); + let _ = done_rx.recv(); } async fn accept_loop( @@ -819,26 +778,16 @@ type AuthenticatedConnection = ( type MatchedReservation = Arc, ConnectionCountGuard)>>>; /// Atomically reserve one slot by incrementing `counter` unless it is already -/// at `limit`, retrying under contention. Connection setup now runs -/// concurrently (one task per accepted connection), so this cannot be a -/// plain load-then-increment: two handshakes could otherwise both observe -/// room for the last slot and both take it. +/// at `limit`, retrying under contention. One task per accepted connection +/// means setups reserve concurrently, so this cannot be a plain +/// load-then-increment: two handshakes could otherwise both observe room for +/// the last slot and both take it. fn try_reserve(counter: &AtomicUsize, limit: usize) -> bool { - let mut current = counter.load(Ordering::Acquire); - loop { - if current >= limit { - return false; - } - match counter.compare_exchange_weak( - current, - current + 1, - Ordering::AcqRel, - Ordering::Acquire, - ) { - Ok(_) => return true, - Err(actual) => current = actual, - } - } + counter + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + (current < limit).then_some(current + 1) + }) + .is_ok() } /// Complete the WebSocket handshake, resolving it against the registry to @@ -959,7 +908,7 @@ async fn connection_lifecycle( let (out_tx, mut out_rx) = mpsc::channel::>(OUTBOUND_QUEUE_CAP); let frame_sink = Arc::new(WsFrameSink::new(out_tx)); let product_runtime = Arc::new(entry.runtime_factory.product_runtime(frame_sink)); - let _dispose_guard = DisposeGuard(product_runtime.clone()); + let dispose_guard = DisposeGuard(product_runtime.clone()); let pump_logger = logger.clone(); let pump = tokio::spawn(async move { @@ -1012,12 +961,18 @@ async fn connection_lifecycle( } // The connection is gone: cancel in-flight dispatches so long-pending - // handlers unwind instead of outliving the connection. `_dispose_guard` - // disposes `product_runtime` when it drops at the end of this function. + // handlers unwind instead of outliving the connection. for task in &in_flight { task.abort(); } + // Dispose and release the runtime before awaiting the pump. The pump ends + // when the last outbound sender drops, and the runtime owns one through its + // frame sink, so holding the runtime here while waiting for the pump would + // wait on something only this function's own return can cause. + drop(dispose_guard); + drop(product_runtime); + let _ = pump.await; logger("truapi.ws_bridge.connection_closed", &peer.to_string()); } @@ -1104,6 +1059,12 @@ mod tests { use crate::frame::{Payload, ProtocolMessage, request_ids}; use crate::test_support::{StubPlatform, test_spawner}; + /// A started bridge with logging discarded, which is all any test here + /// wants from `WsBridge::start`'s deferred-log return. + fn start_test_bridge() -> WsBridge { + WsBridge::start(0, no_log()).expect("start bridge").0 + } + fn test_runtime_factory() -> Arc { runtime_factory_for(Arc::new(StubPlatform::default())) } @@ -1215,7 +1176,7 @@ mod tests { #[test] fn drop_from_shared_executor_does_not_block_worker() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; + let bridge = start_test_bridge(); let (executor, _) = shared_native_executor().expect("shared native executor"); let (dropped_tx, dropped_rx) = std::sync::mpsc::channel(); @@ -1235,7 +1196,7 @@ mod tests { /// `feature_supported` response. #[test] fn round_trip_feature_supported_through_bridge() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; + let bridge = start_test_bridge(); let endpoint = bridge.register(test_runtime_factory()); let url = format!("ws://127.0.0.1:{}/?t={}", endpoint.port, endpoint.token); @@ -1304,7 +1265,7 @@ mod tests { /// execution's runtime, never the other's. #[test] fn two_executions_share_one_port_with_isolated_tokens() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; + let bridge = start_test_bridge(); let first = bridge.register(test_runtime_factory()); let second = bridge.register(test_runtime_factory()); @@ -1323,7 +1284,7 @@ mod tests { /// rejected outright. #[test] fn wrong_or_unknown_token_is_rejected_at_handshake() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; + let bridge = start_test_bridge(); let endpoint = bridge.register(test_runtime_factory()); let _second = bridge.register(test_runtime_factory()); @@ -1348,12 +1309,94 @@ mod tests { /// Revoking one execution's token closes only its own connections and /// rejects future handshakes against it, while a sibling execution on /// the same shared listener keeps working. + /// Hands out the real runtime while keeping the control handle for the + /// connection it was built for, so a test can ask whether that connection's + /// runtime has been disposed. + struct DisposalWatchFactory { + inner: Arc, + control: Mutex>, + } + + impl WsProductRuntimeFactory for DisposalWatchFactory { + fn product_runtime(&self, sink: Arc) -> ProductRuntime { + let runtime = self.inner.product_runtime(sink); + *self.control.lock().expect("disposal watch mutex poisoned") = Some(runtime.control()); + runtime + } + } + + /// Ending a connection disposes the runtime that served it, so the host-core + /// subscriptions and chat state it held are released rather than left live + /// for the rest of the listener's life. + #[test] + fn ending_a_connection_disposes_its_runtime() { + fn is_closed(control: &crate::ProductRuntimeControl) -> bool { + matches!( + control.publish_chat_action(v01::HostChatActionSubscribeItem { + room_id: "support".into(), + peer: "dotli.dot".into(), + payload: v01::ChatActionPayload::ActionTriggered(v01::ActionTrigger { + message_id: "message".into(), + action_id: "vote".into(), + payload: None, + }), + }), + Err(crate::ProductRuntimeError::Closed) + ) + } + + let bridge = start_test_bridge(); + let watch = Arc::new(DisposalWatchFactory { + inner: test_runtime_factory(), + control: Mutex::new(None), + }); + let endpoint = bridge.register(watch.clone()); + + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime"); + let url = format!("ws://127.0.0.1:{}/?t={}", endpoint.port, endpoint.token); + + let control = rt.block_on(async { + let (mut ws, _) = tokio_tungstenite::connect_async(&url).await.expect("dial"); + let control = loop { + if let Some(control) = watch + .control + .lock() + .expect("disposal watch mutex poisoned") + .clone() + { + break control; + } + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + }; + assert!( + !is_closed(&control), + "the runtime should be live while the connection is open" + ); + ws.close(None).await.expect("close client"); + control + }); + + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); + while !is_closed(&control) { + assert!( + std::time::Instant::now() < deadline, + "the connection ended without disposing its runtime" + ); + std::thread::sleep(std::time::Duration::from_millis(5)); + } + + drop(bridge); + } + /// A peer holding every handshake slot open must not be able to lock a /// legitimate connection out of the shared listener: a full backlog evicts /// its oldest entry instead of refusing the newcomer. #[test] fn a_full_handshake_backlog_does_not_lock_out_a_new_connection() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; + let bridge = start_test_bridge(); let endpoint = bridge.register(test_runtime_factory()); // Raw TCP connections that never send a byte, so each one occupies a @@ -1384,9 +1427,12 @@ mod tests { drop(bridge); } + /// Revoking one execution's token tears down that execution's live + /// connection and stops its token authenticating, while a sibling + /// execution on the same listener keeps working. #[test] fn revoking_one_token_leaves_another_operational() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; + let bridge = start_test_bridge(); let revoked = bridge.register(test_runtime_factory()); let survives = bridge.register(test_runtime_factory()); @@ -1449,7 +1495,7 @@ mod tests { /// of the old one. #[test] fn reconnecting_after_revoke_gets_a_fresh_token() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; + let bridge = start_test_bridge(); let first = bridge.register(test_runtime_factory()); bridge.revoke(&first.token); @@ -1476,7 +1522,7 @@ mod tests { /// registered execution's connections, not just one. #[test] fn host_shutdown_closes_every_registered_execution() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; + let bridge = start_test_bridge(); let first = bridge.register(test_runtime_factory()); let second = bridge.register(test_runtime_factory()); @@ -1543,7 +1589,7 @@ mod tests { /// though the shared listener's own total cap has plenty of room left. #[test] fn per_execution_cap_rejects_the_connection_past_the_limit() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; + let bridge = start_test_bridge(); let endpoint = bridge.register(test_runtime_factory()); let url = format!("ws://127.0.0.1:{}/?t={}", endpoint.port, endpoint.token); @@ -1586,7 +1632,7 @@ mod tests { /// this exercises the exact same check with far less real I/O. #[test] fn total_cap_rejects_the_connection_even_for_a_fresh_execution() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; + let bridge = start_test_bridge(); let extra = bridge.register(test_runtime_factory()); bridge .registry @@ -1627,7 +1673,7 @@ mod tests { panic!("intentional test panic: simulating a failing product execution") }); - let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; + let bridge = start_test_bridge(); let failing = bridge.register(panicking_factory); let healthy = bridge.register(test_runtime_factory()); @@ -1671,12 +1717,12 @@ mod tests { /// A peer that opens a TCP connection and never sends the HTTP upgrade /// request — so its handshake never resolves — does not block a /// sibling's connection attempt. Each accepted connection's handshake - /// runs in its own task; before that, everything shared the accept - /// loop's own inline handshake, so a stalled peer there would have - /// blocked every other execution's connections too. + /// runs in its own task, and the connection caps are reserved inside it + /// once a token matches, which is why an unauthenticated socket is bounded + /// by the handshake backlog rather than by those caps. #[test] fn a_stalled_handshake_does_not_block_a_sibling_connection() { - let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; + let bridge = start_test_bridge(); let stalled = bridge.register(test_runtime_factory()); let healthy = bridge.register(test_runtime_factory()); @@ -1723,7 +1769,7 @@ mod tests { }) } - let bridge = WsBridge::start(0, no_log()).expect("start bridge").0; + let bridge = start_test_bridge(); let calls: Vec> = (0..3).map(|_| Arc::new(AtomicUsize::new(0))).collect(); let endpoints: Vec = calls .iter() @@ -1749,8 +1795,10 @@ mod tests { let request_frame = ProtocolMessage { request_id: "p:1".into(), payload: Payload { - id: ids.request_id, - value: HostFeatureSupportedRequest::V1( + trait_id: ids.trait_id, + method_id: ids.method_id, + message_type: crate::frame::MESSAGE_TYPE_REQUEST, + value: truapi::versioned::system::HostFeatureSupportedRequest::V1( v01::HostFeatureSupportedRequest::Chain { genesis_hash: vec![0u8; 32], }, From 165b083362d094cf0054b84fd92aaa546d2f61ca Mon Sep 17 00:00:00 2001 From: Nidish Date: Tue, 15 Sep 2026 16:13:49 +0530 Subject: [PATCH 5/7] fix(truapi-server): keep the reservation loop off a deprecated atomic helper --- .changeset/shared-ws-listener-teardown.md | 9 +++++++++ rust/crates/truapi-server/src/ws_bridge.rs | 23 +++++++++++++++++----- 2 files changed, 27 insertions(+), 5 deletions(-) create mode 100644 .changeset/shared-ws-listener-teardown.md diff --git a/.changeset/shared-ws-listener-teardown.md b/.changeset/shared-ws-listener-teardown.md new file mode 100644 index 000000000..9f8fbcf3b --- /dev/null +++ b/.changeset/shared-ws-listener-teardown.md @@ -0,0 +1,9 @@ +--- +"@parity/truapi-host": patch +--- + +Product executions under one host runtime share a single localhost WebSocket listener, each with its own token. Closing a +connection disposes the runtime that served it, so its host-core subscriptions and chat state are released rather than +held for the life of the listener. Admitted connections are bounded per execution and listener-wide, with the +per-execution cap a share of the listener-wide one; handshakes in flight are bounded separately, and a full backlog +evicts its oldest entry so a stalled peer cannot lock out other executions. diff --git a/rust/crates/truapi-server/src/ws_bridge.rs b/rust/crates/truapi-server/src/ws_bridge.rs index 90da8ba14..2b9125a26 100644 --- a/rust/crates/truapi-server/src/ws_bridge.rs +++ b/rust/crates/truapi-server/src/ws_bridge.rs @@ -783,11 +783,24 @@ type MatchedReservation = Arc, ConnectionCountG /// load-then-increment: two handshakes could otherwise both observe room for /// the last slot and both take it. fn try_reserve(counter: &AtomicUsize, limit: usize) -> bool { - counter - .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { - (current < limit).then_some(current + 1) - }) - .is_ok() + // Spelled out rather than via `fetch_update`, which current nightly + // deprecates in favour of `try_update`, and `try_update` is still unstable + // on the stable toolchain this crate also builds on. + let mut current = counter.load(Ordering::Acquire); + loop { + if current >= limit { + return false; + } + match counter.compare_exchange_weak( + current, + current + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return true, + Err(actual) => current = actual, + } + } } /// Complete the WebSocket handshake, resolving it against the registry to From 5d779b580121ea76fd92967affaac6b2513db389 Mon Sep 17 00:00:00 2001 From: Nidish Date: Wed, 16 Sep 2026 02:00:26 +0530 Subject: [PATCH 6/7] test(truapi-server): pass the pocket callbacks argument in the shared bridge test --- rust/crates/truapi-server/src/native.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/rust/crates/truapi-server/src/native.rs b/rust/crates/truapi-server/src/native.rs index 6341ec6f3..431e6f875 100644 --- a/rust/crates/truapi-server/src/native.rs +++ b/rust/crates/truapi-server/src/native.rs @@ -4054,6 +4054,7 @@ mod tests { .open_product_execution( Arc::new(EventCallbacks::new()), None, + None, native_execution_config("shared.dot", ProductExecutionKind::App), ) .expect("App execution should open"); @@ -4062,6 +4063,7 @@ mod tests { .open_product_execution( chat_host.clone(), Some(chat_host), + None, native_execution_config("shared.dot", ProductExecutionKind::Worker), ) .expect("Chat execution should open"); From d5192a4596b15597a3a35e583e3c7efffb07c690 Mon Sep 17 00:00:00 2001 From: pgherveou Date: Fri, 18 Sep 2026 10:50:06 +0200 Subject: [PATCH 7/7] fix(ws): isolate execution lifecycles --- rust/crates/truapi-server/src/host_core.rs | 25 +- rust/crates/truapi-server/src/native.rs | 149 ++++- rust/crates/truapi-server/src/ws_bridge.rs | 737 +++++++-------------- 3 files changed, 374 insertions(+), 537 deletions(-) diff --git a/rust/crates/truapi-server/src/host_core.rs b/rust/crates/truapi-server/src/host_core.rs index e7e5d5dc7..4f2498891 100644 --- a/rust/crates/truapi-server/src/host_core.rs +++ b/rust/crates/truapi-server/src/host_core.rs @@ -1368,12 +1368,8 @@ impl ProductRuntime { // here, which is exactly the production-host-killing shape the debug tap // above was fixed for. // - // The disposed check above is only a fast path: `dispose` can swap it true - // and drain `in_flight` at any point after that check returns and before - // this insert runs. Re-checking here, under the same lock `dispose` holds - // for its own swap-and-drain, closes that window - whichever runs first is - // what the other observes, so a dispatch that loses the race is turned away - // instead of running past a disposal that already happened. + // Re-check under the disposal lock so a racing dispatch cannot register + // after `dispose` has drained the active requests. { let mut in_flight = self .in_flight @@ -1474,17 +1470,10 @@ impl ProductRuntime { /// futures, and cancels active subscriptions. #[instrument(skip_all, fields(runtime.method = "product_runtime.dispose"))] pub fn dispose(&self) { - // Already-disposed callers answer without taking the lock, which the - // authoritative swap below holds across an abort loop that can run - // arbitrary waker code. + // Aborting under the lock can wake code that re-enters disposal. if self.disposed.load(Ordering::Acquire) { return; } - // The swap and the drain share `in_flight`'s lock with `receive_frame`'s own - // disposed-check-then-insert, so whichever of the two critical sections runs - // first is what the other observes: a dispatch that inserted before this - // drain is caught by it, and one that hasn't inserted yet sees `disposed` - // already true and turns itself away instead of dispatching past disposal. { let mut in_flight = self .in_flight @@ -2747,11 +2736,7 @@ mod tests { assert_eq!(response.payload.value, expected); } - /// A dispatch that passes `receive_frame`'s opening disposed check must not - /// run if `dispose` commits before it registers itself. The debug tap is the - /// only seam that can park a frame between those two points; the sink here - /// blocks on purpose, which its own contract forbids, so that the window is - /// reachable deterministically instead of by chance. + // The debug tap deliberately blocks to expose the registration race reliably. #[test] fn a_dispatch_racing_dispose_does_not_reach_the_platform() { struct ParkingDebugSink { @@ -2821,8 +2806,6 @@ mod tests { }) }; - // The frame is now parked inside the tap, past the opening disposed - // check and before it has registered itself. entered_rx .recv_timeout(std::time::Duration::from_secs(5)) .expect("frame never reached the debug tap"); diff --git a/rust/crates/truapi-server/src/native.rs b/rust/crates/truapi-server/src/native.rs index 431e6f875..00bcb54bf 100644 --- a/rust/crates/truapi-server/src/native.rs +++ b/rust/crates/truapi-server/src/native.rs @@ -623,9 +623,6 @@ pub struct NativeTrUApiHostRuntime { events: Arc, #[cfg(feature = "ws-bridge")] spawner: Spawner, - /// One localhost WebSocket listener shared by every product execution - /// this host runtime opens. Starts lazily on the first execution's - /// `start_ws_bridge` call and lives for as long as this host runtime. #[cfg(feature = "ws-bridge")] ws_bridge: Arc, /// The one Worker execution per product; opening another replaces it. @@ -671,7 +668,9 @@ impl NativeTrUApiHostRuntime { #[cfg(feature = "ws-bridge")] spawner, #[cfg(feature = "ws-bridge")] - ws_bridge: Arc::new(SharedWsBridge::new()), + ws_bridge: Arc::new(SharedWsBridge::new(Arc::new(move |marker, detail| { + callbacks.on_core_log(marker.to_string(), detail.to_string()); + }))), worker_executions: Mutex::new(HashMap::new()), })) } @@ -1036,11 +1035,8 @@ pub struct NativeProductExecution { crate::runtime::ActionChannel, >, closed: AtomicBool, - /// The host runtime's shared listener; every execution under the same - /// host runtime clones this same `Arc`. #[cfg(feature = "ws-bridge")] ws_bridge: Arc, - /// This execution's own registered token, if its bridge is running. #[cfg(feature = "ws-bridge")] bridge_token: Mutex>, #[cfg(feature = "ws-bridge")] @@ -1085,10 +1081,7 @@ impl NativeProductExecution { #[cfg(feature = "ws-bridge")] fn stop_bridge(&self) { - // Taken in its own statement so the guard is released before `revoke`, - // which blocks until this execution's connections have unwound. Held - // across that call, it would make a concurrent `start_ws_bridge` on - // this execution wait out the whole teardown. + // Release the token lock before waiting for connection cancellation. let token = self .bridge_token .lock() @@ -1274,11 +1267,8 @@ impl NativeProductExecution { #[cfg(feature = "ws-bridge")] #[uniffi::export] impl NativeProductExecution { - /// Register this execution against the host runtime's shared localhost - /// bridge, minting an independent authentication token. Every execution - /// under the same host runtime connects through the same port; - /// `bind_port` only has an effect for the first execution to register - /// while the shared listener is still unstarted. + /// Register this execution with its own token on the host's shared listener. + /// `bind_port` applies only when the listener first starts. pub fn start_ws_bridge(&self, bind_port: u16) -> Result { if self.closed.load(Ordering::Acquire) { return Err(WsBridgeStartError::Io( @@ -2259,6 +2249,7 @@ mod tests { } struct EventCallbacks { + logs: Mutex>, chat_room_status: Mutex, chat_created_rooms: Mutex>, chat_bot_status: Mutex, @@ -2293,6 +2284,7 @@ mod tests { fn new() -> Self { Self { + logs: Mutex::new(Vec::new()), chat_room_status: Mutex::new(v01::ChatRoomRegistrationStatus::New), chat_created_rooms: Mutex::new(Vec::new()), chat_bot_status: Mutex::new(v01::ChatBotRegistrationStatus::New), @@ -2323,7 +2315,9 @@ mod tests { #[async_trait::async_trait] impl HostCallbacks for EventCallbacks { - fn on_core_log(&self, _marker: String, _detail: String) {} + fn on_core_log(&self, marker: String, _detail: String) { + self.logs.lock().expect("logs mutex poisoned").push(marker); + } fn worker_demand_changed(&self, product_id: String, transition: WorkerTransition) { self.worker_demand .lock() @@ -4030,11 +4024,113 @@ mod tests { execution.stop_ws_bridge(); } - /// The PR's headline behavior, exercised through the real - /// `NativeProductExecution`/UniFFI-facing surface rather than only - /// `ws_bridge.rs`'s own lower-level unit tests: two executions under one - /// host runtime share the shared bridge's port with isolated tokens, and - /// stopping one leaves the other's live connection untouched. + #[cfg(feature = "ws-bridge")] + #[test] + fn closing_an_execution_releases_its_callbacks_while_the_host_lives() { + let host = native_host_runtime_no_session(); + let callbacks = Arc::new(EventCallbacks::new()); + let weak_callbacks = Arc::downgrade(&callbacks); + let execution = host + .open_product_execution( + callbacks, + None, + None, + native_execution_config("first.dot", ProductExecutionKind::App), + ) + .expect("open execution"); + execution.start_ws_bridge(0).expect("start bridge"); + + execution.shutdown(); + drop(execution); + + assert!( + weak_callbacks.upgrade().is_none(), + "the listener must not retain a closed execution's callbacks" + ); + drop(host); + } + + #[cfg(feature = "ws-bridge")] + #[test] + fn bridge_logs_follow_the_host_and_authenticated_execution() { + use futures::SinkExt; + use tokio_tungstenite::tungstenite::Message as WsMessage; + + let callbacks = [ + Arc::new(EventCallbacks::new()), + Arc::new(EventCallbacks::new()), + Arc::new(EventCallbacks::new()), + ]; + let host = NativeTrUApiHostRuntime::with_runtime_config( + callbacks[0].clone(), + native_host_runtime_config(), + ) + .expect("create host"); + let executions = [(1, "first.dot"), (2, "second.dot")].map(|(index, product_id)| { + host.open_product_execution( + callbacks[index].clone(), + None, + None, + native_execution_config(product_id, ProductExecutionKind::App), + ) + .expect("open execution") + }); + executions[0].start_ws_bridge(0).expect("start first"); + let endpoint = executions[1].start_ws_bridge(0).expect("start second"); + let client = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("client runtime"); + let socket = client.block_on(async { + let (mut socket, _) = tokio_tungstenite::connect_async(format!( + "ws://127.0.0.1:{}/?t={}", + endpoint.port, endpoint.token + )) + .await + .expect("connect second"); + socket + .send(WsMessage::Text("ignored".into())) + .await + .expect("send text"); + socket + }); + crate::test_support::wait_until( + || { + callbacks.iter().any(|callbacks| { + callbacks + .logs + .lock() + .expect("logs mutex poisoned") + .iter() + .any(|marker| marker == "truapi.ws_bridge.text_frame_ignored") + }) + }, + "connection did not process the text frame", + ); + let logs = callbacks.map(|callbacks| { + callbacks + .logs + .lock() + .expect("logs mutex poisoned") + .iter() + .filter(|marker| marker.starts_with("truapi.ws_bridge.")) + .cloned() + .collect::>() + }); + assert_eq!( + logs, + [ + vec!["truapi.ws_bridge.started"], + vec![], + vec![ + "truapi.ws_bridge.connection_open", + "truapi.ws_bridge.text_frame_ignored" + ], + ] + ); + drop(socket); + } + #[cfg(feature = "ws-bridge")] #[test] fn two_executions_share_one_bridge_through_the_native_api() { @@ -4136,8 +4232,6 @@ mod tests { .await .expect("chat dial"); - // Each execution's own token round-trips independently on the - // one shared port. app_ws .send(WsMessage::Binary(round_trip("app:1").encode())) .await @@ -4150,9 +4244,6 @@ mod tests { .expect("send on chat connection"); assert_eq!(answer(&mut chat_ws).await.request_id, "chat:1"); - // Stopping (revoking) App must not touch Chat's live connection, - // and must block until App's own connection has actually finished - // unwinding before returning. app.stop_ws_bridge(); chat_ws @@ -4165,10 +4256,6 @@ mod tests { "Chat's connection must keep answering after a sibling execution stops" ); - // The token is removed from the registry synchronously inside - // `revoke`, so a fresh connect with it is refused with no retry - // loop. This says nothing about whether the aborted connection has - // finished unwinding. assert!( tokio_tungstenite::connect_async(&app_url).await.is_err(), "a revoked token must not still be accepted" diff --git a/rust/crates/truapi-server/src/ws_bridge.rs b/rust/crates/truapi-server/src/ws_bridge.rs index 2b9125a26..02bab8e31 100644 --- a/rust/crates/truapi-server/src/ws_bridge.rs +++ b/rust/crates/truapi-server/src/ws_bridge.rs @@ -5,30 +5,16 @@ //! //! Feature-gated (`ws-bridge`) so wasm32 and no-tokio build paths stay lean. //! -//! Every product execution under one host runtime shares a single -//! [`SharedWsBridge`] listener on the process-wide `tokio` runtime. -//! [`SharedWsBridge::register`] hands each execution its own `{port, token}` -//! endpoint on that one port; [`SharedWsBridge::revoke`] tears down only that -//! execution's connections, leaving the listener and its siblings untouched. +//! Executions under one host share a [`SharedWsBridge`] listener and the +//! process-wide `tokio` runtime, with independent tokens and connections. //! -//! Security model: the listener binds `127.0.0.1` only, and every connection -//! must present its execution's 256-bit token (`?t=`, from the OS -//! CSPRNG) before the upgrade completes. The handshake scans every registered -//! token with a constant-time comparison and exits early on neither a match -//! nor a duplicated `t=` pair, so timing reveals nothing about which token -//! matched. Revocation removes the token before aborting connections, and the -//! revoked flag is read under the same lock that registers a connection, so a -//! revoked token can never carry a frame. Tokens go only to the host's -//! embedded WebView, whose origin is not known a priori, so the `Origin` -//! header is not pinned. +//! Each upgrade requires an execution's random 256-bit token (`?t=`). +//! Token comparisons scan all candidates to avoid revealing which matched. +//! Tokens go only to the host's embedded WebView; its origin is not known in +//! advance, so the `Origin` header is not pinned. //! -//! Bounds, all to contain a misbehaving local peer: inbound message size, each -//! connection's outbound queue, admitted connections per execution and -//! listener-wide, and — since a peer that never presents a token reaches none -//! of those — the number of handshakes in flight -//! ([`MAX_PENDING_HANDSHAKES`]). Worst-case sockets is the sum of the last -//! two. A full handshake backlog evicts its oldest entry rather than refusing -//! the newcomer, so one stalled peer cannot lock the listener. +//! Message size, outbound queues, authenticated connections and pending +//! handshakes are bounded to contain a misbehaving local peer. use std::collections::{HashMap, VecDeque}; use std::io; @@ -36,7 +22,7 @@ use std::net::SocketAddr; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, OnceLock}; -use futures::{SinkExt, StreamExt}; +use futures::{FutureExt, SinkExt, StreamExt}; use rand::RngCore; use tokio::net::TcpListener; use tokio::runtime::{Handle, Runtime}; @@ -45,38 +31,17 @@ use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::tungstenite::Message as WsMessage; use tokio_tungstenite::tungstenite::handshake::server::{ErrorResponse, Request, Response}; use tokio_tungstenite::tungstenite::http::{Response as HttpResponse, StatusCode}; -use tokio_tungstenite::tungstenite::protocol::CloseFrame; use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; -use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; use crate::{FrameSink, ProductRuntime}; -/// Maximum simultaneous connections a single registered execution may hold. -/// An execution uses one connection; the headroom absorbs reconnect churn, -/// and the cap bounds a buggy or hostile local peer opening many sockets -/// against one token. -/// -/// Kept low enough that -/// `MAX_TOTAL_WS_CONNECTIONS / MAX_WS_CONNECTIONS_PER_EXECUTION` executions -/// can each hold their full allowance at once. Were the per-execution cap -/// above that share, one execution could take enough of the shared budget to -/// starve its siblings, which a per-execution listener could not do. +// Allow reconnect overlap without one execution exhausting the shared limit. const MAX_WS_CONNECTIONS_PER_EXECUTION: usize = 8; -/// Maximum simultaneous connections across every execution sharing the -/// listener. Set well above any realistic concurrent-execution count (App, -/// Widget, Chat/Worker) while still bounding total resource use from a -/// misbehaving peer amplifying across many registered tokens. +// Bound resource use even when a peer holds several valid execution tokens. const MAX_TOTAL_WS_CONNECTIONS: usize = 64; -/// Maximum handshakes the listener will carry at once, counted separately -/// from [`MAX_TOTAL_WS_CONNECTIONS`]: a connection occupies a slot here only -/// until its handshake resolves, and one there only once a token has -/// matched. Worst-case sockets is therefore the sum of the two, not either -/// alone. The bound exists because a peer that never presents a token is -/// invisible to the connection caps, which are reserved inside the -/// handshake. Reaching it evicts the oldest handshake rather than turning -/// the newcomer away (see [`accept_loop`]). +// Unauthenticated peers do not count against the connection limits. const MAX_PENDING_HANDSHAKES: usize = 64; /// Bound on the per-connection outbound frame queue. A peer that stops reading @@ -89,13 +54,7 @@ const OUTBOUND_QUEUE_CAP: usize = 4096; /// memory-amplification DoS well below tungstenite's 64 MiB default. const MAX_WS_MESSAGE_BYTES: usize = 8 << 20; -/// Ceiling on how long one connection's own setup task will wait for its -/// handshake (`authenticate_and_upgrade`) to resolve. Each connection's setup -/// runs in its own task, so a peer that never completes the handshake cannot -/// stall any other connection — this bound exists so a stalled setup task -/// (and the socket and file descriptor behind it) cannot linger forever, and -/// so the listener's shutdown never has to wait on one indefinitely. Generous -/// relative to a real localhost upgrade (sub-millisecond in practice). +// Stalled handshakes must eventually release their sockets. const HANDSHAKE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); /// Per-session descriptor returned to the host: product uses `port + token` @@ -113,8 +72,7 @@ pub struct WsBridgeEndpoint { #[derive(Debug, thiserror::Error, uniffi::Error)] #[uniffi(flat_error)] pub enum WsBridgeStartError { - /// This execution already registered a bridge token. The host's shared - /// listener being up is the normal case and not an error. + /// This execution already has a registered bridge token. #[error("ws bridge already running")] AlreadyRunning, /// Anything else (bind failure, runtime spin-up failure, ...). @@ -207,20 +165,14 @@ fn shared_native_executor() -> io::Result<(&'static SharedNativeExecutor, bool)> Ok((executor, initialized)) } -/// One execution's registered runtime factory and live connection state. struct RegistryEntry { runtime_factory: Arc, + logger: BridgeLogger, connection_count: Arc, connections: Mutex, } -/// A connection can finish its handshake against an entry that is revoked a -/// moment later. The `revoked` flag and the handle list share one lock, and -/// `connection_setup` both reads the flag and spawns the connection under a -/// single acquisition of it, so revocation and registration serialize: -/// whichever happens first is what the other observes. A revocation that -/// wins turns the connection away before any task reads from its socket; one -/// that loses finds the handle already recorded and aborts it. +// One lock prevents a completed handshake from registering past revocation. #[derive(Default)] struct EntryConnections { revoked: bool, @@ -228,8 +180,6 @@ struct EntryConnections { } impl EntryConnections { - /// Request cancellation of every tracked connection and hand back sole - /// ownership of their handles, so a caller can join what it aborted. fn abort_and_take(&mut self) -> Vec> { for handle in self.handles.iter() { handle.abort(); @@ -238,8 +188,6 @@ impl EntryConnections { } } -/// Token registry shared by the listener's accept loop and every -/// [`SharedWsBridge::register`]/[`SharedWsBridge::revoke`] call. #[derive(Default)] struct WsBridgeRegistry { entries: Mutex>>, @@ -247,7 +195,12 @@ struct WsBridgeRegistry { } impl WsBridgeRegistry { - fn insert(&self, token: String, runtime_factory: Arc) { + fn insert( + &self, + token: String, + runtime_factory: Arc, + logger: BridgeLogger, + ) { self.entries .lock() .expect("ws bridge registry mutex poisoned") @@ -255,19 +208,13 @@ impl WsBridgeRegistry { token, Arc::new(RegistryEntry { runtime_factory, + logger, connection_count: Arc::new(AtomicUsize::new(0)), connections: Mutex::new(EntryConnections::default()), }), ); } - /// Remove `token`'s entry, mark it revoked, and abort its live - /// connections, handing their handles back so a caller can choose to - /// wait for them to actually finish unwinding. A handshake that matched - /// this token just before the removal, but has not yet registered its - /// handle, observes the `revoked` flag (set under the same lock as the - /// abort loop below) and aborts itself instead of registering. No-op - /// (empty result) for an unknown or already-revoked token. fn revoke(&self, token: &str) -> Vec> { let Some(entry) = self .entries @@ -285,13 +232,9 @@ impl WsBridgeRegistry { state.abort_and_take() } - /// Find the entry whose token matches `path_and_query`'s `?t=` value. - /// Scans every registered token without exiting early on a match, so - /// timing does not reveal which one (if any) matched. + // Scan all candidates so an early match cannot reveal token order. fn find_matching(&self, path_and_query: Option<&str>) -> Option> { - // Snapshot first: the scan below runs over peer-supplied request bytes - // once per registered token, and holding the registry lock across it - // would stall every concurrent handshake, registration and revocation. + // Peer-supplied query processing must not hold up registry updates. let candidates: Vec<(String, Arc)> = self .entries .lock() @@ -308,10 +251,7 @@ impl WsBridgeRegistry { found } - /// Abort and drain every connection under a still-registered execution, - /// handing back the owned handles so the caller can await them. Idempotent. - /// A connection revoked moments earlier is not covered: `revoke` removed - /// its entry and aborted it directly, so this cannot find it to await. + // Revoked entries are already removed; their caller owns the aborted tasks. fn take_all_handles(&self) -> Vec> { let entries = self .entries @@ -329,47 +269,37 @@ impl WsBridgeRegistry { } } -/// Host-runtime-owned shared listener. Every product execution under the -/// same host runtime clones the same `Arc` and registers -/// against it; the underlying listener starts on the first registration and -/// lives for as long as the host runtime does. -#[derive(Default)] +/// Lazy listener shared by a host runtime and its product executions. pub struct SharedWsBridge { inner: Mutex>, + logger: BridgeLogger, } impl SharedWsBridge { - /// Construct an unstarted shared bridge. The listener binds lazily on - /// the first [`Self::register`] call. - pub fn new() -> Self { - Self::default() + /// Construct a lazy listener with host-owned lifecycle logging. + pub fn new(logger: BridgeLogger) -> Self { + Self { + inner: Mutex::new(None), + logger, + } } - /// Ensure the shared listener is running and register a fresh - /// per-execution token against it. + /// Register an execution with its own token and connection logger. /// - /// `bind_port` only takes effect for the first execution to register; - /// once the listener is up, every product connects through that same - /// port regardless of what a later caller requests. A later caller that - /// requested a specific (non-zero) port different from the one already - /// running gets a log line noting the request was ignored, rather than - /// silence. + /// The first registration starts the listener on `bind_port`. Later + /// registrations reuse that port; conflicting nonzero requests are logged. pub fn register( &self, bind_port: u16, runtime_factory: Arc, logger: BridgeLogger, ) -> Result { - // Every logger call below is deferred until after this lock is - // released. A host logger that blocks or calls back into - // `register`/`revoke` would otherwise stall or deadlock every other - // execution sharing this bridge, since this lock is the one thing - // every one of their `register`/`revoke` calls must take too. + // Log after unlocking because host callbacks can re-enter the bridge. let mut pending_logs: Vec<(&'static str, String)> = Vec::new(); let endpoint = { let mut guard = self.inner.lock().expect("shared ws bridge mutex poisoned"); if guard.is_none() { - let (bridge, logs) = WsBridge::start(bind_port, logger.clone())?; + let (bridge, logs) = WsBridge::start(bind_port, self.logger.clone())?; *guard = Some(bridge); pending_logs = logs; } else if bind_port != 0 { @@ -384,10 +314,10 @@ impl SharedWsBridge { guard .as_ref() .expect("shared bridge just inserted") - .register(runtime_factory) + .register(runtime_factory, logger) }; for (event, detail) in &pending_logs { - logger(event, detail); + (self.logger)(event, detail); } Ok(endpoint) } @@ -395,10 +325,9 @@ impl SharedWsBridge { /// Revoke one execution's token. No-op if the listener was never /// started or the token is unknown. /// - /// Off the shared executor this blocks until that execution's connection - /// tasks have been joined, the same wait `stop` performs listener-wide. - /// Joining is not a barrier on the destructors inside those tasks. On the - /// shared executor the wait is skipped to avoid deadlocking that worker. + /// Off the shared executor, waits for connection tasks to release their + /// sockets and capacity. On that executor, cancellation is requested + /// without waiting to avoid deadlocking a worker. pub fn revoke(&self, token: &str) { let aborted = { let guard = self.inner.lock().expect("shared ws bridge mutex poisoned"); @@ -411,15 +340,8 @@ impl SharedWsBridge { } } -/// Running listener handle. Drop or call [`WsBridge::stop`] to shut down. -/// -/// Crate-internal: hosts reach the bridge through [`SharedWsBridge`], which -/// owns the listener and hands out per-execution endpoints. Deliberately not -/// re-exported by the crate root. -/// -/// The listener's tasks run on the process-wide native executor. TrUAPI -/// dispatch futures are `Send`, so connections and independent frames from -/// all registered executions can execute across the shared worker pool. +/// Running listener owned by [`SharedWsBridge`]. Dropping it stops acceptance +/// and cancels its connections without stopping the shared executor. pub(crate) struct WsBridge { shutdown: Option>, stopped: Option>, @@ -430,13 +352,7 @@ pub(crate) struct WsBridge { } impl WsBridge { - /// Bind a localhost listener and start the accept loop on the shared - /// native executor. - /// - /// `logger` is still used for the accept loop's own ongoing (async, - /// lock-free) lifecycle events; the one-time startup events below are - /// instead returned for the caller to emit once it is no longer holding - /// [`SharedWsBridge::inner`]'s lock. + // Return startup logs so the caller can emit them outside its lock. fn start( bind_port: u16, logger: BridgeLogger, @@ -459,10 +375,7 @@ impl WsBridge { TcpListener::from_std(std_listener)? }; - // Queued only past the last fallible step. The executor is a process - // -wide `OnceLock`, so `initialized` is true for exactly one caller in - // the process; dropping its event on an error return would lose it for - // good rather than delay it. + // An error return here would lose the executor's one-time startup event. let mut pending_logs: Vec<(&'static str, String)> = Vec::new(); if initialized { pending_logs.push(( @@ -501,22 +414,22 @@ impl WsBridge { )) } - /// Register one execution's runtime factory, minting a fresh token. - fn register(&self, runtime_factory: Arc) -> WsBridgeEndpoint { + fn register( + &self, + runtime_factory: Arc, + logger: BridgeLogger, + ) -> WsBridgeEndpoint { let mut token_bytes = [0u8; 32]; rand::thread_rng().fill_bytes(&mut token_bytes); let token = hex::encode(token_bytes); - self.registry.insert(token.clone(), runtime_factory); + self.registry.insert(token.clone(), runtime_factory, logger); WsBridgeEndpoint { port: self.port, token, } } - /// Revoke one execution's token and abort its live connections, returning - /// the tasks a caller may join. Empty when there was nothing to abort, or - /// when the caller is itself running on this bridge's executor and must - /// not block on it: those connections still unwind on their own. + // Joining from this executor could block the worker needed for cancellation. fn revoke(&self, token: &str) -> Vec> { let connections = self.registry.revoke(token); if Handle::try_current().is_ok_and(|current| current.id() == self.runtime_id) { @@ -525,23 +438,14 @@ impl WsBridge { connections } - /// Signal the accept loop to exit and abort every tracked connection - /// across every registered execution. - /// - /// Off the shared executor, this blocks until the accept loop has - /// actually drained and awaited every connection, so a caller there gets - /// a real "fully stopped" guarantee. From a task already running on the - /// shared executor, waiting is skipped instead (see below), so this - /// specific caller only gets a best-effort sweep, not that guarantee. + // Off the shared executor, wait for tracked connection tasks to release + // their sockets. Dispatch tasks are cancelled but not joined. fn stop(&mut self) { if let Some(tx) = self.shutdown.take() { let _ = tx.send(()); } - // Blocking is safe only off the shared executor; on one of its own - // workers it could deadlock against the task being waited on. Skipping - // the wait means `stop` can return before the accept loop has even - // observed the signal, which is what the sweep below is for. + // The accept loop needs a worker to observe shutdown and cancel tasks. let called_from_shared_executor = Handle::try_current().is_ok_and(|handle| handle.id() == self.runtime_id); let stopped_cleanly = if called_from_shared_executor { @@ -559,9 +463,7 @@ impl WsBridge { { task.abort(); } - // Finds nothing when the accept loop drained on its own; does real work - // only on the non-blocking path above or after a force-abort. Aborts - // without awaiting, `stop` not being async. + // Cover the non-blocking path or an accept loop that failed to stop. drop(self.registry.take_all_handles()); } } @@ -572,11 +474,8 @@ impl Drop for WsBridge { } } -/// Block the calling thread until every aborted connection's task has been -/// joined. Spawns a small joiner task on the shared executor and blocks on a -/// synchronous channel rather than awaiting directly, since callers are -/// ordinary host threads with no executor of their own. There is no deadline: -/// a connection wedged in non-yielding work holds the caller until it unwedges. +// Synchronous host callers need not have an executor. This wait relies on +// connection tasks yielding so cancellation can complete. fn join_aborted_connections(handles: Vec>) { if handles.is_empty() { return; @@ -600,13 +499,7 @@ async fn accept_loop( logger: BridgeLogger, mut shutdown: oneshot::Receiver<()>, ) { - // Each accepted connection's handshake runs in its own task (see - // `connection_setup`) so a stalled or hostile peer never blocks another - // connection's setup, only its own. This loop tracks those setup tasks so - // shutdown can cancel and await whichever haven't resolved yet, in - // addition to draining every connection already registered, and so a full - // backlog can evict its oldest entry. Accept order is insertion order, so - // the front is the longest-running handshake. + // Independent setup tasks keep a stalled handshake from blocking acceptance. let mut setup_tasks: VecDeque> = VecDeque::new(); loop { tokio::select! { @@ -618,13 +511,6 @@ async fn accept_loop( for task in setup_tasks { let _ = task.await; } - // `take_all_handles` both requests cancellation and hands back - // sole ownership of every handle still under a registered - // execution, so every one of them can genuinely be awaited - // here before this loop (and thus the listener's `stopped` - // signal) returns. It does not see a connection whose token - // was revoked (and thus removed from the registry) moments - // earlier — `revoke` already aborted that one directly. for handle in registry.take_all_handles() { let _ = handle.await; } @@ -639,15 +525,8 @@ async fn accept_loop( } }; setup_tasks.retain(|task| !task.is_finished()); - // A peer that never presents a token holds a task and a file - // descriptor for up to `HANDSHAKE_TIMEOUT` without ever - // reaching the connection caps, which are reserved inside the - // handshake. Bound that backlog by evicting its oldest entry - // rather than refusing this peer: a real localhost handshake - // resolves in well under a millisecond, so the oldest of a - // full backlog is one that has stalled, while refusing the - // newcomer would let a peer holding every slot lock the whole - // shared listener out for every execution. + // Evict the oldest pending handshake so stalled peers cannot + // reserve the entire backlog until their timeouts expire. if setup_tasks.len() >= MAX_PENDING_HANDSHAKES && let Some(oldest) = setup_tasks.pop_front() { @@ -664,19 +543,6 @@ async fn accept_loop( } } -/// Authenticate one accepted connection and, if the handshake succeeds, -/// register it and run its lifecycle. Spawned independently per connection -/// (from `accept_loop`) precisely so a stalled or hostile peer's handshake -/// blocks only this task, never another connection's setup, another -/// execution's connections, or the accept loop's ability to keep accepting. -/// -/// Handshakes therefore resolve concurrently, which is why the -/// connection-count caps inside `authenticate_and_upgrade`'s handshake -/// callback are reserved with a compare-and-swap loop (`try_reserve`) rather -/// than a plain load-then-increment, and why the `revoked` flag below is read -/// under the same lock that registers the connection: a token revoked while -/// this handshake was resolving must not end up with a task reading from its -/// socket, however many handshakes are in flight. async fn connection_setup( stream: tokio::net::TcpStream, peer: SocketAddr, @@ -697,13 +563,7 @@ async fn connection_setup( }) else { return; }; - // Read the revocation flag and spawn the connection under one acquisition - // of the lock `revoke` takes for its own flag-and-abort pass, so the two - // serialize: an execution revoked while this handshake was resolving never - // gets a task reading from the socket, and one still live cannot be - // revoked between the spawn and the handle landing in its list, where it - // would be left running with no owner. `tokio::spawn` only queues the - // task, so nothing host-supplied runs under the lock here. + let logger = entry.logger.clone(); let revoked = { let mut state = entry .connections @@ -713,18 +573,15 @@ async fn connection_setup( if state.revoked { true } else { - let conn_logger = logger.clone(); let conn_entry = entry.clone(); state.handles.push(tokio::spawn(async move { let _guard = guard; - connection_lifecycle(ws, peer, conn_entry, conn_logger).await; + connection_lifecycle(ws, peer, conn_entry).await; })); false } }; - // Logged after the lock is released: the logger is a host callback, and one - // that blocks or re-enters would otherwise stall every `revoke` and - // shutdown pass that needs this entry's lock. + // Host callbacks may re-enter revocation, so logging must stay outside the lock. if revoked { logger( "truapi.ws_bridge.connection_revoked_during_setup", @@ -733,9 +590,6 @@ async fn connection_setup( } } -/// Decrements the global and per-execution connection counts on drop, which -/// runs whether the connection task ends normally or is aborted (e.g. by a -/// revoke), keeping the caps accurate under both. struct ConnectionCountGuard { total: Arc, per_entry: Arc, @@ -748,15 +602,7 @@ impl Drop for ConnectionCountGuard { } } -/// Disposes a connection's `ProductRuntime` when this guard drops, however -/// that happens. `.abort()` — the only mechanism `revoke` and shutdown use to -/// end a connection — drops its task's future at whatever `.await` point it -/// was suspended at (almost always inside the read loop), which skips any -/// code written after that point, including an explicit `dispose()` call at -/// the end of the function. A value still on the stack at that point still -/// gets dropped, though, so holding the runtime here is what actually -/// disposes an aborted connection. `dispose()` is documented as idempotent, -/// so this is safe even alongside a path that also disposes explicitly. +// Task cancellation skips explicit cleanup after an await, but still drops guards. struct DisposeGuard(Arc); impl Drop for DisposeGuard { @@ -765,27 +611,17 @@ impl Drop for DisposeGuard { } } -/// A fully upgraded WebSocket connection, the execution entry its token -/// matched, and the connection-count guard reserved for it. type AuthenticatedConnection = ( WebSocketStream, Arc, ConnectionCountGuard, ); -/// What the handshake callback found and reserved, shared between the -/// callback and the function's own return path. type MatchedReservation = Arc, ConnectionCountGuard)>>>; -/// Atomically reserve one slot by incrementing `counter` unless it is already -/// at `limit`, retrying under contention. One task per accepted connection -/// means setups reserve concurrently, so this cannot be a plain -/// load-then-increment: two handshakes could otherwise both observe room for -/// the last slot and both take it. +// Concurrent handshakes must not both claim the last available slot. fn try_reserve(counter: &AtomicUsize, limit: usize) -> bool { - // Spelled out rather than via `fetch_update`, which current nightly - // deprecates in favour of `try_update`, and `try_update` is still unstable - // on the stable toolchain this crate also builds on. + // `fetch_update` is deprecated on nightly; `try_update` is not yet stable. let mut current = counter.load(Ordering::Acquire); loop { if current >= limit { @@ -803,27 +639,7 @@ fn try_reserve(counter: &AtomicUsize, limit: usize) -> bool { } } -/// Complete the WebSocket handshake, resolving it against the registry to -/// find the matching execution's entry and reserving its connection-count -/// slots for the life of the connection. -/// -/// Called from `connection_setup`, one independent task per accepted -/// connection, under [`HANDSHAKE_TIMEOUT`] — so a peer that never completes -/// its handshake only ever stalls that one task, not the accept loop or any -/// other connection's setup. -/// -/// The connection-count reservation happens inside the synchronous handshake -/// callback via [`try_reserve`], as soon as a token matches and both caps -/// have room, and a [`ConnectionCountGuard`] for it is constructed -/// immediately and stashed alongside the matched entry. Every way this -/// function can end — a successful upgrade, a post-callback handshake I/O -/// failure, or an outer timeout dropping this whole future — drops `matched` -/// and, with it, any reserved-but-unclaimed guard, so a failure after the -/// callback already committed a slot can never leak it. -// `clippy::result_large_err` fires on the handshake callback because -// tokio-tungstenite's `ErrorResponse` type carries the full HTTP response -// (~136 bytes). The closure signature is dictated by tokio-tungstenite's -// API, so the lint can only be silenced at the call site. +// The handshake callback's error type is fixed by tokio-tungstenite. #[allow(clippy::result_large_err)] async fn authenticate_and_upgrade( stream: tokio::net::TcpStream, @@ -844,13 +660,9 @@ async fn authenticate_and_upgrade( return Err(err); }; - // Reserve both caps' slots here, inside the synchronous handshake - // callback, so an over-cap peer is rejected at the HTTP upgrade - // rather than allowed to open a socket that gets dropped right - // after. Each reservation is a compare-and-swap because one task per - // connection means concurrent reservations against the same counter. + // Reject over-cap connections before acknowledging the HTTP upgrade. if !try_reserve(&entry.connection_count, MAX_WS_CONNECTIONS_PER_EXECUTION) { - auth_logger( + (entry.logger)( "truapi.ws_bridge.connection_limit_execution", &peer.to_string(), ); @@ -860,20 +672,14 @@ async fn authenticate_and_upgrade( return Err(err); } if !try_reserve(&auth_registry.total_connections, MAX_TOTAL_WS_CONNECTIONS) { - // Roll back the per-execution reservation above: it was never - // claimed, since the total cap is what ultimately rejected this - // attempt. entry.connection_count.fetch_sub(1, Ordering::AcqRel); - auth_logger("truapi.ws_bridge.connection_limit_total", &peer.to_string()); + (entry.logger)("truapi.ws_bridge.connection_limit_total", &peer.to_string()); let mut err: ErrorResponse = HttpResponse::new(Some("listener at capacity".to_string())); *err.status_mut() = StatusCode::SERVICE_UNAVAILABLE; return Err(err); } - // Built right after both reservations above, so no return from this - // function — success, a later I/O failure, or an outer timeout — can - // drop `matched` without also dropping (and thus balancing) this - // guard. + // Keep the reservation guarded even if the upgrade fails or times out. let guard = ConnectionCountGuard { total: auth_registry.total_connections.clone(), per_entry: entry.connection_count.clone(), @@ -907,7 +713,7 @@ async fn authenticate_and_upgrade( .expect("ws bridge handshake mutex poisoned") .take() .expect("a successful upgrade always resolved a matching registry entry"); - logger("truapi.ws_bridge.connection_open", &peer.to_string()); + (entry.logger)("truapi.ws_bridge.connection_open", &peer.to_string()); Some((ws, entry, guard)) } @@ -915,88 +721,66 @@ async fn connection_lifecycle( ws: WebSocketStream, peer: SocketAddr, entry: Arc, - logger: BridgeLogger, ) { + let logger = &entry.logger; let (mut sink, mut source) = ws.split(); let (out_tx, mut out_rx) = mpsc::channel::>(OUTBOUND_QUEUE_CAP); let frame_sink = Arc::new(WsFrameSink::new(out_tx)); let product_runtime = Arc::new(entry.runtime_factory.product_runtime(frame_sink)); let dispose_guard = DisposeGuard(product_runtime.clone()); - let pump_logger = logger.clone(); - let pump = tokio::spawn(async move { - while let Some(bytes) = out_rx.recv().await { - if let Err(err) = sink.send(WsMessage::Binary(bytes)).await { - pump_logger("truapi.ws_bridge.send_error", &err.to_string()); - break; - } - } - let _ = sink - .send(WsMessage::Close(Some(CloseFrame { - code: CloseCode::Normal, - reason: "bridge closing".into(), - }))) - .await; - let _ = sink.close().await; - }); - // Dispatch each inbound frame on its own `Send` task so a slow request // handler cannot stall the read loop and independent frames can run on // different executor workers. Responses may interleave; the wire protocol // matches them by request id, and `WsFrameSink::emit_frame` is thread-safe. - let mut in_flight: Vec> = Vec::new(); - while let Some(frame) = source.next().await { - match frame { - Ok(WsMessage::Binary(bytes)) => { - in_flight.retain(|task| !task.is_finished()); - let product_runtime = product_runtime.clone(); - let frame_logger = logger.clone(); - in_flight.push(tokio::spawn(async move { - // A frame the runtime cannot decode is a wire mismatch on - // the peer's side. Report it: dropping it unreported is - // indistinguishable from the peer never having sent it, - // and the peer is left waiting for a response forever. - if let Err(err) = product_runtime.receive_frame(bytes.to_vec()).await { - frame_logger("truapi.ws_bridge.frame_error", &err.to_string()); - } - })); - } - Ok(WsMessage::Text(_)) => { - logger("truapi.ws_bridge.text_frame_ignored", ""); + let mut in_flight = tokio::task::JoinSet::new(); + { + let writer = async { + while let Some(bytes) = out_rx.recv().await { + if let Err(err) = sink.send(WsMessage::Binary(bytes)).await { + logger("truapi.ws_bridge.send_error", &err.to_string()); + break; + } } - Ok(WsMessage::Close(_)) => break, - Ok(_) => {} - Err(err) => { - logger("truapi.ws_bridge.read_error", &err.to_string()); - break; + }; + tokio::pin!(writer); + loop { + let frame = tokio::select! { + _ = &mut writer => break, + frame = source.next() => frame, + }; + match frame { + Some(Ok(WsMessage::Binary(bytes))) => { + while in_flight.try_join_next().is_some() {} + let product_runtime = product_runtime.clone(); + let frame_logger = logger.clone(); + in_flight.spawn(async move { + if let Err(err) = product_runtime.receive_frame(bytes.to_vec()).await { + frame_logger("truapi.ws_bridge.frame_error", &err.to_string()); + } + }); + } + Some(Ok(WsMessage::Text(_))) => { + logger("truapi.ws_bridge.text_frame_ignored", ""); + } + None | Some(Ok(WsMessage::Close(_))) => break, + Some(Ok(_)) => {} + Some(Err(err)) => { + logger("truapi.ws_bridge.read_error", &err.to_string()); + break; + } } } } - // The connection is gone: cancel in-flight dispatches so long-pending - // handlers unwind instead of outliving the connection. - for task in &in_flight { - task.abort(); - } - - // Dispose and release the runtime before awaiting the pump. The pump ends - // when the last outbound sender drops, and the runtime owns one through its - // frame sink, so holding the runtime here while waiting for the pump would - // wait on something only this function's own return can cause. + // A slow peer must not retain capacity while its close reply waits. + let _ = sink.close().now_or_never(); + drop(in_flight); drop(dispose_guard); - drop(product_runtime); - - let _ = pump.await; logger("truapi.ws_bridge.connection_closed", &peer.to_string()); } -/// Whether `path_and_query`'s `?t=` value (any occurrence) matches `expected`, -/// compared in constant time. Every `t=` pair present is checked — this does -/// not stop at the first match — so a peer padding the query with several -/// `t=` pairs cannot make a match resolve faster than a non-match and use -/// that timing to test a candidate token. The token length is fixed and -/// public, so a length mismatch may short-circuit; only the value comparison -/// must be constant time. +// Scan duplicate `t=` parameters too, so their order cannot expose an early match. fn path_token_matches(path_and_query: Option<&str>, expected: &str) -> bool { let Some(raw) = path_and_query else { return false; @@ -1072,8 +856,6 @@ mod tests { use crate::frame::{Payload, ProtocolMessage, request_ids}; use crate::test_support::{StubPlatform, test_spawner}; - /// A started bridge with logging discarded, which is all any test here - /// wants from `WsBridge::start`'s deferred-log return. fn start_test_bridge() -> WsBridge { WsBridge::start(0, no_log()).expect("start bridge").0 } @@ -1126,10 +908,6 @@ mod tests { assert!(!path_token_matches(None, "abc")); } - /// A query with more than one `t=` pair is matched if ANY of them equals - /// the expected token, regardless of position — this is what closes the - /// timing differential a peer could otherwise create by padding the - /// query with non-matching `t=` pairs before or after the real one. #[test] fn path_token_matches_every_duplicated_t_pair_not_just_the_first() { assert!(path_token_matches(Some("/?t=wrong&t=abc"), "abc")); @@ -1210,7 +988,7 @@ mod tests { #[test] fn round_trip_feature_supported_through_bridge() { let bridge = start_test_bridge(); - let endpoint = bridge.register(test_runtime_factory()); + let endpoint = bridge.register(test_runtime_factory(), no_log()); let url = format!("ws://127.0.0.1:{}/?t={}", endpoint.port, endpoint.token); // Use a fresh `tokio` runtime on the test thread so the client does @@ -1273,33 +1051,26 @@ mod tests { drop(bridge); } - /// Two executions registered against the same shared listener land on - /// the same port with independent tokens, and a token only opens its own - /// execution's runtime, never the other's. #[test] fn two_executions_share_one_port_with_isolated_tokens() { let bridge = start_test_bridge(); - let first = bridge.register(test_runtime_factory()); - let second = bridge.register(test_runtime_factory()); + let first = bridge.register(test_runtime_factory(), no_log()); + let second = bridge.register(test_runtime_factory(), no_log()); assert_eq!(first.port, second.port); assert_ne!(first.token, second.token); - // Each token independently authenticates against the shared port. connect(first.port, &first.token); connect(second.port, &second.token); drop(bridge); } - /// A handshake presenting one execution's token must not be accepted as - /// belonging to a different execution's entry, and an unknown token is - /// rejected outright. #[test] fn wrong_or_unknown_token_is_rejected_at_handshake() { let bridge = start_test_bridge(); - let endpoint = bridge.register(test_runtime_factory()); - let _second = bridge.register(test_runtime_factory()); + let endpoint = bridge.register(test_runtime_factory(), no_log()); + let _second = bridge.register(test_runtime_factory(), no_log()); let rt = tokio::runtime::Builder::new_current_thread() .enable_all() @@ -1319,12 +1090,6 @@ mod tests { drop(bridge); } - /// Revoking one execution's token closes only its own connections and - /// rejects future handshakes against it, while a sibling execution on - /// the same shared listener keeps working. - /// Hands out the real runtime while keeping the control handle for the - /// connection it was built for, so a test can ask whether that connection's - /// runtime has been disposed. struct DisposalWatchFactory { inner: Arc, control: Mutex>, @@ -1338,11 +1103,8 @@ mod tests { } } - /// Ending a connection disposes the runtime that served it, so the host-core - /// subscriptions and chat state it held are released rather than left live - /// for the rest of the listener's life. #[test] - fn ending_a_connection_disposes_its_runtime() { + fn retained_controls_do_not_keep_ended_connections_alive() { fn is_closed(control: &crate::ProductRuntimeControl) -> bool { matches!( control.publish_chat_action(v01::HostChatActionSubscribeItem { @@ -1358,62 +1120,79 @@ mod tests { ) } - let bridge = start_test_bridge(); - let watch = Arc::new(DisposalWatchFactory { - inner: test_runtime_factory(), - control: Mutex::new(None), - }); - let endpoint = bridge.register(watch.clone()); - - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .expect("test runtime"); - let url = format!("ws://127.0.0.1:{}/?t={}", endpoint.port, endpoint.token); - - let control = rt.block_on(async { - let (mut ws, _) = tokio_tungstenite::connect_async(&url).await.expect("dial"); - let control = loop { - if let Some(control) = watch - .control - .lock() - .expect("disposal watch mutex poisoned") - .clone() - { - break control; - } - tokio::time::sleep(std::time::Duration::from_millis(5)).await; - }; - assert!( - !is_closed(&control), - "the runtime should be live while the connection is open" + for ending in ["close", "revoke", "shutdown"] { + let mut bridge = start_test_bridge(); + let watch = Arc::new(DisposalWatchFactory { + inner: test_runtime_factory(), + control: Mutex::new(None), + }); + let endpoint = bridge.register(watch.clone(), no_log()); + let client = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("client runtime"); + let mut socket = client.block_on(async { + tokio_tungstenite::connect_async(format!( + "ws://127.0.0.1:{}/?t={}", + endpoint.port, endpoint.token + )) + .await + .expect("connect") + .0 + }); + crate::test_support::wait_until( + || { + watch + .control + .lock() + .expect("control mutex poisoned") + .is_some() + }, + "connection did not create its runtime", ); - ws.close(None).await.expect("close client"); - control - }); - - let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); - while !is_closed(&control) { - assert!( - std::time::Instant::now() < deadline, - "the connection ended without disposing its runtime" + let control = watch + .control + .lock() + .expect("control mutex poisoned") + .clone() + .expect("connection control"); + assert!(!is_closed(&control), "connection must begin live"); + + match ending { + "close" => client.block_on(socket.close(None)).expect("close client"), + "revoke" => join_aborted_connections(bridge.revoke(&endpoint.token)), + "shutdown" => bridge.stop(), + _ => unreachable!(), + } + client.block_on(async { + tokio::time::timeout(std::time::Duration::from_secs(2), async { + if ending == "close" { + assert!( + matches!(socket.next().await, Some(Ok(WsMessage::Close(_)))), + "a healthy peer must receive its close acknowledgement" + ); + } + while let Some(Ok(_)) = socket.next().await {} + }) + .await + .unwrap_or_else(|_| { + panic!("{ending} left the socket open with a retained control") + }); + }); + crate::test_support::wait_until( + || bridge.registry.total_connections.load(Ordering::Acquire) == 0, + "ended connection did not release its capacity", ); - std::thread::sleep(std::time::Duration::from_millis(5)); + assert!(is_closed(&control), "{ending} did not dispose the runtime"); } - - drop(bridge); } - /// A peer holding every handshake slot open must not be able to lock a - /// legitimate connection out of the shared listener: a full backlog evicts - /// its oldest entry instead of refusing the newcomer. #[test] fn a_full_handshake_backlog_does_not_lock_out_a_new_connection() { let bridge = start_test_bridge(); - let endpoint = bridge.register(test_runtime_factory()); + let endpoint = bridge.register(test_runtime_factory(), no_log()); - // Raw TCP connections that never send a byte, so each one occupies a - // handshake slot until it is evicted or times out. + // Silent sockets fill the backlog without reaching token authentication. let addr = format!("127.0.0.1:{}", endpoint.port); let mut stalled = Vec::new(); for _ in 0..MAX_PENDING_HANDSHAKES { @@ -1440,14 +1219,11 @@ mod tests { drop(bridge); } - /// Revoking one execution's token tears down that execution's live - /// connection and stops its token authenticating, while a sibling - /// execution on the same listener keeps working. #[test] fn revoking_one_token_leaves_another_operational() { let bridge = start_test_bridge(); - let revoked = bridge.register(test_runtime_factory()); - let survives = bridge.register(test_runtime_factory()); + let revoked = bridge.register(test_runtime_factory(), no_log()); + let survives = bridge.register(test_runtime_factory(), no_log()); let rt = tokio::runtime::Builder::new_current_thread() .enable_all() @@ -1463,10 +1239,7 @@ mod tests { bridge.revoke(&revoked.token); - // The already-open connection is torn down. Aborting the task drops - // the socket without a clean close handshake, so the client - // observes either end-of-stream or a read error, not necessarily a - // `Close` frame. + // Cancellation need not complete a WebSocket close handshake. rt.block_on(async { let deadline = tokio::time::sleep(std::time::Duration::from_secs(2)); tokio::pin!(deadline); @@ -1485,13 +1258,11 @@ mod tests { } }); - // ...and its token no longer authenticates. let err = rt .block_on(async { tokio_tungstenite::connect_async(&revoked_url).await }) .expect_err("revoked token must be rejected"); assert!(format!("{err}").to_lowercase().contains("unauthorized")); - // The sibling execution is unaffected. let survives_url = format!("ws://127.0.0.1:{}/?t={}", survives.port, survives.token); rt.block_on(async { let (mut ws, _) = tokio_tungstenite::connect_async(&survives_url) @@ -1503,16 +1274,13 @@ mod tests { drop(bridge); } - /// After a token is revoked, registering a fresh execution (simulating a - /// reconnect / restart) gets a brand new token that works independently - /// of the old one. #[test] fn reconnecting_after_revoke_gets_a_fresh_token() { let bridge = start_test_bridge(); - let first = bridge.register(test_runtime_factory()); + let first = bridge.register(test_runtime_factory(), no_log()); bridge.revoke(&first.token); - let second = bridge.register(test_runtime_factory()); + let second = bridge.register(test_runtime_factory(), no_log()); assert_ne!(first.token, second.token); assert_eq!(first.port, second.port); @@ -1531,13 +1299,11 @@ mod tests { drop(bridge); } - /// Dropping the shared listener (host runtime shutdown) tears down every - /// registered execution's connections, not just one. #[test] fn host_shutdown_closes_every_registered_execution() { let bridge = start_test_bridge(); - let first = bridge.register(test_runtime_factory()); - let second = bridge.register(test_runtime_factory()); + let first = bridge.register(test_runtime_factory(), no_log()); + let second = bridge.register(test_runtime_factory(), no_log()); let rt = tokio::runtime::Builder::new_current_thread() .enable_all() @@ -1574,11 +1340,9 @@ mod tests { }); } - /// `SharedWsBridge` starts its listener lazily on first registration and - /// hands every subsequent registration the same port. #[test] fn shared_ws_bridge_lazily_starts_and_reuses_its_port() { - let shared = SharedWsBridge::new(); + let shared = SharedWsBridge::new(no_log()); let first = shared .register(0, test_runtime_factory(), no_log()) .expect("first registration starts the listener"); @@ -1593,25 +1357,19 @@ mod tests { connect(second.port, &second.token); shared.revoke(&first.token); - // The second registration is untouched by revoking the first. connect(second.port, &second.token); } - /// Once one execution has `MAX_WS_CONNECTIONS_PER_EXECUTION` live - /// connections, its next connection attempt is refused with a 503 even - /// though the shared listener's own total cap has plenty of room left. #[test] fn per_execution_cap_rejects_the_connection_past_the_limit() { let bridge = start_test_bridge(); - let endpoint = bridge.register(test_runtime_factory()); + let endpoint = bridge.register(test_runtime_factory(), no_log()); let url = format!("ws://127.0.0.1:{}/?t={}", endpoint.port, endpoint.token); let rt = tokio::runtime::Builder::new_current_thread() .enable_all() .build() .expect("test runtime"); - // Keep every socket alive so the execution's connection count never - // drops back below the cap while the next attempt is made. let _sockets = rt.block_on(async { let mut sockets = Vec::new(); for _ in 0..MAX_WS_CONNECTIONS_PER_EXECUTION { @@ -1635,22 +1393,23 @@ mod tests { drop(bridge); } - /// The shared listener's total connection cap is enforced independently - /// of any single execution's own cap: a brand-new execution (one that - /// has never opened a connection before, nowhere near its own - /// per-execution limit) is still refused once the shared budget is - /// gone. Simulates the budget already being exhausted by directly - /// setting the same atomic the real cap check reads, rather than - /// opening `MAX_TOTAL_WS_CONNECTIONS` real sockets just to reach it — - /// this exercises the exact same check with far less real I/O. #[test] - fn total_cap_rejects_the_connection_even_for_a_fresh_execution() { + fn total_capacity_is_reusable_after_a_retained_connection_closes() { let bridge = start_test_bridge(); - let extra = bridge.register(test_runtime_factory()); - bridge - .registry - .total_connections - .store(MAX_TOTAL_WS_CONNECTIONS, Ordering::SeqCst); + let extra = bridge.register(test_runtime_factory(), no_log()); + let controls = Arc::new(Mutex::new(Vec::new())); + let inner = test_runtime_factory(); + let factory: Arc = Arc::new({ + let controls = controls.clone(); + move |sink| { + let runtime = inner.product_runtime(sink); + controls + .lock() + .expect("controls mutex poisoned") + .push(runtime.control()); + runtime + } + }); let rt = tokio::runtime::Builder::new_current_thread() .enable_all() @@ -1658,6 +1417,26 @@ mod tests { .expect("test runtime"); let extra_url = format!("ws://127.0.0.1:{}/?t={}", extra.port, extra.token); + let mut sockets = rt.block_on(async { + let mut sockets = Vec::new(); + for _ in 0..MAX_TOTAL_WS_CONNECTIONS / MAX_WS_CONNECTIONS_PER_EXECUTION { + let endpoint = bridge.register(factory.clone(), no_log()); + for _ in 0..MAX_WS_CONNECTIONS_PER_EXECUTION { + let (socket, _) = tokio_tungstenite::connect_async(format!( + "ws://127.0.0.1:{}/?t={}", + endpoint.port, endpoint.token + )) + .await + .expect("connect within capacity"); + sockets.push(socket); + } + } + sockets + }); + crate::test_support::wait_until( + || controls.lock().expect("controls mutex poisoned").len() == MAX_TOTAL_WS_CONNECTIONS, + "connections did not create their controls", + ); let err = rt .block_on(async { tokio_tungstenite::connect_async(&extra_url).await }) .expect_err("a fresh execution must still be refused once the shared listener is full"); @@ -1667,18 +1446,25 @@ mod tests { "expected a 503 rejection past the total cap, got: {err}", ); + rt.block_on(async { + sockets[0].close(None).await.expect("close one client"); + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + if let Ok((socket, _)) = tokio_tungstenite::connect_async(&extra_url).await { + break socket; + } + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + } + }) + .await + .expect("a disconnected execution must release capacity for another product"); + }); + drop(bridge); } - /// A panic while building one execution's product runtime (e.g. a bug in - /// its own logic) tears down only that connection, not a sibling's — this - /// bridge does not, say, hold a lock across the panicking call that a - /// sibling connection also needs. Tokio's per-task panic containment is - /// what provides the isolation this test observes, but the release - /// profile sets `panic = "abort"`, so that containment — and thus this - /// test's premise — is a property of test builds only; a panic here in - /// an actual release build aborts the whole host process regardless of - /// which task it originated in. + // Tokio contains task panics in test builds. Release builds use panic=abort, + // so this test cannot promise panic isolation in a production host. #[test] fn a_panicking_execution_does_not_affect_a_sibling() { let panicking_factory: Arc = @@ -1687,16 +1473,14 @@ mod tests { }); let bridge = start_test_bridge(); - let failing = bridge.register(panicking_factory); - let healthy = bridge.register(test_runtime_factory()); + let failing = bridge.register(panicking_factory, no_log()); + let healthy = bridge.register(test_runtime_factory(), no_log()); let rt = tokio::runtime::Builder::new_current_thread() .enable_all() .build() .expect("test runtime"); - // The handshake succeeds (the token is valid); the connection is then - // torn down once its factory panics while building the runtime. rt.block_on(async { let failing_url = format!("ws://127.0.0.1:{}/?t={}", failing.port, failing.token); let (mut ws, _) = tokio_tungstenite::connect_async(&failing_url) @@ -1715,7 +1499,6 @@ mod tests { } }); - // The healthy sibling execution is unaffected. let healthy_url = format!("ws://127.0.0.1:{}/?t={}", healthy.port, healthy.token); rt.block_on(async { let (mut ws, _) = tokio_tungstenite::connect_async(&healthy_url) @@ -1727,17 +1510,11 @@ mod tests { drop(bridge); } - /// A peer that opens a TCP connection and never sends the HTTP upgrade - /// request — so its handshake never resolves — does not block a - /// sibling's connection attempt. Each accepted connection's handshake - /// runs in its own task, and the connection caps are reserved inside it - /// once a token matches, which is why an unauthenticated socket is bounded - /// by the handshake backlog rather than by those caps. #[test] fn a_stalled_handshake_does_not_block_a_sibling_connection() { let bridge = start_test_bridge(); - let stalled = bridge.register(test_runtime_factory()); - let healthy = bridge.register(test_runtime_factory()); + let stalled = bridge.register(test_runtime_factory(), no_log()); + let healthy = bridge.register(test_runtime_factory(), no_log()); let rt = tokio::runtime::Builder::new_current_thread() .enable_all() @@ -1745,8 +1522,6 @@ mod tests { .expect("test runtime"); rt.block_on(async { - // A raw TCP connection that never sends any bytes: the accept - // loop sees it, but its handshake never resolves. let _stalled_stream = tokio::net::TcpStream::connect(("127.0.0.1", stalled.port)) .await .expect("open a raw stream to the shared port"); @@ -1768,10 +1543,6 @@ mod tests { drop(bridge); } - /// Registers three tokens with distinguishable runtime factories and - /// confirms connecting through one token invokes only its own factory, - /// never a sibling's — ruling out a routing bug that only happens to - /// work by coincidence with exactly two registry entries. #[test] fn three_tokens_route_to_their_own_factory_only() { fn tracked_factory(called: Arc) -> Arc { @@ -1786,14 +1557,10 @@ mod tests { let calls: Vec> = (0..3).map(|_| Arc::new(AtomicUsize::new(0))).collect(); let endpoints: Vec = calls .iter() - .map(|called| bridge.register(tracked_factory(called.clone()))) + .map(|called| bridge.register(tracked_factory(called.clone()), no_log())) .collect(); - // Connect through the middle token specifically, and wait for a real - // response: the server cannot answer without having already called - // `product_runtime()` on the connection task, which is what makes - // this a reliable synchronization point rather than racing the - // server's own handling of the just-completed handshake. + // A response proves the factory ran; the HTTP upgrade alone does not. let rt = tokio::runtime::Builder::new_current_thread() .enable_all() .build()