diff --git a/crates/mcp/src/server.rs b/crates/mcp/src/server.rs index 75b82d8d8..0d262dbe7 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 / 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) -> std::io::Result> { + 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) -> 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())?; diff --git a/crates/mcp/src/tests.rs b/crates/mcp/src/tests.rs index 71631b5e8..5a5297246 100644 --- a/crates/mcp/src/tests.rs +++ b/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, +} + +impl std::io::Read for HugeLine { + fn read(&mut self, out: &mut [u8]) -> std::io::Result { + 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 = 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();