From 296d1faea6ec1fd3b538eb629e93bd9910725778 Mon Sep 17 00:00:00 2001 From: GTanger <49036591+GTanger@users.noreply.github.com> Date: Sat, 5 Sep 2026 18:34:36 +0800 Subject: [PATCH 1/2] fix(transport): bound buffered Streamable HTTP responses --- crates/rmcp/src/transport/common.rs | 6 + .../common/auth/streamable_http_client.rs | 204 +++++++- .../src/transport/common/client_side_body.rs | 123 +++++ .../common/reqwest/streamable_http_client.rs | 64 ++- .../rmcp/src/transport/common/unix_socket.rs | 81 ++- .../src/transport/streamable_http_client.rs | 126 ++++- ...est_streamable_http_default_error_limit.rs | 78 +++ .../test_streamable_http_limit_lifecycle.rs | 476 ++++++++++++++++++ .../test_streamable_http_response_limits.rs | 352 +++++++++++++ .../tests/test_unix_socket_response_limits.rs | 313 ++++++++++++ 10 files changed, 1776 insertions(+), 47 deletions(-) create mode 100644 crates/rmcp/src/transport/common/client_side_body.rs 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 create mode 100644 crates/rmcp/tests/test_streamable_http_response_limits.rs create mode 100644 crates/rmcp/tests/test_unix_socket_response_limits.rs diff --git a/crates/rmcp/src/transport/common.rs b/crates/rmcp/src/transport/common.rs index 8cf00c02f..0386c4941 100644 --- a/crates/rmcp/src/transport/common.rs +++ b/crates/rmcp/src/transport/common.rs @@ -15,6 +15,12 @@ mod reqwest; #[cfg(feature = "client-side-sse")] pub mod client_side_sse; +#[cfg(any( + feature = "transport-streamable-http-client-reqwest", + all(unix, feature = "transport-streamable-http-client-unix-socket") +))] +pub(crate) mod client_side_body; + #[cfg(feature = "auth")] pub mod auth; 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 97432f90e..c976f4993 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 @@ -208,19 +210,217 @@ 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, Mutex}; + use super::*; use crate::transport::{ auth::{ AuthorizationManager, AuthorizationMetadata, CredentialRefreshGuard, CredentialStore, + InMemoryCredentialStore, OAuthHttpClient, OAuthHttpClientFuture, OAuthHttpRequest, StoredCredentials, }, - 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 { + limit: 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 { limit: 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_body.rs b/crates/rmcp/src/transport/common/client_side_body.rs new file mode 100644 index 000000000..10ed8eba7 --- /dev/null +++ b/crates/rmcp/src/transport/common/client_side_body.rs @@ -0,0 +1,123 @@ +//! Bounded buffering for the built-in Streamable HTTP clients. + +use bytes::Bytes; +use futures::{Stream, StreamExt}; + +use crate::transport::streamable_http_client::StreamableHttpError; + +/// Buffer a response without appending a chunk that would exceed `limit`. +/// +/// A known content length permits early rejection, but is never trusted as the +/// sole bound. The count applies to bytes yielded by the HTTP backend (after +/// decompression, when enabled). The backend may already have allocated a chunk; +/// this helper bounds the accumulated body, not every backend allocation. +pub(crate) async fn read_bounded_body( + stream: S, + content_length: Option, + limit: usize, +) -> Result, StreamableHttpError> +where + S: Stream>, + E: std::error::Error + Send + Sync + 'static, +{ + if content_length.is_some_and(|length| length > limit as u64) { + return Err(StreamableHttpError::ResponseBodyTooLarge { limit }); + } + let mut body = Vec::new(); + futures::pin_mut!(stream); + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(StreamableHttpError::Client)?; + if chunk.len() > limit - body.len() { + return Err(StreamableHttpError::ResponseBodyTooLarge { limit }); + } + body.extend_from_slice(&chunk); + } + Ok(body) +} + +#[cfg(test)] +mod tests { + use std::{ + io, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + }; + + use super::*; + + #[tokio::test] + async fn exact_limit_accepts_multiple_chunks() { + let stream = futures::stream::iter([ + Ok::<_, io::Error>(Bytes::from_static(b"ab")), + Ok(Bytes::from_static(b"cd")), + ]); + assert_eq!(read_bounded_body(stream, None, 4).await.unwrap(), b"abcd"); + } + + #[tokio::test] + async fn declared_oversize_is_rejected_without_polling() { + let stream = + futures::stream::poll_fn(|_| -> std::task::Poll>> { + panic!("a declared oversized body must not be polled") + }); + assert!(matches!( + read_bounded_body(stream, Some(5), 4).await, + Err(StreamableHttpError::ResponseBodyTooLarge { limit: 4 }) + )); + } + + #[tokio::test] + async fn unknown_length_stops_at_first_oversized_chunk() { + let polls = Arc::new(AtomicUsize::new(0)); + let count = polls.clone(); + let stream = futures::stream::poll_fn(move |_| { + let index = count.fetch_add(1, Ordering::SeqCst); + assert!(index < 2, "must not drain an oversized or endless body"); + std::task::Poll::Ready(Some(Ok::<_, io::Error>(Bytes::from_static(b"abc")))) + }); + assert!(matches!( + read_bounded_body(stream, None, 4).await, + Err(StreamableHttpError::ResponseBodyTooLarge { limit: 4 }) + )); + assert_eq!(polls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn understated_length_does_not_bypass_chunk_counting() { + let stream = futures::stream::iter([Ok::<_, io::Error>(Bytes::from_static(b"abcde"))]); + assert!(matches!( + read_bounded_body(stream, Some(1), 4).await, + Err(StreamableHttpError::ResponseBodyTooLarge { limit: 4 }) + )); + } + + #[tokio::test] + async fn zero_limit_accepts_only_empty_bodies() { + assert!( + read_bounded_body( + futures::stream::empty::>(), + None, + 0 + ) + .await + .unwrap() + .is_empty() + ); + let stream = futures::stream::iter([Ok::<_, io::Error>(Bytes::from_static(b"x"))]); + assert!(matches!( + read_bounded_body(stream, None, 0).await, + Err(StreamableHttpError::ResponseBodyTooLarge { limit: 0 }) + )); + } + + #[tokio::test] + async fn stream_failure_remains_a_client_error() { + let stream = futures::stream::iter([Err::(io::Error::other("read failed"))]); + assert!(matches!( + read_bounded_body(stream, None, 8).await, + Err(StreamableHttpError::Client(_)) + )); + } +} 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 b0b5090b5..9125d228b 100644 --- a/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs +++ b/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs @@ -9,6 +9,7 @@ use crate::{ model::{ClientJsonRpcMessage, JsonRpcMessage, ServerJsonRpcMessage}, transport::{ common::{ + client_side_body::read_bounded_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, @@ -176,13 +177,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 } @@ -195,6 +196,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()) @@ -276,10 +300,17 @@ impl StreamableHttpClient for reqwest::Client { // Non-success responses may carry valid JSON-RPC error payloads that // should be surfaced as McpError rather than lost in TransportSend. if !status.is_success() { - let body = response - .text() - .await - .unwrap_or_else(|_| "".to_owned()); + let body = match read_bounded_body( + response.bytes_stream(), + content_length, + limits.max_error_response_size, + ) + .await + { + Ok(body) => String::from_utf8_lossy(&body).into_owned(), + Err(StreamableHttpError::Client(_)) => "".to_owned(), + Err(error) => return Err(error), + }; if content_type .as_deref() .is_some_and(|ct| ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes())) @@ -304,14 +335,31 @@ 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()) => { // Try to parse as a valid JSON-RPC message. If the body is // malformed (e.g. a 200 response to a notification that lacks // an `id` field), treat it as accepted rather than failing. - match response.json::().await { + let body = match read_bounded_body( + response.bytes_stream(), + content_length, + limits.max_json_response_size, + ) + .await + { + Ok(body) => body, + Err(StreamableHttpError::Client(error)) => { + tracing::warn!( + "could not read JSON response, treating as accepted: {error}" + ); + return Ok(StreamableHttpPostResponse::Accepted); + } + Err(error) => return Err(error), + }; + match serde_json::from_slice::(&body) { Ok(message) => Ok(StreamableHttpPostResponse::Json(message, session_id)), Err(e) => { tracing::warn!( diff --git a/crates/rmcp/src/transport/common/unix_socket.rs b/crates/rmcp/src/transport/common/unix_socket.rs index e3b34898f..711458f47 100644 --- a/crates/rmcp/src/transport/common/unix_socket.rs +++ b/crates/rmcp/src/transport/common/unix_socket.rs @@ -1,7 +1,7 @@ use std::{borrow::Cow, collections::HashMap, sync::Arc}; use bytes::Bytes; -use futures::stream::BoxStream; +use futures::{TryStreamExt, stream::BoxStream}; use http::{HeaderName, HeaderValue, Method, Request, StatusCode, header::WWW_AUTHENTICATE}; use http_body_util::{BodyExt, Full}; use hyper::body::Incoming; @@ -13,6 +13,7 @@ use crate::{ model::{ClientJsonRpcMessage, ServerJsonRpcMessage}, transport::{ common::{ + client_side_body::read_bounded_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, @@ -172,13 +173,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 } @@ -191,6 +192,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)))?; @@ -225,6 +249,11 @@ impl StreamableHttpClient for UnixSocketHttpClient { .map_err(StreamableHttpError::Client)?; let status = response.status(); + let content_length = response + .headers() + .get(http::header::CONTENT_LENGTH) + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()); if status == StatusCode::UNAUTHORIZED && let Some(header) = response.headers().get(WWW_AUTHENTICATE) @@ -268,12 +297,20 @@ impl StreamableHttpClient for UnixSocketHttpClient { } if !status.is_success() { - let body = response - .into_body() - .collect() - .await - .map(|c| String::from_utf8_lossy(&c.to_bytes()).into_owned()) - .unwrap_or_else(|_| "".to_owned()); + let body = match read_bounded_body( + response + .into_body() + .into_data_stream() + .map_err(UnixSocketError::Hyper), + content_length, + limits.max_error_response_size, + ) + .await + { + Ok(body) => String::from_utf8_lossy(&body).into_owned(), + Err(StreamableHttpError::Client(_)) => "".to_owned(), + Err(error) => return Err(error), + }; if let Some(response) = legacy_discover_response(&message, session_was_attached, status, &body) { @@ -285,11 +322,6 @@ impl StreamableHttpClient for UnixSocketHttpClient { } let content_type = response.headers().get(http::header::CONTENT_TYPE).cloned(); - let content_length = response - .headers() - .get(http::header::CONTENT_LENGTH) - .and_then(|v| v.to_str().ok()) - .and_then(|v| v.parse::().ok()); let session_id = response .headers() .get(HEADER_SESSION_ID) @@ -310,17 +342,22 @@ 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()) => { - let body = response - .into_body() - .collect() - .await - .map_err(|e| StreamableHttpError::Client(UnixSocketError::Hyper(e)))? - .to_bytes(); + let body = read_bounded_body( + response + .into_body() + .into_data_stream() + .map_err(UnixSocketError::Hyper), + content_length, + limits.max_json_response_size, + ) + .await?; match serde_json::from_slice::(&body) { Ok(message) => Ok(StreamableHttpPostResponse::Json(message, session_id)), Err(e) => { diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index d702bc1ca..d82393468 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -220,6 +220,10 @@ pub enum StreamableHttpError { /// A cancellation or reply POST did not finish in time. #[error("Control POST timed out")] ControlRequestTimeout, + /// A buffered JSON or error response exceeded its configured byte limit. + /// The response body is not included and this error does not trigger retries. + #[error("HTTP response body exceeded {limit} bytes before decoding")] + ResponseBodyTooLarge { limit: usize }, } impl StreamableHttpError { @@ -369,12 +373,40 @@ 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 or JSON response and 64 KiB per HTTP +/// error response. Zero accepts only empty bodies. For compressed responses, +/// body limits count the decompressed bytes yielded by the HTTP backend. +#[derive(Debug, Clone, Copy)] +#[non_exhaustive] +pub struct StreamableHttpResponseLimits { + /// Maximum raw size of an individual SSE event. + pub max_sse_event_size: usize, + /// Maximum buffered JSON success response body. + pub max_json_response_size: usize, + /// Maximum buffered non-success HTTP response body. + 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: 16 * 1024 * 1024, + max_error_response_size: 64 * 1024, + } + } +} + /// HTTP backend used by [`StreamableHttpClientTransport`]. /// /// 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 @@ -414,6 +446,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, @@ -596,7 +655,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 @@ -613,13 +672,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 { @@ -959,7 +1018,7 @@ impl StreamableHttpClientWorker { uri: Arc, auth_header: Option, custom_headers: HashMap, - max_sse_event_size: usize, + response_limits: StreamableHttpResponseLimits, ) -> Result< ( Option>, @@ -969,13 +1028,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::() @@ -1000,13 +1059,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::()?; @@ -1077,13 +1136,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 { @@ -1144,13 +1203,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( @@ -1230,7 +1289,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)), }; @@ -1473,13 +1532,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 { @@ -2025,6 +2084,20 @@ 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 buffered JSON success response body (default: 16 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 buffered HTTP error response body (default: 64 KiB). + /// + /// Enforced before JSON-RPC error parsing or legacy discovery fallback. + /// 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`). /// @@ -2109,6 +2182,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 response body size before decoding. + 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 @@ -2136,6 +2229,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..809ce9878 --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_default_error_limit.rs @@ -0,0 +1,78 @@ +#![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::INTERNAL_SERVER_ERROR, 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_rejects_oversized_error_without_exposing_body() { + const MARKER: &str = "SENSITIVE_TEST_MARKER_DO_NOT_ECHO"; + let mut body = MARKER.to_owned(); + body.push_str(&"x".repeat(65_537 - MARKER.len())); + assert_eq!(body.len(), 65_537); + 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.contains("exceeded"), "expected a size-limit error"); + assert!( + !error.contains(MARKER), + "error must not expose response body" + ); + assert!(error.len() < 512, "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..cd65b8158 --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_limit_lifecycle.rs @@ -0,0 +1,476 @@ +#![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, ClientInfo, 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, + ClientInfo::default().serve_with_lifecycle(transport, auto_lifecycle()), + ) + .await?? + } else { + tokio::time::timeout(TIMEOUT, ClientInfo::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(error_response: bool, 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()); + let (status, content_type, payload) = if error_response { + ( + 422, + "text/plain", + "Unexpected message, expect initialize request".repeat(8), + ) + } else { + ( + 200, + "application/json", + 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(false, false)] +#[case::error(true, false)] +#[case::unfinished_json(false, true)] +#[case::unfinished_error(true, true)] +#[tokio::test] +async fn oversized_discover_fails_without_fallback_or_retry( + #[case] error_response: bool, + #[case] never_finishes: bool, +) { + let server = LoopbackServer::oversized_discover(error_response, 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, + ClientInfo::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 { limit: 64 }) + ), + "unexpected transport error: {error:?}" + ); + assert_eq!(server.posts.lock().unwrap().as_slice(), ["server/discover"]); +} diff --git a/crates/rmcp/tests/test_streamable_http_response_limits.rs b/crates/rmcp/tests/test_streamable_http_response_limits.rs new file mode 100644 index 000000000..1177c222a --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_response_limits.rs @@ -0,0 +1,352 @@ +#![cfg(all( + feature = "client", + feature = "transport-streamable-http-client-reqwest", + not(feature = "local") +))] + +use std::{collections::HashMap, convert::Infallible, sync::Arc, time::Duration}; + +use axum::{Router, body::Body, http::Response, routing::post as post_route}; +use bytes::Bytes; +use futures::StreamExt; +use rmcp::{ + model::{ + ClientJsonRpcMessage, ClientRequest, DiscoverRequest, DiscoverRequestParams, + JsonRpcMessage, PingRequest, RequestId, + }, + transport::streamable_http_client::{ + StreamableHttpClient, StreamableHttpError, StreamableHttpPostResponse, + StreamableHttpResponseLimits, + }, +}; +use rstest::rstest; + +const JSON_RESPONSE: &str = r#"{"jsonrpc":"2.0","id":1,"result":{}}"#; +const JSON_ERROR: &str = + r#"{"jsonrpc":"2.0","id":1,"error":{"code":-32600,"message":"Invalid Request"}}"#; + +struct MockServer { + uri: Arc, + task: tokio::task::JoinHandle<()>, +} + +impl Drop for MockServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +impl MockServer { + async fn start( + status: u16, + content_type: &'static str, + body: impl Into, + chunked: bool, + ) -> Self { + Self::with_headers(status, content_type, body, chunked, Vec::new()).await + } + + async fn with_headers( + status: u16, + content_type: &'static str, + body: impl Into, + chunked: bool, + headers: Vec<(&'static str, &'static str)>, + ) -> Self { + let body = body.into(); + let router = Router::new().route( + "/mcp", + post_route(move || { + let body = body.clone(); + let headers = headers.clone(); + async move { + let mut response = Response::builder() + .status(status) + .header("content-type", content_type) + .header("mcp-session-id", "test-session"); + for (name, value) in headers { + response = response.header(name, value); + } + let body = if chunked { + // Unknown size forces chunked HTTP. Delay each chunk so the + // client must account for bytes across successive reads. + let chunks = body + .chunks((body.len() / 3).max(1)) + .map(Bytes::copy_from_slice) + .collect::>(); + Body::from_stream(futures::stream::iter(chunks).then(|chunk| async move { + tokio::time::sleep(Duration::from_millis(1)).await; + Ok::<_, Infallible>(chunk) + })) + } else { + response = response.header("content-length", body.len()); + Body::from(body) + }; + response.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: Arc::from(format!("http://{address}/mcp")), + task, + } + } +} + +fn client() -> reqwest::Client { + reqwest::Client::builder() + .no_proxy() + .timeout(Duration::from_secs(5)) + .build() + .unwrap() +} + +fn ping() -> ClientJsonRpcMessage { + ClientJsonRpcMessage::request( + ClientRequest::PingRequest(PingRequest::default()), + RequestId::Number(1), + ) +} + +fn discover() -> ClientJsonRpcMessage { + ClientJsonRpcMessage::request( + ClientRequest::DiscoverRequest(DiscoverRequest::new(DiscoverRequestParams {})), + RequestId::Number(1), + ) +} + +fn 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: &MockServer, + message: ClientJsonRpcMessage, + limits: StreamableHttpResponseLimits, +) -> Result> { + client() + .post_message_with_response_limits( + server.uri.clone(), + message, + None, + None, + HashMap::new(), + limits, + ) + .await +} + +#[rstest] +#[case::below_limit(1, false)] +#[case::exact_limit(0, false)] +#[case::over_limit(-1, false)] +#[case::chunked_below_limit(1, true)] +#[case::chunked_exact_limit(0, true)] +#[case::chunked_over_limit(-1, true)] +#[tokio::test] +async fn json_response_limit_is_inclusive(#[case] margin: isize, #[case] chunked: bool) { + let server = MockServer::start(200, "application/json", JSON_RESPONSE, chunked).await; + let limit = JSON_RESPONSE.len().checked_add_signed(margin).unwrap(); + let result = post(&server, ping(), limits(limit, 0, 0)).await; + if margin < 0 { + assert!( + matches!(result, Err(StreamableHttpError::ResponseBodyTooLarge { limit: got }) if got == limit) + ); + } else { + let StreamableHttpPostResponse::Json(JsonRpcMessage::Response(_), session) = + result.unwrap() + else { + panic!("expected JSON-RPC response"); + }; + assert_eq!(session.as_deref(), Some("test-session")); + } +} + +#[rstest] +#[case::fixed(false)] +#[case::chunked(true)] +#[tokio::test] +async fn oversized_malformed_json_is_not_accepted(#[case] chunked: bool) { + let server = MockServer::start(200, "application/json", "not valid json", chunked).await; + assert!(matches!( + post(&server, ping(), limits(4, 100, 100)).await, + Err(StreamableHttpError::ResponseBodyTooLarge { limit: 4 }) + )); + // Preserve the existing malformed-response fallback within the bound. + assert!(matches!( + post(&server, ping(), limits(100, 100, 100)).await, + Ok(StreamableHttpPostResponse::Accepted) + )); +} + +#[rstest] +#[case::below_limit(1, false)] +#[case::exact_limit(0, false)] +#[case::over_limit(-1, false)] +#[case::chunked_below_limit(1, true)] +#[case::chunked_exact_limit(0, true)] +#[case::chunked_over_limit(-1, true)] +#[tokio::test] +async fn json_rpc_error_uses_error_body_limit(#[case] margin: isize, #[case] chunked: bool) { + let server = MockServer::start(400, "application/json", JSON_ERROR, chunked).await; + let limit = JSON_ERROR.len().checked_add_signed(margin).unwrap(); + let result = post(&server, ping(), limits(0, limit, 0)).await; + if margin < 0 { + assert!( + matches!(result, Err(StreamableHttpError::ResponseBodyTooLarge { limit: got }) if got == limit) + ); + } else { + let StreamableHttpPostResponse::Json(JsonRpcMessage::Error(error), session) = + result.unwrap() + else { + panic!("expected JSON-RPC error"); + }; + assert_eq!(error.error.message, "Invalid Request"); + assert_eq!(session.as_deref(), Some("test-session")); + } +} + +#[rstest] +#[case::fixed(false)] +#[case::chunked(true)] +#[tokio::test] +async fn oversized_discovery_rejection_cannot_trigger_legacy_fallback(#[case] chunked: bool) { + let body = "Unexpected message, expect initialize request"; + let server = MockServer::start(422, "text/plain", body, chunked).await; + assert!(matches!( + post(&server, discover(), limits(100, body.len() - 1, 100)).await, + Err(StreamableHttpError::ResponseBodyTooLarge { .. }) + )); + assert!(matches!( + post(&server, discover(), limits(100, body.len(), 100)).await, + Ok(StreamableHttpPostResponse::Json( + JsonRpcMessage::Error(_), + _ + )) + )); +} + +#[rstest] +#[case::fixed(false)] +#[case::chunked(true)] +#[tokio::test] +async fn non_json_error_bodies_are_bounded(#[case] chunked: bool) { + let server = MockServer::start(500, "text/plain", "failure", chunked).await; + assert!(matches!( + post(&server, ping(), limits(100, 6, 100)).await, + Err(StreamableHttpError::ResponseBodyTooLarge { limit: 6 }) + )); + assert!(matches!( + post(&server, ping(), limits(0, 7, 0)).await, + Err(StreamableHttpError::UnexpectedServerResponse(message)) + if message.contains("failure") + )); +} + +#[rstest] +#[case::accepted(202)] +#[case::no_content(204)] +#[tokio::test] +async fn accepted_statuses_do_not_require_a_response_body(#[case] status: u16) { + let server = MockServer::start(status, "application/json", "", false).await; + assert!(matches!( + post(&server, ping(), limits(0, 0, 0)).await, + Ok(StreamableHttpPostResponse::Accepted) + )); +} + +#[tokio::test] +async fn sse_responses_keep_their_independent_event_limit() { + let server = MockServer::start(200, "text/event-stream", "data: example\n\n", true).await; + let StreamableHttpPostResponse::Sse(mut stream, _) = + post(&server, ping(), limits(0, 0, 64)).await.unwrap() + else { + panic!("expected SSE stream"); + }; + assert!(stream.next().await.unwrap().is_ok()); + + let StreamableHttpPostResponse::Sse(mut stream, _) = + post(&server, ping(), limits(100, 100, 4)).await.unwrap() + else { + panic!("expected bounded SSE stream"); + }; + assert!(stream.next().await.unwrap().is_err()); +} + +#[rstest] +#[case::unauthorized(401)] +#[case::forbidden(403)] +#[tokio::test] +async fn authentication_challenges_keep_precedence(#[case] status: u16) { + let server = MockServer::with_headers( + status, + "text/plain", + "authentication required", + false, + vec![("www-authenticate", "Bearer scope=\"read\"")], + ) + .await; + let result = post(&server, ping(), limits(0, 0, 0)).await; + match status { + 401 => assert!(matches!(result, Err(StreamableHttpError::AuthRequired(_)))), + 403 => assert!(matches!( + result, + Err(StreamableHttpError::InsufficientScope(_)) + )), + _ => unreachable!(), + } +} + +#[tokio::test] +async fn expired_sessions_keep_precedence() { + let server = MockServer::start(404, "text/plain", "session missing", false).await; + let result = client() + .post_message_with_response_limits( + server.uri.clone(), + ping(), + Some(Arc::from("expired-session")), + None, + HashMap::new(), + limits(0, 0, 0), + ) + .await; + assert!(matches!(result, Err(StreamableHttpError::SessionExpired))); +} + +#[rstest] +#[case::post_message(false)] +#[case::post_message_with_max_sse_event_size(true)] +#[tokio::test] +async fn existing_entry_points_apply_default_body_limits(#[case] sse_limit_method: bool) { + let limit = StreamableHttpResponseLimits::default().max_error_response_size; + let server = MockServer::start(500, "text/plain", vec![b'x'; limit + 1], true).await; + let client = client(); + let result = if sse_limit_method { + client + .post_message_with_max_sse_event_size( + server.uri.clone(), + ping(), + None, + None, + HashMap::new(), + usize::MAX, + ) + .await + } else { + client + .post_message(server.uri.clone(), ping(), None, None, HashMap::new()) + .await + }; + assert!( + matches!(result, Err(StreamableHttpError::ResponseBodyTooLarge { limit: got }) if got == limit) + ); +} diff --git a/crates/rmcp/tests/test_unix_socket_response_limits.rs b/crates/rmcp/tests/test_unix_socket_response_limits.rs new file mode 100644 index 000000000..cb08e1da3 --- /dev/null +++ b/crates/rmcp/tests/test_unix_socket_response_limits.rs @@ -0,0 +1,313 @@ +#![cfg(all( + unix, + feature = "transport-streamable-http-client-unix-socket", + not(feature = "local") +))] + +use std::{ + collections::HashMap, + future::Future, + path::PathBuf, + sync::atomic::{AtomicUsize, Ordering}, + time::Duration, +}; + +use futures::StreamExt; +use rmcp::{ + model::ClientJsonRpcMessage, + transport::{ + UnixSocketHttpClient, + streamable_http_client::{ + StreamableHttpClient, StreamableHttpError, StreamableHttpPostResponse, + StreamableHttpResponseLimits, + }, + }, +}; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::UnixListener, +}; + +const URI: &str = "http://localhost/mcp"; +const JSON: &str = r#"{"jsonrpc":"2.0","id":1,"result":{}}"#; + +struct TestServer { + path: PathBuf, + task: tokio::task::JoinHandle<()>, +} + +impl TestServer { + fn new(responses: Vec>) -> 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; + } + }); + 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 { + 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 { limit: actual }) if actual == limit) + ); + } else { + assert!( + matches!(result, Ok(StreamableHttpPostResponse::Json(..))), + "{result:?}" + ); + } + } + } +} + +#[tokio::test] +async fn error_thresholds_cannot_be_hidden_by_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 length > limit { + assert!( + matches!(result, Err(StreamableHttpError::ResponseBodyTooLarge { limit: actual }) if actual == limit) + ); + } else { + assert!( + matches!( + result, + Err(StreamableHttpError::UnexpectedServerResponse(_)) + ), + "{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 { limit: 64 }) + )); +} + +#[tokio::test] +async fn legacy_post_methods_apply_default_json_and_error_limits() { + for (status, content_type, limit) in [ + ("200 OK", "application/json", 16 * 1024 * 1024), + ("500 Internal Server Error", "text/plain", 64 * 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 { limit: 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 { limit: actual }) if actual == 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", + ); +} From 4500e81a9a79a37aeea8c153729c0368a7282474 Mon Sep 17 00:00:00 2001 From: GTanger <49036591+GTanger@users.noreply.github.com> Date: Sun, 6 Sep 2026 19:49:00 +0800 Subject: [PATCH 2/2] fix(transport): preserve legacy fallback with bounded diagnostics Apply the complete JSON limit on every HTTP status, raise its default to 64 MiB, and truncate diagnostic prefixes without waiting for EOF. Align Unix socket JSON-RPC error handling with reqwest and add regression coverage for review feedback. --- .../src/transport/common/client_side_body.rs | 91 ++++++++ .../common/reqwest/streamable_http_client.rs | 41 ++-- .../rmcp/src/transport/common/unix_socket.rs | 60 ++++-- .../src/transport/streamable_http_client.rs | 32 ++- ...est_streamable_http_default_error_limit.rs | 22 +- .../test_streamable_http_limit_lifecycle.rs | 84 ++++++-- .../test_streamable_http_response_limits.rs | 115 ++++++++-- .../tests/test_unix_socket_response_limits.rs | 204 +++++++++++++++++- 8 files changed, 544 insertions(+), 105 deletions(-) diff --git a/crates/rmcp/src/transport/common/client_side_body.rs b/crates/rmcp/src/transport/common/client_side_body.rs index 10ed8eba7..75a1b8fde 100644 --- a/crates/rmcp/src/transport/common/client_side_body.rs +++ b/crates/rmcp/src/transport/common/client_side_body.rs @@ -35,6 +35,33 @@ where Ok(body) } +/// Read at most `limit` 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 needed: an oversized declaration must not discard a useful prefix. +/// The returned bytes must never be parsed as a complete JSON-RPC message. +pub(crate) async fn read_truncated_body( + stream: S, + limit: 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() < limit { + let Some(chunk) = stream.next().await else { + break; + }; + let chunk = chunk.map_err(StreamableHttpError::Client)?; + let count = chunk.len().min(limit - body.len()); + body.extend_from_slice(&chunk[..count]); + } + Ok(body) +} + #[cfg(test)] mod tests { use std::{ @@ -120,4 +147,68 @@ mod tests { Err(StreamableHttpError::Client(_)) )); } + + #[tokio::test] + async fn truncated_body_keeps_short_bodies_complete() { + let stream = futures::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::stream::iter([ + Ok::<_, io::Error>(Bytes::from_static(b"ab")), + Ok(Bytes::from_static(b"cd")), + ]) + .chain(futures::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 = futures::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 = + futures::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 = futures::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 = + futures::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 9125d228b..2fbd1a2d0 100644 --- a/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs +++ b/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs @@ -9,7 +9,7 @@ use crate::{ model::{ClientJsonRpcMessage, JsonRpcMessage, ServerJsonRpcMessage}, transport::{ common::{ - client_side_body::read_bounded_body, + client_side_body::{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, @@ -300,22 +300,29 @@ impl StreamableHttpClient for reqwest::Client { // Non-success responses may carry valid JSON-RPC error payloads that // should be surfaced as McpError rather than lost in TransportSend. if !status.is_success() { - let body = match read_bounded_body( - response.bytes_stream(), - content_length, - limits.max_error_response_size, - ) - .await - { - Ok(body) => String::from_utf8_lossy(&body).into_owned(), - Err(StreamableHttpError::Client(_)) => "".to_owned(), + let is_json = content_type + .as_deref() + .is_some_and(|ct| ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes())); + // JSON-RPC errors need the same complete-body limit as success + // messages. Other bodies are diagnostics: retain only a prefix so + // large legacy error pages can still reach discovery fallback. + let body = if is_json { + read_bounded_body( + response.bytes_stream(), + content_length, + limits.max_json_response_size, + ) + .await + } else { + read_truncated_body(response.bytes_stream(), limits.max_error_response_size).await + }; + let body = match body { + Ok(body) => body, + Err(StreamableHttpError::Client(_)) => b"".to_vec(), Err(error) => return Err(error), }; - if content_type - .as_deref() - .is_some_and(|ct| ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes())) - { - match parse_json_rpc_error(&body) { + if is_json { + match parse_json_rpc_error(&String::from_utf8_lossy(&body)) { Some(message) => { return Ok(StreamableHttpPostResponse::Json(message, session_id)); } @@ -324,6 +331,10 @@ impl StreamableHttpClient for reqwest::Client { ), } } + // Even malformed/non-error JSON gets only a diagnostic prefix. + // Truncate raw bytes before lossy UTF-8 conversion, never a String. + 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) { diff --git a/crates/rmcp/src/transport/common/unix_socket.rs b/crates/rmcp/src/transport/common/unix_socket.rs index 711458f47..2a7c036d3 100644 --- a/crates/rmcp/src/transport/common/unix_socket.rs +++ b/crates/rmcp/src/transport/common/unix_socket.rs @@ -10,10 +10,10 @@ use sse_stream::Sse; use tokio::net::UnixStream; use crate::{ - model::{ClientJsonRpcMessage, ServerJsonRpcMessage}, + model::{ClientJsonRpcMessage, JsonRpcMessage, ServerJsonRpcMessage}, transport::{ common::{ - client_side_body::read_bounded_body, + client_side_body::{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, @@ -296,21 +296,46 @@ impl StreamableHttpClient for UnixSocketHttpClient { return Err(StreamableHttpError::SessionExpired); } + 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, - limits.max_error_response_size, - ) - .await - { - Ok(body) => String::from_utf8_lossy(&body).into_owned(), - Err(StreamableHttpError::Client(_)) => "".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(_)) => b"".to_vec(), Err(error) => return Err(error), }; + // Match reqwest: parse only complete JSON bodies, before considering + // legacy discovery fallback or truncating the diagnostic text. + if is_json { + match serde_json::from_str::(&String::from_utf8_lossy(&body)) + { + Ok(message @ JsonRpcMessage::Error(_)) => { + return Ok(StreamableHttpPostResponse::Json(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) { @@ -321,13 +346,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!( diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index d82393468..4b99be4b8 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -220,7 +220,7 @@ pub enum StreamableHttpError { /// A cancellation or reply POST did not finish in time. #[error("Control POST timed out")] ControlRequestTimeout, - /// A buffered JSON or error response exceeded its configured byte limit. + /// A buffered JSON response (including a JSON-RPC error) exceeded its byte limit. /// The response body is not included and this error does not trigger retries. #[error("HTTP response body exceeded {limit} bytes before decoding")] ResponseBodyTooLarge { limit: usize }, @@ -375,17 +375,24 @@ 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 or JSON response and 64 KiB per HTTP -/// error response. Zero accepts only empty bodies. For compressed responses, -/// body limits count the decompressed bytes yielded by the HTTP backend. +/// 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 buffered JSON success response body. + /// Maximum complete JSON response body, including HTTP error responses. pub max_json_response_size: usize, - /// Maximum buffered non-success HTTP response body. + /// 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, } @@ -393,7 +400,7 @@ impl Default for StreamableHttpResponseLimits { fn default() -> Self { Self { max_sse_event_size: DEFAULT_MAX_SSE_EVENT_SIZE, - max_json_response_size: 16 * 1024 * 1024, + max_json_response_size: 64 * 1024 * 1024, max_error_response_size: 64 * 1024, } } @@ -2084,16 +2091,19 @@ 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 buffered JSON success response body (default: 16 MiB). + /// 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 buffered HTTP error response body (default: 64 KiB). + /// Maximum HTTP error diagnostic prefix (default: 64 KiB). /// - /// Enforced before JSON-RPC error parsing or legacy discovery fallback. + /// 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`]. @@ -2188,7 +2198,7 @@ impl StreamableHttpClientTransportConfig { self } - /// Set the maximum HTTP error response body size before decoding. + /// 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 diff --git a/crates/rmcp/tests/test_streamable_http_default_error_limit.rs b/crates/rmcp/tests/test_streamable_http_default_error_limit.rs index 809ce9878..d791c70c3 100644 --- a/crates/rmcp/tests/test_streamable_http_default_error_limit.rs +++ b/crates/rmcp/tests/test_streamable_http_default_error_limit.rs @@ -47,11 +47,10 @@ impl MockServer { /// 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_rejects_oversized_error_without_exposing_body() { - const MARKER: &str = "SENSITIVE_TEST_MARKER_DO_NOT_ECHO"; - let mut body = MARKER.to_owned(); - body.push_str(&"x".repeat(65_537 - MARKER.len())); - assert_eq!(body.len(), 65_537); +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() @@ -69,10 +68,17 @@ async fn default_client_rejects_oversized_error_without_exposing_body() { .expect_err("an oversized HTTP error must fail") .to_string(); - assert!(error.contains("exceeded"), "expected a size-limit error"); + assert!( + error.starts_with("unexpected server response: HTTP 500"), + "{error}" + ); + assert!( + error.contains(&"x".repeat(65_536)), + "expected diagnostic prefix" + ); assert!( !error.contains(MARKER), - "error must not expose response body" + "error must not expose the discarded tail" ); - assert!(error.len() < 512, "error must remain bounded"); + 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 index cd65b8158..d06079bf1 100644 --- a/crates/rmcp/tests/test_streamable_http_limit_lifecycle.rs +++ b/crates/rmcp/tests/test_streamable_http_limit_lifecycle.rs @@ -369,7 +369,11 @@ impl Drop for LoopbackServer { } impl LoopbackServer { - async fn oversized_discover(error_response: bool, never_finishes: bool) -> Self { + 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( @@ -382,22 +386,33 @@ impl LoopbackServer { .lock() .unwrap() .push(request["method"].as_str().unwrap().to_owned()); - let (status, content_type, payload) = if error_response { - ( - 422, - "text/plain", - "Unexpected message, expect initialize request".repeat(8), - ) + 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 { - ( - 200, - "application/json", - json!({ - "jsonrpc": "2.0", "id": request["id"], - "error": { "code": -32601, "message": "legacy".repeat(80) }, - }) - .to_string(), - ) + 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. @@ -430,16 +445,17 @@ impl LoopbackServer { } #[rstest] -#[case::json(false, false)] -#[case::error(true, false)] -#[case::unfinished_json(false, true)] -#[case::unfinished_error(true, true)] +#[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] error_response: bool, + #[case] status: u16, #[case] never_finishes: bool, ) { - let server = LoopbackServer::oversized_discover(error_response, never_finishes).await; + 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()) @@ -474,3 +490,27 @@ async fn oversized_discover_fails_without_fallback_or_retry( ); 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, + ClientInfo::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 1177c222a..a7804a7a2 100644 --- a/crates/rmcp/tests/test_streamable_http_response_limits.rs +++ b/crates/rmcp/tests/test_streamable_http_response_limits.rs @@ -196,10 +196,14 @@ async fn oversized_malformed_json_is_not_accepted(#[case] chunked: bool) { #[case::chunked_exact_limit(0, true)] #[case::chunked_over_limit(-1, true)] #[tokio::test] -async fn json_rpc_error_uses_error_body_limit(#[case] margin: isize, #[case] chunked: bool) { - let server = MockServer::start(400, "application/json", JSON_ERROR, chunked).await; +async fn json_rpc_error_uses_json_limit_for_every_status( + #[case] margin: isize, + #[case] chunked: bool, + #[values(200, 400)] status: u16, +) { + let server = MockServer::start(status, "application/json", JSON_ERROR, chunked).await; let limit = JSON_ERROR.len().checked_add_signed(margin).unwrap(); - let result = post(&server, ping(), limits(0, limit, 0)).await; + let result = post(&server, ping(), limits(limit, 0, 0)).await; if margin < 0 { assert!( matches!(result, Err(StreamableHttpError::ResponseBodyTooLarge { limit: got }) if got == limit) @@ -219,20 +223,27 @@ async fn json_rpc_error_uses_error_body_limit(#[case] margin: isize, #[case] chu #[case::fixed(false)] #[case::chunked(true)] #[tokio::test] -async fn oversized_discovery_rejection_cannot_trigger_legacy_fallback(#[case] chunked: bool) { +async fn oversized_discovery_rejection_preserves_legacy_fallback(#[case] chunked: bool) { let body = "Unexpected message, expect initialize request"; let server = MockServer::start(422, "text/plain", body, chunked).await; - assert!(matches!( - post(&server, discover(), limits(100, body.len() - 1, 100)).await, - Err(StreamableHttpError::ResponseBodyTooLarge { .. }) - )); - assert!(matches!( - post(&server, discover(), limits(100, body.len(), 100)).await, - Ok(StreamableHttpPostResponse::Json( - JsonRpcMessage::Error(_), - _ - )) - )); + for limit in [0, body.len() - 1, body.len()] { + let StreamableHttpPostResponse::Json(JsonRpcMessage::Error(error), session) = + post(&server, discover(), limits(0, limit, 0)) + .await + .unwrap() + else { + panic!("expected legacy discovery rejection"); + }; + assert_eq!(error.id, Some(RequestId::Number(1))); + assert!(session.is_none()); + assert_eq!( + error.error.message, + format!( + "server/discover rejected with HTTP 422 Unprocessable Entity: {}", + &body[..limit] + ) + ); + } } #[rstest] @@ -243,7 +254,8 @@ async fn non_json_error_bodies_are_bounded(#[case] chunked: bool) { let server = MockServer::start(500, "text/plain", "failure", chunked).await; assert!(matches!( post(&server, ping(), limits(100, 6, 100)).await, - Err(StreamableHttpError::ResponseBodyTooLarge { limit: 6 }) + Err(StreamableHttpError::UnexpectedServerResponse(message)) + if message == format!("HTTP 500 Internal Server Error: {}", &"failure"[..6]) )); assert!(matches!( post(&server, ping(), limits(0, 7, 0)).await, @@ -346,7 +358,74 @@ async fn existing_entry_points_apply_default_body_limits(#[case] sse_limit_metho .post_message(server.uri.clone(), ping(), None, None, HashMap::new()) .await }; - assert!( - matches!(result, Err(StreamableHttpError::ResponseBodyTooLarge { limit: got }) if got == limit) + let Err(StreamableHttpError::UnexpectedServerResponse(message)) = result else { + panic!("expected a bounded HTTP diagnostic"); + }; + assert_eq!( + message, + format!("HTTP 500 Internal Server Error: {}", "x".repeat(limit)) ); } + +#[test] +fn defaults_allow_larger_json_without_changing_sse_or_diagnostics() { + let limits = StreamableHttpResponseLimits::default(); + assert_eq!(limits.max_json_response_size, 64 * 1024 * 1024); + assert_eq!(limits.max_sse_event_size, 16 * 1024 * 1024); + assert_eq!(limits.max_error_response_size, 64 * 1024); +} + +#[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 = MockServer::start(200, "application/json", body, false).await; + assert!(matches!( + client() + .post_message(server.uri.clone(), ping(), None, None, HashMap::new()) + .await, + Ok(StreamableHttpPostResponse::Json( + JsonRpcMessage::Response(_), + _ + )) + )); +} + +#[rstest] +#[case::fixed(false)] +#[case::chunked(true)] +#[tokio::test] +async fn truncated_diagnostics_cannot_be_parsed_as_json_rpc(#[case] chunked: bool) { + // The prefix is valid JSON, but the complete body is not a protocol message. + let body = format!("{JSON_ERROR}not-json"); + for content_type in ["text/plain", "application/json"] { + let server = MockServer::start(400, content_type, body.clone(), chunked).await; + assert!(matches!( + post(&server, ping(), limits(body.len(), JSON_ERROR.len(), 0)).await, + Err(StreamableHttpError::UnexpectedServerResponse(_)) + )); + if content_type == "application/json" { + assert!(matches!( + post(&server, discover(), limits(JSON_ERROR.len(), body.len(), 0)).await, + Err(StreamableHttpError::ResponseBodyTooLarge { .. }) + )); + } + } +} + +#[rstest] +#[case::unauthorized(401)] +#[case::forbidden(403)] +#[case::server_error(500)] +#[tokio::test] +async fn truncated_diagnostics_do_not_expand_discovery_fallback(#[case] status: u16) { + let server = MockServer::start(status, "text/html", "large error page", true).await; + assert!(matches!( + post(&server, discover(), limits(0, 2, 0)).await, + Err(StreamableHttpError::UnexpectedServerResponse(_)) + )); +} diff --git a/crates/rmcp/tests/test_unix_socket_response_limits.rs b/crates/rmcp/tests/test_unix_socket_response_limits.rs index cb08e1da3..556566a02 100644 --- a/crates/rmcp/tests/test_unix_socket_response_limits.rs +++ b/crates/rmcp/tests/test_unix_socket_response_limits.rs @@ -14,7 +14,10 @@ use std::{ use futures::StreamExt; use rmcp::{ - model::ClientJsonRpcMessage, + model::{ + ClientJsonRpcMessage, ClientRequest, DiscoverRequest, DiscoverRequestParams, + JsonRpcMessage, RequestId, + }, transport::{ UnixSocketHttpClient, streamable_http_client::{ @@ -38,6 +41,10 @@ struct TestServer { 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", @@ -69,6 +76,10 @@ impl TestServer { } // 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 } @@ -106,6 +117,12 @@ fn response(status: &str, content_type: &str, body: &[u8], chunked: bool) -> Vec } 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 })) @@ -165,7 +182,7 @@ async fn json_thresholds_cover_content_length_chunked_and_recovery() { } #[tokio::test] -async fn error_thresholds_cannot_be_hidden_by_discovery_fallback() { +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]; @@ -197,15 +214,18 @@ async fn error_thresholds_cannot_be_hidden_by_discovery_fallback() { limits(1, limit), )) .await; - if length > limit { - assert!( - matches!(result, Err(StreamableHttpError::ResponseBodyTooLarge { limit: actual }) if actual == limit) - ); + 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(_)) + Err(StreamableHttpError::UnexpectedServerResponse(ref body)) + if body == &format!("HTTP 400 Bad Request: {}", "x".repeat(length.min(limit))) ), "{result:?}" ); @@ -238,10 +258,10 @@ async fn oversized_invalid_json_is_not_accepted() { } #[tokio::test] -async fn legacy_post_methods_apply_default_json_and_error_limits() { +async fn legacy_post_methods_apply_default_json_limit_on_every_status() { for (status, content_type, limit) in [ - ("200 OK", "application/json", 16 * 1024 * 1024), - ("500 Internal Server Error", "text/plain", 64 * 1024), + ("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", @@ -276,6 +296,170 @@ async fn legacy_post_methods_apply_default_json_and_error_limits() { } } +#[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(