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..75a1b8fde --- /dev/null +++ b/crates/rmcp/src/transport/common/client_side_body.rs @@ -0,0 +1,214 @@ +//! 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) +} + +/// 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::{ + 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(_)) + )); + } + + #[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 b0b5090b5..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,6 +9,7 @@ use crate::{ model::{ClientJsonRpcMessage, JsonRpcMessage, ServerJsonRpcMessage}, transport::{ common::{ + 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, @@ -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,15 +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 = response - .text() - .await - .unwrap_or_else(|_| "".to_owned()); - if content_type + let is_json = content_type .as_deref() - .is_some_and(|ct| ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes())) - { - match parse_json_rpc_error(&body) { + .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 is_json { + match parse_json_rpc_error(&String::from_utf8_lossy(&body)) { Some(message) => { return Ok(StreamableHttpPostResponse::Json(message, session_id)); } @@ -293,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) { @@ -304,14 +346,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..2a7c036d3 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; @@ -10,9 +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, 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, @@ -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) @@ -267,13 +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 = response + 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() - .collect() - .await - .map(|c| String::from_utf8_lossy(&c.to_bytes()).into_owned()) - .unwrap_or_else(|_| "".to_owned()); + .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) { @@ -284,18 +346,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) - .and_then(|v| v.to_str().ok()) - .map(|s| s.to_string()); - if status.is_success() && content_length == Some(0) && matches!( @@ -310,17 +360,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..4b99be4b8 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 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 }, } impl StreamableHttpError { @@ -369,12 +373,47 @@ 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`]. /// /// 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 +453,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 +662,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 +679,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 +1025,7 @@ impl StreamableHttpClientWorker { uri: Arc, auth_header: Option, custom_headers: HashMap, - max_sse_event_size: usize, + response_limits: StreamableHttpResponseLimits, ) -> Result< ( Option>, @@ -969,13 +1035,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 +1066,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 +1143,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 +1210,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 +1296,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 +1539,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 +2091,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`). /// @@ -2109,6 +2192,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 @@ -2136,6 +2239,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..d791c70c3 --- /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::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_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 500"), + "{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..d06079bf1 --- /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, 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( + 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, + 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"]); +} + +#[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 new file mode 100644 index 000000000..a7804a7a2 --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_response_limits.rs @@ -0,0 +1,431 @@ +#![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_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(limit, 0, 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_preserves_legacy_fallback(#[case] chunked: bool) { + let body = "Unexpected message, expect initialize request"; + let server = MockServer::start(422, "text/plain", body, chunked).await; + 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] +#[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::UnexpectedServerResponse(message)) + if message == format!("HTTP 500 Internal Server Error: {}", &"failure"[..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 + }; + 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 new file mode 100644 index 000000000..556566a02 --- /dev/null +++ b/crates/rmcp/tests/test_unix_socket_response_limits.rs @@ -0,0 +1,497 @@ +#![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, ClientRequest, DiscoverRequest, DiscoverRequestParams, + JsonRpcMessage, RequestId, + }, + 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 { + 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 { limit: 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 { limit: 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 { 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 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", + ); +}