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/android/truapi-host/README.md b/android/truapi-host/README.md index 23b6fb6c2..569ee8a6e 100644 --- a/android/truapi-host/README.md +++ b/android/truapi-host/README.md @@ -67,7 +67,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/android/truapi-host/src/main/kotlin/io/parity/truapi/TrUAPIHost.kt b/android/truapi-host/src/main/kotlin/io/parity/truapi/TrUAPIHost.kt index f19341999..5c0e37925 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 @@ -974,11 +974,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/host_core.rs b/rust/crates/truapi-server/src/host_core.rs index c139a6569..d4211d3eb 100644 --- a/rust/crates/truapi-server/src/host_core.rs +++ b/rust/crates/truapi-server/src/host_core.rs @@ -1408,10 +1408,19 @@ 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); + // + // 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 + .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; @@ -1502,16 +1511,21 @@ 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) { + // Aborting under the lock can wake code that re-enters disposal. + if self.disposed.load(Ordering::Acquire) { return; } - for (_, handle) in self - .in_flight - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .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(); + } } self.admin.product_runtime.detach_chat(); self.admin.product_runtime.detach_renderer(); @@ -1628,7 +1642,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, @@ -2765,6 +2779,92 @@ mod tests { assert_eq!(response.payload.value, expected); } + // The debug tap deliberately blocks to expose the registration race reliably. + #[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"); + }) + }; + + 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)); diff --git a/rust/crates/truapi-server/src/native.rs b/rust/crates/truapi-server/src/native.rs index d31596b3b..89b0a4a45 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`. @@ -697,6 +697,8 @@ pub struct NativeTrUApiHostRuntime { events: Arc, #[cfg(feature = "ws-bridge")] spawner: Spawner, + #[cfg(feature = "ws-bridge")] + ws_bridge: Arc, /// The one Worker execution per product; opening another replaces it. worker_executions: Mutex>>, } @@ -739,6 +741,10 @@ impl NativeTrUApiHostRuntime { events, #[cfg(feature = "ws-bridge")] spawner, + #[cfg(feature = "ws-bridge")] + 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()), })) } @@ -789,7 +795,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)), }); @@ -1181,7 +1189,9 @@ pub struct NativeProductExecution { >, closed: AtomicBool, #[cfg(feature = "ws-bridge")] - bridge: Mutex>, + ws_bridge: Arc, + #[cfg(feature = "ws-bridge")] + bridge_token: Mutex>, #[cfg(feature = "ws-bridge")] product_control: Arc>>, } @@ -1224,13 +1234,14 @@ impl NativeProductExecution { #[cfg(feature = "ws-bridge")] fn stop_bridge(&self) { - if let Some(mut bridge) = self - .bridge + // Release the token lock before waiting for connection cancellation. + let token = self + .bridge_token .lock() .expect("native product bridge mutex poisoned") - .take() - { - bridge.stop(); + .take(); + if let Some(token) = token { + self.ws_bridge.revoke(&token); } *self .product_control @@ -1409,7 +1420,8 @@ impl NativeProductExecution { #[cfg(feature = "ws-bridge")] #[uniffi::export] impl NativeProductExecution { - /// Start this execution's independently authenticated localhost bridge. + /// 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( @@ -1417,7 +1429,7 @@ impl NativeProductExecution { )); } let mut guard = self - .bridge + .bridge_token .lock() .expect("native product bridge mutex poisoned"); if guard.is_some() { @@ -1441,12 +1453,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(); } @@ -2427,6 +2441,7 @@ mod tests { } struct EventCallbacks { + logs: Mutex>, chat_room_status: Mutex, chat_created_rooms: Mutex>, chat_bot_status: Mutex, @@ -2461,6 +2476,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), @@ -2491,7 +2507,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() @@ -4276,6 +4294,247 @@ mod tests { execution.stop_ws_bridge(); } + #[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() { + 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, + 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), + None, + 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 { + 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], + }) + .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"); + + 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"); + + 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" + ); + + 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 1ddc6b886..2b552bef5 100644 --- a/rust/crates/truapi-server/src/ws_bridge.rs +++ b/rust/crates/truapi-server/src/ws_bridge.rs @@ -5,41 +5,44 @@ //! //! 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. +//! Executions under one host share a [`SharedWsBridge`] listener and the +//! process-wide `tokio` runtime, with independent tokens and connections. //! -//! 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. +//! 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. +//! +//! 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; 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}; 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}; -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 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; +// Allow reconnect overlap without one execution exhausting the shared limit. +const MAX_WS_CONNECTIONS_PER_EXECUTION: usize = 8; + +// Bound resource use even when a peer holds several valid execution tokens. +const MAX_TOTAL_WS_CONNECTIONS: usize = 64; + +// 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 /// cannot make the core buffer responses without limit; once the queue fills @@ -51,6 +54,9 @@ const OUTBOUND_QUEUE_CAP: usize = 4096; /// memory-amplification DoS well below tungstenite's 64 MiB default. const MAX_WS_MESSAGE_BYTES: usize = 8 << 20; +// 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` /// to build its WebSocket URL (e.g. `ws://127.0.0.1:/?t=`). #[derive(Clone, Debug, uniffi::Record)] @@ -66,7 +72,7 @@ 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 has a registered bridge token. #[error("ws bridge already running")] AlreadyRunning, /// Anything else (bind failure, runtime spin-up failure, ...). @@ -159,33 +165,200 @@ fn shared_native_executor() -> io::Result<(&'static SharedNativeExecutor, bool)> Ok((executor, initialized)) } -/// Running bridge 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 { +struct RegistryEntry { + runtime_factory: Arc, + logger: BridgeLogger, + connection_count: Arc, + connections: Mutex, +} + +// One lock prevents a completed handshake from registering past revocation. +#[derive(Default)] +struct EntryConnections { + revoked: bool, + handles: Vec>, +} + +impl EntryConnections { + fn abort_and_take(&mut self) -> Vec> { + for handle in self.handles.iter() { + handle.abort(); + } + std::mem::take(&mut self.handles) + } +} + +#[derive(Default)] +struct WsBridgeRegistry { + entries: Mutex>>, + total_connections: Arc, +} + +impl WsBridgeRegistry { + fn insert( + &self, + token: String, + runtime_factory: Arc, + logger: BridgeLogger, + ) { + self.entries + .lock() + .expect("ws bridge registry mutex poisoned") + .insert( + token, + Arc::new(RegistryEntry { + runtime_factory, + logger, + connection_count: Arc::new(AtomicUsize::new(0)), + connections: Mutex::new(EntryConnections::default()), + }), + ); + } + + fn revoke(&self, token: &str) -> Vec> { + let Some(entry) = self + .entries + .lock() + .expect("ws bridge registry mutex poisoned") + .remove(token) + else { + return Vec::new(); + }; + let mut state = entry + .connections + .lock() + .expect("ws bridge registry entry mutex poisoned"); + state.revoked = true; + state.abort_and_take() + } + + // Scan all candidates so an early match cannot reveal token order. + fn find_matching(&self, path_and_query: Option<&str>) -> Option> { + // Peer-supplied query processing must not hold up registry updates. + let candidates: Vec<(String, Arc)> = self + .entries + .lock() + .expect("ws bridge registry mutex poisoned") + .iter() + .map(|(token, entry)| (token.clone(), entry.clone())) + .collect(); + let mut found = None; + for (token, entry) in candidates.iter() { + if path_token_matches(path_and_query, token) { + found = Some(entry.clone()); + } + } + found + } + + // Revoked entries are already removed; their caller owns the aborted tasks. + 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"); + all.append(&mut state.abort_and_take()); + } + all + } +} + +/// Lazy listener shared by a host runtime and its product executions. +pub struct SharedWsBridge { + inner: Mutex>, + logger: BridgeLogger, +} + +impl SharedWsBridge { + /// Construct a lazy listener with host-owned lifecycle logging. + pub fn new(logger: BridgeLogger) -> Self { + Self { + inner: Mutex::new(None), + logger, + } + } + + /// Register an execution with its own token and connection logger. + /// + /// 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 { + // 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, self.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, logger) + }; + for (event, detail) in &pending_logs { + (self.logger)(event, detail); + } + Ok(endpoint) + } + + /// Revoke one execution's token. No-op if the listener was never + /// started or the token is unknown. + /// + /// 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"); + match guard.as_ref() { + Some(bridge) => bridge.revoke(token), + None => return, + } + }; + join_aborted_connections(aborted); + } +} + +/// 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>, 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( + // Return startup logs so the caller can emit them outside its lock. + 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); - + ) -> io::Result<(Self, Vec<(&'static str, String)>)> { // 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)?; @@ -194,15 +367,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. @@ -210,29 +374,32 @@ impl WsBridge { let _entered = handle.enter(); TcpListener::from_std(std_listener)? }; + + // 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(( + "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::<()>(); - 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( + pending_logs.push(( "truapi.ws_bridge.started", - &format!( - "port={port} token_len={} runtime_id={runtime_id}", - token.len() - ), - ); + format!("port={port} runtime_id={runtime_id}"), + )); Ok(( Self { @@ -240,23 +407,45 @@ impl WsBridge { stopped: Some(stopped_rx), accept_task: Some(accept_task), runtime_id, + registry, + port, }, - WsBridgeEndpoint { port, token }, + pending_logs, )) } - /// Signal this bridge's accept loop to exit without stopping the shared - /// native executor used by other products. - pub fn stop(&mut self) { + 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, logger); + WsBridgeEndpoint { + port: self.port, + token, + } + } + + // 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) { + return Vec::new(); + } + connections + } + + // 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(()); } - // 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; the shutdown signal still lets the task clean itself up. + // 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 { @@ -274,6 +463,8 @@ impl WsBridge { { task.abort(); } + // Cover the non-blocking path or an accept loop that failed to stop. + drop(self.registry.take_all_handles()); } } @@ -283,23 +474,45 @@ impl Drop for WsBridge { } } +// 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; + } + 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, - runtime_factory: Arc, - expected_token: String, + registry: Arc, logger: BridgeLogger, mut shutdown: oneshot::Receiver<()>, ) { - let mut handles: Vec> = Vec::new(); + // Independent setup tasks keep a stalled handshake from blocking acceptance. + let mut setup_tasks: VecDeque> = VecDeque::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 task in setup_tasks { + let _ = task.await; } - for h in handles { - let _ = h.await; + for handle in registry.take_all_handles() { + let _ = handle.await; } break; } @@ -311,48 +524,171 @@ 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; + setup_tasks.retain(|task| !task.is_finished()); + // 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() + { + oldest.abort(); + logger("truapi.ws_bridge.handshake_backlog_evicted", &peer.to_string()); } - let runtime_factory = runtime_factory.clone(); + 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_back(tokio::spawn(async move { + connection_setup(stream, peer, registry, logger).await; })); } } } } -// `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 connection_setup( stream: tokio::net::TcpStream, peer: SocketAddr, - runtime_factory: Arc, - expected_token: String, + registry: Arc, logger: BridgeLogger, ) { - 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) + 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 logger = entry.logger.clone(); + 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_entry = entry.clone(); + state.handles.push(tokio::spawn(async move { + let _guard = guard; + connection_lifecycle(ws, peer, conn_entry).await; + })); + false + } + }; + // Host callbacks may re-enter revocation, so logging must stay outside the lock. + if revoked { + logger( + "truapi.ws_bridge.connection_revoked_during_setup", + &peer.to_string(), + ); + } +} + +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); + } +} + +// Task cancellation skips explicit cleanup after an await, but still drops guards. +struct DisposeGuard(Arc); + +impl Drop for DisposeGuard { + fn drop(&mut self) { + self.0.dispose(); + } +} + +type AuthenticatedConnection = ( + WebSocketStream, + Arc, + ConnectionCountGuard, +); + +type MatchedReservation = Arc, ConnectionCountGuard)>>>; + +// Concurrent handshakes must not both claim the last available slot. +fn try_reserve(counter: &AtomicUsize, limit: usize) -> bool { + // `fetch_update` is deprecated on nightly; `try_update` is not yet stable. + 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, + } + } +} + +// 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, + peer: SocketAddr, + 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 = 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); + }; + + // Reject over-cap connections before acknowledging the HTTP upgrade. + if !try_reserve(&entry.connection_count, MAX_WS_CONNECTIONS_PER_EXECUTION) { + (entry.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) { + entry.connection_count.fetch_sub(1, Ordering::AcqRel); + (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); + } + // 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(), + }; + + *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,77 +704,83 @@ async fn handle_connection( Ok(ws) => ws, Err(err) => { logger("truapi.ws_bridge.handshake_error", &err.to_string()); - return; + return None; } }; - logger("truapi.ws_bridge.connection_open", &peer.to_string()); + let (entry, guard) = matched + .lock() + .expect("ws bridge handshake mutex poisoned") + .take() + .expect("a successful upgrade always resolved a matching registry entry"); + (entry.logger)("truapi.ws_bridge.connection_open", &peer.to_string()); + Some((ws, entry, guard)) +} + +async fn connection_lifecycle( + ws: WebSocketStream, + peer: SocketAddr, + entry: Arc, +) { + 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(runtime_factory.product_runtime(frame_sink)); - - 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; - }); + let product_runtime = Arc::new(entry.runtime_factory.product_runtime(frame_sink)); + let dispose_guard = DisposeGuard(product_runtime.clone()); // 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(); - } - - product_runtime.dispose(); - let _ = pump.await; + // A slow peer must not retain capacity while its close reply waits. + let _ = sink.close().now_or_never(); + drop(in_flight); + drop(dispose_guard); logger("truapi.ws_bridge.connection_closed", &peer.to_string()); } +// 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; @@ -447,16 +789,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 @@ -513,6 +856,10 @@ mod tests { use crate::frame::{Payload, ProtocolMessage, request_ids}; use crate::test_support::{StubPlatform, test_spawner}; + 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())) } @@ -538,6 +885,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")); @@ -548,6 +909,13 @@ mod tests { assert!(!path_token_matches(None, "abc")); } + #[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"); @@ -600,8 +968,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 = start_test_bridge(); let (executor, _) = shared_native_executor().expect("shared native executor"); let (dropped_tx, dropped_rx) = std::sync::mpsc::channel(); @@ -615,15 +982,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 = start_test_bridge(); + 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 @@ -683,97 +1049,573 @@ 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. #[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 = start_test_bridge(); + 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); - let ids = runtime_ids.lock().unwrap().clone(); - assert_eq!(ids.len(), 2); - assert_eq!(ids[0], ids[1]); + connect(first.port, &first.token); + connect(second.port, &second.token); - first.stop(); + drop(bridge); + } + + #[test] + fn wrong_or_unknown_token_is_rejected_at_handshake() { + let bridge = start_test_bridge(); + 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() + .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); + } + + 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 + } + } + + #[test] + fn retained_controls_do_not_keep_ended_connections_alive() { + 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) + ) + } + + 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", + ); + 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", + ); + assert!(is_closed(&control), "{ending} did not dispose the runtime"); + } + } + + #[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(), no_log()); + + // 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 { + 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); - let client = tokio::runtime::Builder::new_current_thread() + 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 = start_test_bridge(); + 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() .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); + + // 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); + 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, + } + } + } + } + }); + + 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")); + + 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. #[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 = start_test_bridge(); + let first = bridge.register(test_runtime_factory(), no_log()); + bridge.revoke(&first.token); + + let second = bridge.register(test_runtime_factory(), no_log()); + 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); + } + + #[test] + fn host_shutdown_closes_every_registered_execution() { + let bridge = start_test_bridge(); + 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() + .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() {} + } => {} + } + }); + } + + #[test] + fn shared_ws_bridge_lazily_starts_and_reuses_its_port() { + let shared = SharedWsBridge::new(no_log()); + 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); + connect(second.port, &second.token); + } + + #[test] + fn per_execution_cap_rejects_the_connection_past_the_limit() { + let bridge = start_test_bridge(); + 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"); + 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. #[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_capacity_is_reusable_after_a_retained_connection_closes() { + let bridge = start_test_bridge(); + 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() + .build() + .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"); + let msg = format!("{err}").to_lowercase(); + assert!( + msg.contains("503") || msg.contains("service unavailable"), + "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); + } + + // 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 = + Arc::new(|_sink: Arc| -> ProductRuntime { + panic!("intentional test panic: simulating a failing product execution") + }); + + let bridge = start_test_bridge(); + 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"); + + 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, + } + } + } + }); + + 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); + } + + #[test] + fn a_stalled_handshake_does_not_block_a_sibling_connection() { + let bridge = start_test_bridge(); + 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() + .build() + .expect("test runtime"); + + rt.block_on(async { + 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); + } + + #[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) + }) + } - // 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 = start_test_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()), no_log())) + .collect(); + + // A response proves the factory ran; the HTTP upgrade alone does not. + 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 { + 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], + }, + ) + .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); } }