Merge PR #1722

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<String> {
+    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<Value> {
+    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<Value, String>)) {
+        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();
                 }
             }
         });

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.