MCP: the stdio server keeps at most 16 MB of a message (#2164)

`serve` read a whole line before checking it against `MAX_LINE`, so a client sending a
newline-free payload grew the buffer without limit. The line is now read through `take`, as the
control server does; the rest of an over-long line is read and dropped, the client gets the
"message too large" error, and the next message is served (stdio has one client, so the stream
stays open rather than closing as the control server's connections do).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

crates/mcp/src/server.rs

@@ -1,6 +1,6 @@
 //! JSON-RPC 2.0 framing and the MCP lifecycle / tools / resources / prompts methods.
 
-use std::io::{BufRead, Write};
+use std::io::{BufRead, Read, Write};
 
 use serde_json::{Value, json};
 
@@ -45,6 +45,45 @@ fn response(id: Value, result: Value) -> Value {
     json!({"jsonrpc": "2.0", "id": id, "result": result})
 }
 
+/// Reads the next newline-delimited message into `buf`, keeping at most `MAX_LINE` bytes of it:
+/// a client sending no newline can't make the server allocate more (#2164). `None` at the end of
+/// `input`; `Some(false)` for a line over the limit, whose rest is read and dropped so the next
+/// message is read whole.
+pub(crate) fn read_message(input: &mut impl BufRead, buf: &mut Vec<u8>) -> std::io::Result<Option<bool>> {
+    buf.clear();
+    // One byte past the limit tells an over-long line from one of exactly `MAX_LINE` bytes.
+    let limit = u64::try_from(MAX_LINE).unwrap_or(u64::MAX).saturating_add(1);
+    if input.take(limit).read_until(b'\n', buf)? == 0 {
+        return Ok(None);
+    }
+    if buf.len() > MAX_LINE && buf.last() != Some(&b'\n') {
+        skip_line(input)?;
+        buf.clear();
+        return Ok(Some(false));
+    }
+    Ok(Some(true))
+}
+
+/// Reads up to and including the next newline (or to the end of `input`) without keeping it.
+fn skip_line(input: &mut impl BufRead) -> std::io::Result<()> {
+    loop {
+        let chunk = input.fill_buf()?;
+        if chunk.is_empty() {
+            return Ok(());
+        }
+        match chunk.iter().position(|&b| b == b'\n') {
+            Some(i) => {
+                input.consume(i + 1);
+                return Ok(());
+            }
+            None => {
+                let n = chunk.len();
+                input.consume(n);
+            }
+        }
+    }
+}
+
 fn error(id: Value, code: i64, message: impl Into<String>) -> Value {
     json!({"jsonrpc": "2.0", "id": id, "error": {"code": code, "message": message.into()}})
 }
@@ -68,14 +107,13 @@ impl Server {
     pub fn serve(&mut self, mut input: impl BufRead, mut output: impl Write) -> std::io::Result<()> {
         let mut buf = Vec::new();
         loop {
-            buf.clear();
-            let n = input.read_until(b'\n', &mut buf)?;
-            if n == 0 {
-                return Ok(());
-            }
-            let reply = match std::str::from_utf8(&buf) {
-                Ok(line) => self.handle_line(line),
-                Err(_) => Some(error(Value::Null, PARSE_ERROR, "message is not valid UTF-8").to_string()),
+            let reply = match read_message(&mut input, &mut buf)? {
+                None => return Ok(()),
+                Some(false) => Some(error(Value::Null, INVALID_REQUEST, "message too large (16 MB max)").to_string()),
+                Some(true) => match std::str::from_utf8(&buf) {
+                    Ok(line) => self.handle_line(line),
+                    Err(_) => Some(error(Value::Null, PARSE_ERROR, "message is not valid UTF-8").to_string()),
+                },
             };
             if let Some(reply) = reply {
                 output.write_all(reply.as_bytes())?;

crates/mcp/src/tests.rs

@@ -359,6 +359,71 @@ fn serve_loop_over_buffers() {
     assert!(lines[2]["result"]["content"][0]["text"].as_str().unwrap().contains("42"));
 }
 
+/// Streams `len` bytes of `a` from one reused buffer, then `tail`: a huge line costs no memory to
+/// produce, so only the server can make it expensive.
+struct HugeLine {
+    left: u64,
+    tail: &'static [u8],
+    chunk: Vec<u8>,
+}
+
+impl std::io::Read for HugeLine {
+    fn read(&mut self, out: &mut [u8]) -> std::io::Result<usize> {
+        let buf = std::io::BufRead::fill_buf(self)?;
+        let n = buf.len().min(out.len());
+        out[..n].copy_from_slice(&buf[..n]);
+        std::io::BufRead::consume(self, n);
+        Ok(n)
+    }
+}
+
+impl std::io::BufRead for HugeLine {
+    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
+        if self.left > 0 {
+            let n = self.chunk.len().min(usize::try_from(self.left).unwrap_or(usize::MAX));
+            return Ok(&self.chunk[..n]);
+        }
+        Ok(self.tail)
+    }
+    fn consume(&mut self, n: usize) {
+        if self.left > 0 {
+            self.left -= n as u64;
+        } else {
+            self.tail = &self.tail[n..];
+        }
+    }
+}
+
+#[test]
+fn a_line_without_newline_keeps_at_most_the_limit_in_memory() {
+    // #2164: the whole line was buffered before its size was checked (4 GiB here; any size a
+    // client cares to send).
+    let mut input =
+        HugeLine { left: 4 << 30, tail: concat!("\n", r#"{"jsonrpc":"2.0","id":7,"method":"ping"}"#, "\n").as_bytes(), chunk: vec![b'a'; 1 << 20] };
+    let mut buf = Vec::new();
+    assert_eq!(crate::server::read_message(&mut input, &mut buf).unwrap(), Some(false));
+    assert!(buf.capacity() <= 2 * (16 << 20) + 1, "kept {} bytes", buf.capacity());
+    assert_eq!(crate::server::read_message(&mut input, &mut buf).unwrap(), Some(true));
+    assert!(buf.starts_with(br#"{"jsonrpc""#));
+    assert_eq!(crate::server::read_message(&mut input, &mut buf).unwrap(), None);
+}
+
+#[test]
+fn serve_loop_bounds_a_line_without_newline() {
+    // #2164: a 4 GiB line was read whole before its size was checked. Now the server keeps at
+    // most 16 MB of it, answers that it is too large, and reads the next message normally.
+    let input =
+        HugeLine { left: 4 << 30, tail: concat!("\n", r#"{"jsonrpc":"2.0","id":7,"method":"ping"}"#, "\n").as_bytes(), chunk: vec![b'a'; 1 << 20] };
+    let mut out = Vec::new();
+    server().serve(input, &mut out).unwrap();
+    let lines: Vec<Value> = String::from_utf8(out).unwrap().lines().map(|l| serde_json::from_str(l).unwrap()).collect();
+    assert_eq!(lines.len(), 2);
+    assert_eq!(lines[0]["error"]["code"], -32600);
+    assert!(lines[0]["error"]["message"].as_str().unwrap().contains("too large"));
+    assert_eq!(lines[1]["id"], 7);
+    assert!(lines[1].get("result").is_some(), "{}", lines[1]);
+}
+
 #[test]
 fn cell_set_without_input_is_error() {
     let mut s = server();