shifathosn
/
cadcraft
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

control server: bound request lines and concurrent connections (#678) 

Squash of PR #678 by @SteamSmith.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
SteamSmith 2026-10-12 03:12:34 -04:00 committed by GitHub
parent 9a0d1b3a62
commit 61c018845a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 149 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

156
apps/cadcraft/src/control_server.rs
View File

@ -1,14 +1,25 @@
//! Loopback JSON-lines control server: one request per line, one reply per line.
//! This is the transport the MCP server (`cadcraft-cli mcp --connect`) wraps.
//!
//! Lines are bounded ([`MAX_LINE`]): a longer one is refused and the connection closed, without
//! buffering the rest. So is the number of connections served at once ([`MAX_CONNECTIONS`]).
use std::io::{BufRead, BufReader, Write};
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 cadcraft_ui_egui::ControlRequest;
use serde_json::{Value, json};
/// Longest accepted request line (bytes, without the newline). Generous, so requests that carry
/// whole drawings or images as base64 still fit.
const MAX_LINE: usize = 256 * 1024 * 1024;
/// Connections served at once; further ones get an error line and are closed.
const MAX_CONNECTIONS: usize = 16;
pub fn start(port: u16, ctx: egui::Context) -> Receiver<ControlRequest> {
let (tx, rx) = channel::<ControlRequest>();
let listener = match TcpListener::bind(("127.0.0.1", port)) {
@ -19,25 +30,82 @@ pub fn start(port: u16, ctx: egui::Context) -> Receiver<ControlRequest> {
}
};
log::info!("control server listening on 127.0.0.1:{port}");
std::thread::spawn(move || {
for stream in listener.incoming().flatten() {
std::thread::spawn(move || accept_loop(listener, tx, ctx));
rx
}
/// A connection's place among [`MAX_CONNECTIONS`], given back when the connection ends (also when
/// its thread can't be started, or panics).
struct Slot(Arc<AtomicUsize>);
impl Drop for Slot {
fn drop(&mut self) {
self.0.fetch_sub(1, Ordering::SeqCst);
}
}
fn accept_loop(listener: TcpListener, tx: Sender<ControlRequest>, ctx: egui::Context) {
let open = Arc::new(AtomicUsize::new(0));
for mut stream in listener.incoming().flatten() {
if open.fetch_add(1, Ordering::SeqCst) >= MAX_CONNECTIONS {
open.fetch_sub(1, Ordering::SeqCst);
let _ = writeln!(stream, "{}", json!({"ok": false, "error": "too many control connections"}));
continue;
}
let slot = Slot(Arc::clone(&open));
let tx = tx.clone();
let ctx = ctx.clone();
std::thread::spawn(move || serve(stream, tx, ctx));
}
let _ = std::thread::Builder::new().spawn(move || {
let _slot = slot;
serve(stream, tx, ctx);
});
rx
}
}
/// 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<u8>, max: usize) -> Result<bool, String> {
buf.clear();
// `take` bounds the bytes buffered for one line: a peer can't make us allocate more than `max`.
let limit = u64::try_from(max).unwrap_or(u64::MAX).saturating_add(1);
let n = reader.take(limit).read_until(b'\n', buf).map_err(|e| e.to_string())?;
if n == 0 {
return Ok(false);
}
if buf.last() == Some(&b'\n') {
buf.pop();
if buf.last() == Some(&b'\r') {
buf.pop();
}
return Ok(true);
}
if buf.len() > max {
return Err(format!("request line longer than {max} bytes; closing the connection"));
}
// Last line without a newline, then end of stream.
Ok(true)
}
fn serve(stream: TcpStream, tx: Sender<ControlRequest>, ctx: egui::Context) {
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);
let mut buf = Vec::new();
loop {
match read_bounded_line(&mut reader, &mut buf, MAX_LINE) {
Ok(true) => {}
Ok(false) => break,
Err(e) => {
let _ = writeln!(out, "{}", json!({"ok": false, "error": e}));
break;
}
}
// Bytes that aren't UTF-8 end the connection, as `BufRead::lines` did.
let Ok(line) = std::str::from_utf8(&buf) else { break };
if line.trim().is_empty() {
continue;
}
let reply = match serde_json::from_str::<Value>(&line) {
let reply = match serde_json::from_str::<Value>(line) {
Ok(msg) => {
let id = msg.get("id").cloned().unwrap_or(Value::Null);
let method = msg.get("method").and_then(Value::as_str).unwrap_or("").to_string();
@ -62,3 +130,73 @@ fn serve(stream: TcpStream, tx: Sender<ControlRequest>, ctx: egui::Context) {
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn lines_are_read_up_to_the_limit() {
let mut buf = Vec::new();
let mut input: &[u8] = b"aaaaaaaaaa\nbb\r\ncc";
assert_eq!(read_bounded_line(&mut input, &mut buf, 10), Ok(true));
assert_eq!(buf, b"aaaaaaaaaa", "a line exactly at the limit is accepted");
assert_eq!(read_bounded_line(&mut input, &mut buf, 10), Ok(true));
assert_eq!(buf, b"bb", "CRLF is stripped");
assert_eq!(read_bounded_line(&mut input, &mut buf, 10), Ok(true));
assert_eq!(buf, b"cc", "a last line without a newline is still read");
assert_eq!(read_bounded_line(&mut input, &mut buf, 10), Ok(false), "then end of stream");
}
#[test]
fn an_overlong_line_is_refused_without_reading_the_rest() {
let mut buf = Vec::new();
let long = [b'x'; 100];
let mut input: &[u8] = &long;
assert!(read_bounded_line(&mut input, &mut buf, 10).is_err());
assert!(buf.len() <= 11, "at most limit + 1 bytes are buffered, got {}", buf.len());
assert_eq!(input.len(), 100 - 11, "the rest of the line is left unread");
}
/// The control server on a loopback port, with a stand-in UI that echoes each method.
fn serve_echo() -> std::net::SocketAddr {
let listener = TcpListener::bind(("127.0.0.1", 0)).unwrap();
let addr = listener.local_addr().unwrap();
let (tx, rx) = channel::<ControlRequest>();
std::thread::spawn(move || accept_loop(listener, tx, egui::Context::default()));
std::thread::spawn(move || {
for req in rx {
let _ = req.reply.send(json!({"ok": true, "method": req.method}));
}
});
addr
}
#[test]
fn requests_are_answered_with_their_ids() {
let addr = serve_echo();
let mut c = TcpStream::connect(addr).unwrap();
c.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
c.write_all(b"{\"id\":3,\"method\":\"ping\"}\r\n\n{\"id\":4,\"method\":\"pong\"}\n").unwrap();
let mut r = BufReader::new(c);
for (id, method) in [(3, "ping"), (4, "pong")] {
let mut line = String::new();
r.read_line(&mut line).unwrap();
let v: Value = serde_json::from_str(&line).unwrap();
assert_eq!((v["id"].as_i64(), v["method"].as_str()), (Some(id), Some(method)), "{line}");
}
}
#[test]
fn extra_connections_are_refused() {
let addr = serve_echo();
// Idle connections are accepted in order and held open.
let held: Vec<TcpStream> = (0..MAX_CONNECTIONS).map(|_| TcpStream::connect(addr).unwrap()).collect();
let mut extra = TcpStream::connect(addr).unwrap();
extra.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
let mut line = String::new();
BufReader::new(&mut extra).read_line(&mut line).unwrap();
assert!(line.contains("too many control connections"), "{line}");
drop(held);
}
}