Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
onatozmenn wants to merge 1 commit into
base:main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
154 changes: 140 additions & 14 deletions crates/rmcp/src/transport/async_rw.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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,
Expand All@@ -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")]
Expand DownExpand Up@@ -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;
}
Expand Down
205 changes: 205 additions & 0 deletions crates/rmcp/tests/test_async_rw_max_line_length.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -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"
);
}
, 'i'); if (__m === '*' || __re.test(location.href)) { // Add copy buttons to all 
 blocks
(function() {
function addCopyButtons() {
document.querySelectorAll('pre code').forEach(function(codeBlock) {
if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;
codeBlock.parentElement.setAttribute('data-copy-added', 'true');
var btn = document.createElement('button');
btn.textContent = 'Copy';
btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';
btn.onmouseover = function() { this.style.opacity = '1'; };
btn.onmouseout = function() { this.style.opacity = '0.7'; };
btn.onclick = function() {
navigator.clipboard.writeText(codeBlock.textContent).then(function() {
btn.textContent = 'Copied!';
setTimeout(function() { btn.textContent = 'Copy'; }, 1500);
});
};
codeBlock.parentElement.style.position = 'relative';
codeBlock.parentElement.appendChild(btn);
});
}
addCopyButtons();
// Re-run on dynamic content
var observer = new MutationObserver(addCopyButtons);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
fix(transport): bound the incoming line buffer in AsyncRwTransport by onatozmenn · Pull Request #1049 · modelcontextprotocol/rust-sdk · GitHub
Skip to content
fix(transport): bound the incoming line buffer in AsyncRwTransport#1049

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Open

onatozmenn
wants to merge
1
commit into

base:main
Choose a base branch

from

 Open


fix(transport): bound the incoming line buffer in AsyncRwTransport
#1049
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


154 changes: 140 additions & 14 deletions
crates/rmcp/src/transport/async_rw.rs

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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,
Expand All@@ -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")]
Expand DownExpand Up@@ -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;
 }
Expand Down

205 changes: 205 additions & 0 deletions
crates/rmcp/tests/test_async_rw_max_line_length.rs

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -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"
 );
}

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Force GitHub README to respect dark mode
(function() {
var style = document.createElement('style');
style.textContent = '
.markdown-body {
color-scheme: dark light;
}
.markdown-body pre { background: #161b22 !important; }
.markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }
.markdown-body table th, .markdown-body table td { border-color: #30363d !important; }
.markdown-body img { background: #0d1117; }
.markdown-body blockquote { border-left-color: #8b949e; }
.markdown-body hr { border-color: #30363d; }
';
document.head.appendChild(style);
})();
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix(transport): bound the incoming line buffer in AsyncRwTransport by onatozmenn · Pull Request #1049 · modelcontextprotocol/rust-sdk · GitHub
Skip to content
fix(transport): bound the incoming line buffer in AsyncRwTransport#1049

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Open

onatozmenn
wants to merge
1
commit into

base:main
Choose a base branch

from

 Open


fix(transport): bound the incoming line buffer in AsyncRwTransport
#1049
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


154 changes: 140 additions & 14 deletions
crates/rmcp/src/transport/async_rw.rs

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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,
Expand All@@ -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")]
Expand DownExpand Up@@ -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;
 }
Expand Down

205 changes: 205 additions & 0 deletions
crates/rmcp/tests/test_async_rw_max_line_length.rs

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -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"
 );
}

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Highlight search terms from Google/DuckDuckGo/Bing referrer
(function() {
var ref = document.referrer;
var terms = [];
if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {
var url = new URL(ref);
var q = url.searchParams.get('q') || url.searchParams.get('p');
if (q) {
terms = q.split(/\s+/).filter(function(t) { return t.length > 2; });
}
}
if (terms.length === 0) return;
var style = document.createElement('style');
style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';
document.head.appendChild(style);
function highlight(node) {
if (node.nodeType === 3) { // text node
var text = node.textContent;
var found = false;
terms.forEach(function(term) {
var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\]\\]/g, '\\') + ')', 'gi');
if (regex.test(text)) {
found = true;
var frag = document.createDocumentFragment();
var parts = text.split(regex);
parts.forEach(function(part, i) {
if (i % 2 === 0) {
frag.appendChild(document.createTextNode(part));
} else {
var span = document.createElement('span');
span.className = 'userscript-highlight';
span.textContent = part;
frag.appendChild(span);
}
});
node.parentNode.replaceChild(frag, node);
}
});
} else if (node.nodeType === 1 && node.childNodes) { // element
var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];
if (!skipTags.includes(node.tagName)) {
Array.from(node.childNodes).forEach(highlight);
}
}
}
highlight(document.body);
// Re-highlight on dynamic content
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1 || node.nodeType === 3) highlight(node);
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix(transport): bound the incoming line buffer in AsyncRwTransport by onatozmenn · Pull Request #1049 · modelcontextprotocol/rust-sdk · GitHub
Skip to content
fix(transport): bound the incoming line buffer in AsyncRwTransport#1049

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Open

