diff --git a/crates/mcp/src/backend.rs b/crates/mcp/src/backend.rs index 73a02624e..0c59afa6a 100644 --- a/crates/mcp/src/backend.rs +++ b/crates/mcp/src/backend.rs @@ -1,11 +1,14 @@ //! Where MCP tool calls end up: a control-channel method call. -use std::io::{BufRead, BufReader, Write}; +use std::io::{BufRead, BufReader, Read, Write}; use std::net::{TcpStream, ToSocketAddrs}; use std::time::Duration; use serde_json::{Value, json}; +/// Reply payloads include base64-encoded previews and screenshots, unlike the small requests. +const MAX_REPLY: usize = 64 * 1024 * 1024; + /// Something that answers control-channel methods (`engine.execute`, `engine.commands`, /// `ui.render`, `ui.screenshot`, `ui.key`, …). See `docs/control-protocol.md`. pub trait Backend { @@ -84,11 +87,7 @@ impl Remote { data.push(b'\n'); writer.write_all(&data).map_err(|e| (e, false))?; writer.flush().map_err(|e| (e, true))?; - let mut reply = String::new(); - if reader.read_line(&mut reply).map_err(|e| (e, true))? == 0 { - return Err((std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "control channel closed"), true)); - } - Ok(reply) + read_reply(reader, MAX_REPLY).map_err(|e| (e, true)) } fn unknown_outcome(&self, method: &str, e: &std::io::Error) -> String { @@ -97,6 +96,36 @@ impl Remote { } } +fn read_reply(reader: &mut impl BufRead, max: usize) -> std::io::Result { + let mut reply = String::new(); + let limit = u64::try_from(max).unwrap_or(u64::MAX).saturating_add(1); + reader.take(limit).read_line(&mut reply)?; + if reply.ends_with('\n') { + Ok(reply) + } else if reply.len() > max { + Err(std::io::Error::new(std::io::ErrorKind::InvalidData, format!("control reply longer than {max} bytes"))) + } else { + Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "control channel closed before a complete reply line")) + } +} + +fn parse_reply(reply: &str, id: u64) -> std::io::Result { + let invalid = |message| std::io::Error::new(std::io::ErrorKind::InvalidData, message); + let v: Value = serde_json::from_str(reply).map_err(|e| invalid(format!("bad reply from app: {e}")))?; + let ok = v.get("ok").and_then(Value::as_bool); + match v.get("id") { + Some(reply_id) if reply_id.as_u64() == Some(id) => {} + // The server refuses connections or malformed input before dispatch without an id. + None if ok == Some(false) => {} + _ => return Err(invalid(format!("control reply id does not match request {id}"))), + } + match ok { + Some(true) if v.get("result").is_some() => Ok(v), + Some(false) if v.get("error").and_then(Value::as_str).is_some() => Ok(v), + _ => Err(invalid("control reply must contain boolean `ok` and a result or string error".to_string())), + } +} + /// True when the app has closed `stream` (or it is unusable), checked without blocking or consuming data. fn peer_closed(stream: &TcpStream) -> bool { if stream.set_nonblocking(true).is_err() { @@ -133,7 +162,13 @@ impl Backend for Remote { return Err(self.unknown_outcome(method, &e)); } }; - let v: Value = serde_json::from_str(reply.trim()).map_err(|e| format!("bad reply from app: {e}"))?; + let v = parse_reply(&reply, id).map_err(|e| { + self.conn = None; + self.unknown_outcome(method, &e) + })?; + if v.get("id").is_none() { + self.conn = None; + } if v.get("ok").and_then(Value::as_bool) == Some(true) { Ok(v.get("result").cloned().unwrap_or(Value::Null)) } else { @@ -157,6 +192,143 @@ mod tests { use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; + fn with_reply(reply: &[u8], check: impl FnOnce(&mut Remote, Result)) { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap().to_string(); + let reply = reply.to_vec(); + let worker = std::thread::spawn(move || { + let (stream, _) = listener.accept().unwrap(); + stream.set_read_timeout(Some(Duration::from_secs(2))).unwrap(); + let mut writer = stream.try_clone().unwrap(); + let mut reader = BufReader::new(stream); + let mut request = String::new(); + reader.read_line(&mut request).unwrap(); + let request: Value = serde_json::from_str(&request).unwrap(); + assert_eq!(request["id"], 1); + writer.write_all(&reply).unwrap(); + writer.shutdown(std::net::Shutdown::Write).unwrap(); + let mut extra = String::new(); + assert_eq!(reader.read_line(&mut extra).unwrap(), 0, "request was resent on the same connection"); + listener + }); + let mut remote = Remote::lazy(&addr); + let result = remote.call("engine.execute", json!({"command": "edit.undo"})); + check(&mut remote, result); + drop(remote); + let listener = worker.join().unwrap(); + listener.set_nonblocking(true).unwrap(); + assert_eq!(listener.accept().unwrap_err().kind(), std::io::ErrorKind::WouldBlock, "request was resent on a new connection"); + } + + #[test] + fn mismatched_reply_id_is_rejected_without_resending() { + with_reply(b"{\"id\":999,\"ok\":true,\"result\":\"unrelated\"}\n", |remote, result| { + let err = result.unwrap_err(); + assert!(err.contains("outcome is unknown"), "{err}"); + assert!(remote.conn.is_none(), "discard the untrustworthy connection"); + }); + } + + #[test] + fn unterminated_reply_is_rejected_without_resending() { + with_reply(b"{\"id\":1,\"ok\":true,\"result\":7}", |remote, result| { + let err = result.unwrap_err(); + assert!(err.contains("outcome is unknown"), "{err}"); + assert!(remote.conn.is_none(), "discard the incomplete response"); + }); + } + + #[test] + fn malformed_replies_are_rejected_without_resending() { + for reply in [ + "not JSON\n", + "[]\n", + "{\"id\":1,\"ok\":\"true\",\"result\":7}\n", + "{\"id\":1,\"result\":7}\n", + "{\"id\":1,\"ok\":true}\n", + "{\"ok\":true,\"result\":7}\n", + "{\"id\":null,\"ok\":true,\"result\":7}\n", + "{\"id\":\"1\",\"ok\":true,\"result\":7}\n", + "{\"id\":1.0,\"ok\":true,\"result\":7}\n", + "{\"id\":999,\"ok\":false,\"error\":\"unrelated\"}\n", + "{\"id\":1,\"ok\":false}\n", + "{\"id\":1,\"ok\":false,\"error\":7}\n", + "{\"ok\":false}\n", + ] { + with_reply(reply.as_bytes(), |remote, result| { + let err = result.unwrap_err(); + assert!(err.contains("outcome is unknown"), "{reply}: {err}"); + assert!(remote.conn.is_none()); + }); + } + with_reply(b"\xff\n", |remote, result| { + assert!(result.unwrap_err().contains("outcome is unknown")); + assert!(remote.conn.is_none()); + }); + } + + #[test] + fn valid_replies_and_errors_are_preserved() { + for (reply, expected) in [ + ("{\"id\":1,\"ok\":true,\"result\":7}\n", Ok(json!(7))), + ("{\"id\":1,\"ok\":true,\"result\":null}\r\n", Ok(Value::Null)), + ("{\"id\":1,\"ok\":false,\"error\":\"unknown command\"}\n", Err("unknown command".to_string())), + ("{\"ok\":false,\"error\":\"too many control connections\"}\n", Err("too many control connections".to_string())), + ] { + with_reply(reply.as_bytes(), |remote, result| { + assert_eq!(result, expected); + assert_eq!(remote.conn.is_none(), !reply.contains("\"id\"")); + }); + } + } + + #[test] + fn reply_line_is_bounded_and_requires_a_newline() { + let reply = b"{\"id\":1,\"ok\":true,\"result\":7}\n"; + let payload_len = reply.len() - 1; + assert_eq!(read_reply(&mut reply.as_slice(), payload_len).unwrap().as_bytes(), reply); + assert_eq!(read_reply(&mut reply.as_slice(), payload_len - 1).unwrap_err().kind(), std::io::ErrorKind::InvalidData); + assert_eq!(read_reply(&mut &reply[..payload_len], payload_len).unwrap_err().kind(), std::io::ErrorKind::UnexpectedEof); + assert_eq!(read_reply(&mut &b""[..], payload_len).unwrap_err().kind(), std::io::ErrorKind::UnexpectedEof); + let mut oversized = std::io::Cursor::new(vec![b'x'; 100]); + assert_eq!(read_reply(&mut oversized, 10).unwrap_err().kind(), std::io::ErrorKind::InvalidData); + assert_eq!(oversized.position(), 11, "read only the bound plus one byte"); + } + + #[test] + fn preview_sized_replies_use_the_documented_limit() { + let input = std::io::repeat(b'x').take(MAX_REPLY as u64).chain(std::io::Cursor::new(b"\n")); + let reply = read_reply(&mut BufReader::new(input), MAX_REPLY).unwrap(); + assert_eq!(reply.len(), MAX_REPLY + 1); + assert!(reply.ends_with('\n')); + let mut oversized = BufReader::new(std::io::repeat(b'x')); + assert_eq!(read_reply(&mut oversized, MAX_REPLY).unwrap_err().kind(), std::io::ErrorKind::InvalidData); + } + + #[test] + fn consecutive_replies_match_their_requests_on_one_connection() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap().to_string(); + let worker = std::thread::spawn(move || { + let (stream, _) = listener.accept().unwrap(); + stream.set_read_timeout(Some(Duration::from_secs(2))).unwrap(); + let mut writer = stream.try_clone().unwrap(); + let mut reader = BufReader::new(stream); + for id in 1..=2 { + let mut line = String::new(); + reader.read_line(&mut line).unwrap(); + let request: Value = serde_json::from_str(&line).unwrap(); + assert_eq!(request["id"], id); + writeln!(writer, "{}", json!({"id": id, "ok": true, "result": id})).unwrap(); + } + }); + let mut remote = Remote::lazy(&addr); + for id in 1..=2 { + assert_eq!(remote.call("engine.commands", json!({})).unwrap(), json!(id)); + } + worker.join().unwrap(); + } + #[test] fn lost_reply_is_not_resent() { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); @@ -192,7 +364,8 @@ mod tests { let mut line = String::new(); if BufReader::new(s).read_line(&mut line).unwrap_or(0) > 0 { seen2.fetch_add(1, Ordering::SeqCst); - writer.write_all(b"{\"ok\": true, \"result\": 7}\n").unwrap(); + let request: Value = serde_json::from_str(&line).unwrap(); + writeln!(writer, "{}", json!({"id": request["id"], "ok": true, "result": 7})).unwrap(); } } }); diff --git a/docs/mcp.md b/docs/mcp.md index 6e243dfc5..5a8a124ae 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -24,7 +24,12 @@ it open, `--library` on the same folder fails with "This library is already open the app is not running yet and connects on the first call (and reconnects if the app restarts). A request is sent again only when it never reached the app. If the connection drops after a request was sent but before its reply, the tool fails with "the outcome is unknown" and nothing is resent, so a -command such as undo can't run twice: check the app's state before retrying. +command such as undo can't run twice: check the app's state before retrying. Replies must be +newline-terminated, at most 64 MiB before the newline (including encoded previews/screenshots), +and carry the request's integer `id`, a boolean `ok`, and either a `result` or a string `error`. +A malformed, incomplete, oversized or mismatched reply also closes the connection and reports +an unknown outcome without resending. An idless `ok: false` with a string `error` is preserved +as a pre-dispatch refusal (for example, too many control connections); the next call reconnects. Logs go to stderr; stdout carries only protocol messages.