diff --git a/crates/mcp/src/backend.rs b/crates/mcp/src/backend.rs index 78cec12b2..a07c89391 100644 --- a/crates/mcp/src/backend.rs +++ b/crates/mcp/src/backend.rs @@ -1,12 +1,17 @@ //! Where MCP tool calls end up: a control-channel method call. -use std::io::{BufRead, BufReader, Write}; +use std::io::{BufReader, Write}; use std::net::{IpAddr, SocketAddr, TcpStream, ToSocketAddrs}; use std::path::Path; use std::time::Duration; use serde_json::{Value, json}; +use crate::server::{BoundedLine, read_bounded_line}; + +/// Longest reply line accepted from the app (screenshots and `document.inspect` are large). +const MAX_REPLY_LINE: usize = 256 << 20; + /// Something that answers control-channel methods (`engine.execute`, `document.inspect`, /// `ui.pointer`, `ui.render`, …). See `wordcraft_ui_egui::control` for the full list. pub trait Backend { @@ -161,11 +166,11 @@ impl Remote { writer.write_all(line.as_bytes())?; writer.write_all(b"\n")?; writer.flush()?; - let mut reply = String::new(); - if reader.read_line(&mut reply)? == 0 { - return Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "control channel closed")); + match read_bounded_line(reader, MAX_REPLY_LINE)? { + BoundedLine::Line(reply) => Ok(reply), + BoundedLine::Eof => Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "control channel closed")), + BoundedLine::TooLong => Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "reply from app too long")), } - Ok(reply) } /// One request and its reply, with one retry on a new connection (the app may have @@ -189,6 +194,11 @@ impl Remote { Ok(r) => break r, Err(e) => { self.conn = None; + // A reply that arrived but can't be read (too long, not UTF-8): the app has + // run the request, so sending it again could run a state-changing one twice. + if e.kind() == std::io::ErrorKind::InvalidData { + return Err(format!("bad reply from app: {e}")); + } if !first { return Err(unreachable(&self.addr, e, self.connected)); } @@ -246,6 +256,7 @@ impl Backend for Remote { #[cfg(test)] mod tests { use super::*; + use std::io::BufRead; use std::net::{SocketAddr, TcpListener}; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -335,6 +346,37 @@ mod tests { (port, rx) } + #[test] + fn an_unreadable_reply_is_not_retried() { + // The app answered something that can't be read (here not UTF-8; an over-long reply + // takes the same path): the request ran, so it must not be sent again. + let listener = TcpListener::bind(("127.0.0.1", 0)).unwrap(); + let port = listener.local_addr().unwrap().port(); + let seen = Arc::new(AtomicUsize::new(0)); + let n = seen.clone(); + std::thread::spawn(move || { + for stream in listener.incoming().flatten() { + let n = n.clone(); + std::thread::spawn(move || { + let mut out = stream.try_clone().unwrap(); + for line in BufReader::new(stream).lines() { + if line.is_err() { + return; + } + n.fetch_add(1, Ordering::SeqCst); + if out.write_all(b"\xff\xfe\n").is_err() { + return; + } + } + }); + } + }); + let mut r = Remote::connect_with(&format!("127.0.0.1:{port}"), fixed_key("k1")).unwrap(); + let e = r.call("engine.execute", json!({"command": "text.insert", "params": {"text": "x"}})).unwrap_err(); + assert!(e.contains("bad reply"), "{e}"); + assert_eq!(seen.load(Ordering::SeqCst), 1, "the request was sent once"); + } + fn fixed_key(key: &'static str) -> KeyFinder { Box::new(move |_| Ok(key.to_string())) } diff --git a/crates/mcp/src/server.rs b/crates/mcp/src/server.rs index 63f6764c6..0bca29e64 100644 --- a/crates/mcp/src/server.rs +++ b/crates/mcp/src/server.rs @@ -1,6 +1,6 @@ //! JSON-RPC 2.0 framing and the MCP lifecycle / tools / resources methods. -use std::io::{BufRead, Write}; +use std::io::{BufRead, Read, Write}; use serde_json::{Value, json}; @@ -28,6 +28,49 @@ execute format.bold → inspect_document to verify → render_page to look at th pub const DOC_URI: &str = "wordcraft://document"; pub const COMMANDS_URI: &str = "wordcraft://commands"; +/// Longest request line accepted on stdin. +pub const MAX_REQUEST_LINE: usize = 1 << 20; + +/// One bounded line read. +pub(crate) enum BoundedLine { + Eof, + Line(String), + /// The line exceeded the limit; the rest of it has been discarded. + TooLong, +} + +/// Read one line of at most `max` bytes (excluding the newline) without buffering more than +/// `max + 1` bytes; an over-long line is skipped up to its newline. +pub(crate) fn read_bounded_line(reader: &mut impl BufRead, max: usize) -> std::io::Result { + let mut buf = Vec::new(); + reader.by_ref().take(max as u64 + 1).read_until(b'\n', &mut buf)?; + if buf.is_empty() { + return Ok(BoundedLine::Eof); + } + if buf.last() != Some(&b'\n') && buf.len() > max { + loop { + let chunk = reader.fill_buf()?; + if chunk.is_empty() { + break; + } + let found = chunk.iter().position(|&b| b == b'\n'); + let used = found.map_or(chunk.len(), |i| i + 1); + reader.consume(used); + if found.is_some() { + break; + } + } + return Ok(BoundedLine::TooLong); + } + if buf.last() == Some(&b'\n') { + buf.pop(); + if buf.last() == Some(&b'\r') { + buf.pop(); + } + } + String::from_utf8(buf).map(BoundedLine::Line).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e)) +} + /// An MCP server bound to one backend. pub struct Server { backend: Box, @@ -58,10 +101,14 @@ impl Server { /// Serve newline-delimited JSON-RPC until `input` closes. Logs go to stderr only (stdout is /// the protocol stream). - pub fn serve(&mut self, input: impl BufRead, mut output: impl Write) -> std::io::Result<()> { - for line in input.lines() { - let line = line?; - if let Some(reply) = self.handle_line(&line) { + pub fn serve(&mut self, mut input: impl BufRead, mut output: impl Write) -> std::io::Result<()> { + loop { + let reply = match read_bounded_line(&mut input, MAX_REQUEST_LINE)? { + BoundedLine::Eof => break, + BoundedLine::Line(line) => self.handle_line(&line), + BoundedLine::TooLong => Some(error(Value::Null, INVALID_REQUEST, "request line too long").to_string()), + }; + if let Some(reply) = reply { output.write_all(reply.as_bytes())?; output.write_all(b"\n")?; output.flush()?; diff --git a/crates/mcp/src/tests.rs b/crates/mcp/src/tests.rs index 12ea973f1..f23f72f75 100644 --- a/crates/mcp/src/tests.rs +++ b/crates/mcp/src/tests.rs @@ -228,3 +228,20 @@ fn type_text_paragraphs_reports_a_refusal_and_types_nothing_after_it() { assert_eq!(calls[0].0, "text.insert"); assert_eq!(calls[0].1["text"], "one\rtwo", "the line break is a paragraph break"); } + +/// #3685: an over-long request line gets an error reply and the server keeps serving. +#[test] +fn over_long_request_line_is_rejected_and_serving_continues() { + let mut s = Server::new(Box::new(Headless::default())); + let mut input = vec![b'a'; (1 << 20) + 10]; + input.push(b'\n'); + input.extend_from_slice(br#"{"jsonrpc":"2.0","id":5,"method":"ping"}"#); + input.push(b'\n'); + let mut out = Vec::new(); + s.serve(std::io::Cursor::new(input), &mut out).unwrap(); + let text = String::from_utf8(out).unwrap(); + let lines: Vec<&str> = text.lines().collect(); + assert_eq!(lines.len(), 2, "{text}"); + assert!(lines[0].contains("too long"), "{}", lines[0]); + assert!(lines[1].contains(r#""id":5"#), "{}", lines[1]); +}