onatozmenn
wants to merge
1
commit into

base:main
Choose a base branch

from

 Open


fix(transport): bound the incoming line buffer in AsyncRwTransport
#1049
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


154 changes: 140 additions & 14 deletions
crates/rmcp/src/transport/async_rw.rs

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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,
Expand All@@ -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")]
Expand DownExpand Up@@ -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;
 }
Expand Down

205 changes: 205 additions & 0 deletions
crates/rmcp/tests/test_async_rw_max_line_length.rs

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -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"
 );
}

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Strip utm_, fbclid, gclid, etc. from all links on page
(function() {
var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',
'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',
'ref', 'ref_src', 'source', 'medium', 'campaign'];
function cleanUrl(url) {
try {
var u = new URL(url, window.location.origin);
var changed = false;
trackingParams.forEach(function(p) {
if (u.searchParams.has(p)) {
u.searchParams.delete(p);
changed = true;
}
});
return changed ? u.toString() : url;
} catch (e) {
return url;
}
}
function cleanLinks() {
document.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
cleanLinks();
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1) {
if (node.tagName === 'A') cleanLinks();
node.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '
fix(transport): bound the incoming line buffer in AsyncRwTransport by onatozmenn · Pull Request #1049 · modelcontextprotocol/rust-sdk · GitHub
Skip to content
fix(transport): bound the incoming line buffer in AsyncRwTransport#1049

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Open

onatozmenn
wants to merge
1
commit into

base:main
Choose a base branch

from

 Open


fix(transport): bound the incoming line buffer in AsyncRwTransport
#1049
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


154 changes: 140 additions & 14 deletions
crates/rmcp/src/transport/async_rw.rs

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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,
Expand All@@ -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")]
Expand DownExpand Up@@ -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;
 }
Expand Down

205 changes: 205 additions & 0 deletions
crates/rmcp/tests/test_async_rw_max_line_length.rs

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -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"
 );
}

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Auto-enable theater mode on YouTube
(function() {
function tryTheater() {
var btn = document.querySelector('button[aria-label="Theater mode"], ytd-player #player button[title="Theater mode"]');
if (btn && !btn.classList.contains('activated')) {
btn.click();
}
}
// Try immediately
tryTheater();
// Try after navigation (SPA)
var lastUrl = location.href;
setInterval(function() {
if (location.href !== lastUrl) {
lastUrl = location.href;
setTimeout(tryTheater, 500);
}
}, 1000);
// Also try on player load
var observer = new MutationObserver(tryTheater);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix(transport): bound the incoming line buffer in AsyncRwTransport by onatozmenn · Pull Request #1049 · modelcontextprotocol/rust-sdk · GitHub
Skip to content
fix(transport): bound the incoming line buffer in AsyncRwTransport#1049

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Open

onatozmenn
wants to merge
1
commit into

base:main
Choose a base branch

from

 Open


fix(transport): bound the incoming line buffer in AsyncRwTransport
#1049
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


154 changes: 140 additions & 14 deletions
crates/rmcp/src/transport/async_rw.rs

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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,
Expand All@@ -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")]
Expand DownExpand Up@@ -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;
 }
Expand Down

205 changes: 205 additions & 0 deletions
crates/rmcp/tests/test_async_rw_max_line_length.rs

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -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"
 );
}

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Remove or un-stick sticky/fixed headers that block content
(function() {
function unstick() {
document.querySelectorAll('header, nav, [role="banner"], .header, .navbar, .sticky, .fixed-top, [style*="position: fixed"], [style*="position:sticky"]').forEach(function(el) {
if (el.style.position === 'fixed' || el.style.position === 'sticky' || getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {
el.style.position = 'static';
el.style.top = 'auto';
el.style.zIndex = 'auto';
}
});
}
unstick();
var observer = new MutationObserver(unstick);
observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });
})();
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix(transport): bound the incoming line buffer in AsyncRwTransport by onatozmenn · Pull Request #1049 · modelcontextprotocol/rust-sdk · GitHub
Skip to content
fix(transport): bound the incoming line buffer in AsyncRwTransport#1049

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Open

