diff --git a/Cargo.toml b/Cargo.toml index 4db8013..682adc0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,6 +43,7 @@ uuid = { version = "1", features = ["v4"] } futures = "0.3" openapiv3 = "2" http = "1" +http-body-util = "0.1" url = "2" bytes = "1" jsonschema = { version = "0.46", default-features = false } diff --git a/src/server/adapter.rs b/src/server/adapter.rs index 2b6af0f..6953479 100644 --- a/src/server/adapter.rs +++ b/src/server/adapter.rs @@ -160,6 +160,7 @@ fn build_router(state: RouterState, extra_routes: Option) -> Router { "/mcp", crate::adapters::to_mcp_service(std::sync::Arc::new(dispatch)), ) + .layer(axum::middleware::from_fn(mcp_body_limit)) .layer(from_fn_with_state( auth_state.clone(), bearer_auth_middleware, @@ -240,6 +241,115 @@ async fn rejected_reserved_path() -> axum::response::Response { use axum::middleware::from_fn_with_state; +/// Cap the `/mcp` body at 8 MiB (feature `mcp`). +/// +/// The nested rmcp `StreamableHttpService` collects the raw request body +/// itself (`expect_json` → `body.collect()`), bypassing axum's +/// extractor-based default limit: `DefaultBodyLimit` works by inserting +/// an extension that `FromRequest` extractors consult, so it has no +/// effect on a service that reads the body directly (rmcp 1.8 +/// `server_side_http::expect_json` never checks it). This middleware is +/// both the cap and the status source: it wraps the body in a counting +/// stream that stops at [`MCP_BODY_LIMIT`] with an explicit error and +/// post-checks a exceedance flag to answer `413 Payload Too Large`, +/// replacing whatever the inner service answered (rmcp maps body-read +/// errors to `500`). +/// +/// The limit is 8 MiB — headroom over the gateway's 2 MiB whole-body +/// default for JSON-RPC batch payloads on the MCP surface. +#[cfg(feature = "mcp")] +const MCP_BODY_LIMIT: usize = 8 * 1024 * 1024; + +#[cfg(feature = "mcp")] +const MCP_BODY_LIMIT_EXCEEDED: &str = "mcp body limit exceeded"; + +#[cfg(feature = "mcp")] +async fn mcp_body_limit( + req: axum::extract::Request, + next: axum::middleware::Next, +) -> axum::response::Response { + use axum::response::IntoResponse; + + let (parts, body) = req.into_parts(); + + if let Some(len) = parts + .headers + .get(http::header::CONTENT_LENGTH) + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()) + { + if len > MCP_BODY_LIMIT { + return ( + axum::http::StatusCode::PAYLOAD_TOO_LARGE, + "Payload Too Large: /mcp body exceeds the 8 MiB limit", + ) + .into_response(); + } + } + + let exceeded = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let counting = CountingBody { + inner: body.into_data_stream(), + remaining: MCP_BODY_LIMIT, + exceeded: Arc::clone(&exceeded), + }; + + let mut limited_req = + axum::extract::Request::from_parts(parts, axum::body::Body::from_stream(counting)); + limited_req.extensions_mut().insert(exceeded.clone()); + + let response = next.run(limited_req).await; + + if exceeded.load(std::sync::atomic::Ordering::Relaxed) { + return ( + axum::http::StatusCode::PAYLOAD_TOO_LARGE, + "Payload Too Large: /mcp body exceeds the 8 MiB limit", + ) + .into_response(); + } + response +} + +#[cfg(feature = "mcp")] +struct CountingBody { + inner: axum::body::BodyDataStream, + remaining: usize, + exceeded: Arc, +} + +#[cfg(feature = "mcp")] +impl futures::Stream for CountingBody { + type Item = Result; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + let this = &mut *self; + match std::pin::Pin::new(&mut this.inner).poll_next(cx) { + std::task::Poll::Ready(Some(Ok(data))) => { + let len = data.len(); + if len > this.remaining { + this.remaining = 0; + this.exceeded + .store(true, std::sync::atomic::Ordering::Relaxed); + return std::task::Poll::Ready(Some(Err(std::io::Error::other( + MCP_BODY_LIMIT_EXCEEDED, + )))); + } + this.remaining -= len; + std::task::Poll::Ready(Some(Ok(data))) + } + std::task::Poll::Ready(Some(Err(e))) => { + let _ = e; + std::task::Poll::Ready(Some(Err(std::io::Error::other(MCP_BODY_LIMIT_EXCEEDED)))) + } + std::task::Poll::Pending => std::task::Poll::Pending, + std::task::Poll::Ready(None) => std::task::Poll::Ready(None), + } + } +} + #[async_trait] impl alkcall::core::types::ProtocolHandler for HttpAdapter { fn alpn(&self) -> &'static [u8] { @@ -576,6 +686,94 @@ mod tests { let _ = server_task.await; } + #[cfg(feature = "mcp")] + #[tokio::test] + async fn mcp_rejects_oversized_body_declared_content_length_with_413() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let oversized = MCP_BODY_LIMIT + 1; + let head = format!( + "POST /mcp HTTP/1.1\r\nHost: localhost\r\nContent-Type: application/json\r\nAccept: application/json, text/event-stream\r\nContent-Length: {oversized}\r\nConnection: close\r\n\r\n" + ); + + let adapter = HttpAdapter::new(provider(), empty_registry()); + let (client, server) = tokio::io::duplex(64 * 1024); + let conn = Connection::from_bidi(server, b"http/1.1".to_vec(), None); + let auth = AuthContext::anonymous(b"http/1.1"); + let server_task = tokio::spawn(async move { + let _ = ProtocolHandler::handle(&adapter, conn, &auth).await; + }); + + let mut client = client; + client.write_all(head.as_bytes()).await.unwrap(); + + let mut response = Vec::new(); + let _ = tokio::time::timeout( + std::time::Duration::from_secs(5), + client.read_to_end(&mut response), + ) + .await + .expect("read timed out") + .unwrap(); + + let text = String::from_utf8_lossy(&response); + assert!( + text.starts_with("HTTP/1.1 413 Payload Too Large"), + "declared oversized body got: {text}" + ); + + let _ = server_task.await; + } + + #[cfg(feature = "mcp")] + #[tokio::test] + async fn mcp_rejects_oversized_chunked_body_with_413() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let adapter = HttpAdapter::new(provider(), empty_registry()); + let (client, server) = tokio::io::duplex(64 * 1024); + let conn = Connection::from_bidi(server, b"http/1.1".to_vec(), None); + let auth = AuthContext::anonymous(b"http/1.1"); + let server_task = tokio::spawn(async move { + let _ = ProtocolHandler::handle(&adapter, conn, &auth).await; + }); + + let mut client = client; + client + .write_all( + b"POST /mcp HTTP/1.1\r\nHost: localhost\r\nContent-Type: application/json\r\nAccept: application/json, text/event-stream\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n", + ) + .await + .unwrap(); + + let chunk_size = 64 * 1024; + for _ in 0..(MCP_BODY_LIMIT / chunk_size) + 1 { + let header = format!("{chunk_size:x}\r\n"); + client.write_all(header.as_bytes()).await.unwrap(); + let chunk = vec![b'a'; chunk_size]; + client.write_all(&chunk).await.unwrap(); + client.write_all(b"\r\n").await.unwrap(); + } + client.write_all(b"0\r\n\r\n").await.unwrap(); + + let mut response = Vec::new(); + let _ = tokio::time::timeout( + std::time::Duration::from_secs(10), + client.read_to_end(&mut response), + ) + .await + .expect("read timed out") + .unwrap(); + + let text = String::from_utf8_lossy(&response); + assert!( + text.starts_with("HTTP/1.1 413 Payload Too Large"), + "chunked oversized body got: {text}" + ); + + let _ = server_task.await; + } + struct StaticProvider; impl IdentityProvider for StaticProvider { fn resolve_from_fingerprint(&self, _: &str) -> Option {