diff --git a/crates/ui-egui/src/control.rs b/crates/ui-egui/src/control.rs index eaa102b56..753971678 100644 --- a/crates/ui-egui/src/control.rs +++ b/crates/ui-egui/src/control.rs @@ -780,27 +780,113 @@ pub struct Endpoint { pub token: String, } +/// What a connection may do before it authenticates, so an unauthenticated local client cannot +/// hold a thread for ever or exhaust them (#1806). +#[cfg(not(target_arch = "wasm32"))] +#[derive(Clone, Copy)] +struct Limits { + /// How long a connection has to authenticate, counted from its acceptance: an absolute + /// deadline, so a client feeding one byte at a time cannot extend it. + auth_deadline: std::time::Duration, + /// How many connections may wait for authentication at once; further ones are refused. + max_unauthenticated: usize, + /// How long a write may block on a client that never reads its replies. + write_timeout: std::time::Duration, +} + +#[cfg(not(target_arch = "wasm32"))] +const LIMITS: Limits = + Limits { auth_deadline: std::time::Duration::from_secs(10), max_unauthenticated: 8, write_timeout: std::time::Duration::from_secs(30) }; + +/// One of the [`Limits::max_unauthenticated`] places, given back when dropped: when the +/// connection authenticates, and on every way its thread can end. +#[cfg(not(target_arch = "wasm32"))] +struct Slot(std::sync::Arc); + +#[cfg(not(target_arch = "wasm32"))] +impl Slot { + fn take(waiting: &std::sync::Arc, max: usize) -> Option { + use std::sync::atomic::Ordering; + let mut n = waiting.load(Ordering::Acquire); + loop { + if n >= max { + return None; + } + match waiting.compare_exchange_weak(n, n.saturating_add(1), Ordering::AcqRel, Ordering::Acquire) { + Ok(_) => return Some(Slot(waiting.clone())), + Err(now) => n = now, + } + } + } +} + +#[cfg(not(target_arch = "wasm32"))] +impl Drop for Slot { + fn drop(&mut self) { + self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel); + } +} + +/// The reading half of a connection, with the authentication deadline: every read before it +/// gets the time left as its timeout, and none at all once it has passed. +#[cfg(not(target_arch = "wasm32"))] +struct Reading { + stream: std::net::TcpStream, + until: Option, +} + +#[cfg(not(target_arch = "wasm32"))] +impl std::io::Read for Reading { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + if let Some(until) = self.until { + let left = until.saturating_duration_since(std::time::Instant::now()); + if left.is_zero() { + return Err(std::io::Error::new(std::io::ErrorKind::TimedOut, "the authentication deadline passed")); + } + self.stream.set_read_timeout(Some(left))?; + } + self.stream.read(buf) + } +} + /// Listen on `127.0.0.1` (random port) and forward authenticated requests to `client`. /// /// Protocol: newline-delimited JSON-RPC 2.0. The first request on a connection must be /// `{"method": "auth", "params": {"token": …}}`; anything else closes the connection, and so does -/// a line longer than [`MAX_AUTH_LINE`] before that. +/// a line longer than [`MAX_AUTH_LINE`] before that, or [`LIMITS`] passing without it. #[cfg(not(target_arch = "wasm32"))] pub fn serve(client: ControlClient) -> std::io::Result { + serve_with(client, LIMITS) +} + +#[cfg(not(target_arch = "wasm32"))] +fn serve_with(client: ControlClient, limits: Limits) -> std::io::Result { use std::io::{BufReader, Write}; let listener = std::net::TcpListener::bind(("127.0.0.1", 0))?; let port = listener.local_addr()?.port(); let token = random_token()?; let expected = token.clone(); + let waiting = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); std::thread::Builder::new().name("pdfcraft-control".into()).spawn(move || { for stream in listener.incoming().flatten() { + let _ = stream.set_write_timeout(Some(limits.write_timeout)); + // Admitted only while fewer than the limit wait to authenticate: refused on the + // accepting thread, so a refused connection costs no thread at all. + let Some(slot) = Slot::take(&waiting, limits.max_unauthenticated) else { + let mut write = stream; + let e = rpc_error(Value::Null, -32003, "too many connections waiting to authenticate; try again later"); + let _ = writeln!(write, "{e}").and_then(|_| write.flush()); + continue; + }; let client = client.clone(); let expected = expected.clone(); let _ = std::thread::Builder::new().name("pdfcraft-control-conn".into()).spawn(move || { + // Held until `auth` succeeds; given back however this thread ends. + let mut slot = Some(slot); let Ok(read) = stream.try_clone() else { return }; let mut write = stream; let mut authed = false; - let mut reader = BufReader::new(read); + let mut reader = BufReader::new(Reading { stream: read, until: Some(std::time::Instant::now() + limits.auth_deadline) }); let mut buf = Vec::new(); loop { // Until `auth` succeeds, lines are capped at MAX_AUTH_LINE; afterwards they @@ -808,6 +894,13 @@ pub fn serve(client: ControlClient) -> std::io::Result { let max = if authed { usize::MAX } else { MAX_AUTH_LINE }; match read_line_limited(&mut reader, &mut buf, max) { Ok(LineRead::Line) => {} + // The deadline passed before `auth`: say so and close. + Err(e) if !authed && matches!(e.kind(), std::io::ErrorKind::TimedOut | std::io::ErrorKind::WouldBlock) => { + let secs = limits.auth_deadline.as_secs_f64(); + let e = rpc_error(Value::Null, -32004, &format!("not authenticated within {secs} s")); + let _ = writeln!(write, "{e}").and_then(|_| write.flush()); + break; + } // End of stream or a failed read ends the connection, as before. Ok(LineRead::End) | Err(_) => break, Ok(LineRead::TooLong) => { @@ -833,6 +926,10 @@ pub fn serve(client: ControlClient) -> std::io::Result { method == "auth" && params.get("token").and_then(Value::as_str).is_some_and(|t| constant_time_eq(t, &expected)); authed = ok; if ok { + // Authenticated: no deadline from here on, and its place is free. + reader.get_mut().until = None; + let _ = reader.get_mut().stream.set_read_timeout(None); + drop(slot.take()); (json!({ "jsonrpc": "2.0", "id": id, "result": { "ok": true } }), false) } else { (rpc_error(id, -32001, "authenticate first: auth {token}"), true) @@ -985,6 +1082,109 @@ mod transport_tests { assert!(closed, "the connection stayed open: {next:?}"); } + /// Small limits for the tests below: a 300 ms deadline, two places to wait in. + #[cfg(not(target_arch = "wasm32"))] + const TEST_LIMITS: super::Limits = super::Limits { + auth_deadline: std::time::Duration::from_millis(300), + max_unauthenticated: 2, + write_timeout: std::time::Duration::from_secs(5), + }; + + /// The next line the server sends, as JSON, or `None` when it closed the connection instead. + #[cfg(not(target_arch = "wasm32"))] + fn next_line(r: &mut std::io::BufReader) -> Option { + use std::io::BufRead; + let mut line = String::new(); + match r.read_line(&mut line) { + Ok(0) | Err(_) => None, + Ok(_) => Some(serde_json::from_str(&line).unwrap()), + } + } + + #[cfg(not(target_arch = "wasm32"))] + fn connect(port: u16) -> std::net::TcpStream { + let s = std::net::TcpStream::connect(("127.0.0.1", port)).unwrap(); + s.set_read_timeout(Some(std::time::Duration::from_secs(5))).unwrap(); + s + } + + /// #1806: a connection that never authenticates is closed at the deadline, counted from its + /// acceptance, and a client feeding a byte at a time cannot push that deadline back. + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn an_unauthenticated_connection_is_closed_at_the_deadline() { + use std::io::{BufReader, Write}; + let (_control, client) = super::attach(&egui::Context::default()); + let ep = super::serve_with(client, TEST_LIMITS).unwrap(); + // Idle. + let started = std::time::Instant::now(); + let mut r = BufReader::new(connect(ep.port)); + let reply = next_line(&mut r).expect("the server says why before closing"); + assert_eq!(reply["error"]["code"], -32004, "{reply}"); + let took = started.elapsed(); + assert!(took >= std::time::Duration::from_millis(250) && took < std::time::Duration::from_secs(3), "{took:?}"); + assert!(next_line(&mut r).is_none(), "then the connection is closed"); + // Dripping: a byte every 50 ms for as long as the server keeps the connection. + let started = std::time::Instant::now(); + let s = connect(ep.port); + let mut w = s.try_clone().unwrap(); + let drip = std::thread::spawn(move || { + while w.write_all(b"x").is_ok() { + std::thread::sleep(std::time::Duration::from_millis(50)); + if started.elapsed() > std::time::Duration::from_secs(3) { + break; + } + } + }); + let mut r = BufReader::new(s); + let reply = next_line(&mut r).expect("the server says why before closing"); + assert_eq!(reply["error"]["code"], -32004, "{reply}"); + let took = started.elapsed(); + assert!(took < std::time::Duration::from_secs(2), "the deadline is absolute: {took:?}"); + drop(r); + drip.join().unwrap(); + } + + /// #1806: only so many connections may wait to authenticate at once. One more is refused + /// on the spot; once a waiting one authenticates or times out, its place is free again. + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn only_so_many_connections_may_wait_to_authenticate() { + use std::io::{BufReader, Write}; + let (_control, client) = super::attach(&egui::Context::default()); + let ep = super::serve_with(client, TEST_LIMITS).unwrap(); + let auth = serde_json::json!({ "jsonrpc": "2.0", "id": 1, "method": "auth", "params": { "token": ep.token } }); + let first = connect(ep.port); + let second = connect(ep.port); + // The kernel accepts a connection before the listener admits it: give the listener a + // moment to admit both, well within their deadline. + std::thread::sleep(std::time::Duration::from_millis(100)); + // The third finds both places taken. + let mut r3 = BufReader::new(connect(ep.port)); + let reply = next_line(&mut r3).expect("refused with a reason"); + assert_eq!(reply["error"]["code"], -32003, "{reply}"); + assert!(next_line(&mut r3).is_none(), "and closed"); + // The second authenticates: its place is free, so a new connection is admitted. + let mut r2 = BufReader::new(second.try_clone().unwrap()); + let mut w2 = second.try_clone().unwrap(); + writeln!(w2, "{auth}").unwrap(); + assert_eq!(next_line(&mut r2).unwrap()["result"]["ok"], true); + let fourth = connect(ep.port); + let mut r4 = BufReader::new(fourth.try_clone().unwrap()); + let mut w4 = fourth.try_clone().unwrap(); + writeln!(w4, "{auth}").unwrap(); + assert_eq!(next_line(&mut r4).unwrap()["result"]["ok"], true, "admitted and authenticated"); + // The first was closed at the deadline (its place freed on that way out too): after it, + // another connection is admitted even though the first never authenticated. + drop(first); + std::thread::sleep(TEST_LIMITS.auth_deadline + std::time::Duration::from_millis(100)); + let fifth = connect(ep.port); + let mut r5 = BufReader::new(fifth.try_clone().unwrap()); + let mut w5 = fifth.try_clone().unwrap(); + writeln!(w5, "{auth}").unwrap(); + assert_eq!(next_line(&mut r5).unwrap()["result"]["ok"], true); + } + /// After `auth` the limit no longer applies: a long line is read as before (here it isn't /// JSON, so it gets a parse error and the connection stays open). #[cfg(not(target_arch = "wasm32"))]