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