diff --git a/crates/rmcp/Cargo.toml b/crates/rmcp/Cargo.toml index 5ab429abc..a963c156d 100644 --- a/crates/rmcp/Cargo.toml +++ b/crates/rmcp/Cargo.toml @@ -153,7 +153,7 @@ server-side-http = [ transport-worker = ["dep:tokio-stream"] # SSE stream parsing utilities (used by streamable HTTP client for SSE-formatted responses) -client-side-sse = ["dep:sse-stream", "dep:http", "base64"] +client-side-sse = ["dep:sse-stream", "dep:http", "dep:bytes", "base64"] # Streamable HTTP client transport-streamable-http-client = ["client-side-sse", "transport-worker"] 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 47f08f13e..f0a6211af 100644 --- a/crates/rmcp/src/transport/common/auth/streamable_http_client.rs +++ b/crates/rmcp/src/transport/common/auth/streamable_http_client.rs @@ -47,6 +47,33 @@ where .await } + async fn get_stream_with_max_sse_event_size( + &self, + uri: std::sync::Arc, + session_id: std::sync::Arc, + last_event_id: Option, + mut auth_token: Option, + custom_headers: HashMap, + max_sse_event_size: usize, + ) -> Result< + futures::stream::BoxStream<'static, Result>, + crate::transport::streamable_http_client::StreamableHttpError, + > { + if auth_token.is_none() { + auth_token = Some(self.get_access_token().await?); + } + self.http_client + .get_stream_with_max_sse_event_size( + uri, + session_id, + last_event_id, + auth_token, + custom_headers, + max_sse_event_size, + ) + .await + } + async fn post_message( &self, uri: std::sync::Arc, @@ -65,4 +92,31 @@ where .post_message(uri, message, session_id, auth_token, custom_headers) .await } + + async fn post_message_with_max_sse_event_size( + &self, + uri: std::sync::Arc, + message: crate::model::ClientJsonRpcMessage, + session_id: Option>, + mut auth_token: Option, + custom_headers: HashMap, + max_sse_event_size: usize, + ) -> Result< + crate::transport::streamable_http_client::StreamableHttpPostResponse, + StreamableHttpError, + > { + if auth_token.is_none() { + auth_token = Some(self.get_access_token().await?); + } + self.http_client + .post_message_with_max_sse_event_size( + uri, + message, + session_id, + auth_token, + custom_headers, + max_sse_event_size, + ) + .await + } } diff --git a/crates/rmcp/src/transport/common/client_side_sse.rs b/crates/rmcp/src/transport/common/client_side_sse.rs index fc9e15eb7..6408514f4 100644 --- a/crates/rmcp/src/transport/common/client_side_sse.rs +++ b/crates/rmcp/src/transport/common/client_side_sse.rs @@ -5,13 +5,166 @@ use std::{ time::Duration, }; -use futures::{Stream, stream::BoxStream}; -use sse_stream::{Error as SseError, Sse}; +use bytes::Bytes; +use futures::{Stream, StreamExt, stream::BoxStream}; +use sse_stream::{Error as SseError, Sse, SseStream}; +use thiserror::Error; use crate::model::ServerJsonRpcMessage; pub type BoxedSseResponse = BoxStream<'static, Result>; +/// Maximum raw size of one SSE event accepted from a remote server. +pub const DEFAULT_MAX_SSE_EVENT_SIZE: usize = 16 * 1024 * 1024; + +#[derive(Debug, Error)] +enum BoundedSseStreamError { + #[error(transparent)] + Source(Box), + #[error("SSE event exceeded the maximum size of {max_size} bytes")] + EventTooLarge { max_size: usize }, +} + +#[derive(Debug)] +struct SseEventSizeLimiter { + max_size: usize, + retained_size: usize, + line_size: usize, + line_is_comment: bool, + previous_was_cr: bool, +} + +impl SseEventSizeLimiter { + fn new(max_size: usize) -> Self { + Self { + max_size, + retained_size: 0, + line_size: 0, + line_is_comment: false, + previous_was_cr: false, + } + } + + fn observe(&mut self, chunk: &[u8]) -> Result<(), ()> { + for &byte in chunk { + if self.previous_was_cr { + self.previous_was_cr = false; + if byte == b'\n' { + continue; + } + } + + match byte { + b'\r' => { + self.finish_line()?; + self.previous_was_cr = true; + } + b'\n' => self.finish_line()?, + _ => { + if self.line_size == 0 { + self.line_is_comment = byte == b':'; + } + self.line_size = self.line_size.saturating_add(1); + self.check_limit()?; + } + } + } + Ok(()) + } + + fn finish_line(&mut self) -> Result<(), ()> { + if self.line_size == 0 { + self.retained_size = 0; + } else if !self.line_is_comment { + // The SSE parser inserts a newline when joining multiple data fields. + self.retained_size = self + .retained_size + .saturating_add(self.line_size) + .saturating_add(1); + } + self.line_size = 0; + self.line_is_comment = false; + self.check_limit() + } + + fn check_limit(&self) -> Result<(), ()> { + if self.retained_size.saturating_add(self.line_size) > self.max_size { + Err(()) + } else { + Ok(()) + } + } +} + +pin_project_lite::pin_project! { + struct BoundedSseByteStream { + #[pin] + inner: S, + limiter: SseEventSizeLimiter, + failed: bool, + } +} + +impl Stream for BoundedSseByteStream +where + S: Stream>, + E: std::error::Error + Send + Sync + 'static, +{ + type Item = Result; + + fn poll_next( + self: Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll> { + let mut this = self.project(); + if *this.failed { + return Poll::Ready(None); + } + + match ready!(this.inner.as_mut().poll_next(cx)) { + Some(Ok(chunk)) => { + if this.limiter.observe(&chunk).is_err() { + *this.failed = true; + Poll::Ready(Some(Err(BoundedSseStreamError::EventTooLarge { + max_size: this.limiter.max_size, + }))) + } else { + Poll::Ready(Some(Ok(chunk))) + } + } + Some(Err(error)) => { + *this.failed = true; + Poll::Ready(Some(Err(BoundedSseStreamError::Source(Box::new(error))))) + } + None => Poll::Ready(None), + } + } +} + +pub(crate) fn bounded_sse_stream(stream: S, max_event_size: usize) -> BoxedSseResponse +where + S: Stream> + Send + 'static, + E: std::error::Error + Send + Sync + 'static, +{ + let stream = BoundedSseByteStream { + inner: stream, + limiter: SseEventSizeLimiter::new(max_event_size), + failed: false, + }; + SseStream::from_bytes_stream(stream).boxed() +} + +fn is_event_too_large_error(error: &SseError) -> bool { + matches!( + error, + SseError::Body(error) + if matches!( + error.downcast_ref::(), + Some(BoundedSseStreamError::EventTooLarge { .. }) + ) + ) +} + pub trait SseRetryPolicy: std::fmt::Debug + Send + Sync { fn retry(&self, current_times: usize) -> Option; } @@ -124,6 +277,10 @@ pub(crate) trait SseStreamReconnect { tracing::warn!("sse stream error: {error}"); } } + fn map_fatal_stream_error(&mut self, error: SseError) -> Option { + tracing::warn!("fatal sse stream error: {error}"); + None + } } pin_project_lite::pin_project! { @@ -248,6 +405,10 @@ where } } Some(Err(e)) => { + if is_event_too_large_error(&e) { + this.state.set(SseAutoReconnectStreamState::Terminated); + return Poll::Ready(this.connector.map_fatal_stream_error(e).map(Err)); + } this.connector .handle_stream_error(&e, this.last_event_id.as_deref()); let retrying = this @@ -327,3 +488,202 @@ where self.poll_next(cx) } } + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use super::*; + + #[derive(Debug, Error)] + enum TestReconnectError { + #[error("SSE stream error: {0}")] + Sse(SseError), + #[error("unexpected reconnect")] + Reconnect, + } + + struct CountingReconnect { + attempts: Arc, + } + + impl SseStreamReconnect for CountingReconnect { + type Error = TestReconnectError; + type Future = futures::future::Ready>; + + fn retry_connection(&mut self, _last_event_id: Option<&str>) -> Self::Future { + self.attempts.fetch_add(1, Ordering::Relaxed); + futures::future::ready(Err(TestReconnectError::Reconnect)) + } + + fn map_fatal_stream_error(&mut self, error: SseError) -> Option { + Some(TestReconnectError::Sse(error)) + } + } + + #[tokio::test] + async fn bounded_sse_stream_rejects_unterminated_event_over_limit() { + let source = futures::stream::iter([ + Ok::<_, std::io::Error>(Bytes::from_static(b"data: aaaaa")), + Ok(Bytes::from_static(b"aaaaaa")), + ]); + let mut stream = bounded_sse_stream(source, 16); + + let error = stream.next().await.unwrap().unwrap_err(); + + assert!( + error.to_string().contains("maximum size of 16 bytes"), + "unexpected error: {error}" + ); + } + + #[tokio::test] + async fn bounded_sse_stream_resets_limit_after_event_terminator() { + let source = futures::stream::iter([ + Ok::<_, std::io::Error>(Bytes::from_static(b"data: a\r")), + Ok(Bytes::from_static(b"\n\r\ndata: b\n\n")), + ]); + let mut stream = bounded_sse_stream(source, 8); + + let first = stream.next().await.unwrap().unwrap(); + let second = stream.next().await.unwrap().unwrap(); + + assert_eq!( + (first.data.as_deref(), second.data.as_deref()), + (Some("a"), Some("b")) + ); + } + + #[tokio::test] + async fn bounded_sse_stream_passes_event_at_exact_limit() { + let source = + futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(b"data: ab\n\n"))]); + let mut stream = bounded_sse_stream(source, 9); + + let event = stream.next().await.unwrap().unwrap(); + assert_eq!(event.data.as_deref(), Some("ab")); + } + + #[tokio::test] + async fn bounded_sse_stream_rejects_oversize_split_across_many_chunks() { + let source = futures::stream::iter( + std::iter::repeat_with(|| Ok::<_, std::io::Error>(Bytes::from_static(b"data: x"))) + .take(20), + ); + let mut stream = bounded_sse_stream(source, 32); + + let mut found_error = false; + while let Some(item) = stream.next().await { + if item.is_err() { + found_error = true; + break; + } + } + assert!(found_error, "expected oversize error"); + } + + #[tokio::test] + async fn bounded_sse_stream_handles_crlf_split_across_chunks() { + let source = futures::stream::iter([ + Ok::<_, std::io::Error>(Bytes::from_static(b"data: hello\r")), + Ok(Bytes::from_static(b"\n\ndata: world\n\n")), + ]); + let mut stream = bounded_sse_stream(source, 64); + + let first = stream.next().await.unwrap().unwrap(); + let second = stream.next().await.unwrap().unwrap(); + + assert_eq!(first.data.as_deref(), Some("hello")); + assert_eq!(second.data.as_deref(), Some("world")); + } + + #[tokio::test] + async fn bounded_sse_stream_discards_completed_comment_lines_from_limit() { + let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static( + b": ping\n: ping\n: ping\n", + ))]); + let mut stream = bounded_sse_stream(source, 6); + + assert!(stream.next().await.is_none()); + } + + #[tokio::test] + async fn bounded_sse_stream_comments_do_not_reset_accumulated_data() { + let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static( + b"data: a\n: ping\ndata: b\n", + ))]); + let mut stream = bounded_sse_stream(source, 14); + + assert!(stream.next().await.unwrap().is_err()); + } + + #[tokio::test] + async fn bounded_sse_stream_counts_multiline_data_join_newlines() { + let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static( + b"data: aaa\ndata: bbb\n\n", + ))]); + let mut stream = bounded_sse_stream(source, 18); + + let error = stream.next().await.unwrap().unwrap_err(); + assert!( + error.to_string().contains("maximum size"), + "unexpected error: {error}" + ); + } + + #[tokio::test] + async fn bounded_sse_stream_propagates_source_error() { + let source = futures::stream::iter([Err::(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + "connection reset", + ))]); + let mut stream = bounded_sse_stream(source, 1024); + + let error = stream.next().await.unwrap().unwrap_err(); + assert!( + error.to_string().contains("connection reset"), + "unexpected error: {error}" + ); + } + + #[tokio::test] + async fn is_event_too_large_error_detects_oversize() { + let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from(vec![b'A'; 100]))]); + let mut stream = bounded_sse_stream(source, 8); + + let error = stream.next().await.unwrap().unwrap_err(); + assert!(is_event_too_large_error(&error)); + } + + #[tokio::test] + async fn is_event_too_large_error_rejects_other_errors() { + let source = + futures::stream::iter([Err::(std::io::Error::other("something else"))]); + let mut stream = bounded_sse_stream(source, 1024); + + let error = stream.next().await.unwrap().unwrap_err(); + assert!(!is_event_too_large_error(&error)); + } + + #[tokio::test] + async fn oversized_event_returns_error_without_reconnecting() { + let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from(vec![b'A'; 100]))]); + let attempts = Arc::new(AtomicUsize::new(0)); + let connector = CountingReconnect { + attempts: attempts.clone(), + }; + let stream = SseAutoReconnectStream::new( + bounded_sse_stream(source, 8), + connector, + Arc::new(NeverRetry), + ); + let mut stream = std::pin::pin!(stream); + + let result = stream.next().await; + + assert!( + matches!(result, Some(Err(TestReconnectError::Sse(_)))) + && attempts.load(Ordering::Relaxed) == 0 + ); + } +} 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 57c12f1ca..7032e1a87 100644 --- a/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs +++ b/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs @@ -1,16 +1,19 @@ use std::{borrow::Cow, collections::HashMap, sync::Arc}; -use futures::{StreamExt, stream::BoxStream}; +use futures::stream::BoxStream; use http::{HeaderName, HeaderValue, header::WWW_AUTHENTICATE}; use reqwest::header::ACCEPT; -use sse_stream::{Sse, SseStream}; +use sse_stream::Sse; use crate::{ model::{ClientJsonRpcMessage, JsonRpcMessage, ServerJsonRpcMessage}, transport::{ - common::http_header::{ - EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE, - extract_scope_from_header, validate_custom_header, + common::{ + 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, + extract_scope_from_header, validate_custom_header, + }, }, streamable_http_client::*, }, @@ -53,6 +56,26 @@ impl StreamableHttpClient for reqwest::Client { last_event_id: Option, auth_token: Option, custom_headers: HashMap, + ) -> Result>, StreamableHttpError> { + self.get_stream_with_max_sse_event_size( + uri, + session_id, + last_event_id, + auth_token, + custom_headers, + DEFAULT_MAX_SSE_EVENT_SIZE, + ) + .await + } + + async fn get_stream_with_max_sse_event_size( + &self, + uri: Arc, + session_id: Arc, + last_event_id: Option, + auth_token: Option, + custom_headers: HashMap, + max_sse_event_size: usize, ) -> Result>, StreamableHttpError> { let mut request_builder = self .get(uri.as_ref()) @@ -84,7 +107,7 @@ impl StreamableHttpClient for reqwest::Client { return Err(StreamableHttpError::UnexpectedContentType(None)); } } - let event_stream = SseStream::from_bytes_stream(response.bytes_stream()).boxed(); + let event_stream = bounded_sse_stream(response.bytes_stream(), max_sse_event_size); Ok(event_stream) } @@ -119,6 +142,26 @@ impl StreamableHttpClient for reqwest::Client { session_id: Option>, auth_token: Option, custom_headers: HashMap, + ) -> Result> { + self.post_message_with_max_sse_event_size( + uri, + message, + session_id, + auth_token, + custom_headers, + DEFAULT_MAX_SSE_EVENT_SIZE, + ) + .await + } + + async fn post_message_with_max_sse_event_size( + &self, + uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + auth_token: Option, + custom_headers: HashMap, + max_sse_event_size: usize, ) -> Result> { let mut request = self .post(uri.as_ref()) @@ -223,7 +266,7 @@ 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 = SseStream::from_bytes_stream(response.bytes_stream()).boxed(); + let event_stream = bounded_sse_stream(response.bytes_stream(), max_sse_event_size); Ok(StreamableHttpPostResponse::Sse(event_stream, session_id)) } Some(ct) if ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes()) => { @@ -364,6 +407,57 @@ mod tests { assert!(parse_json_rpc_error(body).is_none()); } + #[tokio::test] + async fn post_sse_response_honors_configured_event_limit() -> anyhow::Result<()> { + use std::{collections::HashMap, net::SocketAddr, sync::Arc}; + + use axum::{Router, routing::post}; + use futures::StreamExt; + + use crate::transport::streamable_http_client::{ + StreamableHttpClient, StreamableHttpPostResponse, + }; + + let listener = tokio::net::TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0))).await?; + let addr = listener.local_addr()?; + let server = tokio::spawn(async move { + let app = Router::new().route( + "/mcp", + post(|| async { + ( + [(http::header::CONTENT_TYPE, "text/event-stream")], + "data: this event is too large\n", + ) + }), + ); + axum::serve(listener, app).await + }); + let message = ClientJsonRpcMessage::request( + ClientRequest::PingRequest(PingRequest::default()), + RequestId::Number(1), + ); + let client = reqwest::Client::new(); + + let response = client + .post_message_with_max_sse_event_size( + Arc::::from(format!("http://{addr}/mcp")), + message, + None, + None, + HashMap::new(), + 16, + ) + .await?; + let StreamableHttpPostResponse::Sse(mut stream, _) = response else { + anyhow::bail!("expected SSE response"); + }; + let error = stream.next().await.unwrap().unwrap_err(); + + server.abort(); + assert!(error.to_string().contains("maximum size of 16 bytes")); + Ok(()) + } + #[tokio::test] async fn default_http_client_does_not_leak_custom_headers_to_redirect_target() -> anyhow::Result<()> { diff --git a/crates/rmcp/src/transport/common/server_side_http.rs b/crates/rmcp/src/transport/common/server_side_http.rs index 39a321f9b..32609bb75 100644 --- a/crates/rmcp/src/transport/common/server_side_http.rs +++ b/crates/rmcp/src/transport/common/server_side_http.rs @@ -170,31 +170,61 @@ pub(crate) fn unexpected_message_response(expect: &str) -> Response( body: B, + max_bytes: usize, ) -> Result>> where B: Body + Send + 'static, B::Error: Display, { - match body.collect().await { - Ok(bytes) => { - match serde_json::from_reader::<_, ClientJsonRpcMessage>(bytes.aggregate().reader()) { - Ok(message) => Ok(message), - Err(e) => { + let mut collected = bytes::BytesMut::new(); + let mut body = std::pin::pin!(body); + loop { + let frame = futures::future::poll_fn(|cx| body.as_mut().poll_frame(cx)).await; + match frame { + None => break, + Some(Ok(frame)) => { + let Ok(mut data) = frame.into_data() else { + continue; + }; + if data.remaining() > max_bytes.saturating_sub(collected.len()) { let response = Response::builder() - .status(http::StatusCode::UNSUPPORTED_MEDIA_TYPE) + .status(http::StatusCode::PAYLOAD_TOO_LARGE) .body( - Full::new(Bytes::from(format!("fail to deserialize request body {e}"))) - .boxed(), + Full::new(Bytes::from(format!( + "Payload Too Large: request body exceeds {max_bytes} bytes" + ))) + .boxed(), ) .expect("valid response"); - Err(response) + return Err(response); } + while data.has_remaining() { + let chunk = data.chunk(); + let chunk_len = chunk.len(); + collected.extend_from_slice(chunk); + data.advance(chunk_len); + } + } + Some(Err(e)) => { + let response = Response::builder() + .status(http::StatusCode::INTERNAL_SERVER_ERROR) + .body( + Full::new(Bytes::from(format!("Failed to read request body: {e}"))).boxed(), + ) + .expect("valid response"); + return Err(response); } } + } + + match serde_json::from_slice::(&collected) { + Ok(message) => Ok(message), Err(e) => { let response = Response::builder() - .status(http::StatusCode::INTERNAL_SERVER_ERROR) - .body(Full::new(Bytes::from(format!("Failed to read request body: {e}"))).boxed()) + .status(http::StatusCode::UNSUPPORTED_MEDIA_TYPE) + .body( + Full::new(Bytes::from(format!("fail to deserialize request body {e}"))).boxed(), + ) .expect("valid response"); Err(response) } @@ -203,9 +233,17 @@ where #[cfg(test)] mod tests { + use std::convert::Infallible; + + use futures::stream; + use http_body::Frame; + use http_body_util::StreamBody; + use super::*; use crate::model::{EmptyResult, JsonRpcResponse, JsonRpcVersion2_0, RequestId, ServerResult}; + const INITIALIZE_REQUEST: &str = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"0.1"}}}"#; + fn dummy_message() -> ServerJsonRpcMessage { ServerJsonRpcMessage::Response(JsonRpcResponse { jsonrpc: JsonRpcVersion2_0, @@ -245,4 +283,62 @@ mod tests { assert!(msg.message.is_none()); assert_eq!(msg.retry, Some(Duration::from_secs(5))); } + + #[tokio::test] + async fn expect_json_accepts_body_under_limit() { + let body = Full::new(Bytes::from_static(INITIALIZE_REQUEST.as_bytes())); + let result = expect_json(body, 4 * 1024 * 1024).await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn expect_json_accepts_non_contiguous_body_data() { + let split = INITIALIZE_REQUEST.len() / 2; + let data = Bytes::copy_from_slice(&INITIALIZE_REQUEST.as_bytes()[..split]).chain( + Bytes::copy_from_slice(&INITIALIZE_REQUEST.as_bytes()[split..]), + ); + let result = expect_json(Full::new(data), 4 * 1024 * 1024).await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn expect_json_accepts_multi_frame_body() { + let split = INITIALIZE_REQUEST.len() / 2; + let body = StreamBody::new(stream::iter([ + Ok::<_, Infallible>(Frame::data(Bytes::copy_from_slice( + &INITIALIZE_REQUEST.as_bytes()[..split], + ))), + Ok(Frame::data(Bytes::copy_from_slice( + &INITIALIZE_REQUEST.as_bytes()[split..], + ))), + ])); + let result = expect_json(body, 4 * 1024 * 1024).await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn expect_json_rejects_oversized_body() { + let big_body = Full::new(Bytes::from(vec![b'x'; 128])); + let result = expect_json(big_body, 64).await; + let response = result.unwrap_err(); + assert_eq!(response.status(), http::StatusCode::PAYLOAD_TOO_LARGE); + } + + #[tokio::test] + async fn expect_json_rejects_oversized_multi_frame_body() { + let body = StreamBody::new(stream::iter([ + Ok::<_, Infallible>(Frame::data(Bytes::from_static(b"12345678"))), + Ok(Frame::data(Bytes::from_static(b"9"))), + ])); + let response = expect_json(body, 8).await.unwrap_err(); + assert_eq!(response.status(), http::StatusCode::PAYLOAD_TOO_LARGE); + } + + #[tokio::test] + async fn expect_json_returns_415_for_invalid_json_under_limit() { + let body = Full::new(Bytes::from("not valid json")); + let result = expect_json(body, 4 * 1024 * 1024).await; + let response = result.unwrap_err(); + assert_eq!(response.status(), http::StatusCode::UNSUPPORTED_MEDIA_TYPE); + } } diff --git a/crates/rmcp/src/transport/common/unix_socket.rs b/crates/rmcp/src/transport/common/unix_socket.rs index 8ea30f57f..899548313 100644 --- a/crates/rmcp/src/transport/common/unix_socket.rs +++ b/crates/rmcp/src/transport/common/unix_socket.rs @@ -1,20 +1,23 @@ use std::{borrow::Cow, collections::HashMap, sync::Arc}; use bytes::Bytes; -use futures::{StreamExt, stream::BoxStream}; +use futures::stream::BoxStream; use http::{HeaderName, HeaderValue, Method, Request, StatusCode, header::WWW_AUTHENTICATE}; use http_body_util::{BodyExt, Full}; use hyper::body::Incoming; use hyper_util::rt::TokioIo; -use sse_stream::{Sse, SseStream}; +use sse_stream::Sse; use tokio::net::UnixStream; use crate::{ model::{ClientJsonRpcMessage, ServerJsonRpcMessage}, transport::{ - common::http_header::{ - EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE, - extract_scope_from_header, validate_custom_header, + common::{ + 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, + extract_scope_from_header, validate_custom_header, + }, }, streamable_http_client::*, }, @@ -168,6 +171,26 @@ impl StreamableHttpClient for UnixSocketHttpClient { session_id: Option>, auth_token: Option, custom_headers: HashMap, + ) -> Result> { + self.post_message_with_max_sse_event_size( + uri, + message, + session_id, + auth_token, + custom_headers, + DEFAULT_MAX_SSE_EVENT_SIZE, + ) + .await + } + + async fn post_message_with_max_sse_event_size( + &self, + uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + auth_token: Option, + custom_headers: HashMap, + max_sse_event_size: usize, ) -> Result> { let json_body = serde_json::to_string(&message) .map_err(|e| StreamableHttpError::Client(UnixSocketError::Json(e)))?; @@ -282,7 +305,8 @@ 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 = SseStream::new(response.into_body()).boxed(); + let sse_stream = + bounded_sse_stream(response.into_body().into_data_stream(), 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()) => { @@ -357,6 +381,27 @@ impl StreamableHttpClient for UnixSocketHttpClient { auth_token: Option, custom_headers: HashMap, ) -> Result>, StreamableHttpError> + { + self.get_stream_with_max_sse_event_size( + uri, + session_id, + last_event_id, + auth_token, + custom_headers, + DEFAULT_MAX_SSE_EVENT_SIZE, + ) + .await + } + + async fn get_stream_with_max_sse_event_size( + &self, + uri: Arc, + session_id: Arc, + last_event_id: Option, + auth_token: Option, + custom_headers: HashMap, + max_sse_event_size: usize, + ) -> Result>, StreamableHttpError> { let mut builder = Request::builder() .method(Method::GET) @@ -444,7 +489,10 @@ impl StreamableHttpClient for UnixSocketHttpClient { } } - Ok(SseStream::new(response.into_body()).boxed()) + Ok(bounded_sse_stream( + response.into_body().into_data_stream(), + max_sse_event_size, + )) } } diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index 2379baf95..a83cd24ff 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -13,7 +13,9 @@ use thiserror::Error; use tokio_util::sync::CancellationToken; use tracing::debug; -use super::common::client_side_sse::{ExponentialBackoff, SseRetryPolicy, SseStreamReconnect}; +use super::common::client_side_sse::{ + DEFAULT_MAX_SSE_EVENT_SIZE, ExponentialBackoff, SseRetryPolicy, SseStreamReconnect, +}; use crate::{ RoleClient, model::{ @@ -266,6 +268,12 @@ impl StreamableHttpPostResponse { } } +/// 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. pub trait StreamableHttpClient: Clone + Send + 'static { type Error: std::error::Error + Send + Sync + 'static; fn post_message( @@ -278,6 +286,23 @@ pub trait StreamableHttpClient: Clone + Send + 'static { ) -> impl Future>> + Send + '_; + /// Send a message while enforcing a maximum raw SSE event size. + /// + /// Custom clients that parse SSE responses should override this method. + /// The default implementation delegates to [`Self::post_message`]. + 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, + ) -> impl Future>> + + Send + + '_ { + self.post_message(uri, message, session_id, auth_header, custom_headers) + } fn delete_session( &self, uri: Arc, @@ -299,6 +324,27 @@ pub trait StreamableHttpClient: Clone + Send + 'static { >, > + Send + '_; + /// Open an SSE stream while enforcing a maximum raw event size. + /// + /// Custom clients that parse SSE responses should override this method. + /// The default implementation delegates to [`Self::get_stream`]. + fn get_stream_with_max_sse_event_size( + &self, + uri: Arc, + session_id: Arc, + last_event_id: Option, + auth_header: Option, + custom_headers: HashMap, + _max_sse_event_size: usize, + ) -> impl Future< + Output = Result< + BoxStream<'static, Result>, + StreamableHttpError, + >, + > + Send + + '_ { + self.get_stream(uri, session_id, last_event_id, auth_header, custom_headers) + } } #[non_exhaustive] @@ -313,6 +359,7 @@ struct StreamableHttpClientReconnect { pub uri: Arc, pub auth_header: Option, pub custom_headers: HashMap, + pub max_sse_event_size: usize, } impl SseStreamReconnect for StreamableHttpClientReconnect { @@ -324,13 +371,25 @@ impl SseStreamReconnect for StreamableHttpClientReconne let session_id = self.session_id.clone(); let auth_header = self.auth_header.clone(); let custom_headers = self.custom_headers.clone(); + let max_sse_event_size = self.max_sse_event_size; let last_event_id = last_event_id.map(|s| s.to_owned()); Box::pin(async move { client - .get_stream(uri, session_id, last_event_id, auth_header, custom_headers) + .get_stream_with_max_sse_event_size( + uri, + session_id, + last_event_id, + auth_header, + custom_headers, + max_sse_event_size, + ) .await }) } + + fn map_fatal_stream_error(&mut self, error: SseError) -> Option { + Some(StreamableHttpError::Sse(error)) + } } /// Info retained for cleaning up the session when the worker exits. @@ -485,6 +544,7 @@ impl StreamableHttpClientWorker { uri: Arc, auth_header: Option, custom_headers: HashMap, + max_sse_event_size: usize, retry_config: Arc, ) -> impl Stream>> + Send + 'static { @@ -496,6 +556,7 @@ impl StreamableHttpClientWorker { uri, auth_header, custom_headers, + max_sse_event_size, }, retry_config, ) @@ -514,6 +575,7 @@ impl StreamableHttpClientWorker { uri: Arc, auth_header: Option, custom_headers: HashMap, + max_sse_event_size: usize, retry_config: Arc, ) -> BoxStream<'static, Result>> { match session_id { @@ -524,6 +586,7 @@ impl StreamableHttpClientWorker { uri, auth_header, custom_headers, + max_sse_event_size, retry_config, ) .boxed(), @@ -591,15 +654,17 @@ impl StreamableHttpClientWorker { uri: Arc, auth_header: Option, custom_headers: HashMap, + max_sse_event_size: usize, ) -> Result<(Option>, HashMap), StreamableHttpError> { let (init_msg, new_session_id_str) = client - .post_message( + .post_message_with_max_sse_event_size( uri.clone(), saved_init_request, None, auth_header.clone(), custom_headers.clone(), + max_sse_event_size, ) .await? .expect_initialized::() @@ -624,12 +689,13 @@ impl StreamableHttpClientWorker { &negotiated_version, ); client - .post_message( + .post_message_with_max_sse_event_size( uri, initialized_notification, new_session_id.clone(), auth_header, initialized_headers, + max_sse_event_size, ) .await? .expect_accepted_or_json::()?; @@ -670,12 +736,13 @@ impl Worker for StreamableHttpClientWorker { let saved_init_request = initialize_request.clone(); let (message, session_id) = match self .client - .post_message( + .post_message_with_max_sse_event_size( config.uri.clone(), initialize_request, None, config.auth_header.clone(), config.custom_headers.clone(), + config.max_sse_event_size, ) .await { @@ -730,12 +797,13 @@ impl Worker for StreamableHttpClientWorker { &negotiated_version, ); self.client - .post_message( + .post_message_with_max_sse_event_size( config.uri.clone(), initialized_notification.message, session_id.clone(), config.auth_header.clone(), initialized_headers, + config.max_sse_event_size, ) .await .map_err(WorkerQuitReason::fatal_context( @@ -765,15 +833,17 @@ impl Worker for StreamableHttpClientWorker { let config_uri = config.uri.clone(); let config_auth_header = config.auth_header.clone(); let spawn_headers = protocol_headers.clone(); + let max_sse_event_size = config.max_sse_event_size; streams.spawn(async move { match client - .get_stream( + .get_stream_with_max_sse_event_size( uri.clone(), session_id.clone(), None, auth_header.clone(), spawn_headers.clone(), + max_sse_event_size, ) .await { @@ -786,6 +856,7 @@ impl Worker for StreamableHttpClientWorker { uri: config_uri, auth_header: config_auth_header, custom_headers: spawn_headers, + max_sse_event_size, }, retry_config, ); @@ -855,12 +926,13 @@ impl Worker for StreamableHttpClientWorker { ); let response = self .client - .post_message( + .post_message_with_max_sse_event_size( config.uri.clone(), message.clone(), session_id.clone(), config.auth_header.clone(), request_headers, + config.max_sse_event_size, ) .await; let send_result = match response { @@ -879,6 +951,7 @@ impl Worker for StreamableHttpClientWorker { config.uri.clone(), config.auth_header.clone(), config.custom_headers.clone(), + config.max_sse_event_size, ) .await { @@ -926,14 +999,16 @@ impl Worker for StreamableHttpClientWorker { let config_uri = config.uri.clone(); let config_auth = config.auth_header.clone(); let spawn_headers = protocol_headers.clone(); + let max_sse_event_size = config.max_sse_event_size; streams.spawn(async move { match client - .get_stream( + .get_stream_with_max_sse_event_size( uri, new_sid.clone(), None, auth_header.clone(), spawn_headers.clone(), + max_sse_event_size, ) .await { @@ -946,6 +1021,7 @@ impl Worker for StreamableHttpClientWorker { uri: config_uri, auth_header: config_auth, custom_headers: spawn_headers, + max_sse_event_size, }, retry_config, ); @@ -981,12 +1057,13 @@ impl Worker for StreamableHttpClientWorker { ); let retry_response = self .client - .post_message( + .post_message_with_max_sse_event_size( config.uri.clone(), message, session_id.clone(), config.auth_header.clone(), retry_headers, + config.max_sse_event_size, ) .await; match retry_response { @@ -1021,6 +1098,7 @@ impl Worker for StreamableHttpClientWorker { config.uri.clone(), config.auth_header.clone(), protocol_headers.clone(), + config.max_sse_event_size, self.config.retry_config.clone(), ); streams.spawn(Self::execute_sse_stream( @@ -1064,6 +1142,7 @@ impl Worker for StreamableHttpClientWorker { config.uri.clone(), config.auth_header.clone(), protocol_headers.clone(), + config.max_sse_event_size, self.config.retry_config.clone(), ); streams.spawn(Self::execute_sse_stream( @@ -1339,6 +1418,12 @@ pub struct StreamableHttpClientTransportConfig { pub auth_header: Option, /// Custom HTTP headers to include with every request pub custom_headers: HashMap, + /// Maximum raw size of one SSE event accepted from the server. + /// + /// The built-in reqwest and Unix socket clients enforce this value. Custom + /// [`StreamableHttpClient`] implementations must override the corresponding + /// `*_with_max_sse_event_size` methods to enforce it. + pub max_sse_event_size: usize, /// Enables transparent recovery when the server reports an expired session (`HTTP 404`). /// /// When enabled, the transport performs one automatic recovery attempt: @@ -1397,6 +1482,12 @@ impl StreamableHttpClientTransportConfig { self } + /// Set the maximum raw size of one SSE event accepted from the server. + pub fn max_sse_event_size(mut self, bytes: usize) -> Self { + self.max_sse_event_size = bytes; + self + } + /// Set whether the transport should attempt transparent re-initialization on session expiration /// See [`Self::reinit_on_expired_session`] for details. /// # Example @@ -1420,6 +1511,7 @@ impl Default for StreamableHttpClientTransportConfig { allow_stateless: true, auth_header: None, custom_headers: HashMap::new(), + max_sse_event_size: DEFAULT_MAX_SSE_EVENT_SIZE, reinit_on_expired_session: true, } } diff --git a/crates/rmcp/src/transport/streamable_http_server.rs b/crates/rmcp/src/transport/streamable_http_server.rs index df1945ab2..4655f32f3 100644 --- a/crates/rmcp/src/transport/streamable_http_server.rs +++ b/crates/rmcp/src/transport/streamable_http_server.rs @@ -3,4 +3,6 @@ pub mod session; pub mod tower; pub use session::{RestoreOutcome, SessionId, SessionManager, SessionRestoreMarker}; #[cfg(all(feature = "transport-streamable-http-server", not(feature = "local")))] -pub use tower::{StreamableHttpServerConfig, StreamableHttpService}; +pub use tower::{ + DEFAULT_MAX_REQUEST_BODY_BYTES, StreamableHttpServerConfig, StreamableHttpService, +}; diff --git a/crates/rmcp/src/transport/streamable_http_server/tower.rs b/crates/rmcp/src/transport/streamable_http_server/tower.rs index d8fe2d426..97cd5ddfc 100644 --- a/crates/rmcp/src/transport/streamable_http_server/tower.rs +++ b/crates/rmcp/src/transport/streamable_http_server/tower.rs @@ -38,6 +38,9 @@ use crate::{ }, }; +/// Default maximum POST request body size (4 MiB). +pub const DEFAULT_MAX_REQUEST_BODY_BYTES: usize = 4 * 1024 * 1024; + #[non_exhaustive] #[derive(Debug, Clone)] pub struct StreamableHttpServerConfig { @@ -99,6 +102,12 @@ pub struct StreamableHttpServerConfig { /// }; /// ``` pub session_store: Option>, + /// Maximum POST request body size in bytes. + /// + /// Enforced while streaming the body, independent of `Content-Length`, + /// chunked transfer encoding, or HTTP version. Oversized payloads receive + /// a `413 Payload Too Large` response. + pub max_request_body_bytes: usize, } impl std::fmt::Debug for dyn SessionStore { @@ -118,6 +127,7 @@ impl Default for StreamableHttpServerConfig { allowed_hosts: vec!["localhost".into(), "127.0.0.1".into(), "::1".into()], allowed_origins: vec![], session_store: None, + max_request_body_bytes: DEFAULT_MAX_REQUEST_BODY_BYTES, } } } @@ -171,6 +181,12 @@ impl StreamableHttpServerConfig { self.cancellation_token = token; self } + + /// Set the maximum POST request body size in bytes. + pub fn with_max_request_body_bytes(mut self, bytes: usize) -> Self { + self.max_request_body_bytes = bytes; + self + } } #[expect( @@ -1140,7 +1156,7 @@ where // json deserialize request body let (part, body) = request.into_parts(); - let mut message = match expect_json(body).await { + let mut message = match expect_json(body, self.config.max_request_body_bytes).await { Ok(message) => message, Err(response) => return Ok(response), };