diff --git a/crates/obscura-cdp/src/lib.rs b/crates/obscura-cdp/src/lib.rs index 9e0f7e8..1bda909 100644 --- a/crates/obscura-cdp/src/lib.rs +++ b/crates/obscura-cdp/src/lib.rs @@ -7,5 +7,6 @@ pub(crate) mod util; pub use server::{ start, start_with_full_options, start_with_full_serve_options, start_with_host, - start_with_host_and_security, start_with_options, + start_with_host_and_security, start_with_options, start_with_serve_options_and_limit, + DEFAULT_MAX_CONNECTIONS, }; diff --git a/crates/obscura-cdp/src/server.rs b/crates/obscura-cdp/src/server.rs index 21907d6..e52b356 100644 --- a/crates/obscura-cdp/src/server.rs +++ b/crates/obscura-cdp/src/server.rs @@ -1,6 +1,6 @@ use std::collections::HashMap; use std::net::SocketAddr; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::Arc; use futures_util::{SinkExt, StreamExt}; @@ -27,6 +27,25 @@ const MAX_DEFERRED_MESSAGES: usize = 256; // backlog still absorbs short-term spikes, but a long-term stall now // fails loudly at accept time rather than silently piling up FDs. const MAX_PENDING_WS_HANDOFFS: usize = 128; + +// Cap on *live* CDP connections, each of which costs one OS thread and its own +// V8 isolates. `MAX_PENDING_WS_HANDOFFS` above bounds only the handoff queue — +// connections that have already been handed off are unbounded without this. +// +// 128 matches the handoff bound and is well above any real client fan-out +// (Playwright/Puppeteer use one connection per browser). Threads are what this +// actually bounds: with arenas capped by `cap_malloc_arenas`, 128 idle +// connections cost 146 threads, 33.2 GiB of reserved address space and 51 MiB +// resident -- and nearly all of that 33.2 GiB is V8's process-wide sandbox, +// which is there at zero connections. Override with `--max-connections`. +pub const DEFAULT_MAX_CONNECTIONS: usize = 128; + +// Sent to a client that arrives while the server is at `max_connections`, in +// place of dropping the socket unexplained. The client sees a refusal it can +// retry rather than a bare connection reset. +const CONNECTION_LIMIT_RESPONSE: &str = "HTTP/1.1 503 Service Unavailable\r\n\ + Content-Length: 0\r\nConnection: close\r\n\ + X-Obscura-Reason: max-connections\r\n\r\n"; use crate::types::CdpRequest; struct CdpMessage { @@ -116,6 +135,35 @@ pub async fn start_with_full_serve_options( allow_file_access: bool, storage_dir: Option, allow_private_network: bool, +) -> anyhow::Result<()> { + start_with_serve_options_and_limit( + port, + host, + proxy, + stealth, + user_agent, + allow_file_access, + storage_dir, + allow_private_network, + DEFAULT_MAX_CONNECTIONS, + ) + .await +} + +/// As `start_with_full_serve_options`, with an explicit cap on live CDP +/// connections. Each connection owns an OS thread and its pages' V8 isolates, +/// so this is what bounds the server's thread and memory footprint. +#[allow(clippy::too_many_arguments)] +pub async fn start_with_serve_options_and_limit( + port: u16, + host: &str, + proxy: Option, + stealth: bool, + user_agent: Option, + allow_file_access: bool, + storage_dir: Option, + allow_private_network: bool, + max_connections: usize, ) -> anyhow::Result<()> { let ip: std::net::IpAddr = host .parse() @@ -243,6 +291,11 @@ pub async fn start_with_full_serve_options( cap_malloc_arenas(); + // Live CDP connections, incremented on accept and decremented when a + // connection thread exits (see `run_connection`). + let live_connections = Arc::new(AtomicUsize::new(0)); + info!("Connection limit: {}", max_connections); + // Accept loop: hand each WebSocket connection to its own OS thread so its // pages' isolates live on a dedicated thread. loop { @@ -266,7 +319,28 @@ pub async fn start_with_full_serve_options( .set_nodelay(true) .map_err(|e| error!("set_nodelay on WS stream: {}", e)) .ok(); - run_connection(stream, shared_ctx.clone(), shutdown_notify.clone()); + // Reserve a slot before spawning. `fetch_update` (rather than a load + // then a store) keeps the check atomic against the accept thread + // handing off the next stream concurrently. + let reserved = live_connections + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |n| { + (n < max_connections).then_some(n + 1) + }) + .is_ok(); + if !reserved { + warn!( + "refusing CDP connection: at --max-connections ({})", + max_connections + ); + refuse_connection(stream); + continue; + } + run_connection( + stream, + shared_ctx.clone(), + shutdown_notify.clone(), + live_connections.clone(), + ); } // Server is shutting down: persist the shared cookie jar once. @@ -327,10 +401,24 @@ fn run_connection( std_stream: std::net::TcpStream, default_context: Arc, shutdown_notify: Arc, + live_connections: Arc, ) { - std::thread::Builder::new() + // Releases the slot reserved by the accept loop when the thread unwinds, + // however it exits — clean close, error return, or panic. A plain + // decrement at the end of the closure would leak slots on the early + // returns below until the cap wedged the server shut. + struct SlotGuard(Arc); + impl Drop for SlotGuard { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::AcqRel); + } + } + + let slot = live_connections.clone(); + let spawned = std::thread::Builder::new() .name("obscura-cdp-conn".into()) .spawn(move || { + let _slot = SlotGuard(slot); let rt = match tokio::runtime::Builder::new_current_thread() .enable_all() .build() @@ -363,8 +451,27 @@ fn run_connection( // processor so the thread can exit. processor.abort(); }); - }) - .ok(); + }); + + // The closure never ran, so its `SlotGuard` never existed: release the + // reserved slot here or the cap drifts down on every failed spawn. + if let Err(e) = spawned { + error!("connection thread spawn failed: {}", e); + live_connections.fetch_sub(1, Ordering::AcqRel); + } +} + +/// Turn away a connection that arrived while the server was at its limit. +/// +/// Best-effort: the socket is going away either way, so a failed write just +/// means the client sees a reset instead of the 503. +fn refuse_connection(stream: std::net::TcpStream) { + use std::io::Write; + let mut stream = stream; + let _ = stream.set_nonblocking(false); + let _ = stream.write_all(CONNECTION_LIMIT_RESPONSE.as_bytes()); + let _ = stream.flush(); + let _ = stream.shutdown(std::net::Shutdown::Both); } const HTTP_PEEK_BUF: usize = 4096; diff --git a/crates/obscura-cdp/tests/max_connections_cap.rs b/crates/obscura-cdp/tests/max_connections_cap.rs new file mode 100644 index 0000000..4d78c4a --- /dev/null +++ b/crates/obscura-cdp/tests/max_connections_cap.rs @@ -0,0 +1,162 @@ +//! `--max-connections` bounds the thread-per-connection server. +//! +//! Each CDP connection owns an OS thread and its pages' V8 isolates, so without +//! a cap a client can grow the server's thread count and memory without limit. +//! The cap must do three things, and this test pins all three: +//! +//! 1. connections up to the limit are accepted and usable; +//! 2. the one past the limit is refused with an explicit `503` carrying +//! `X-Obscura-Reason: max-connections`, not dropped with a bare reset; +//! 3. closing a connection frees its slot, so the server recovers rather than +//! wedging shut once it has ever been full. +//! +//! (3) is the one that actually bites: a slot leaked on any connection-teardown +//! path would only show up as a server that stops accepting after N lifetime +//! connections, which no other test would catch. +//! +//! Run with `cargo nextest run -p obscura-cdp -E 'test(max_connections_cap)'`. + +use std::time::Duration; + +use futures_util::{SinkExt, StreamExt}; +use serde_json::json; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpListener; +use tokio_tungstenite::{connect_async, tungstenite::Message}; + +const LIMIT: usize = 2; + +async fn pick_port() -> u16 { + let l = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = l.local_addr().unwrap().port(); + drop(l); + port +} + +/// A CDP connection that is actually driven, not just opened: `Target.getTargets` +/// round-trips through the connection's own processor, proving the connection +/// holds a live slot rather than a socket the server has forgotten about. +async fn open_and_use( + ws_port: u16, + id: u64, +) -> Result>, String> +{ + let url = format!("ws://127.0.0.1:{}/devtools/browser", ws_port); + let (mut ws, _) = connect_async(&url).await.map_err(|e| e.to_string())?; + ws.send(Message::Text( + json!({"id": id, "method": "Target.getTargets"}).to_string().into(), + )) + .await + .map_err(|e| e.to_string())?; + + let deadline = tokio::time::Instant::now() + Duration::from_secs(10); + loop { + if tokio::time::Instant::now() >= deadline { + return Err("timeout waiting for Target.getTargets reply".into()); + } + let msg = tokio::time::timeout(Duration::from_secs(5), ws.next()) + .await + .map_err(|_| "timeout".to_string())? + .ok_or("ws closed")? + .map_err(|e| e.to_string())?; + if let Message::Text(t) = msg { + let v: serde_json::Value = serde_json::from_str(&t).map_err(|e| e.to_string())?; + if v.get("id").and_then(|i| i.as_u64()) == Some(id) { + return Ok(ws); + } + } + } +} + +/// Raw handshake, returning the server's status line. Deliberately not +/// `connect_async`: a refusal is an HTTP response, and we want to read it +/// rather than have the websocket client turn it into an opaque error. +async fn raw_handshake_status(ws_port: u16) -> String { + let mut sock = tokio::net::TcpStream::connect(("127.0.0.1", ws_port)) + .await + .expect("connect"); + let req = format!( + "GET /devtools/browser HTTP/1.1\r\nHost: 127.0.0.1:{}\r\n\ + Upgrade: websocket\r\nConnection: Upgrade\r\n\ + Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\r\n", + ws_port + ); + sock.write_all(req.as_bytes()).await.expect("write"); + let mut buf = vec![0u8; 1024]; + let n = tokio::time::timeout(Duration::from_secs(5), sock.read(&mut buf)) + .await + .expect("read timed out") + .expect("read"); + String::from_utf8_lossy(&buf[..n]).to_string() +} + +#[test] +fn max_connections_refuses_then_recovers() { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let local = tokio::task::LocalSet::new(); + + local.block_on(&rt, async move { + let ws_port = pick_port().await; + tokio::task::spawn_local(async move { + let _ = obscura_cdp::start_with_serve_options_and_limit( + ws_port, + "127.0.0.1", + None, + false, + None, + false, + None, + true, + LIMIT, + ) + .await; + }); + tokio::time::sleep(Duration::from_millis(300)).await; + + // 1. Fill the server to its limit. + let mut held = Vec::new(); + for i in 0..LIMIT { + held.push( + open_and_use(ws_port, 100 + i as u64) + .await + .unwrap_or_else(|e| panic!("connection {} within the limit must be accepted: {}", i, e)), + ); + } + + // 2. One past the limit is refused, and says why. + let refused = raw_handshake_status(ws_port).await; + assert!( + refused.starts_with("HTTP/1.1 503"), + "over-limit connection must be refused with 503, got: {:?}", + refused.lines().next() + ); + assert!( + refused.contains("X-Obscura-Reason: max-connections"), + "refusal must name the reason so a client can tell it apart from a crash: {:?}", + refused + ); + + // 3. Free one slot and confirm the server accepts again. Closing is + // asynchronous (the connection thread unwinds and drops its guard), + // so poll rather than assume an instant handover. + let mut freed = held.pop().unwrap(); + let _ = freed.close(None).await; + drop(freed); + + let mut reconnected = None; + for _ in 0..40 { + tokio::time::sleep(Duration::from_millis(100)).await; + if let Ok(ws) = open_and_use(ws_port, 900).await { + reconnected = Some(ws); + break; + } + } + assert!( + reconnected.is_some(), + "a closed connection must release its slot; the server is wedged at the cap" + ); + }); +} diff --git a/crates/obscura-cli/src/main.rs b/crates/obscura-cli/src/main.rs index 1210131..5b8d63f 100644 --- a/crates/obscura-cli/src/main.rs +++ b/crates/obscura-cli/src/main.rs @@ -78,6 +78,13 @@ enum Command { #[arg(long, default_value_t = 1)] workers: u16, + /// Maximum live CDP connections. Each connection runs on its own OS + /// thread with its own V8 isolates, so this bounds the server's thread + /// and memory footprint. Connections beyond the limit are refused with + /// a 503 rather than queued. + #[arg(long, default_value_t = obscura_cdp::DEFAULT_MAX_CONNECTIONS)] + max_connections: usize, + /// Allow CDP clients to navigate to file:// URLs. Off by /// default so a CDP connection cannot read arbitrary local /// files. Enable only when serving local HTML for testing @@ -328,7 +335,7 @@ async fn main() -> anyhow::Result<()> { let stealth = args.stealth; match args.command { - Some(Command::Serve { port, host, proxy, user_agent, workers, allow_file_access, storage_dir, quiet: _ }) => { + Some(Command::Serve { port, host, proxy, user_agent, workers, max_connections, allow_file_access, storage_dir, quiet: _ }) => { // Fall back to OBSCURA_PROXY so a proxy can be supplied without // putting credentials on the command line. The multi-worker load // balancer passes the proxy to each worker this way (issue #366). @@ -357,9 +364,9 @@ async fn main() -> anyhow::Result<()> { tracing::info!("{} worker processes", workers); run_multi_worker_serve(port, host, workers, proxy, stealth, user_agent).await?; } else { - obscura_cdp::start_with_full_serve_options( + obscura_cdp::start_with_serve_options_and_limit( port, &host, proxy, stealth, user_agent, allow_file_access, storage_dir, - args.allow_private_network, + args.allow_private_network, max_connections, ).await?; } }