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

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"
+    );
+}