Merge pull request #2326 from nihilist-narwhal/fix/2164-mcp-stdio-line-bound

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

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();