diff --git a/apps/deckcraft/src/control_server.rs b/apps/deckcraft/src/control_server.rs index f50de025..6dbf0740 100644 --- a/apps/deckcraft/src/control_server.rs +++ b/apps/deckcraft/src/control_server.rs @@ -1,39 +1,146 @@ //! Loopback JSON-lines control server: one request per line, one reply per line. //! This is the transport the MCP server (`deckcraft-cli mcp --connect`) wraps. +//! +//! Every resource a client can hold is bounded ([`Limits`]): simultaneous connections, request +//! line length, queued requests, and how long a socket may sit idle or blocked on a write. -use std::io::{BufRead, BufReader, Write}; +use std::io::{BufRead, BufReader, ErrorKind, Read, Write}; use std::net::{TcpListener, TcpStream}; -use std::sync::mpsc::{Receiver, Sender, channel}; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::mpsc::{Receiver, SyncSender, TrySendError, sync_channel}; use std::time::Duration; use deckcraft_ui_egui::ControlRequest; use serde_json::{Value, json}; +/// Resource bounds for the control server. +#[derive(Clone, Copy, Debug)] +pub struct Limits { + /// Connections served at once; more are told `busy` and closed. + pub max_clients: usize, + /// Longest request line in bytes (newline excluded); a longer one closes the connection. + pub max_line: usize, + /// Requests waiting for the UI at once; more are answered `busy`. + pub queue: usize, + /// How long a connection may go without sending a byte before it is closed. + pub idle_timeout: Duration, + /// How long a reply write may block before the connection is closed. + pub write_timeout: Duration, + /// How long a request waits for the UI's answer. + pub reply_timeout: Duration, +} + +impl Default for Limits { + fn default() -> Self { + Self { + max_clients: 16, + max_line: 16 * 1024 * 1024, + queue: 64, + idle_timeout: Duration::from_secs(600), + write_timeout: Duration::from_secs(30), + reply_timeout: Duration::from_secs(60), + } + } +} + pub fn start(port: u16, ctx: egui::Context) -> Receiver { - let (tx, rx) = channel::(); - let listener = match TcpListener::bind(("127.0.0.1", port)) { - Ok(l) => l, + match TcpListener::bind(("127.0.0.1", port)) { + Ok(l) => { + log::info!("control server listening on 127.0.0.1:{port}"); + serve_listener(l, Limits::default(), ctx) + } Err(e) => { log::error!("control server failed to bind 127.0.0.1:{port}: {e}"); - return rx; + sync_channel::(0).1 } - }; - log::info!("control server listening on 127.0.0.1:{port}"); + } +} + +/// Accept connections on `listener` under `limits`; requests arrive on the returned receiver. +fn serve_listener(listener: TcpListener, limits: Limits, ctx: egui::Context) -> Receiver { + let (tx, rx) = sync_channel::(limits.queue.max(1)); + let active = Arc::new(AtomicUsize::new(0)); std::thread::spawn(move || { - for stream in listener.incoming().flatten() { + for mut stream in listener.incoming().flatten() { + let _ = stream.set_write_timeout(Some(limits.write_timeout)); + if active.fetch_add(1, Ordering::SeqCst) >= limits.max_clients { + active.fetch_sub(1, Ordering::SeqCst); + let _ = writeln!(stream, "{}", json!({"ok": false, "error": "busy: too many control connections"})); + continue; + } + let slot = Slot(Arc::clone(&active)); let tx = tx.clone(); let ctx = ctx.clone(); - std::thread::spawn(move || serve(stream, tx, ctx)); + let spawned = std::thread::Builder::new().name("control-client".into()).spawn(move || { + let _slot = slot; + serve(stream, &tx, &ctx, limits); + }); + if spawned.is_err() { + log::warn!("control server could not start a client thread"); + } } }); rx } -fn serve(stream: TcpStream, tx: Sender, ctx: egui::Context) { +/// Holds one connection slot; gives it back when the client thread ends, however it ends. +struct Slot(Arc); + +impl Drop for Slot { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::SeqCst); + } +} + +enum Line { + Text(Vec), + TooLong, + Closed, +} + +/// Read one `\n`-terminated line of at most `max` bytes without buffering more than `max + 1`. +fn read_line(r: &mut impl BufRead, max: usize) -> Line { + let mut buf = Vec::new(); + let limit = u64::try_from(max).unwrap_or(u64::MAX).saturating_add(1); + match r.by_ref().take(limit).read_until(b'\n', &mut buf) { + Ok(0) | Err(_) => Line::Closed, + Ok(_) => { + if buf.last() == Some(&b'\n') { + buf.pop(); + Line::Text(buf) + } else if buf.len() > max { + Line::TooLong + } else { + // EOF inside a line: serve what arrived, the next read reports Closed. + Line::Text(buf) + } + } + } +} + +fn serve(stream: TcpStream, tx: &SyncSender, ctx: &egui::Context, limits: Limits) { + if stream.set_read_timeout(Some(limits.idle_timeout)).is_err() || stream.set_write_timeout(Some(limits.write_timeout)).is_err() { + return; + } let Ok(read) = stream.try_clone() else { return }; let mut out = stream; - for line in BufReader::new(read).lines() { - let Ok(line) = line else { break }; + let mut reader = BufReader::new(read); + loop { + let line = match read_line(&mut reader, limits.max_line) { + Line::Text(t) => t, + Line::TooLong => { + let _ = writeln!(out, "{}", json!({"ok": false, "error": format!("request line longer than {} bytes", limits.max_line)})); + break; + } + Line::Closed => break, + }; + let Ok(line) = String::from_utf8(line) else { + if writeln!(out, "{}", json!({"ok": false, "error": "bad JSON: not UTF-8"})).is_err() { + break; + } + continue; + }; if line.trim().is_empty() { continue; } @@ -43,11 +150,14 @@ fn serve(stream: TcpStream, tx: Sender, ctx: egui::Context) { let method = msg.get("method").and_then(Value::as_str).unwrap_or("").to_string(); let params = msg.get("params").cloned().unwrap_or(json!({})); let (req, rrx) = ControlRequest::new(method, params); - if tx.send(req).is_err() { - break; - } - ctx.request_repaint(); - let mut r = rrx.recv_timeout(Duration::from_secs(60)).unwrap_or_else(|_| json!({"ok": false, "error": "timeout"})); + let mut r = match tx.try_send(req) { + Ok(()) => { + ctx.request_repaint(); + rrx.recv_timeout(limits.reply_timeout).unwrap_or_else(|_| json!({"ok": false, "error": "timeout"})) + } + Err(TrySendError::Full(_)) => json!({"ok": false, "error": "busy: too many queued requests"}), + Err(TrySendError::Disconnected(_)) => break, + }; if let Some(o) = r.as_object_mut() { o.insert("id".into(), id); } @@ -55,8 +165,134 @@ fn serve(stream: TcpStream, tx: Sender, ctx: egui::Context) { } Err(e) => json!({"ok": false, "error": format!("bad JSON: {e}")}), }; - if writeln!(out, "{reply}").is_err() { + if let Err(e) = writeln!(out, "{reply}") { + if matches!(e.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) { + log::debug!("control client stopped reading replies; closing"); + } break; } } } + +#[cfg(test)] +mod tests { + use super::*; + use std::net::SocketAddr; + use std::time::Instant; + + fn small() -> Limits { + Limits { + max_clients: 2, + max_line: 64, + queue: 1, + idle_timeout: Duration::from_millis(300), + write_timeout: Duration::from_millis(300), + reply_timeout: Duration::from_millis(300), + } + } + + fn server(limits: Limits) -> (SocketAddr, Receiver) { + let l = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = l.local_addr().unwrap(); + (addr, serve_listener(l, limits, egui::Context::default())) + } + + fn client(addr: SocketAddr) -> (BufReader, TcpStream) { + let s = TcpStream::connect(addr).unwrap(); + s.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); + (BufReader::new(s.try_clone().unwrap()), s) + } + + fn reply(r: &mut BufReader) -> Option { + let mut s = String::new(); + match r.read_line(&mut s) { + Ok(0) | Err(_) => None, + Ok(_) => serde_json::from_str(&s).ok(), + } + } + + /// Answer every request with its method name until the server goes away. + fn ui(rx: Receiver) { + std::thread::spawn(move || { + for req in rx { + let _ = req.reply.send(json!({"ok": true, "result": req.method})); + } + }); + } + + #[test] + fn request_round_trip() { + let (addr, rx) = server(small()); + ui(rx); + let (mut r, mut w) = client(addr); + writeln!(w, r#"{{"id":7,"method":"ui.inspect"}}"#).unwrap(); + let v = reply(&mut r).unwrap(); + assert_eq!(v["id"], 7); + assert_eq!(v["result"], "ui.inspect"); + } + + #[test] + fn oversized_line_closes_connection() { + let (addr, rx) = server(small()); + ui(rx); + let (mut r, mut w) = client(addr); + w.write_all(&[b'x'; 200]).unwrap(); + let v = reply(&mut r).unwrap(); + assert!(v["error"].as_str().unwrap().contains("longer than 64")); + assert!(reply(&mut r).is_none()); + } + + #[test] + fn idle_connection_is_closed() { + let (addr, rx) = server(small()); + ui(rx); + let (mut r, mut w) = client(addr); + w.write_all(b"{\"id\":1,").unwrap(); // unterminated, then silence + let t = Instant::now(); + assert!(reply(&mut r).is_none()); + assert!(t.elapsed() < Duration::from_secs(4)); + } + + #[test] + fn connection_limit_refuses_then_frees_slot() { + let (addr, rx) = server(small()); + ui(rx); + let a = client(addr); + let (mut br, mut bw) = client(addr); + // Make sure both are being served before the third arrives. + writeln!(bw, r#"{{"id":1,"method":"m"}}"#).unwrap(); + assert!(reply(&mut br).is_some()); + let (mut cr, _cw) = client(addr); + let v = reply(&mut cr).unwrap(); + assert!(v["error"].as_str().unwrap().starts_with("busy")); + drop(a); + // The freed slot is reusable once the server notices the disconnect. + let deadline = Instant::now() + Duration::from_secs(4); + loop { + let (mut dr, mut dw) = client(addr); + writeln!(dw, r#"{{"id":2,"method":"m"}}"#).unwrap(); + if reply(&mut dr).is_some_and(|v| v["ok"] == true) { + break; + } + assert!(Instant::now() < deadline, "slot never freed"); + std::thread::sleep(Duration::from_millis(50)); + } + } + + #[test] + fn blocked_ui_times_out_and_full_queue_is_busy() { + let (addr, rx) = server(small()); + // No UI: the first request fills the one-slot queue and times out; the second finds it full. + let (mut ar, mut aw) = client(addr); + let (mut br, mut bw) = client(addr); + writeln!(aw, r#"{{"id":1,"method":"m"}}"#).unwrap(); + std::thread::sleep(Duration::from_millis(50)); + writeln!(bw, r#"{{"id":2,"method":"m"}}"#).unwrap(); + let b = reply(&mut br).unwrap(); + assert!(b["error"].as_str().unwrap().starts_with("busy")); + assert_eq!(b["id"], 2); + let a = reply(&mut ar).unwrap(); + assert_eq!(a["error"], "timeout"); + drop(rx); + } +} diff --git a/docs/control-protocol.md b/docs/control-protocol.md index 07741fde..41f890c0 100644 --- a/docs/control-protocol.md +++ b/docs/control-protocol.md @@ -3,6 +3,11 @@ Start the app with `--control `. The server listens on `127.0.0.1` only and speaks JSON lines: one request object per line, one reply per line. +The server bounds what a client can hold: 16 connections at once (more get `{"ok":false,"error":"busy: …"}` +and are closed), 16 MiB per request line (longer closes the connection), 64 queued requests (more are +answered `busy`), 10 minutes idle and 30 seconds per blocked reply write before the socket is closed, and +60 seconds waiting for the UI (`"error":"timeout"`). + ```json {"id": 1, "method": "engine.execute", "params": {"command": "shape.insert", "params": {"preset": "roundRect", "rect": [100, 100, 240, 120], "text": "Hello"}}} {"id": 1, "ok": true, "result": {"id": 258}}