onatozmenn
wants to merge
1
commit into

base:main
Choose a base branch

from

 Open


fix(transport): bound the incoming line buffer in AsyncRwTransport
#1049
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


154 changes: 140 additions & 14 deletions
crates/rmcp/src/transport/async_rw.rs

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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,
Expand All@@ -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")]
Expand DownExpand Up@@ -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;
 }
Expand Down

205 changes: 205 additions & 0 deletions
crates/rmcp/tests/test_async_rw_max_line_length.rs

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -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"
 );
}

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Universal Dark Mode - works on any site
(function() {
var enabled = true;
function applyDarkMode() {
if (!enabled) return;
// Create style element if it doesn't exist
var style = document.getElementById('universal-dark-mode-style');
if (!style) {
style = document.createElement('style');
style.id = 'universal-dark-mode-style';
document.head.appendChild(style);
}
// Dark mode CSS - inverts colors but preserves images/video
style.textContent = '
/* Invert everything except media */
html {
filter: invert(1) hue-rotate(180deg) !important;
background: #1a1a2e !important;
}
/* Restore images, videos, iframes, canvas */
img, video, iframe, canvas, svg, picture, [style*="background-image"] {
filter: invert(1) hue-rotate(180deg) !important;
}
/* Preserve specific elements that should not be inverted */
.no-dark-mode, .no-dark-mode *,
[data-theme="light"], [data-theme="light"],
.ace_editor, .ace_editor *,
.CodeMirror, .CodeMirror *,
.monaco-editor, .monaco-editor *,
.markdown-body pre, .markdown-body pre *,
.highlight, .highlight *,
pre code, pre code * {
filter: none !important;
}
/* Fix common UI elements */
.modal, .popup, .dropdown-menu, .tooltip, .popover {
filter: invert(1) hue-rotate(180deg) !important;
background: #2d2d44 !important;
border-color: #444 !important;
}
/* Scrollbars */
::-webkit-scrollbar { background: #1a1a2e !important; }
::-webkit-scrollbar-thumb { background: #444 !important; }
::-webkit-scrollbar-thumb:hover { background: #555 !important; }
/* Selection */
::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
';
}
function removeDarkMode() {
var style = document.getElementById('universal-dark-mode-style');
if (style) style.remove();
}
// Toggle with Alt+Shift+D
document.addEventListener('keydown', function(e) {
if (e.altKey && e.shiftKey && e.key === 'D') {
e.preventDefault();
enabled = !enabled;
if (enabled) {
applyDarkMode();
console.log('[Universal Dark Mode] Enabled');
} else {
removeDarkMode();
console.log('[Universal Dark Mode] Disabled');
}
}
});
// Apply on load
applyDarkMode();
// Re-apply on dynamic content
var observer = new MutationObserver(function(mutations) {
if (enabled && !document.getElementById('universal-dark-mode-style')) {
applyDarkMode();
}
});
observer.observe(document.head, { childList: true });
console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');
})();
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();
fix(transport): bound the incoming line buffer in AsyncRwTransport by onatozmenn · Pull Request #1049 · modelcontextprotocol/rust-sdk · GitHub
Skip to content
fix(transport): bound the incoming line buffer in AsyncRwTransport#1049

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Open

onatozmenn
wants to merge
1
commit into

base:main
Choose a base branch

from

 Open


fix(transport): bound the incoming line buffer in AsyncRwTransport
#1049
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


154 changes: 140 additions & 14 deletions
crates/rmcp/src/transport/async_rw.rs

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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,
Expand All@@ -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")]
Expand DownExpand Up@@ -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;
 }
Expand Down

205 changes: 205 additions & 0 deletions
crates/rmcp/tests/test_async_rw_max_line_length.rs

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -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"
 );
}