From b9e0ad48394f83e5772e35a15e794bbe83f1909c Mon Sep 17 00:00:00 2001 From: GTanger <49036591+GTanger@users.noreply.github.com> Date: Wed, 11 Nov 2026 16:43:04 +0800 Subject: [PATCH] fix(transport): bound buffered Streamable HTTP responses Expose configurable StreamableHttpResponseLimits (defaulting to 64 MiB for complete JSON bodies on any HTTP status, 64 KiB for HTTP diagnostic prefixes, and 16 MiB per SSE event) while preserving the v3.9.0 ResponseBodyTooLarge { max_size } error variant, HTTP 5xx status metadata fallback, and reqwest charset/BOM decoding. Truncate non-JSON diagnostic error bodies without waiting for EOF so oversized legacy 4xx pages still reach server/discover fallback, align Unix socket JSON-RPC error handling with reqwest, and forward limits across all worker POST sites and AuthClient token refreshes. --- .../common/auth/streamable_http_client.rs | 203 ++++++- .../src/transport/common/client_side_http.rs | 108 +++- .../common/reqwest/streamable_http_client.rs | 194 +++++-- .../rmcp/src/transport/common/unix_socket.rs | 117 +++- .../src/transport/streamable_http_client.rs | 145 ++++- ...est_streamable_http_default_error_limit.rs | 84 +++ .../test_streamable_http_limit_lifecycle.rs | 516 +++++++++++++++++ .../test_streamable_http_response_limits.rs | 416 +++++++++++--- .../tests/test_unix_socket_response_limits.rs | 522 +++++++++++++++++- 9 files changed, 2099 insertions(+), 206 deletions(-) create mode 100644 crates/rmcp/tests/test_streamable_http_default_error_limit.rs create mode 100644 crates/rmcp/tests/test_streamable_http_limit_lifecycle.rs diff --git a/crates/rmcp/src/transport/common/auth/streamable_http_client.rs b/crates/rmcp/src/transport/common/auth/streamable_http_client.rs index 47b6caf55..3e0f86551 100644 --- a/crates/rmcp/src/transport/common/auth/streamable_http_client.rs +++ b/crates/rmcp/src/transport/common/auth/streamable_http_client.rs @@ -5,7 +5,9 @@ use tracing::debug; use crate::transport::{ auth::{AuthClient, AuthError}, - streamable_http_client::{StreamableHttpClient, StreamableHttpError}, + streamable_http_client::{ + StreamableHttpClient, StreamableHttpError, StreamableHttpResponseLimits, + }, }; impl AuthClient @@ -213,11 +215,44 @@ where }) .await } + + async fn post_message_with_response_limits( + &self, + uri: std::sync::Arc, + message: crate::model::ClientJsonRpcMessage, + session_id: Option>, + auth_token: Option, + custom_headers: HashMap, + limits: StreamableHttpResponseLimits, + ) -> Result< + crate::transport::streamable_http_client::StreamableHttpPostResponse, + StreamableHttpError, + > { + self.call_reacting_to_challenges(auth_token, |token| { + let uri = uri.clone(); + let message = message.clone(); + let session_id = session_id.clone(); + let custom_headers = custom_headers.clone(); + async move { + self.http_client + .post_message_with_response_limits( + uri, + message, + session_id, + token, + custom_headers, + limits, + ) + .await + } + }) + .await + } } #[cfg(all(test, feature = "transport-streamable-http-client-reqwest"))] mod tests { - use std::sync::Arc; + use std::sync::{Arc, Mutex}; use oauth2::{AccessToken, RefreshToken, basic::BasicTokenType}; @@ -228,9 +263,171 @@ mod tests { InMemoryCredentialStore, OAuthHttpClient, OAuthHttpClientFuture, OAuthHttpRequest, OAuthTokenResponse, StoredCredentials, VendorExtraTokenFields, }, - streamable_http_client::AuthRequiredError, + streamable_http_client::{AuthRequiredError, StreamableHttpPostResponse}, }; + struct RefreshClient; + + impl OAuthHttpClient for RefreshClient { + fn execute(&self, request: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> { + Box::pin(async move { + assert_eq!(request.request.uri(), "https://auth.example.com/token"); + assert_eq!(request.request.method(), http::Method::POST); + Ok(oauth2::http::Response::builder() + .status(200) + .header("content-type", "application/json") + .body(br#"{"access_token":"fresh-token","token_type":"Bearer"}"#.to_vec()) + .unwrap()) + }) + } + } + + type RecordedPostCall = (Option, StreamableHttpResponseLimits); + + #[derive(Clone, Default)] + struct RecordingLimitsClient { + calls: Arc>>, + } + + impl StreamableHttpClient for RecordingLimitsClient { + type Error = std::io::Error; + + async fn post_message( + &self, + _: Arc, + _: crate::model::ClientJsonRpcMessage, + _: Option>, + _: Option, + _: HashMap, + ) -> Result> { + panic!("response limits must not use the legacy POST method"); + } + + async fn delete_session( + &self, + _: Arc, + _: Arc, + _: Option, + _: HashMap, + ) -> Result<(), StreamableHttpError> { + unreachable!("this test only sends POST requests") + } + + async fn get_stream( + &self, + _: Arc, + _: Option>, + _: Option, + _: Option, + _: HashMap, + ) -> Result< + futures::stream::BoxStream<'static, Result>, + StreamableHttpError, + > { + unreachable!("this test only sends POST requests") + } + + async fn post_message_with_response_limits( + &self, + uri: Arc, + message: crate::model::ClientJsonRpcMessage, + session_id: Option>, + auth_token: Option, + custom_headers: HashMap, + limits: StreamableHttpResponseLimits, + ) -> Result> { + assert_eq!(uri.as_ref(), "https://mcp.example.com/mcp"); + assert_eq!(session_id.as_deref(), Some("test-session")); + assert_eq!(serde_json::to_value(message).unwrap()["method"], "ping"); + assert_eq!( + custom_headers.get(&HeaderName::from_static("x-test-header")), + Some(&HeaderValue::from_static("preserved")), + ); + let mut calls = self.calls.lock().unwrap(); + calls.push((auth_token, limits)); + if calls.len() == 1 { + Err(StreamableHttpError::AuthRequired(AuthRequiredError::new( + "Bearer".into(), + ))) + } else { + Err(StreamableHttpError::ResponseBodyTooLarge { + max_size: limits.max_json_response_size, + }) + } + } + } + + #[tokio::test] + async fn response_limits_survive_auth_refresh_and_size_errors_do_not_retry() { + let mut manager = AuthorizationManager::new_with_oauth_http_client( + "https://mcp.example.com/mcp", + Arc::new(RefreshClient), + ) + .await + .unwrap(); + manager.set_metadata(AuthorizationMetadata { + authorization_endpoint: "https://auth.example.com/authorize".into(), + token_endpoint: "https://auth.example.com/token".into(), + ..Default::default() + }); + manager.configure_client_id("client").unwrap(); + let store = InMemoryCredentialStore::new(); + let credentials = serde_json::from_value(serde_json::json!({ + "access_token": "old-token", + "refresh_token": "refresh-token", + "token_type": "Bearer", + })) + .unwrap(); + store + .save(StoredCredentials::new( + "client".into(), + Some(credentials), + vec![], + None, + )) + .await + .unwrap(); + manager.set_credential_store(store); + let recording = RecordingLimitsClient::default(); + let client = AuthClient::new(recording.clone(), manager); + let mut headers = HashMap::new(); + headers.insert( + HeaderName::from_static("x-test-header"), + HeaderValue::from_static("preserved"), + ); + let limits = StreamableHttpResponseLimits { + max_sse_event_size: 17, + max_json_response_size: 31, + max_error_response_size: 13, + }; + let result = client + .post_message_with_response_limits( + "https://mcp.example.com/mcp".into(), + serde_json::from_value( + serde_json::json!({"jsonrpc":"2.0", "id":1, "method":"ping"}), + ) + .unwrap(), + Some("test-session".into()), + None, + headers, + limits, + ) + .await; + assert!(matches!( + result, + Err(StreamableHttpError::ResponseBodyTooLarge { max_size: 31 }) + )); + let calls = recording.calls.lock().unwrap(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].0.as_deref(), Some("old-token")); + assert_eq!(calls[1].0.as_deref(), Some("fresh-token")); + for (_, observed) in calls.iter() { + assert_eq!(observed.max_sse_event_size, 17); + assert_eq!(observed.max_json_response_size, 31); + assert_eq!(observed.max_error_response_size, 13); + } + } + struct UnavailableStore; #[async_trait::async_trait] diff --git a/crates/rmcp/src/transport/common/client_side_http.rs b/crates/rmcp/src/transport/common/client_side_http.rs index 33d9e2aaa..0ddf0c687 100644 --- a/crates/rmcp/src/transport/common/client_side_http.rs +++ b/crates/rmcp/src/transport/common/client_side_http.rs @@ -5,9 +5,6 @@ use futures::{Stream, TryStreamExt}; use crate::transport::streamable_http_client::StreamableHttpError; -pub(super) const MAX_JSON_RESPONSE_SIZE: usize = 16 * 1024 * 1024; -pub(super) const MAX_ERROR_RESPONSE_SIZE: usize = 64 * 1024; - /// Count bytes before retaining them; a length hint can reject a body early, /// but never replaces the streamed check or determines allocation size. pub(super) async fn read_bounded_body( @@ -38,9 +35,45 @@ where Ok(body) } +/// Read at most `max_size` bytes for a diagnostic, dropping the rest of the stream. +/// +/// Unlike a protocol message, a diagnostic need not be complete. Stop as soon as +/// the prefix is full, without polling for another chunk or EOF. `Content-Length` +/// is not consulted: an oversized declaration must not discard a useful prefix. +/// The returned bytes must never be parsed as a complete JSON-RPC message. +pub(super) async fn read_truncated_body( + stream: S, + max_size: usize, +) -> Result, StreamableHttpError> +where + S: Stream>, + E: std::error::Error + Send + Sync + 'static, +{ + let mut body = Vec::new(); + futures::pin_mut!(stream); + while body.len() < max_size { + let Some(chunk) = stream + .try_next() + .await + .map_err(StreamableHttpError::Client)? + else { + break; + }; + let count = chunk.len().min(max_size - body.len()); + body.extend_from_slice(&chunk[..count]); + } + Ok(body) +} + #[cfg(test)] mod tests { - use std::io; + use std::{ + io, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + }; use futures::stream; use rstest::rstest; @@ -119,5 +152,72 @@ mod tests { .unwrap() .is_empty() ); + let stream = stream::iter([Ok::<_, io::Error>(Bytes::from_static(b"x"))]); + assert!(matches!( + read_bounded_body(stream, None, 0).await, + Err(StreamableHttpError::ResponseBodyTooLarge { max_size: 0 }) + )); + } + + #[tokio::test] + async fn truncated_body_keeps_short_bodies_complete() { + let stream = stream::iter([ + Ok::<_, io::Error>(Bytes::from_static(b"ab")), + Ok(Bytes::from_static(b"c")), + ]); + assert_eq!(read_truncated_body(stream, 4).await.unwrap(), b"abc"); + } + + #[tokio::test] + async fn truncated_body_stops_at_exact_limit_without_polling_for_eof() { + let stream = futures::StreamExt::chain( + stream::iter([ + Ok::<_, io::Error>(Bytes::from_static(b"ab")), + Ok(Bytes::from_static(b"cd")), + ]), + stream::poll_fn(|_| panic!("a full diagnostic prefix must not poll for EOF")), + ); + assert_eq!(read_truncated_body(stream, 4).await.unwrap(), b"abcd"); + } + + #[tokio::test] + async fn truncated_body_stops_in_the_middle_of_an_oversized_chunk() { + let polls = Arc::new(AtomicUsize::new(0)); + let count = polls.clone(); + let stream = stream::poll_fn(move |_| { + assert!( + count.fetch_add(1, Ordering::SeqCst) < 2, + "must not drain the body" + ); + std::task::Poll::Ready(Some(Ok::<_, io::Error>(Bytes::from_static(b"abc")))) + }); + assert_eq!(read_truncated_body(stream, 4).await.unwrap(), b"abca"); + assert_eq!(polls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn truncated_body_zero_limit_does_not_poll() { + let stream = stream::poll_fn(|_| -> std::task::Poll>> { + panic!("zero diagnostic limit must not read the body") + }); + assert!(read_truncated_body(stream, 0).await.unwrap().is_empty()); + } + + #[tokio::test] + async fn truncated_body_preserves_io_errors_before_the_limit() { + let stream = stream::iter([ + Ok(Bytes::from_static(b"a")), + Err(io::Error::other("read failed")), + ]); + assert!(matches!( + read_truncated_body(stream, 4).await, + Err(StreamableHttpError::Client(_)) + )); + } + + #[tokio::test] + async fn truncated_body_counts_bytes_not_utf8_characters() { + let stream = stream::iter([Ok::<_, io::Error>(Bytes::from_static("é".as_bytes()))]); + assert_eq!(read_truncated_body(stream, 1).await.unwrap(), [0xc3]); } } diff --git a/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs b/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs index 1b08b152d..4787db9bb 100644 --- a/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs +++ b/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs @@ -12,9 +12,7 @@ use crate::{ model::{ClientJsonRpcMessage, JsonRpcMessage, ServerJsonRpcMessage}, transport::{ common::{ - client_side_http::{ - MAX_ERROR_RESPONSE_SIZE, MAX_JSON_RESPONSE_SIZE, read_bounded_body, - }, + client_side_http::{read_bounded_body, read_truncated_body}, client_side_sse::{DEFAULT_MAX_SSE_EVENT_SIZE, bounded_sse_stream}, http_header::{ EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE, @@ -31,6 +29,35 @@ impl From for StreamableHttpError { } } +const BODY_READ_FAILURE: &str = ""; + +fn truncated_read_failure(max_error_response_size: usize) -> String { + BODY_READ_FAILURE[..BODY_READ_FAILURE.len().min(max_error_response_size)].to_owned() +} + +async fn decode_error_body( + body: Vec, + content_type: Option, + max_error_response_size: usize, +) -> String { + if body.is_empty() { + return String::new(); + } + // Decode only the bounded, in-memory body, using reqwest's + // feature-dependent charset and BOM handling. No network + // stream or Content-Encoding is carried into this response. + let mut buffered = http::Response::new(body); + if let Some(content_type) = content_type { + buffered + .headers_mut() + .insert(reqwest::header::CONTENT_TYPE, content_type); + } + reqwest::Response::from(buffered) + .text() + .await + .unwrap_or_else(|_| truncated_read_failure(max_error_response_size)) +} + /// Applies custom headers to a request builder, rejecting reserved headers. fn apply_custom_headers( mut builder: reqwest::RequestBuilder, @@ -198,13 +225,13 @@ impl StreamableHttpClient for reqwest::Client { auth_token: Option, custom_headers: HashMap, ) -> Result> { - self.post_message_with_max_sse_event_size( + self.post_message_with_response_limits( uri, message, session_id, auth_token, custom_headers, - DEFAULT_MAX_SSE_EVENT_SIZE, + StreamableHttpResponseLimits::default(), ) .await } @@ -217,6 +244,29 @@ impl StreamableHttpClient for reqwest::Client { auth_token: Option, custom_headers: HashMap, max_sse_event_size: usize, + ) -> Result> { + self.post_message_with_response_limits( + uri, + message, + session_id, + auth_token, + custom_headers, + StreamableHttpResponseLimits { + max_sse_event_size, + ..Default::default() + }, + ) + .await + } + + async fn post_message_with_response_limits( + &self, + uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + auth_token: Option, + custom_headers: HashMap, + limits: StreamableHttpResponseLimits, ) -> Result> { let mut request = self .post(uri.as_ref()) @@ -280,7 +330,7 @@ impl StreamableHttpClient for reqwest::Client { .as_deref() .is_some_and(|ct| ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes())); // Rate limits and non-JSON failures return without reading the body. JSON - // server failures are read within MAX_ERROR_RESPONSE_SIZE to preserve + // server failures are read within max_json_response_size to preserve // their JSON-RPC error; an oversized body falls back to status metadata. let status_error = match ( http_status_error(status, response.headers()), @@ -319,40 +369,6 @@ impl StreamableHttpClient for reqwest::Client { .headers() .get(reqwest::header::CONTENT_TYPE) .cloned(); - let body = match read_bounded_body( - response.bytes_stream(), - content_length, - MAX_ERROR_RESPONSE_SIZE, - ) - .await - { - Ok(body) => { - // Decode only the bounded, in-memory body, using reqwest's - // feature-dependent charset and BOM handling. No network - // stream or Content-Encoding is carried into this response. - let mut buffered = http::Response::new(body); - if let Some(content_type) = content_type_header { - buffered - .headers_mut() - .insert(reqwest::header::CONTENT_TYPE, content_type); - } - reqwest::Response::from(buffered) - .text() - .await - .unwrap_or_else(|_| "".to_owned()) - } - Err(error @ StreamableHttpError::ResponseBodyTooLarge { .. }) => { - return Err(status_error.map(StreamableHttpError::from).unwrap_or(error)); - } - Err(_) => "".to_owned(), - }; - // Must precede the JSON-RPC branch below, which would forward a - // discover rejection with an id the lifecycle cannot correlate. - if let Some(response) = - legacy_discover_response(&message, session_was_attached, status, &body) - { - return Ok(response); - } // A 401/403 without a challenge (those with one returned above) // must keep its status rather than pass as an ordinary JSON-RPC // error, matching `legacy_discover_response`. @@ -360,15 +376,94 @@ impl StreamableHttpClient for reqwest::Client { status, reqwest::StatusCode::UNAUTHORIZED | reqwest::StatusCode::FORBIDDEN ); - if !is_auth_failure && is_json_response { - match parse_json_rpc_error(&body) { - Some(message) => { - return Ok(StreamableHttpPostResponse::Json(message, session_id)); + // JSON-RPC errors need the same complete-body limit as success + // messages. Non-JSON error bodies are diagnostics: retain only a + // prefix so large legacy error pages still reach discovery fallback. + let body = if is_json_response { + match read_bounded_body( + response.bytes_stream(), + content_length, + limits.max_json_response_size, + ) + .await + { + Ok(mut raw_body) if is_auth_failure => { + raw_body.truncate(limits.max_error_response_size); + decode_error_body( + raw_body, + content_type_header, + limits.max_error_response_size, + ) + .await } - None => tracing::warn!( - "HTTP {status}: could not parse JSON body as a JSON-RPC error" - ), + Ok(raw_body) => { + let prefix_bytes = (raw_body.len() > limits.max_error_response_size) + .then(|| raw_body[..limits.max_error_response_size].to_vec()); + let full_body = decode_error_body( + raw_body, + content_type_header.clone(), + limits.max_error_response_size, + ) + .await; + if let Some(error_message) = parse_json_rpc_error(&full_body) { + // Must precede returning the raw JSON-RPC error, which + // would forward a discover rejection with an id the + // lifecycle cannot correlate. + if let Some(response) = legacy_discover_response( + &message, + session_was_attached, + status, + &full_body, + ) { + return Ok(response); + } + return Ok(StreamableHttpPostResponse::Json(error_message, session_id)); + } + tracing::warn!( + "HTTP {status}: could not parse JSON body as a JSON-RPC error" + ); + match prefix_bytes { + Some(prefix) => { + decode_error_body( + prefix, + content_type_header, + limits.max_error_response_size, + ) + .await + } + None => full_body, + } + } + Err(error @ StreamableHttpError::ResponseBodyTooLarge { .. }) => { + return Err(status_error.map(StreamableHttpError::from).unwrap_or(error)); + } + Err(StreamableHttpError::Client(_)) => { + truncated_read_failure(limits.max_error_response_size) + } + Err(error) => return Err(error), } + } else { + match read_truncated_body(response.bytes_stream(), limits.max_error_response_size) + .await + { + Ok(prefix) => { + decode_error_body( + prefix, + content_type_header, + limits.max_error_response_size, + ) + .await + } + Err(StreamableHttpError::Client(_)) => { + truncated_read_failure(limits.max_error_response_size) + } + Err(error) => return Err(error), + } + }; + if let Some(response) = + legacy_discover_response(&message, session_was_attached, status, &body) + { + return Ok(response); } if let Some(error) = status_error { return Err(error.into()); @@ -379,7 +474,8 @@ impl StreamableHttpClient for reqwest::Client { } match content_type.as_deref() { Some(ct) if ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes()) => { - let event_stream = bounded_sse_stream(response.bytes_stream(), max_sse_event_size); + let event_stream = + bounded_sse_stream(response.bytes_stream(), limits.max_sse_event_size); Ok(StreamableHttpPostResponse::Sse(event_stream, session_id)) } Some(ct) if ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes()) => { @@ -390,7 +486,7 @@ impl StreamableHttpClient for reqwest::Client { let body = read_bounded_body( response.bytes_stream(), content_length, - MAX_JSON_RESPONSE_SIZE, + limits.max_json_response_size, ) .await?; match serde_json::from_slice::(&body) { diff --git a/crates/rmcp/src/transport/common/unix_socket.rs b/crates/rmcp/src/transport/common/unix_socket.rs index 441b66421..c55cf2c17 100644 --- a/crates/rmcp/src/transport/common/unix_socket.rs +++ b/crates/rmcp/src/transport/common/unix_socket.rs @@ -10,12 +10,10 @@ use sse_stream::Sse; use tokio::net::UnixStream; use crate::{ - model::{ClientJsonRpcMessage, ServerJsonRpcMessage}, + model::{ClientJsonRpcMessage, JsonRpcMessage, ServerJsonRpcMessage}, transport::{ common::{ - client_side_http::{ - MAX_ERROR_RESPONSE_SIZE, MAX_JSON_RESPONSE_SIZE, read_bounded_body, - }, + client_side_http::{read_bounded_body, read_truncated_body}, client_side_sse::{DEFAULT_MAX_SSE_EVENT_SIZE, bounded_sse_stream}, http_header::{ EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE, @@ -164,6 +162,12 @@ fn apply_custom_headers( Ok(builder) } +const BODY_READ_FAILURE: &str = ""; + +fn truncated_read_failure(max_error_response_size: usize) -> Vec { + BODY_READ_FAILURE.as_bytes()[..BODY_READ_FAILURE.len().min(max_error_response_size)].to_vec() +} + impl StreamableHttpClient for UnixSocketHttpClient { type Error = UnixSocketError; @@ -175,13 +179,13 @@ impl StreamableHttpClient for UnixSocketHttpClient { auth_token: Option, custom_headers: HashMap, ) -> Result> { - self.post_message_with_max_sse_event_size( + self.post_message_with_response_limits( uri, message, session_id, auth_token, custom_headers, - DEFAULT_MAX_SSE_EVENT_SIZE, + StreamableHttpResponseLimits::default(), ) .await } @@ -194,6 +198,29 @@ impl StreamableHttpClient for UnixSocketHttpClient { auth_token: Option, custom_headers: HashMap, max_sse_event_size: usize, + ) -> Result> { + self.post_message_with_response_limits( + uri, + message, + session_id, + auth_token, + custom_headers, + StreamableHttpResponseLimits { + max_sse_event_size, + ..Default::default() + }, + ) + .await + } + + async fn post_message_with_response_limits( + &self, + uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + auth_token: Option, + custom_headers: HashMap, + limits: StreamableHttpResponseLimits, ) -> Result> { let json_body = serde_json::to_string(&message) .map_err(|e| StreamableHttpError::Client(UnixSocketError::Json(e)))?; @@ -275,21 +302,58 @@ impl StreamableHttpClient for UnixSocketHttpClient { .get(http::header::CONTENT_LENGTH) .and_then(|v| v.to_str().ok()) .and_then(|v| v.parse::().ok()); + let content_type = response.headers().get(http::header::CONTENT_TYPE).cloned(); + let session_id = response + .headers() + .get(HEADER_SESSION_ID) + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()); + if !status.is_success() { - let body = match read_bounded_body( - response - .into_body() - .into_data_stream() - .map_err(UnixSocketError::Hyper), - content_length, - MAX_ERROR_RESPONSE_SIZE, - ) - .await - { - Ok(body) => String::from_utf8_lossy(&body).into_owned(), - Err(error @ StreamableHttpError::ResponseBodyTooLarge { .. }) => return Err(error), - Err(_) => "".to_owned(), + let is_json = content_type + .as_ref() + .is_some_and(|ct| ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes())); + let stream = response + .into_body() + .into_data_stream() + .map_err(UnixSocketError::Hyper); + let body = if is_json { + read_bounded_body(stream, content_length, limits.max_json_response_size).await + } else { + read_truncated_body(stream, limits.max_error_response_size).await + }; + let body = match body { + Ok(body) => body, + Err(StreamableHttpError::Client(_)) => { + truncated_read_failure(limits.max_error_response_size) + } + Err(error) => return Err(error), }; + // Match reqwest: preserve 401/403 status and re-correlate discover + // rejections before returning a complete JSON-RPC error. + let is_auth_failure = + matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN); + if !is_auth_failure && is_json { + let full_body = String::from_utf8_lossy(&body); + match serde_json::from_str::(&full_body) { + Ok(error_message @ JsonRpcMessage::Error(_)) => { + if let Some(response) = legacy_discover_response( + &message, + session_was_attached, + status, + &full_body, + ) { + return Ok(response); + } + return Ok(StreamableHttpPostResponse::Json(error_message, session_id)); + } + _ => tracing::warn!( + "HTTP {status}: could not parse JSON body as a JSON-RPC error" + ), + } + } + let body = + String::from_utf8_lossy(&body[..body.len().min(limits.max_error_response_size)]); if let Some(response) = legacy_discover_response(&message, session_was_attached, status, &body) { @@ -300,13 +364,6 @@ impl StreamableHttpClient for UnixSocketHttpClient { ))); } - let content_type = response.headers().get(http::header::CONTENT_TYPE).cloned(); - let session_id = response - .headers() - .get(HEADER_SESSION_ID) - .and_then(|v| v.to_str().ok()) - .map(|s| s.to_string()); - if status.is_success() && content_length == Some(0) && matches!( @@ -321,8 +378,10 @@ impl StreamableHttpClient for UnixSocketHttpClient { match content_type { Some(ref ct) if ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes()) => { - let sse_stream = - bounded_sse_stream(response.into_body().into_data_stream(), max_sse_event_size); + let sse_stream = bounded_sse_stream( + response.into_body().into_data_stream(), + limits.max_sse_event_size, + ); Ok(StreamableHttpPostResponse::Sse(sse_stream, session_id)) } Some(ref ct) if ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes()) => { @@ -332,7 +391,7 @@ impl StreamableHttpClient for UnixSocketHttpClient { .into_data_stream() .map_err(UnixSocketError::Hyper), content_length, - MAX_JSON_RESPONSE_SIZE, + limits.max_json_response_size, ) .await?; match serde_json::from_slice::(&body) { diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index 7b1981ff0..ab55a03f3 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -453,20 +453,56 @@ pub(super) fn legacy_discover_response( )) } +/// Byte limits enforced by the built-in Streamable HTTP clients before parsing. +/// +/// The defaults are 16 MiB per SSE event, 64 MiB per JSON response (regardless of +/// HTTP status), and 64 KiB per HTTP diagnostic prefix. JSON bodies above their +/// limit are rejected; non-JSON error bodies are truncated, preserving legacy +/// discovery fallback. A zero JSON limit accepts only empty bodies, while a zero +/// diagnostic limit skips the body. For compressed responses, body limits count +/// the decompressed bytes yielded by the HTTP backend, not total memory usage. +#[derive(Debug, Clone, Copy)] +#[non_exhaustive] +pub struct StreamableHttpResponseLimits { + /// Maximum raw size of an individual SSE event. + pub max_sse_event_size: usize, + /// Maximum complete JSON response body, including HTTP error responses. + pub max_json_response_size: usize, + /// Maximum diagnostic prefix of a non-success HTTP response body. + /// + /// Non-JSON bodies are read only up to this limit. JSON bodies are first + /// bounded by `max_json_response_size` and parsed in full; only when they are + /// not JSON-RPC errors is their diagnostic text truncated to this limit. + pub max_error_response_size: usize, +} + +impl Default for StreamableHttpResponseLimits { + fn default() -> Self { + Self { + max_sse_event_size: DEFAULT_MAX_SSE_EVENT_SIZE, + max_json_response_size: 64 * 1024 * 1024, + max_error_response_size: 64 * 1024, + } + } +} + /// HTTP backend used by [`StreamableHttpClientTransport`]. /// -/// The built-in reqwest and Unix socket clients limit successful JSON response -/// bodies to 16 MiB and non-success response bodies to 64 KiB. These fixed limits -/// count bytes before parsing (after decompression, when enabled in reqwest), -/// independently of the SSE event-size setting. Oversized bodies are rejected, -/// not truncated and parsed; reqwest JSON 5xx failures retain HTTP status metadata. -/// Custom clients must enforce their own limits on non-streaming response bodies. +/// The built-in reqwest and Unix socket clients enforce [`StreamableHttpResponseLimits`] +/// (defaulting to 64 MiB per complete JSON response regardless of HTTP status, +/// 64 KiB per HTTP diagnostic prefix, and 16 MiB per SSE event). Limits count +/// bytes before parsing (after decompression, when enabled in reqwest). Oversized +/// JSON bodies are rejected rather than truncated and parsed (reqwest JSON 5xx +/// failures retain HTTP status metadata), while non-JSON error bodies retain at +/// most the diagnostic prefix so legacy discovery fallback still succeeds. /// These are size limits, not timeouts for servers that stop sending data. /// /// Custom implementations that parse SSE responses must override /// [`Self::post_message_with_max_sse_event_size`] and /// [`Self::get_stream_with_max_sse_event_size`] to enforce the transport's /// configured event-size limit. +/// Implementations that buffer JSON or HTTP error responses must also override +/// [`Self::post_message_with_response_limits`] to enforce the body-size limits. /// /// For legacy http, the transport keeps an open response stream alive until /// its cancellation send finishes or is dropped. This lets a custom client @@ -506,6 +542,33 @@ pub trait StreamableHttpClient: Clone + Send + 'static { + '_ { self.post_message(uri, message, session_id, auth_header, custom_headers) } + /// Send a message with transport-wide limits applied before response parsing. + /// + /// The built-in reqwest and Unix socket clients enforce all three limits. + /// Custom clients that buffer responses must override this method. The + /// default preserves compatibility with existing clients and forwards only + /// the SSE limit to [`Self::post_message_with_max_sse_event_size`]; it cannot + /// impose byte limits on responses that a custom client has already decoded. + fn post_message_with_response_limits( + &self, + uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + auth_header: Option, + custom_headers: HashMap, + limits: StreamableHttpResponseLimits, + ) -> impl Future>> + + Send + + '_ { + self.post_message_with_max_sse_event_size( + uri, + message, + session_id, + auth_header, + custom_headers, + limits.max_sse_event_size, + ) + } fn delete_session( &self, uri: Arc, @@ -698,7 +761,7 @@ impl StreamableHttpClientWorker { ) -> BoxFuture<'static, PostResult> { let uri = config.uri.clone(); let auth_header = config.auth_header.clone(); - let max_sse_event_size = config.max_sse_event_size; + let response_limits = config.response_limits(); let control_request_timeout = config.control_request_timeout; let is_control = Self::is_control_message(&send_request.message); let cancellation = send_request @@ -715,13 +778,13 @@ impl StreamableHttpClientWorker { _ = tokio::time::sleep(control_request_timeout), if is_control => { Some(Err(StreamableHttpError::ControlRequestTimeout)) }, - response = client.post_message_with_max_sse_event_size( + response = client.post_message_with_response_limits( uri, send_request.message.clone(), session.id, auth_header, session.headers, - max_sse_event_size, + response_limits, ) => Some(response), }; PostResult { @@ -1062,7 +1125,7 @@ impl StreamableHttpClientWorker { uri: Arc, auth_header: Option, custom_headers: HashMap, - max_sse_event_size: usize, + response_limits: StreamableHttpResponseLimits, ) -> Result< ( Option>, @@ -1072,13 +1135,13 @@ impl StreamableHttpClientWorker { StreamableHttpError, > { let (init_msg, new_session_id_str) = client - .post_message_with_max_sse_event_size( + .post_message_with_response_limits( uri.clone(), saved_init_request, None, auth_header.clone(), custom_headers.clone(), - max_sse_event_size, + response_limits, ) .await? .expect_initialized::() @@ -1106,13 +1169,13 @@ impl StreamableHttpClientWorker { &negotiated_version, ); client - .post_message_with_max_sse_event_size( + .post_message_with_response_limits( uri, initialized_notification, new_session_id.clone(), auth_header, initialized_headers, - max_sse_event_size, + response_limits, ) .await? .expect_accepted_or_json::()?; @@ -1183,13 +1246,13 @@ impl Worker for StreamableHttpClientWorker { }; let (message, session_id) = match self .client - .post_message_with_max_sse_event_size( + .post_message_with_response_limits( config.uri.clone(), startup_request, None, config.auth_header.clone(), bootstrap_headers.clone(), - config.max_sse_event_size, + config.response_limits(), ) .await { @@ -1250,13 +1313,13 @@ impl Worker for StreamableHttpClientWorker { &negotiated_version, ); self.client - .post_message_with_max_sse_event_size( + .post_message_with_response_limits( config.uri.clone(), initialized_notification.message, session_id.clone(), config.auth_header.clone(), initialized_headers, - config.max_sse_event_size, + config.response_limits(), ) .await .map_err(WorkerQuitReason::fatal_context( @@ -1336,7 +1399,7 @@ impl Worker for StreamableHttpClientWorker { config.uri.clone(), config.auth_header.clone(), config.custom_headers.clone(), - config.max_sse_event_size, + config.response_limits(), ), ) => result.unwrap_or(Err(StreamableHttpError::SessionRecoveryTimeout)), }; @@ -1579,13 +1642,13 @@ impl Worker for StreamableHttpClientWorker { let response = self .client - .post_message_with_max_sse_event_size( + .post_message_with_response_limits( config.uri.clone(), message, None, config.auth_header.clone(), config.custom_headers.clone(), - config.max_sse_event_size, + config.response_limits(), ) .await; let response = match response { @@ -2131,6 +2194,23 @@ pub struct StreamableHttpClientTransportConfig { /// [`StreamableHttpClient`] implementations must override the corresponding /// `*_with_max_sse_event_size` methods to enforce it. pub max_sse_event_size: usize, + /// Maximum complete JSON response body, regardless of HTTP status (default: 64 MiB). + /// + /// Built-in clients enforce this before decoding, including initialization, + /// discovery, control requests and session recovery. Custom clients must + /// override [`StreamableHttpClient::post_message_with_response_limits`]. + /// Zero accepts only an empty body; increase this for larger tool results. + pub max_json_response_size: usize, + /// Maximum HTTP error diagnostic prefix (default: 64 KiB). + /// + /// Non-JSON error bodies are truncated, without draining the stream, so a + /// large diagnostic page does not prevent legacy discovery fallback. JSON + /// errors instead use `max_json_response_size` before parsing; only malformed + /// or non-error JSON is truncated for diagnostics. Zero skips diagnostic text. + /// Authentication challenges and session-expired responses are returned + /// without buffering their bodies. Custom clients must override + /// [`StreamableHttpClient::post_message_with_response_limits`]. + pub max_error_response_size: usize, /// Automatically creates a new session when the server reports an expired /// session (`http 404`). /// @@ -2215,6 +2295,26 @@ impl StreamableHttpClientTransportConfig { self } + /// Set the maximum JSON response body size before decoding. + pub fn max_json_response_size(mut self, bytes: usize) -> Self { + self.max_json_response_size = bytes; + self + } + + /// Set the maximum HTTP error diagnostic prefix size (not the JSON limit). + pub fn max_error_response_size(mut self, bytes: usize) -> Self { + self.max_error_response_size = bytes; + self + } + + fn response_limits(&self) -> StreamableHttpResponseLimits { + StreamableHttpResponseLimits { + max_sse_event_size: self.max_sse_event_size, + max_json_response_size: self.max_json_response_size, + max_error_response_size: self.max_error_response_size, + } + } + /// Set whether the transport should attempt transparent re-initialization on session expiration /// See [`Self::reinit_on_expired_session`] for details. /// # Example @@ -2242,6 +2342,9 @@ impl Default for StreamableHttpClientTransportConfig { auth_header: None, custom_headers: HashMap::new(), max_sse_event_size: DEFAULT_MAX_SSE_EVENT_SIZE, + max_json_response_size: StreamableHttpResponseLimits::default().max_json_response_size, + max_error_response_size: StreamableHttpResponseLimits::default() + .max_error_response_size, reinit_on_expired_session: true, } } diff --git a/crates/rmcp/tests/test_streamable_http_default_error_limit.rs b/crates/rmcp/tests/test_streamable_http_default_error_limit.rs new file mode 100644 index 000000000..cf4bf4944 --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_default_error_limit.rs @@ -0,0 +1,84 @@ +#![cfg(all( + feature = "client", + feature = "transport-streamable-http-client-reqwest", + not(feature = "local") +))] + +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use axum::{Router, http::StatusCode, routing::post}; +use rmcp::{ + model::{ClientJsonRpcMessage, ClientRequest, PingRequest, RequestId}, + transport::streamable_http_client::StreamableHttpClient, +}; + +struct MockServer { + uri: Arc, + task: tokio::task::JoinHandle<()>, +} + +impl Drop for MockServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +impl MockServer { + async fn start(body: String) -> Self { + let router = Router::new().route( + "/mcp", + post(move || { + let body = body.clone(); + async move { (StatusCode::BAD_REQUEST, body) } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let task = tokio::spawn(async move { + axum::serve(listener, router).await.unwrap(); + }); + Self { + uri: Arc::from(format!("http://{address}/mcp")), + task, + } + } +} + +/// Exercise only the existing public API so this regression also compiles +/// against an SDK that still buffers HTTP error responses without a bound. +#[tokio::test] +async fn default_client_truncates_error_body_without_exposing_the_tail() { + const MARKER: &str = "TAIL_TEST_MARKER_DO_NOT_ECHO"; + let mut body = "x".repeat(65_536); + body.push_str(MARKER); + let server = MockServer::start(body).await; + let client = reqwest::Client::builder() + .no_proxy() + .timeout(Duration::from_secs(5)) + .build() + .unwrap(); + let message = ClientJsonRpcMessage::request( + ClientRequest::PingRequest(PingRequest::default()), + RequestId::Number(1), + ); + + let error = client + .post_message(server.uri.clone(), message, None, None, HashMap::new()) + .await + .expect_err("an oversized HTTP error must fail") + .to_string(); + + assert!( + error.starts_with("unexpected server response: HTTP 400"), + "{error}" + ); + assert!( + error.contains(&"x".repeat(65_536)), + "expected diagnostic prefix" + ); + assert!( + !error.contains(MARKER), + "error must not expose the discarded tail" + ); + assert!(error.len() < 65_636, "error must remain bounded"); +} diff --git a/crates/rmcp/tests/test_streamable_http_limit_lifecycle.rs b/crates/rmcp/tests/test_streamable_http_limit_lifecycle.rs new file mode 100644 index 000000000..9016fbb78 --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_limit_lifecycle.rs @@ -0,0 +1,516 @@ +#![cfg(all( + feature = "client", + feature = "transport-streamable-http-client-reqwest", + not(feature = "local") +))] + +use std::{ + collections::HashMap, + convert::Infallible, + io, + sync::{Arc, Mutex}, + time::Duration, +}; + +use axum::{Router, body::Body, http::Response, routing::post}; +use bytes::Bytes; +use futures::{StreamExt, stream::BoxStream}; +use http::{HeaderName, HeaderValue}; +use rmcp::{ + ServiceExt, + model::{ + CancelledNotificationParam, ClientConfig, ClientJsonRpcMessage, ClientRequest, + DiscoverResult, PingRequest, ProtocolVersion, RequestId, ServerJsonRpcMessage, + }, + service::{ClientInitializeError, ClientLifecycleMode, ClientServiceExt}, + transport::streamable_http_client::{ + StreamableHttpClient, StreamableHttpClientTransport, StreamableHttpClientTransportConfig, + StreamableHttpError, StreamableHttpPostResponse, StreamableHttpResponseLimits, + }, +}; +use rstest::rstest; +use serde_json::{Value, json}; + +const TIMEOUT: Duration = Duration::from_secs(2); +const LIMITS: (usize, usize, usize) = (111, 222, 333); + +#[derive(Debug)] +struct RecordedPost { + method: String, + id: Value, + session: Option>, + limits: (usize, usize, usize), +} + +#[derive(Default)] +struct RecordingState { + posts: Vec, + initializations: usize, + pings: usize, +} + +#[derive(Clone)] +struct RecordingClient { + modern: bool, + state: Arc>, +} + +impl StreamableHttpClient for RecordingClient { + type Error = io::Error; + + async fn post_message( + &self, + _uri: Arc, + _message: ClientJsonRpcMessage, + _session_id: Option>, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result> { + panic!("transport bypassed the configured response limits") + } + + async fn post_message_with_response_limits( + &self, + _uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + _auth_header: Option, + _custom_headers: HashMap, + limits: StreamableHttpResponseLimits, + ) -> Result> { + let message = serde_json::to_value(message).unwrap(); + let method = message["method"].as_str().unwrap(); + let id = message["id"].clone(); + let mut state = self.state.lock().unwrap(); + state.posts.push(RecordedPost { + method: method.to_owned(), + id: id.clone(), + session: session_id, + limits: ( + limits.max_sse_event_size, + limits.max_json_response_size, + limits.max_error_response_size, + ), + }); + let (response, session) = match method { + "server/discover" if self.modern => ( + json!({ + "jsonrpc": "2.0", "id": id, + "result": DiscoverResult::new( + vec![ProtocolVersion::V_2026_07_28], Default::default() + ), + }), + None, + ), + "server/discover" => ( + json!({ + "jsonrpc": "2.0", "id": id, + "error": { "code": -32601, "message": "legacy server" }, + }), + None, + ), + "initialize" => { + state.initializations += 1; + ( + json!({ + "jsonrpc": "2.0", "id": id, + "result": { + "protocolVersion": "2025-11-25", "capabilities": {}, + "serverInfo": { "name": "limit-recorder", "version": "1" }, + }, + }), + Some(format!("session-{}", state.initializations)), + ) + } + "notifications/initialized" | "notifications/cancelled" => { + return Ok(StreamableHttpPostResponse::Accepted); + } + "ping" => { + state.pings += 1; + if !self.modern && state.pings == 1 { + return Err(StreamableHttpError::SessionExpired); + } + (json!({ "jsonrpc": "2.0", "id": id, "result": {} }), None) + } + other => panic!("unexpected POST method: {other}"), + }; + Ok(StreamableHttpPostResponse::Json( + serde_json::from_value::(response).unwrap(), + session, + )) + } + + async fn delete_session( + &self, + _uri: Arc, + _session_id: Arc, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result<(), StreamableHttpError> { + Ok(()) + } + + async fn get_stream( + &self, + _uri: Arc, + _session_id: Option>, + _last_event_id: Option, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result< + BoxStream<'static, Result>, + StreamableHttpError, + > { + Ok(Box::pin(futures::stream::pending())) + } +} + +fn auto_lifecycle() -> ClientLifecycleMode { + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + } +} + +/// An existing client implements only the pre-existing SSE-limited entry point. +#[derive(Clone, Default)] +struct LegacySseOnlyClient(Arc>>); + +impl StreamableHttpClient for LegacySseOnlyClient { + type Error = io::Error; + + async fn post_message( + &self, + _uri: Arc, + _message: ClientJsonRpcMessage, + _session_id: Option>, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result> { + panic!("new limits entry point bypassed the old SSE override") + } + + async fn post_message_with_max_sse_event_size( + &self, + _uri: Arc, + _message: ClientJsonRpcMessage, + _session_id: Option>, + _auth_header: Option, + _custom_headers: HashMap, + max_sse_event_size: usize, + ) -> Result> { + self.0.lock().unwrap().push(max_sse_event_size); + Ok(StreamableHttpPostResponse::Accepted) + } + + async fn delete_session( + &self, + _uri: Arc, + _session_id: Arc, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result<(), StreamableHttpError> { + unreachable!("compatibility test does not create a session") + } + + async fn get_stream( + &self, + _uri: Arc, + _session_id: Option>, + _last_event_id: Option, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result< + BoxStream<'static, Result>, + StreamableHttpError, + > { + unreachable!("compatibility test does not open a stream") + } +} + +#[tokio::test] +async fn new_limits_entry_point_preserves_legacy_sse_override() { + let client = LegacySseOnlyClient::default(); + let mut limits = StreamableHttpResponseLimits::default(); + limits.max_sse_event_size = LIMITS.0; + limits.max_json_response_size = LIMITS.1; + limits.max_error_response_size = LIMITS.2; + let response = client + .post_message_with_response_limits( + Arc::from("http://127.0.0.1/record-only"), + ClientJsonRpcMessage::request( + ClientRequest::PingRequest(PingRequest::default()), + RequestId::Number(1), + ), + None, + None, + HashMap::new(), + limits, + ) + .await + .unwrap(); + assert!(matches!(response, StreamableHttpPostResponse::Accepted)); + assert_eq!(client.0.lock().unwrap().as_slice(), [LIMITS.0]); +} + +#[rstest] +#[case::legacy(false, false)] +#[case::discover_fallback(true, false)] +#[case::modern_discover(true, true)] +#[tokio::test] +async fn response_limits_reach_every_lifecycle_post( + #[case] auto: bool, + #[case] modern: bool, +) -> anyhow::Result<()> { + let recorder = RecordingClient { + modern, + state: Arc::default(), + }; + let transport = StreamableHttpClientTransport::with_client( + recorder.clone(), + StreamableHttpClientTransportConfig::with_uri("http://127.0.0.1/record-only") + .max_sse_event_size(LIMITS.0) + .max_json_response_size(LIMITS.1) + .max_error_response_size(LIMITS.2) + .reinit_on_expired_session(true), + ); + let client = if auto { + tokio::time::timeout( + TIMEOUT, + ClientConfig::default().serve_with_lifecycle(transport, auto_lifecycle()), + ) + .await?? + } else { + tokio::time::timeout(TIMEOUT, ClientConfig::default().serve(transport)).await?? + }; + + tokio::time::timeout( + TIMEOUT, + client.send_request(ClientRequest::PingRequest(PingRequest::default())), + ) + .await??; + if !modern { + tokio::time::timeout( + TIMEOUT, + client.notify_cancelled(CancelledNotificationParam::new( + Some(RequestId::Number(999)), + None, + )), + ) + .await??; + } + tokio::time::timeout(TIMEOUT, client.cancel()).await??; + + let state = recorder.state.lock().unwrap(); + let expected = if modern { + vec!["server/discover", "ping"] + } else { + let mut expected = vec![ + "initialize", + "notifications/initialized", + "ping", + "initialize", + "notifications/initialized", + "ping", + "notifications/cancelled", + ]; + if auto { + expected.insert(0, "server/discover"); + } + expected + }; + assert_eq!( + state + .posts + .iter() + .map(|post| post.method.as_str()) + .collect::>(), + expected, + ); + for post in &state.posts { + assert_eq!(post.limits, LIMITS, "limits changed for {post:?}"); + } + if !modern { + let pings = state + .posts + .iter() + .filter(|post| post.method == "ping") + .collect::>(); + assert_eq!(pings[0].session.as_deref(), Some("session-1")); + assert_eq!(pings[1].session.as_deref(), Some("session-2")); + assert_eq!( + pings[0].id, pings[1].id, + "recovery must retry the original request" + ); + for initialize in state + .posts + .iter() + .filter(|post| post.method == "initialize") + { + assert!( + initialize.session.is_none(), + "initialize must start a new session" + ); + } + } + Ok(()) +} + +struct LoopbackServer { + uri: String, + posts: Arc>>, + task: tokio::task::JoinHandle<()>, +} + +impl Drop for LoopbackServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +impl LoopbackServer { + async fn oversized_discover( + status: u16, + content_type: &'static str, + never_finishes: bool, + ) -> Self { + let posts = Arc::new(Mutex::new(Vec::new())); + let recorded_posts = posts.clone(); + let router = Router::new().route( + "/mcp", + post(move |request: Bytes| { + let posts = recorded_posts.clone(); + async move { + let request: Value = serde_json::from_slice(&request).unwrap(); + posts + .lock() + .unwrap() + .push(request["method"].as_str().unwrap().to_owned()); + if request["method"] == "initialize" { + return Response::builder() + .status(200) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "jsonrpc": "2.0", "id": request["id"], + "result": { + "protocolVersion": "2025-11-25", "capabilities": {}, + "serverInfo": {"name": "legacy", "version": "1.0"} + } + }) + .to_string(), + )) + .unwrap(); + } + if request["method"] == "notifications/initialized" { + return Response::builder().status(202).body(Body::empty()).unwrap(); + } + let payload = if content_type == "text/plain" { + "Unexpected message, expect initialize request".repeat(2048) + } else { + json!({ + "jsonrpc": "2.0", "id": request["id"], + "error": { "code": -32601, "message": "legacy".repeat(80) }, + }) + .to_string() + }; + let body = if never_finishes { + // The limit must abort reading without waiting for EOF. + Body::from_stream( + futures::stream::once(async move { + Ok::<_, Infallible>(Bytes::from(payload)) + }) + .chain(futures::stream::pending()), + ) + } else { + Body::from(payload) + }; + Response::builder() + .status(status) + .header("content-type", content_type) + .body(body) + .unwrap() + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let task = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); + Self { + uri: format!("http://{address}/mcp"), + posts, + task, + } + } +} + +#[rstest] +#[case::json(200, false)] +#[case::json_error(400, false)] +#[case::unfinished_json(200, true)] +#[case::unfinished_json_error(400, true)] +#[tokio::test] +async fn oversized_discover_fails_without_fallback_or_retry( + #[case] status: u16, + #[case] never_finishes: bool, +) { + let server = + LoopbackServer::oversized_discover(status, "application/json", never_finishes).await; + let transport = StreamableHttpClientTransport::with_client( + reqwest::Client::builder().no_proxy().build().unwrap(), + StreamableHttpClientTransportConfig::with_uri(server.uri.clone()) + .max_json_response_size(64) + .max_error_response_size(64) + .reinit_on_expired_session(true), + ); + let result = tokio::time::timeout( + TIMEOUT, + ClientConfig::default().serve_with_lifecycle(transport, auto_lifecycle()), + ) + .await + .expect("oversized discovery must fail before EOF or the auto fallback timeout"); + let error = match result { + Ok(client) => { + client.cancel().await.unwrap(); + panic!("oversized discovery was accepted"); + } + Err(error) => error, + }; + let ClientInitializeError::TransportError { error, .. } = error else { + panic!("expected the transport size error without lifecycle fallback: {error:?}"); + }; + assert!( + matches!( + error + .error + .downcast_ref::>(), + Some(StreamableHttpError::ResponseBodyTooLarge { max_size: 64 }) + ), + "unexpected transport error: {error:?}" + ); + assert_eq!(server.posts.lock().unwrap().as_slice(), ["server/discover"]); +} + +#[rstest] +#[case::fixed(false)] +#[case::unfinished_chunked(true)] +#[tokio::test] +async fn oversized_error_page_completes_legacy_handshake(#[case] never_finishes: bool) { + let server = LoopbackServer::oversized_discover(422, "text/plain", never_finishes).await; + let transport = StreamableHttpClientTransport::with_client( + reqwest::Client::builder().no_proxy().build().unwrap(), + StreamableHttpClientTransportConfig::with_uri(server.uri.clone()), + ); + let client = tokio::time::timeout( + TIMEOUT, + ClientConfig::default().serve_with_lifecycle(transport, auto_lifecycle()), + ) + .await + .expect("legacy startup must not wait for the error page to finish") + .expect("a truncated diagnostic must preserve legacy startup"); + assert_eq!( + server.posts.lock().unwrap().as_slice(), + ["server/discover", "initialize", "notifications/initialized"] + ); + client.cancel().await.unwrap(); +} diff --git a/crates/rmcp/tests/test_streamable_http_response_limits.rs b/crates/rmcp/tests/test_streamable_http_response_limits.rs index adf40e18a..6973a11e5 100644 --- a/crates/rmcp/tests/test_streamable_http_response_limits.rs +++ b/crates/rmcp/tests/test_streamable_http_response_limits.rs @@ -3,11 +3,11 @@ feature = "transport-streamable-http-client-reqwest", not(feature = "local") ))] -//! The reqwest backend reads non-SSE POST bodies only up to a fixed size: JSON -//! success bodies up to 16 MiB, error bodies up to 64 KiB. An oversized body -//! fails closed: no prefix is parsed and a notification is not reported as -//! accepted. The client must also stop reading a body that is still open -//! once the limit is exceeded. +//! The reqwest backend enforces `StreamableHttpResponseLimits` on non-SSE POST +//! responses: complete JSON bodies (regardless of HTTP status) default to +//! 64 MiB and fail closed when exceeded, while non-JSON error diagnostics +//! default to 64 KiB and stop reading once that prefix is full so large legacy +//! 4xx pages still reach `server/discover` fallback. //! //! A raw TCP fixture controls the framing exactly, including `Content-Length` //! values that do not match the bytes sent, chunked bodies without a final @@ -15,6 +15,7 @@ use std::{collections::HashMap, fmt::Debug, io, sync::Arc, time::Duration}; +use futures::StreamExt; use rmcp::{ model::{ ClientJsonRpcMessage, ClientNotification, ClientRequest, DiscoverRequest, @@ -22,6 +23,7 @@ use rmcp::{ }, transport::streamable_http_client::{ StreamableHttpClient, StreamableHttpError, StreamableHttpPostResponse, + StreamableHttpResponseLimits, }, }; use tokio::{ @@ -30,13 +32,14 @@ use tokio::{ task::{JoinHandle, JoinSet}, }; -/// Must match the private limits in the reqwest backend. -const MAX_JSON_BODY: usize = 16 * 1024 * 1024; +/// Default limits enforced by the reqwest backend. +const MAX_JSON_BODY: usize = 64 * 1024 * 1024; const MAX_ERROR_BODY: usize = 64 * 1024; +/// Smaller custom JSON bound for parameterized framing matrices. +const CUSTOM_JSON_BODY: usize = 128 * 1024; /// Far above any limit; never sent, only declared. const HUGE_DECLARED: u64 = 1 << 40; const JSON_OVER: u64 = MAX_JSON_BODY as u64 + 1; -const ERROR_OVER: u64 = MAX_ERROR_BODY as u64 + 1; /// Generous; a bounded read finishes in well under a second. const DEADLINE: Duration = Duration::from_secs(10); const WRITE_CHUNK: usize = 64 * 1024; @@ -77,6 +80,8 @@ enum Payload { Error, /// Plain text. Text, + /// Raw SSE event stream. + Sse, } impl Payload { @@ -90,6 +95,7 @@ impl Payload { Payload::Result => (RESULT_JSON, b' '), Payload::Error => (ERROR_JSON, b' '), Payload::Text => (TEXT, b'x'), + Payload::Sse => ("data: example\n\n", b' '), }; assert!(prefix.len() <= len); let mut body = Vec::with_capacity(len); @@ -254,6 +260,14 @@ async fn write_reply(stream: &mut TcpStream, reply: &Reply) -> io::Result<()> { stream.flush().await } +fn custom_limits(json: usize, error: usize, sse: usize) -> StreamableHttpResponseLimits { + let mut limits = StreamableHttpResponseLimits::default(); + limits.max_json_response_size = json; + limits.max_error_response_size = error; + limits.max_sse_event_size = sse; + limits +} + async fn post(server: &RawServer, message: Message, session: Option<&str>) -> PostResult { tokio::time::timeout( DEADLINE, @@ -269,6 +283,27 @@ async fn post(server: &RawServer, message: Message, session: Option<&str>) -> Po .expect("the POST must finish without reading past the body limit") } +async fn post_with_limits( + server: &RawServer, + message: Message, + session: Option<&str>, + limits: StreamableHttpResponseLimits, +) -> PostResult { + tokio::time::timeout( + DEADLINE, + reqwest::Client::new().post_message_with_response_limits( + server.uri.clone(), + message.build(), + session.map(Arc::from), + None, + HashMap::new(), + limits, + ), + ) + .await + .expect("the POST must finish without reading past the body limit") +} + /// Bounded `Debug` output, so a failure never dumps a large body. fn summary(value: &impl Debug) -> String { const MAX: usize = 300; @@ -293,6 +328,7 @@ enum Outcome { Accepted, SessionExpired, AuthRequired, + InsufficientScope, } fn assert_outcome(result: PostResult, expected: Outcome) { @@ -323,11 +359,40 @@ fn assert_outcome(result: PostResult, expected: Outcome) { (Outcome::Unexpected, Err(StreamableHttpError::UnexpectedServerResponse(_))) | (Outcome::Accepted, Ok(StreamableHttpPostResponse::Accepted)) | (Outcome::SessionExpired, Err(StreamableHttpError::SessionExpired)) - | (Outcome::AuthRequired, Err(StreamableHttpError::AuthRequired(_))) => {} + | (Outcome::AuthRequired, Err(StreamableHttpError::AuthRequired(_))) + | (Outcome::InsufficientScope, Err(StreamableHttpError::InsufficientScope(_))) => {} (expected, _) => panic!("expected {expected:?}, got {rendered}"), } } +#[test] +fn defaults_allow_larger_json_without_changing_sse_or_diagnostics() { + let limits = StreamableHttpResponseLimits::default(); + assert_eq!(limits.max_json_response_size, MAX_JSON_BODY); + assert_eq!(limits.max_sse_event_size, 16 * 1024 * 1024); + assert_eq!(limits.max_error_response_size, MAX_ERROR_BODY); +} + +#[tokio::test] +async fn default_json_limit_accepts_base64_tool_results_above_sixteen_mib() { + let data = "A".repeat(17 * 1024 * 1024); + let body = serde_json::json!({ + "jsonrpc": "2.0", "id": 1, + "result": {"content": [{"type": "image", "data": data, "mimeType": "image/png"}]} + }) + .to_string(); + let server = serve(Reply::new( + 200, + "application/json", + body.into_bytes().into(), + )) + .await; + assert_outcome( + post(&server, Message::Request, None).await, + Outcome::JsonResult, + ); +} + /// Valid JSON followed by whitespace past the limit is rejected, so the /// limit is enforced on bytes read rather than on what parses. Held-open /// bodies prove the read stops at the limit. @@ -366,21 +431,30 @@ async fn oversized_json_success_fails_closed( #[case] framing: Framing, #[case] end: End, ) { - let body = payload.bytes(MAX_JSON_BODY + 1); + let body = payload.bytes(CUSTOM_JSON_BODY + 1); let server = serve(Reply::new(200, "application/json", body).framed(framing, end)).await; assert_outcome( - post(&server, message, None).await, - Outcome::TooLarge(MAX_JSON_BODY), + post_with_limits( + &server, + message, + None, + custom_limits(CUSTOM_JSON_BODY, MAX_ERROR_BODY, 1024), + ) + .await, + Outcome::TooLarge(CUSTOM_JSON_BODY), ); } #[tokio::test] async fn raising_the_sse_limit_does_not_disable_the_json_body_limit() { - let server = serve(Reply::new( - 200, - "application/json", - Payload::Result.bytes(MAX_JSON_BODY + 1), - )) + let server = serve( + Reply::new( + 200, + "application/json", + Payload::Result.bytes(MAX_JSON_BODY + 1), + ) + .framed(Framing::Chunked, End::Hold), + ) .await; let result = tokio::time::timeout( DEADLINE, @@ -399,10 +473,10 @@ async fn raising_the_sse_limit_does_not_disable_the_json_body_limit() { } #[rstest::rstest] -#[case::exact_content_length(Message::Request, MAX_JSON_BODY, Framing::Length)] -#[case::exact_chunked(Message::Request, MAX_JSON_BODY, Framing::Chunked)] -#[case::under_content_length(Message::Request, MAX_JSON_BODY - 1, Framing::Length)] -#[case::notification_exact(Message::Notification, MAX_JSON_BODY, Framing::Chunked)] +#[case::exact_content_length(Message::Request, CUSTOM_JSON_BODY, Framing::Length)] +#[case::exact_chunked(Message::Request, CUSTOM_JSON_BODY, Framing::Chunked)] +#[case::under_content_length(Message::Request, CUSTOM_JSON_BODY - 1, Framing::Length)] +#[case::notification_exact(Message::Notification, CUSTOM_JSON_BODY, Framing::Chunked)] #[tokio::test] async fn json_success_at_or_under_the_limit_is_parsed( #[case] message: Message, @@ -411,72 +485,181 @@ async fn json_success_at_or_under_the_limit_is_parsed( ) { let body = Payload::Result.bytes(len); let server = serve(Reply::new(200, "application/json", body).framed(framing, End::Close)).await; - assert_outcome(post(&server, message, None).await, Outcome::JsonResult); + assert_outcome( + post_with_limits( + &server, + message, + None, + custom_limits(CUSTOM_JSON_BODY, 0, 0), + ) + .await, + Outcome::JsonResult, + ); } #[rstest::rstest] -#[case::json_rpc_error( - Message::Request, - "application/json", - Payload::Error, - Framing::Length, - End::Close -)] -#[case::json_rpc_error_chunked_held_open( +#[case::json_rpc_error(Message::Request, Framing::Length, End::Close)] +#[case::json_rpc_error_chunked_held_open(Message::Request, Framing::Chunked, End::Hold)] +#[case::discover_json_rpc_error(Message::Discover, Framing::Length, End::Close)] +#[case::notification_json_rpc_error(Message::Notification, Framing::Length, End::Close)] +#[tokio::test] +async fn oversized_json_client_error_bodies_fail_closed( + #[case] message: Message, + #[case] framing: Framing, + #[case] end: End, +) { + let body = Payload::Error.bytes(CUSTOM_JSON_BODY + 1); + let server = serve(Reply::new(400, "application/json", body).framed(framing, end)).await; + assert_outcome( + post_with_limits( + &server, + message, + None, + custom_limits(CUSTOM_JSON_BODY, MAX_ERROR_BODY, 0), + ) + .await, + Outcome::TooLarge(CUSTOM_JSON_BODY), + ); +} + +#[rstest::rstest] +#[case::below_limit(1, Framing::Length)] +#[case::exact_limit(0, Framing::Length)] +#[case::over_limit(-1, Framing::Length)] +#[case::chunked_below_limit(1, Framing::Chunked)] +#[case::chunked_exact_limit(0, Framing::Chunked)] +#[case::chunked_over_limit(-1, Framing::Chunked)] +#[tokio::test] +async fn json_rpc_error_uses_json_limit_for_every_status( + #[case] margin: isize, + #[case] framing: Framing, + #[values(200, 400)] status: u16, +) { + let body: Arc<[u8]> = Arc::from(ERROR_JSON.as_bytes()); + let limit = ERROR_JSON.len().checked_add_signed(margin).unwrap(); + let server = + serve(Reply::new(status, "application/json", body).framed(framing, End::Close)).await; + let result = + post_with_limits(&server, Message::Request, None, custom_limits(limit, 0, 0)).await; + if margin < 0 { + assert_outcome(result, Outcome::TooLarge(limit)); + } else { + assert_outcome(result, Outcome::JsonError(SERVER_ERROR)); + } +} + +#[rstest::rstest] +#[case::request_length(Message::Request, Framing::Length, End::Close, Outcome::Unexpected)] +#[case::request_chunked_held_open( Message::Request, - "application/json", - Payload::Error, Framing::Chunked, - End::Hold -)] -#[case::plain_text( - Message::Request, - "text/plain", - Payload::Text, - Framing::Length, - End::Close + End::Hold, + Outcome::Unexpected )] -#[case::plain_text_chunked_held_open( +#[case::request_huge_declared( Message::Request, - "text/plain", - Payload::Text, - Framing::Chunked, - End::Hold + Framing::Declared(HUGE_DECLARED), + End::Hold, + Outcome::Unexpected )] -#[case::discover_json_rpc_error( +#[case::discover_length( Message::Discover, - "application/json", - Payload::Error, Framing::Length, - End::Close + End::Close, + Outcome::JsonError(INVALID_REQUEST) )] -#[case::discover_plain_text( +#[case::discover_chunked_held_open( Message::Discover, - "text/plain", - Payload::Text, Framing::Chunked, - End::Hold + End::Hold, + Outcome::JsonError(INVALID_REQUEST) )] -#[case::notification_json_rpc_error( - Message::Notification, - "application/json", - Payload::Error, - Framing::Length, - End::Close +#[case::discover_huge_declared( + Message::Discover, + Framing::Declared(HUGE_DECLARED), + End::Hold, + Outcome::JsonError(INVALID_REQUEST) )] #[tokio::test] -async fn oversized_client_error_bodies_fail_closed( +async fn oversized_non_json_error_truncates_diagnostic_and_preserves_legacy_discover( #[case] message: Message, - #[case] content_type: &'static str, - #[case] payload: Payload, #[case] framing: Framing, #[case] end: End, + #[case] expected: Outcome, ) { - let body = payload.bytes(MAX_ERROR_BODY + 1); - let server = serve(Reply::new(400, content_type, body).framed(framing, end)).await; + let mut bytes = vec![b'x'; MAX_ERROR_BODY]; + bytes.extend_from_slice(b"DISCARDED_TAIL"); + let server = serve(Reply::new(400, "text/plain", bytes.into()).framed(framing, end)).await; + let result = post(&server, message, None).await; + match (&result, expected) { + (Err(StreamableHttpError::UnexpectedServerResponse(actual)), Outcome::Unexpected) => { + assert!(actual.ends_with(&"x".repeat(MAX_ERROR_BODY))); + assert!(!actual.contains("DISCARDED_TAIL")); + } + ( + Ok(StreamableHttpPostResponse::Json(JsonRpcMessage::Error(error), _)), + Outcome::JsonError(code), + ) => { + assert_eq!(error.error.code.0, code); + assert_eq!(error.id, Some(RequestId::Number(1))); + assert!(error.error.message.ends_with(&"x".repeat(MAX_ERROR_BODY))); + assert!(!error.error.message.contains("DISCARDED_TAIL")); + } + _ => panic!("expected {expected:?}, got {}", summary(&result)), + } +} + +#[rstest::rstest] +#[case::fixed(Framing::Length)] +#[case::chunked(Framing::Chunked)] +#[tokio::test] +async fn truncated_diagnostics_cannot_be_parsed_as_json_rpc(#[case] framing: Framing) { + let body: Arc<[u8]> = format!("{ERROR_JSON}not-json").into_bytes().into(); + for content_type in ["text/plain", "application/json"] { + let server = + serve(Reply::new(400, content_type, body.clone()).framed(framing, End::Close)).await; + assert_outcome( + post_with_limits( + &server, + Message::Request, + None, + custom_limits(body.len(), ERROR_JSON.len(), 0), + ) + .await, + Outcome::Unexpected, + ); + if content_type == "application/json" { + assert_outcome( + post_with_limits( + &server, + Message::Discover, + None, + custom_limits(ERROR_JSON.len(), body.len(), 0), + ) + .await, + Outcome::TooLarge(ERROR_JSON.len()), + ); + } + } +} + +#[rstest::rstest] +#[case::unauthorized(401, Outcome::Unexpected)] +#[case::forbidden(403, Outcome::Unexpected)] +#[case::server_error(500, Outcome::Status(500))] +#[tokio::test] +async fn truncated_diagnostics_do_not_expand_discovery_fallback( + #[case] status: u16, + #[case] expected: Outcome, +) { + let server = serve( + Reply::new(status, "text/html", Arc::from(&b"large error page"[..])) + .framed(Framing::Chunked, End::Close), + ) + .await; assert_outcome( - post(&server, message, None).await, - Outcome::TooLarge(MAX_ERROR_BODY), + post_with_limits(&server, Message::Discover, None, custom_limits(0, 2, 0)).await, + expected, ); } @@ -491,10 +674,16 @@ async fn oversized_json_server_error_keeps_status_metadata( #[case] framing: Framing, #[case] end: End, ) { - let body = Payload::Error.bytes(MAX_ERROR_BODY + 1); + let body = Payload::Error.bytes(CUSTOM_JSON_BODY + 1); let server = serve(Reply::new(status, "application/json", body).framed(framing, end)).await; assert_outcome( - post(&server, Message::Request, None).await, + post_with_limits( + &server, + Message::Request, + None, + custom_limits(CUSTOM_JSON_BODY, MAX_ERROR_BODY, 0), + ) + .await, Outcome::Status(status), ); } @@ -527,15 +716,8 @@ async fn oversized_json_server_error_keeps_status_metadata( 400, "application/json", Message::Request, - ERROR_OVER, - Outcome::TooLarge(MAX_ERROR_BODY) -)] -#[case::client_error_text( - 400, - "text/plain", - Message::Request, - HUGE_DECLARED, - Outcome::TooLarge(MAX_ERROR_BODY) + JSON_OVER, + Outcome::TooLarge(MAX_JSON_BODY) )] #[case::server_error_json( 500, @@ -629,6 +811,7 @@ async fn error_bodies_at_the_limit_keep_their_meaning( /// so a huge declared body that never arrives does not matter. #[rstest::rstest] #[case::accepted(202, "application/json", None, None, Outcome::Accepted)] +#[case::no_content(204, "application/json", None, None, Outcome::Accepted)] #[case::session_not_found( 404, "application/json", @@ -637,6 +820,13 @@ async fn error_bodies_at_the_limit_keep_their_meaning( Outcome::SessionExpired )] #[case::auth_challenge(401, "application/json", Some(CHALLENGE), None, Outcome::AuthRequired)] +#[case::insufficient_scope( + 403, + "application/json", + Some(CHALLENGE), + None, + Outcome::InsufficientScope +)] #[case::rate_limited_json(429, "application/json", None, None, Outcome::Status(429))] #[case::server_error_text(503, "text/plain", None, None, Outcome::Status(503))] #[tokio::test] @@ -654,6 +844,32 @@ async fn body_free_responses_ignore_the_declared_size( assert_outcome(post(&server, Message::Request, session).await, expected); } +#[tokio::test] +async fn sse_responses_keep_their_independent_event_limit() { + let server = serve( + Reply::new(200, "text/event-stream", Payload::Sse.bytes(16)) + .framed(Framing::Chunked, End::Close), + ) + .await; + let StreamableHttpPostResponse::Sse(mut stream, _) = + post_with_limits(&server, Message::Request, None, custom_limits(0, 0, 64)) + .await + .unwrap() + else { + panic!("expected SSE stream"); + }; + assert!(stream.next().await.unwrap().is_ok()); + + let StreamableHttpPostResponse::Sse(mut stream, _) = + post_with_limits(&server, Message::Request, None, custom_limits(100, 100, 4)) + .await + .unwrap() + else { + panic!("expected bounded SSE stream"); + }; + assert!(stream.next().await.unwrap().is_err()); +} + // Run this suite with and without `reqwest/charset`: Cargo can enable the // dependency feature through another consumer without an rmcp feature flag. #[rstest::rstest] @@ -715,17 +931,47 @@ async fn error_body_decoding_matches_reqwest_text( #[case::length(Framing::Length, End::Close)] #[case::chunked(Framing::Chunked, End::Hold)] #[tokio::test] -async fn oversized_charset_body_is_rejected_before_decoding( +async fn oversized_charset_diagnostic_truncates_input_bytes_without_waiting_for_eof( #[case] framing: Framing, #[case] end: End, ) { - let body: Arc<[u8]> = vec![0x80; MAX_ERROR_BODY + 1].into(); - let server = - serve(Reply::new(400, "text/plain; charset=windows-1252", body).framed(framing, end)).await; - assert_outcome( - post(&server, Message::Request, None).await, - Outcome::TooLarge(MAX_ERROR_BODY), - ); + let mut prefix_plus_tail = vec![0x80; MAX_ERROR_BODY]; + prefix_plus_tail.extend_from_slice(b"TAIL_MARKER"); + let expected_prefix = serve(Reply::new( + 400, + "text/plain; charset=windows-1252", + vec![0x80; MAX_ERROR_BODY].into(), + )) + .await; + let expected = reqwest::Client::new() + .post(expected_prefix.uri.as_ref()) + .json(&Message::Request.build()) + .send() + .await + .unwrap() + .text() + .await + .unwrap(); + + let server = serve( + Reply::new( + 400, + "text/plain; charset=windows-1252", + prefix_plus_tail.into(), + ) + .framed(framing, end), + ) + .await; + match post(&server, Message::Request, None).await { + Err(StreamableHttpError::UnexpectedServerResponse(actual)) => { + assert_eq!(actual, format!("HTTP 400 Bad Request: {expected}")); + assert!(!actual.contains("TAIL_MARKER")); + } + other => panic!( + "expected truncated decoded HTTP error, got {}", + summary(&other) + ), + } } #[tokio::test] diff --git a/crates/rmcp/tests/test_unix_socket_response_limits.rs b/crates/rmcp/tests/test_unix_socket_response_limits.rs index e9264bcd9..c62a569aa 100644 --- a/crates/rmcp/tests/test_unix_socket_response_limits.rs +++ b/crates/rmcp/tests/test_unix_socket_response_limits.rs @@ -5,6 +5,8 @@ ))] use std::{ + collections::HashMap, + future::Future, path::PathBuf, sync::{ Arc, @@ -13,11 +15,18 @@ use std::{ time::Duration, }; +use futures::StreamExt; use rmcp::{ - model::{ClientJsonRpcMessage, ClientRequest, PingRequest, RequestId}, + model::{ + ClientJsonRpcMessage, ClientRequest, DiscoverRequest, DiscoverRequestParams, + JsonRpcMessage, PingRequest, RequestId, + }, transport::{ UnixSocketHttpClient, - streamable_http_client::{StreamableHttpClient, StreamableHttpError}, + streamable_http_client::{ + StreamableHttpClient, StreamableHttpError, StreamableHttpPostResponse, + StreamableHttpResponseLimits, + }, }, }; use rstest::rstest; @@ -27,6 +36,9 @@ use tokio::{ task::JoinHandle, }; +const URI: &str = "http://localhost/mcp"; +const JSON: &str = r#"{"jsonrpc":"2.0","id":1,"result":{}}"#; + struct Server { path: PathBuf, task: JoinHandle<()>, @@ -39,7 +51,12 @@ impl Drop for Server { } } -async fn oversized_response(status: u16, limit: usize, known_length: bool) -> Server { +async fn oversized_response( + status: u16, + content_type: &'static str, + limit: usize, + known_length: bool, +) -> Server { static NEXT: AtomicUsize = AtomicUsize::new(0); let path = std::env::temp_dir().join(format!( "rmcp-body-limit-{}-{}.sock", @@ -62,7 +79,7 @@ async fn oversized_response(status: u16, limit: usize, known_length: bool) -> Se "Transfer-Encoding: chunked\r\n".to_owned() }; let headers = - format!("HTTP/1.1 {status} Test\r\nContent-Type: application/json\r\n{framing}\r\n"); + format!("HTTP/1.1 {status} Test\r\nContent-Type: {content_type}\r\n{framing}\r\n"); socket.write_all(headers.as_bytes()).await.unwrap(); if !known_length { let chunk = vec![b' '; 4096]; @@ -82,18 +99,21 @@ async fn oversized_response(status: u16, limit: usize, known_length: bool) -> Se } #[rstest] -#[case::json_declared(200, 16 * 1024 * 1024, true)] -#[case::json_streamed(200, 16 * 1024 * 1024, false)] -#[case::error_declared(400, 64 * 1024, true)] -#[case::error_streamed(400, 64 * 1024, false)] -#[case::server_error_streamed(500, 64 * 1024, false)] +#[case::json_declared(200, "application/json", 64 * 1024 * 1024, true, true)] +#[case::json_streamed(200, "application/json", 64 * 1024 * 1024, false, true)] +#[case::json_error_declared(400, "application/json", 64 * 1024 * 1024, true, true)] +#[case::json_error_streamed(400, "application/json", 64 * 1024 * 1024, false, true)] +#[case::text_error_streamed(400, "text/plain", 64 * 1024, false, false)] +#[case::server_error_streamed(500, "text/plain", 64 * 1024, false, false)] #[tokio::test] -async fn unix_socket_rejects_oversized_response_before_eof( +async fn unix_socket_bounds_oversized_response_before_eof( #[case] status: u16, + #[case] content_type: &'static str, #[case] limit: usize, #[case] known_length: bool, + #[case] rejected_as_too_large: bool, ) { - let server = oversized_response(status, limit, known_length).await; + let server = oversized_response(status, content_type, limit, known_length).await; let uri: Arc = Arc::from("http://localhost/mcp"); let client = UnixSocketHttpClient::new(server.path.to_str().unwrap(), &uri); let message = ClientJsonRpcMessage::request( @@ -105,9 +125,481 @@ async fn unix_socket_rejects_oversized_response_before_eof( client.post_message(uri, message, None, None, Default::default()), ) .await - .expect("oversized response must be rejected before EOF"); - assert!( - matches!(result, Err(StreamableHttpError::ResponseBodyTooLarge { max_size }) if max_size == limit), - "expected response size error, got {result:?}" + .expect("oversized response must stop before EOF"); + if rejected_as_too_large { + assert!( + matches!(result, Err(StreamableHttpError::ResponseBodyTooLarge { max_size }) if max_size == limit), + "expected response size error, got {result:?}" + ); + } else { + assert!( + matches!( + result, + Err(StreamableHttpError::UnexpectedServerResponse(_)) + ), + "expected truncated diagnostic response, got {result:?}" + ); + } +} + +struct TestServer { + path: PathBuf, + task: JoinHandle<()>, +} + +impl TestServer { + fn new(responses: Vec>) -> Self { + Self::with_eof(responses, true) + } + + fn with_eof(responses: Vec>, finish_body: bool) -> Self { + static NEXT_SOCKET: AtomicUsize = AtomicUsize::new(0); + let path = std::env::temp_dir().join(format!( + "rmcp-body-limits-{}-{}.sock", + std::process::id(), + NEXT_SOCKET.fetch_add(1, Ordering::Relaxed), + )); + let listener = UnixListener::bind(&path).unwrap(); + let task = tokio::spawn(async move { + for response in responses { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + let mut buffer = [0; 1024]; + loop { + let count = stream.read(&mut buffer).await.unwrap(); + assert_ne!(count, 0, "client closed before sending its request"); + request.extend_from_slice(&buffer[..count]); + if let Some(end) = request.windows(4).position(|part| part == b"\r\n\r\n") { + let headers = String::from_utf8_lossy(&request[..end]); + let length = headers + .lines() + .filter_map(|line| line.split_once(':')) + .find(|(name, _)| name.eq_ignore_ascii_case("content-length")) + .map(|(_, value)| value.trim().parse::().unwrap()) + .unwrap_or(0); + if request.len() >= end + 4 + length { + break; + } + } + } + // An oversized response may be rejected before all bytes are sent. + let _ = stream.write_all(&response).await; + if !finish_body { + // Keep the body open until the client drops the response. + let _ = stream.read_u8().await; + } + } + }); + Self { path, task } + } + + fn client(&self) -> UnixSocketHttpClient { + UnixSocketHttpClient::new(self.path.to_str().unwrap(), URI) + } +} + +impl Drop for TestServer { + fn drop(&mut self) { + self.task.abort(); + let _ = std::fs::remove_file(&self.path); + } +} + +fn response(status: &str, content_type: &str, body: &[u8], chunked: bool) -> Vec { + let mut wire = + format!("HTTP/1.1 {status}\r\nContent-Type: {content_type}\r\nConnection: close\r\n") + .into_bytes(); + if chunked { + wire.extend_from_slice(b"Transfer-Encoding: chunked\r\n\r\n"); + for chunk in body.chunks(7) { + wire.extend_from_slice(format!("{:x}\r\n", chunk.len()).as_bytes()); + wire.extend_from_slice(chunk); + wire.extend_from_slice(b"\r\n"); + } + wire.extend_from_slice(b"0\r\n\r\n"); + } else { + wire.extend_from_slice(format!("Content-Length: {}\r\n\r\n", body.len()).as_bytes()); + wire.extend_from_slice(body); + } + wire +} + +fn message(method: &str) -> ClientJsonRpcMessage { + if method == "server/discover" { + return ClientJsonRpcMessage::request( + ClientRequest::DiscoverRequest(DiscoverRequest::new(DiscoverRequestParams {})), + RequestId::Number(1), + ); + } + serde_json::from_value(serde_json::json!({ + "jsonrpc": "2.0", "id": 1, "method": method + })) + .unwrap() +} + +fn limits(json: usize, error: usize) -> StreamableHttpResponseLimits { + let mut limits = StreamableHttpResponseLimits::default(); + limits.max_json_response_size = json; + limits.max_error_response_size = error; + limits +} + +async fn within_deadline(future: F) -> F::Output { + tokio::time::timeout(Duration::from_secs(2), future) + .await + .expect("Unix socket response operation exceeded its two-second deadline") +} + +#[tokio::test] +async fn json_thresholds_cover_content_length_chunked_and_recovery() { + let limit = JSON.len() + 8; + for chunked in [false, true] { + let lengths = [limit - 1, limit, limit + 1, JSON.len()]; + let server = TestServer::new( + lengths + .iter() + .map(|&length| { + let body = format!("{JSON}{}", " ".repeat(length - JSON.len())); + response("200 OK", "application/json", body.as_bytes(), chunked) + }) + .collect(), + ); + let client = server.client(); + for length in lengths { + let result = within_deadline(client.post_message_with_response_limits( + URI.into(), + message("ping"), + None, + None, + HashMap::new(), + limits(limit, 1), + )) + .await; + if length > limit { + assert!( + matches!(result, Err(StreamableHttpError::ResponseBodyTooLarge { max_size: actual }) if actual == limit) + ); + } else { + assert!( + matches!(result, Ok(StreamableHttpPostResponse::Json(..))), + "{result:?}" + ); + } + } + } +} + +#[tokio::test] +async fn diagnostic_prefixes_preserve_http_errors_and_discovery_fallback() { + let limit = 32; + for chunked in [false, true] { + let lengths = [limit - 1, limit, limit + 1, limit + 1, 1]; + let server = TestServer::new( + lengths + .iter() + .map(|&length| { + response( + "400 Bad Request", + "text/plain", + &vec![b'x'; length], + chunked, + ) + }) + .collect(), + ); + let client = server.client(); + for (index, length) in lengths.into_iter().enumerate() { + let result = within_deadline(client.post_message_with_response_limits( + URI.into(), + message(if index == 3 { + "server/discover" + } else { + "ping" + }), + None, + None, + HashMap::new(), + limits(1, limit), + )) + .await; + if index == 3 { + let Ok(StreamableHttpPostResponse::Json(JsonRpcMessage::Error(error), _)) = result + else { + panic!("expected legacy discovery fallback"); + }; + assert!(error.error.message.ends_with(&"x".repeat(limit))); + } else { + assert!( + matches!( + result, + Err(StreamableHttpError::UnexpectedServerResponse(ref body)) + if body == &format!("HTTP 400 Bad Request: {}", "x".repeat(length.min(limit))) + ), + "{result:?}" + ); + } + } + } +} + +#[tokio::test] +async fn oversized_invalid_json_is_not_accepted() { + let server = TestServer::new(vec![response( + "200 OK", + "application/json", + &[b'x'; 65], + true, + )]); + let result = within_deadline(server.client().post_message_with_response_limits( + URI.into(), + message("ping"), + None, + None, + HashMap::new(), + limits(64, 64), + )) + .await; + assert!(matches!( + result, + Err(StreamableHttpError::ResponseBodyTooLarge { max_size: 64 }) + )); +} + +#[tokio::test] +async fn legacy_post_methods_apply_default_json_limit_on_every_status() { + for (status, content_type, limit) in [ + ("200 OK", "application/json", 64 * 1024 * 1024), + ("400 Bad Request", "application/json", 64 * 1024 * 1024), + ] { + let header = format!( + "HTTP/1.1 {status}\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\n\r\n", + limit + 1, + ) + .into_bytes(); + let server = TestServer::new(vec![header.clone(), header]); + let client = server.client(); + let first = within_deadline(client.post_message( + URI.into(), + message("ping"), + None, + None, + HashMap::new(), + )) + .await; + assert!( + matches!(first, Err(StreamableHttpError::ResponseBodyTooLarge { max_size: actual }) if actual == limit) + ); + let second = within_deadline(client.post_message_with_max_sse_event_size( + URI.into(), + message("ping"), + None, + None, + HashMap::new(), + 128, + )) + .await; + assert!( + matches!(second, Err(StreamableHttpError::ResponseBodyTooLarge { max_size: actual }) if actual == limit) + ); + } +} + +#[tokio::test] +async fn legacy_post_methods_truncate_diagnostics_and_keep_auth_session_precedence() { + let limit = 64 * 1024; + let body = vec![b'x'; limit + 1]; + let server = TestServer::new(vec![ + response("500 Internal Server Error", "text/plain", &body, false), + response("500 Internal Server Error", "text/plain", &body, true), + ]); + let client = server.client(); + let results = [ + within_deadline(client.post_message( + URI.into(), + message("ping"), + None, + None, + HashMap::new(), + )) + .await, + within_deadline(client.post_message_with_max_sse_event_size( + URI.into(), + message("ping"), + None, + None, + HashMap::new(), + 0, + )) + .await, + ]; + for result in results { + assert!( + matches!(result, Err(StreamableHttpError::UnexpectedServerResponse(body)) + if body == format!("HTTP 500 Internal Server Error: {}", "x".repeat(limit))) + ); + } + + for status in [ + "401 Unauthorized", + "403 Forbidden", + "404 Not Found", + "500 Internal Server Error", + ] { + let server = TestServer::new(vec![response(status, "text/html", &body, true)]); + let result = within_deadline(server.client().post_message_with_response_limits( + URI.into(), + message("server/discover"), + if status.starts_with("404") { + Some("expired".into()) + } else { + None + }, + None, + HashMap::new(), + limits(0, 0), + )) + .await; + if status.starts_with("404") { + assert!(matches!(result, Err(StreamableHttpError::SessionExpired))); + } else { + assert!(matches!( + result, + Err(StreamableHttpError::UnexpectedServerResponse(_)) + )); + } + } + for status in ["401 Unauthorized", "403 Forbidden"] { + let wire = format!("HTTP/1.1 {status}\r\nWWW-Authenticate: Bearer scope=\"read\"\r\nContent-Length: 999999\r\n\r\n").into_bytes(); + let server = TestServer::with_eof(vec![wire], false); + let result = within_deadline(server.client().post_message_with_response_limits( + URI.into(), + message("server/discover"), + None, + None, + HashMap::new(), + limits(0, 0), + )) + .await; + match status { + "401 Unauthorized" => { + assert!(matches!(result, Err(StreamableHttpError::AuthRequired(_)))) + } + _ => assert!(matches!( + result, + Err(StreamableHttpError::InsufficientScope(_)) + )), + } + } +} + +#[tokio::test] +async fn json_rpc_errors_use_json_limit_on_success_and_failure() { + let body = br#"{"jsonrpc":"2.0","id":1,"error":{"code":-32600,"message":"real error"}}"#; + for status in ["200 OK", "400 Bad Request"] { + for chunked in [false, true] { + for limit in [body.len() - 1, body.len(), body.len() + 1] { + let server = + TestServer::new(vec![response(status, "application/json", body, chunked)]); + let result = within_deadline(server.client().post_message_with_response_limits( + URI.into(), + message("server/discover"), + None, + None, + HashMap::new(), + limits(limit, 0), + )) + .await; + if limit < body.len() { + assert!(matches!( + result, + Err(StreamableHttpError::ResponseBodyTooLarge { .. }) + )); + } else { + let Ok(StreamableHttpPostResponse::Json(JsonRpcMessage::Error(error), _)) = + result + else { + panic!( + "expected the complete JSON-RPC error, not a diagnostic or synthetic fallback" + ); + }; + assert_eq!(error.error.message, "real error"); + } + } + } + } +} + +#[tokio::test] +async fn diagnostic_prefix_does_not_wait_for_eof_or_parse_partial_json() { + let prefix = br#"{"jsonrpc":"2.0","id":1,"error":{"code":-32600,"message":"prefix"}}"#; + let mut body = prefix.to_vec(); + body.extend_from_slice(b"not-json"); + for content_type in ["text/plain", "application/json"] { + let server = TestServer::new(vec![response("400 Bad Request", content_type, &body, true)]); + let result = within_deadline(server.client().post_message_with_response_limits( + URI.into(), + message("ping"), + None, + None, + HashMap::new(), + limits(body.len(), prefix.len()), + )) + .await; + assert!(matches!( + result, + Err(StreamableHttpError::UnexpectedServerResponse(_)) + )); + } + let limit = 64 * 1024; + let mut wire = b"HTTP/1.1 422 Unprocessable Entity\r\nContent-Type: text/html\r\nTransfer-Encoding: chunked\r\n\r\n".to_vec(); + wire.extend_from_slice(format!("{:x}\r\n{}\r\n", limit, "x".repeat(limit)).as_bytes()); + let server = TestServer::with_eof(vec![wire], false); + let result = within_deadline(server.client().post_message( + URI.into(), + message("server/discover"), + None, + None, + HashMap::new(), + )) + .await; + let Ok(StreamableHttpPostResponse::Json(JsonRpcMessage::Error(error), _)) = result else { + panic!("expected legacy fallback without waiting for another chunk or EOF"); + }; + assert!(error.error.message.ends_with(&"x".repeat(limit))); +} + +#[tokio::test] +async fn response_limits_keep_sse_event_limit_independent() { + let server = TestServer::new(vec![response( + "200 OK", + "text/event-stream", + b"data: a response larger than the event limit\n\n", + true, + )]); + let mut limits = limits(1, 1); + limits.max_sse_event_size = 16; + let result = within_deadline(server.client().post_message_with_response_limits( + URI.into(), + message("ping"), + None, + None, + HashMap::new(), + limits, + )) + .await + .unwrap(); + let StreamableHttpPostResponse::Sse(mut stream, _) = result else { + panic!("expected SSE response"); + }; + let error = within_deadline(stream.next()) + .await + .expect("oversized SSE event must produce an error") + .unwrap_err(); + let sse_stream::Error::Body(error) = error else { + panic!("expected an SSE body-size error, got {error}"); + }; + assert_eq!( + error.to_string(), + "SSE event exceeded the maximum size of 16 bytes", ); }