diff --git a/crates/rmcp/src/transport/async_rw.rs b/crates/rmcp/src/transport/async_rw.rs index 46d5deaa..cca65092 100644 --- a/crates/rmcp/src/transport/async_rw.rs +++ b/crates/rmcp/src/transport/async_rw.rs @@ -48,13 +48,38 @@ where pub type TransportWriter = FramedWrite>>; +/// Default cap on the size of a single incoming line (one JSON-RPC message) for +/// [`AsyncRwTransport`]. +/// +/// Without a cap, a peer that never sends a newline makes the read buffer grow +/// until the process runs out of memory. 16 MiB matches the streamable HTTP +/// client's `DEFAULT_MAX_SSE_EVENT_SIZE`, so a payload that is acceptable over +/// HTTP is also acceptable over stdio. +pub const DEFAULT_MAX_LINE_LENGTH: usize = 16 * 1024 * 1024; + pub struct AsyncRwTransport { read: BufReader, line_buf: Vec, + max_line_length: usize, + /// Set once an oversized line has been rejected, until its terminating + /// newline is seen. Lives on the struct rather than in `read_line` because + /// `receive` is polled inside a `select!` and can be cancelled mid-discard. + discarding: bool, write: Arc>>>, _role: PhantomData Role>, } +/// Outcome of a single bounded line read. +enum LineRead { + /// A whole `\n`-terminated line is in `line_buf`. + Line, + /// The line exceeded `max_line_length`; it has been dropped. + Oversized, + /// The peer closed the stream. + Eof, + Io(std::io::Error), +} + impl AsyncRwTransport where R: Send + AsyncRead + Unpin, @@ -69,10 +94,105 @@ where Self { read, line_buf: Vec::new(), + max_line_length: DEFAULT_MAX_LINE_LENGTH, + discarding: false, write, _role: PhantomData, } } + + /// Override the maximum size of a single incoming line. + /// + /// Defaults to [`DEFAULT_MAX_LINE_LENGTH`]. Lines longer than this are + /// dropped and logged rather than buffered, and the connection stays open. + /// `usize::MAX` restores the previous unbounded behaviour. + pub fn with_max_line_length(mut self, max_line_length: usize) -> Self { + self.max_line_length = max_line_length; + self + } + + /// Read one `\n`-terminated line into `self.line_buf` without ever letting + /// it grow past `self.max_line_length`. + /// + /// This replaces `read_until`, which cannot be bounded: it does not return + /// until it reaches a delimiter or EOF, so by the time its length could be + /// inspected the memory has already been committed. Reading through + /// `fill_buf`/`consume` lets the limit be checked before each append. + /// + /// Cancellation safety is preserved. `fill_buf` consumes nothing if the + /// future is dropped, and the copy into `line_buf` and the matching + /// `consume` are synchronous with no await between them, so a cancelled + /// read leaves a partial line in `line_buf` for the next call to resume, + /// exactly as `read_until` did. + async fn read_line(&mut self) -> LineRead { + loop { + // The borrow of `self.read` taken by `fill_buf` has to end before + // `consume` can be called, so the decision is made in this block + // and applied after it. + let (consumed, step) = { + let available = match self.read.fill_buf().await { + Ok(available) => available, + Err(e) => return LineRead::Io(e), + }; + if available.is_empty() { + return LineRead::Eof; + } + + let newline = available.iter().position(|b| *b == b'\n'); + let take = newline.map_or(available.len(), |idx| idx + 1); + + if self.discarding { + // Still dropping the tail of a line already rejected. + ( + take, + newline.map(|_| Step::EndDiscard).unwrap_or(Step::More), + ) + } else if self.line_buf.len().saturating_add(take) > self.max_line_length { + self.line_buf.clear(); + ( + take, + if newline.is_some() { + // The oversized line ends here, nothing left to skip. + Step::Oversized + } else { + Step::OversizedNeedsDiscard + }, + ) + } else { + self.line_buf.extend_from_slice(&available[..take]); + ( + take, + if newline.is_some() { + Step::Line + } else { + Step::More + }, + ) + } + }; + self.read.consume(consumed); + + match step { + Step::Line => return LineRead::Line, + Step::More => {} + Step::EndDiscard => self.discarding = false, + Step::Oversized => return LineRead::Oversized, + Step::OversizedNeedsDiscard => { + self.discarding = true; + return LineRead::Oversized; + } + } + } + } +} + +/// What to do once the `fill_buf` borrow has been released. +enum Step { + Line, + More, + EndDiscard, + Oversized, + OversizedNeedsDiscard, } #[cfg(feature = "client")] @@ -124,22 +244,28 @@ where async fn receive(&mut self) -> Option> { loop { - // `read_until` is not cancellation-safe on its own, and `receive` is - // polled inside a `select!` in the service loop: an in-progress line - // read is dropped whenever another branch (e.g. an outgoing response) - // becomes ready. We rely on `read_until` appending into `self.line_buf` - // and only returning at a delimiter or EOF, so a cancelled read leaves - // its partial bytes in `self.line_buf`. Keeping that buffer across - // calls lets the next read resume the same line; it is cleared only - // after a whole line has been consumed. Clearing at the top of the - // loop (the previous behaviour) discarded the partial read and so - // dropped incoming requests under concurrent response load. - match self.read.read_until(b'\n', &mut self.line_buf).await { + // `receive` is polled inside a `select!` in the service loop, so an + // in-progress line read is dropped whenever another branch (e.g. an + // outgoing response) becomes ready. `read_line` appends into + // `self.line_buf` and only reports a line once it hits a delimiter, + // so a cancelled read leaves its partial bytes there and the next + // call resumes the same line; the buffer is cleared only after a + // whole line has been consumed. Clearing at the top of the loop (the + // behaviour before #947) discarded the partial read and so dropped + // incoming requests under concurrent response load. + match self.read_line().await { + LineRead::Line => {} + LineRead::Oversized => { + tracing::error!( + max_line_length = self.max_line_length, + "Incoming message exceeded the maximum line length, discarding it" + ); + continue; + } // EOF. Any bytes still in `line_buf` are an incomplete trailing // message with no delimiter, so there is nothing to deliver. - Ok(0) => return None, - Ok(_) => {} - Err(e) => { + LineRead::Eof => return None, + LineRead::Io(e) => { tracing::error!("Error reading from stream: {}", e); return None; } diff --git a/crates/rmcp/tests/test_async_rw_max_line_length.rs b/crates/rmcp/tests/test_async_rw_max_line_length.rs new file mode 100644 index 00000000..f4292025 --- /dev/null +++ b/crates/rmcp/tests/test_async_rw_max_line_length.rs @@ -0,0 +1,205 @@ +//! Regression tests for the line-length bound on `AsyncRwTransport`. +//! +//! The stdio server transport and the `TokioChildProcess` client transport both +//! read through `AsyncRwTransport`. Before this bound existed, the read side +//! buffered an incoming line with no ceiling, so a peer that sent an +//! unterminated or oversized line could grow the process's memory until it was +//! killed. See https://github.com/modelcontextprotocol/rust-sdk/issues/1030. + +use rmcp::{ + RoleServer, + transport::{ + Transport, + async_rw::{AsyncRwTransport, DEFAULT_MAX_LINE_LENGTH}, + }, +}; +use tokio::io::{AsyncWriteExt, DuplexStream}; + +const MAX: usize = 4 * 1024; + +/// A single-line JSON-RPC request padded out to at least `total` bytes. +fn padded_request(id: u64, total: usize) -> String { + let base = format!(r#"{{"jsonrpc":"2.0","id":{id},"method":"ping","params":{{"pad":""}}}}"#); + let pad = total.saturating_sub(base.len()); + format!( + r#"{{"jsonrpc":"2.0","id":{id},"method":"ping","params":{{"pad":"{}"}}}}"#, + "x".repeat(pad) + ) +} + +fn small_request(id: u64) -> String { + format!(r#"{{"jsonrpc":"2.0","id":{id},"method":"ping"}}"#) +} + +fn transport(max_line_length: usize) -> (DuplexStream, impl Transport) { + let (peer, ours) = tokio::io::duplex(64 * 1024); + let transport = AsyncRwTransport::::new(ours, tokio::io::sink()) + .with_max_line_length(max_line_length); + (peer, transport) +} + +fn received_id(msg: &rmcp::service::RxJsonRpcMessage) -> serde_json::Value { + serde_json::to_value(msg).expect("received message is serializable")["id"].clone() +} + +/// A *valid* message over the limit must be dropped rather than delivered, and +/// the stream must keep working afterwards. +/// +/// Using a well-formed message matters: unparsable junk is discarded by the +/// existing error handling either way, so only a valid oversized message +/// distinguishes a bounded read from an unbounded one. +#[tokio::test] +async fn oversized_message_is_dropped_and_stream_recovers() { + let (mut peer, mut transport) = transport(MAX); + + tokio::spawn(async move { + peer.write_all(padded_request(1, MAX * 4).as_bytes()) + .await + .unwrap(); + peer.write_all(b"\n").await.unwrap(); + peer.write_all(small_request(2).as_bytes()).await.unwrap(); + peer.write_all(b"\n").await.unwrap(); + }); + + let msg = transport.receive().await.expect("second message delivered"); + assert_eq!( + received_id(&msg), + serde_json::json!(2), + "the oversized message should have been dropped, not delivered" + ); +} + +/// The DoS shape from the report: bytes keep arriving with no newline at all. +/// The transport must not accumulate them, and must recover once a delimiter +/// finally shows up. +#[tokio::test] +async fn unterminated_flood_is_discarded_and_stream_recovers() { + let (mut peer, mut transport) = transport(MAX); + + tokio::spawn(async move { + let chunk = vec![b'A'; 8 * 1024]; + // Well past the limit, still no newline. + for _ in 0..16 { + peer.write_all(&chunk).await.unwrap(); + } + peer.write_all(b"\n").await.unwrap(); + peer.write_all(small_request(7).as_bytes()).await.unwrap(); + peer.write_all(b"\n").await.unwrap(); + }); + + let msg = transport + .receive() + .await + .expect("message after the flood is delivered"); + assert_eq!(received_id(&msg), serde_json::json!(7)); +} + +/// A message that fits must still be delivered, including right at the boundary. +#[tokio::test] +async fn message_within_the_limit_is_delivered() { + let (mut peer, mut transport) = transport(MAX); + + // `MAX` counts the trailing newline too, so this is the largest line that fits. + let line = padded_request(3, MAX - 1); + assert_eq!(line.len(), MAX - 1); + + tokio::spawn(async move { + peer.write_all(line.as_bytes()).await.unwrap(); + peer.write_all(b"\n").await.unwrap(); + }); + + let msg = transport.receive().await.expect("message delivered"); + assert_eq!(received_id(&msg), serde_json::json!(3)); +} + +/// The default must be generous enough for real payloads, e.g. an embedded +/// image, so existing callers are not broken by the bound being introduced. +#[tokio::test] +async fn default_limit_accepts_a_large_realistic_message() { + let (peer, ours) = tokio::io::duplex(64 * 1024); + let mut transport = AsyncRwTransport::::new(ours, tokio::io::sink()); + let mut peer = peer; + + assert_eq!(DEFAULT_MAX_LINE_LENGTH, 16 * 1024 * 1024); + + tokio::spawn(async move { + peer.write_all(padded_request(4, 1024 * 1024).as_bytes()) + .await + .unwrap(); + peer.write_all(b"\n").await.unwrap(); + }); + + let msg = transport.receive().await.expect("1 MiB message delivered"); + assert_eq!(received_id(&msg), serde_json::json!(4)); +} + +/// The bounded read replaced `read_until`, which used to be what carried a +/// partially read line across cancellations. A line arriving in several chunks +/// must still be reassembled. +#[tokio::test] +async fn line_split_across_many_reads_is_reassembled() { + let (mut peer, mut transport) = transport(MAX); + + let line = padded_request(5, MAX / 2); + + tokio::spawn(async move { + for chunk in line.as_bytes().chunks(97) { + peer.write_all(chunk).await.unwrap(); + tokio::task::yield_now().await; + } + peer.write_all(b"\n").await.unwrap(); + }); + + let msg = transport.receive().await.expect("reassembled message"); + assert_eq!(received_id(&msg), serde_json::json!(5)); +} + +/// Two oversized messages in a row must not wedge the transport. +#[tokio::test] +async fn consecutive_oversized_messages_still_recover() { + let (mut peer, mut transport) = transport(MAX); + + tokio::spawn(async move { + for id in [1, 2] { + peer.write_all(padded_request(id, MAX * 3).as_bytes()) + .await + .unwrap(); + peer.write_all(b"\n").await.unwrap(); + } + peer.write_all(small_request(9).as_bytes()).await.unwrap(); + peer.write_all(b"\n").await.unwrap(); + }); + + let msg = transport + .receive() + .await + .expect("message after two oversized"); + assert_eq!(received_id(&msg), serde_json::json!(9)); +} + +/// Negative control for the bound itself. +/// +/// `usize::MAX` is the pre-fix behaviour, and with it the very same oversized +/// message is delivered instead of dropped. This is what makes +/// `oversized_message_is_dropped_and_stream_recovers` meaningful: the outcome +/// changes only because of the limit. +#[tokio::test] +async fn unbounded_limit_still_delivers_an_oversized_message() { + let (mut peer, mut transport) = transport(usize::MAX); + + tokio::spawn(async move { + peer.write_all(padded_request(1, MAX * 4).as_bytes()) + .await + .unwrap(); + peer.write_all(b"\n").await.unwrap(); + peer.write_all(small_request(2).as_bytes()).await.unwrap(); + peer.write_all(b"\n").await.unwrap(); + }); + + let msg = transport.receive().await.expect("message delivered"); + assert_eq!( + received_id(&msg), + serde_json::json!(1), + "with no bound the oversized message is buffered and delivered, which is the behaviour the bound removes" + ); +}