diff --git a/apps/gridcraft/src/control_server.rs b/apps/gridcraft/src/control_server.rs --- a/apps/gridcraft/src/control_server.rs +++ b/apps/gridcraft/src/control_server.rs @@ -12,16 +12,47 @@ use std::io::{BufRead, BufReader, Read, Write}; use std::net::{TcpListener, TcpStream}; use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::mpsc::{Receiver, Sender, channel}; -use std::time::Duration; +use std::time::{Duration, Instant}; use gridcraft_mcp::control_auth; use gridcraft_ui_egui::ControlRequest; use serde_json::{Value as Json, json}; const MAX_LINE: usize = 64 * 1024 * 1024; -/// How long a new connection may take to send its token. +/// How long a new connection may take to send its token, in total. const AUTH_TIMEOUT: Duration = Duration::from_secs(10); +/// How many connections may be waiting to authenticate at once. Anyone on the machine can +/// connect without the token, so past this a new connection is closed before anything is read: +/// each waiting one holds a thread and a socket. +const MAX_PENDING_AUTH: usize = 16; + +/// How the server treats connections that haven't authenticated yet. +#[derive(Clone, Copy)] +struct AuthLimits { + timeout: Duration, + max_pending: usize, +} + +const AUTH_LIMITS: AuthLimits = AuthLimits { timeout: AUTH_TIMEOUT, max_pending: MAX_PENDING_AUTH }; + +/// A connection that is still waiting to authenticate; dropping it frees its place. +struct Pending(Arc); + +impl Pending { + /// A place among the `max` waiting connections, or `None` when they are all taken. + fn claim(waiting: &Arc, max: usize) -> Option { + waiting.try_update(Ordering::AcqRel, Ordering::Acquire, |n| (n < max).then_some(n + 1)).ok()?; + Some(Pending(Arc::clone(waiting))) + } +} + +impl Drop for Pending { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::AcqRel); + } +} /// The secret of this run: the one the user chose (`--control-token`, `GRIDCRAFT_CONTROL_TOKEN`), /// else 256 random bits from the operating system, written to a file only the user can read. @@ -51,35 +82,75 @@ pub fn start(port: u16, chosen_token: Option, ctx: egui::Context) -> std /// Accept connections on `listener`; requests arrive on the returned receiver. fn serve_listener(listener: TcpListener, token: Arc, ctx: egui::Context) -> std::io::Result> { + serve_listener_with(listener, token, ctx, AUTH_LIMITS) +} + +fn serve_listener_with(listener: TcpListener, token: Arc, ctx: egui::Context, limits: AuthLimits) -> std::io::Result> { let (tx, rx) = channel(); + let waiting = Arc::new(AtomicUsize::new(0)); std::thread::Builder::new().name("control".into()).spawn(move || { - for stream in listener.incoming().flatten() { + for stream in listener.incoming() { + let Ok(stream) = stream else { + // Out of sockets or handles: wait for some to free instead of spinning on the error. + std::thread::sleep(Duration::from_millis(50)); + continue; + }; + // Dropping the stream closes it before anything is read. + let Some(place) = Pending::claim(&waiting, limits.max_pending) else { continue }; let tx = tx.clone(); let ctx = ctx.clone(); let token = Arc::clone(&token); - let _ = std::thread::Builder::new().name("control-conn".into()).spawn(move || serve(stream, &token, tx, ctx)); + let _ = + std::thread::Builder::new().name("control-conn".into()).spawn(move || serve(stream, &token, tx, ctx, limits.timeout, Some(place))); } })?; Ok(rx) } /// Read the connection's first line (at most [`control_auth::MAX_AUTH_LINE`] bytes, within -/// [`AUTH_TIMEOUT`]) and check the token. `true` when the peer may go on. -fn authenticate(reader: &mut BufReader, out: &mut TcpStream, token: &str) -> bool { - let _ = out.set_read_timeout(Some(AUTH_TIMEOUT)); - let mut first = String::new(); - let read = reader.by_ref().take(control_auth::MAX_AUTH_LINE as u64).read_line(&mut first); - let verdict = match read { - Ok(n) if n > 0 && first.ends_with('\n') => control_auth::check_first_line(&first, token), +/// `timeout` in total) and check the token. `true` when the peer may go on. +#[allow(clippy::disallowed_methods)] // the desktop app (native only): wasm never runs it +fn authenticate(reader: &mut BufReader, out: &mut TcpStream, token: &str, timeout: Duration) -> bool { + let deadline = Instant::now().checked_add(timeout).unwrap_or_else(Instant::now); + let mut first = Vec::new(); + let read = read_auth_line(reader, &mut first, deadline); + let verdict = match (read, std::str::from_utf8(&first)) { + (Ok(()), Ok(line)) if line.ends_with('\n') => control_auth::check_first_line(line, token), _ => Err(json!({"ok": false, "error": "authentication required"})), }; let allowed = verdict.is_ok(); let reply = verdict.unwrap_or_else(|e| e); let _ = writeln!(out, "{reply}"); - let _ = out.set_read_timeout(None); + let _ = reader.get_ref().set_read_timeout(None); allowed } +/// Read the first line, at most [`control_auth::MAX_AUTH_LINE`] bytes, by `deadline`. +/// +/// The timeout goes on the handle the reader reads: on Windows a `try_clone` of the stream keeps +/// its own, so one set on the writing half never fired (a silent peer held its thread forever). +/// It is shortened before every read, so a peer sending one byte at a time can't stretch it. +#[allow(clippy::disallowed_methods)] // the desktop app (native only): wasm never runs it +fn read_auth_line(reader: &mut BufReader, line: &mut Vec, deadline: Instant) -> std::io::Result<()> { + let max = control_auth::MAX_AUTH_LINE; + while line.len() < max && line.last() != Some(&b'\n') { + let left = deadline.saturating_duration_since(Instant::now()); + if left.is_zero() { + return Err(std::io::ErrorKind::TimedOut.into()); + } + reader.get_ref().set_read_timeout(Some(left))?; + let buf = reader.fill_buf()?; + if buf.is_empty() { + break; + } + let room = max.saturating_sub(line.len()); + let take = buf.iter().take(room).position(|b| *b == b'\n').map_or(buf.len().min(room), |i| i + 1); + line.extend(buf.iter().take(take)); + reader.consume(take); + } + Ok(()) +} + /// Read one `\n`-terminated line of at most `max` bytes into `buf` (newline stripped). /// `Ok(false)` at end of stream; `Err` when the line is too long or the read fails. fn read_bounded_line(reader: &mut impl BufRead, buf: &mut Vec, max: usize) -> Result { @@ -148,11 +219,14 @@ fn is_http_request_line(text: &str) -> bool { || text.rsplit_once(' ').is_some_and(|(_, version)| version.starts_with("HTTP/")) } -fn serve(stream: TcpStream, token: &str, tx: Sender, ctx: egui::Context) { +fn serve(stream: TcpStream, token: &str, tx: Sender, ctx: egui::Context, auth_timeout: Duration, place: Option) { let Ok(read) = stream.try_clone() else { return }; let mut out = stream; let mut reader = BufReader::new(read); - if !authenticate(&mut reader, &mut out, token) { + let allowed = authenticate(&mut reader, &mut out, token, auth_timeout); + // Authenticated or refused, the connection no longer counts as waiting. + drop(place); + if !allowed { return; } let mut buf = Vec::new(); @@ -294,7 +368,7 @@ mod tests { std::thread::spawn(move || { for stream in listener.incoming().flatten() { let tx = tx.clone(); - std::thread::spawn(move || serve(stream, TOKEN, tx, egui::Context::default())); + std::thread::spawn(move || serve(stream, TOKEN, tx, egui::Context::default(), AUTH_TIMEOUT, None)); } }); addr @@ -346,6 +420,81 @@ mod tests { assert_eq!(got[1]["ok"], json!(false)); } + /// Whether the server closed `c` (end of stream, or a refusal and then end of stream) within + /// `within`, and how long it took. + #[allow(clippy::disallowed_methods)] // a native-only test of the desktop app + fn closed_within(c: TcpStream, within: Duration) -> (bool, Duration) { + let start = Instant::now(); + c.set_read_timeout(Some(within)).unwrap(); + let mut r = BufReader::new(c); + let mut line = String::new(); + loop { + line.clear(); + match r.read_line(&mut line) { + Ok(0) => return (true, start.elapsed()), + Ok(_) => continue, + Err(e) if matches!(e.kind(), std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut) => return (false, start.elapsed()), + Err(_) => return (true, start.elapsed()), + } + } + } + + /// The token deadline is for the whole first line: neither a silent peer nor one sending a + /// byte at a time keeps the connection (and its thread) past it. + #[test] + fn a_silent_or_slow_peer_is_closed_at_the_auth_deadline() { + let listener = TcpListener::bind(("127.0.0.1", 0)).unwrap(); + let addr = listener.local_addr().unwrap(); + let limits = AuthLimits { timeout: Duration::from_millis(300), max_pending: MAX_PENDING_AUTH }; + let _rx = serve_listener_with(listener, TOKEN.into(), egui::Context::default(), limits).unwrap(); + + let silent = TcpStream::connect(addr).unwrap(); + let mut slow = TcpStream::connect(addr).unwrap(); + let mut drip = slow.try_clone().unwrap(); + std::thread::spawn(move || { + for b in control_auth::auth_line(1, TOKEN).bytes() { + if drip.write_all(&[b]).is_err() { + break; + } + std::thread::sleep(Duration::from_millis(50)); + } + }); + let (closed, took) = closed_within(silent, Duration::from_secs(5)); + assert!(closed && took < Duration::from_secs(3), "silent peer kept for {took:?}"); + let (closed, took) = closed_within(slow.try_clone().unwrap(), Duration::from_secs(5)); + assert!(closed && took < Duration::from_secs(3), "byte-at-a-time peer kept for {took:?}"); + let _ = slow.flush(); + } + + /// Peers without the token can't pile up threads: past the cap a new connection is closed at + /// once, and a waiting one that authenticates frees its place. + #[test] + fn connections_waiting_to_authenticate_are_capped() { + let listener = TcpListener::bind(("127.0.0.1", 0)).unwrap(); + let addr = listener.local_addr().unwrap(); + let limits = AuthLimits { timeout: Duration::from_secs(20), max_pending: 2 }; + let _rx = serve_listener_with(listener, TOKEN.into(), egui::Context::default(), limits).unwrap(); + + let mut first = TcpStream::connect(addr).unwrap(); + let _second = TcpStream::connect(addr).unwrap(); + std::thread::sleep(Duration::from_millis(200)); + let (closed, took) = closed_within(TcpStream::connect(addr).unwrap(), Duration::from_secs(5)); + assert!(closed && took < Duration::from_secs(3), "a third waiting connection was kept"); + + first.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); + first.write_all(format!("{}\n", control_auth::auth_line(1, TOKEN)).as_bytes()).unwrap(); + let mut reply = String::new(); + BufReader::new(first.try_clone().unwrap()).read_line(&mut reply).unwrap(); + assert_eq!(serde_json::from_str::(&reply).unwrap()["ok"], json!(true), "{reply}"); + std::thread::sleep(Duration::from_millis(200)); + let mut next = TcpStream::connect(addr).unwrap(); + next.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); + next.write_all(format!("{}\n", control_auth::auth_line(1, TOKEN)).as_bytes()).unwrap(); + let mut reply = String::new(); + BufReader::new(next).read_line(&mut reply).unwrap(); + assert_eq!(serde_json::from_str::(&reply).unwrap()["ok"], json!(true), "a freed place was not reused: {reply}"); + } + #[test] fn a_chosen_token_must_be_long_enough() { assert!(make_token(1, Some("short".into())).is_err()); diff --git a/docs/control-protocol.md b/docs/control-protocol.md --- a/docs/control-protocol.md +++ b/docs/control-protocol.md @@ -23,7 +23,8 @@ The control port is reachable by every program on the machine, and by a web page `{"id":0,"ok":true,"result":{"authenticated":true}}`. Any other first line (a wrong or missing token, another method, text that isn't JSON such as an HTTP request, no line within 10 s, or one longer than 4 KiB) gets `{"ok":false,"error":"authentication required"}` and the - connection is closed. After that, a line that isn't JSON also closes the connection. + connection is closed. At most 16 connections may be waiting to send their token at once; past that a new + connection is closed before anything is read. After that, a line that isn't JSON also closes the connection. - `gridcraft-cli send`/`mcp --connect` and `gridcraft_mcp::Remote` do this themselves: they use `GRIDCRAFT_CONTROL_TOKEN` if set, else the token file of that port (so they work for the same user on the same machine without setup). Your own client reads the file or the variable.