fix(transport): bound the incoming line buffer in AsyncRwTransport

`AsyncRwTransport::receive` read a line with `read_until`, which has no
ceiling, so a peer that never sends a newline makes `line_buf` grow until
the process runs out of memory. This affects every stdio server built on
`rmcp::transport::io::stdio()` and every client using `TokioChildProcess`.

`JsonRpcMessageCodec` already implements this bound correctly, but that
logic lives in `Decoder::decode`, which this transport never calls.

Reading through `fill_buf`/`consume` lets the limit be checked before each
append, which `read_until` cannot do: it does not return until it reaches a
delimiter or EOF, so by the time its result could be inspected the memory is
already committed. An oversized line is dropped and logged and the
connection stays open, matching the discard behaviour the codec's decoder
already uses.

The default is 16 MiB, the same bound the streamable HTTP client applies to
a single SSE event, so a payload that is acceptable over HTTP is also
acceptable over stdio. `with_max_line_length` overrides it and `usize::MAX`
restores the previous unbounded behaviour.

Cancellation safety is preserved. `fill_buf` consumes nothing if the future
is dropped, and the copy into `line_buf` and the matching `consume` have no
await between them, so a partially read line still survives to the next
call, which is what #947 relies on.

Fixes #1030

Author: onatozmenn

crates/rmcp/src/transport/async_rw.rs

@@ -48,13 +48,38 @@ where
 
 pub type TransportWriter<Role, W> = FramedWrite<W, JsonRpcMessageCodec<TxJsonRpcMessage<Role>>>;
 
+/// 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<Role: ServiceRole, R: AsyncRead, W: AsyncWrite> {
     read: BufReader<R>,
     line_buf: Vec<u8>,
+    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<Mutex<Option<TransportWriter<Role, W>>>>,
     _role: PhantomData<fn() -> 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<Role: ServiceRole, R, W> AsyncRwTransport<Role, R, W>
 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<RxJsonRpcMessage<Role>> {
         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;
                 }

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<RoleServer>) {
+    let (peer, ours) = tokio::io::duplex(64 * 1024);
+    let transport = AsyncRwTransport::<RoleServer, _, _>::new(ours, tokio::io::sink())
+        .with_max_line_length(max_line_length);
+    (peer, transport)
+}
+
+fn received_id(msg: &rmcp::service::RxJsonRpcMessage<RoleServer>) -> 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::<RoleServer, _, _>::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"
+    );
+}