From bd4bfd112b2f6e40cebd8d0e3d34a84a8ed50989 Mon Sep 17 00:00:00 2001 From: Raymond Khalife Date: Thu, 12 Nov 2026 15:17:53 +0000 Subject: [PATCH] Cap the MCP request body behind the bearer gate The streamable-HTTP transport reads each request body whole into memory before parsing it as JSON-RPC, and axum's default extractor limit never reaches it (the transport is a tower service, not an extractor), so the hosted /mcp route had no bound on what one request could make the process buffer. Caddy in front sets a read timeout but no size cap. A middleware layer inside the bearer gate now refuses a body over the cap with 413: a declared Content-Length over it is refused unread, and an undeclared (chunked) body is read up to the cap and no further. The gate still answers an unauthenticated request on its headers alone, so only an authenticated caller's body is read at all. Default 4 MiB (an IC ingress message is at most 2 MiB, so any one canister call fits even spelled out as JSON); embedders set it with McpServer::with_max_request_body, the binary from $MCP_MAX_REQUEST_BODY. Co-Authored-By: Claude Fable 5.1 --- Cargo.lock | 1 + Cargo.toml | 7 ++- README.md | 3 +- src/auth.rs | 19 ++++++ src/body_limit.rs | 151 ++++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 74 ++++++++++++++++++++++- src/main.rs | 23 ++++++- 7 files changed, 273 insertions(+), 5 deletions(-) create mode 100644 src/body_limit.rs diff --git a/Cargo.lock b/Cargo.lock index 9087355..ea0302b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1798,6 +1798,7 @@ dependencies = [ "axum", "base64 0.22.1", "candid", + "futures-util", "getrandom 0.3.4", "hex", "http-body-util", diff --git a/Cargo.toml b/Cargo.toml index 7d0dab9..1ab1fd7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -102,6 +102,9 @@ tracing-subscriber = { workspace = true } tokio-util = "0.7" uuid = { workspace = true } tower-http = { version = "0.7", features = ["cors"] } +# `LengthLimitError`, to tell a body over the MCP request cap from a transport +# error (src/body_limit.rs). Already in the tree via axum. +http-body-util = "0.1" # Metrics exposition at /metrics. Pinned to the version dfinity/ic uses: these # series are destined for the same Prometheus / Victoria Metrics clusters that # scrape the IC, so matching the org's client keeps one exposition dialect in the @@ -130,7 +133,9 @@ e2e = ["dep:pocket-ic"] # Router-level contract tests (tests/routers.rs) drive the composed app with # tower::ServiceExt::oneshot and read bodies via http-body-util. tower = { version = "0.5", features = ["util"] } -http-body-util = "0.1" +# Streamed (chunked) request bodies for the body-cap tests. Already in the +# tree via tokio and hyper. +futures-util = "0.3" # `test-util` is what `#[tokio::test(start_paused = true)]` needs: the pending # connect tests in `auth.rs` walk a connect past its TTL and retention on a # paused clock instead of waiting them out. A feature on a dependency the crate diff --git a/README.md b/README.md index dd9013f..d35d246 100644 --- a/README.md +++ b/README.md @@ -569,7 +569,8 @@ cargo run # serves http://0.0.0.0:8000 (MCP at /mcp against production II, OAuth under it; / redirects to the landing site) # honours $PORT (default 8000), $PUBLIC_URL (default http://localhost:8000), # $MCP_SERVE_BETA (set it to also serve the beta II instance at /mcp-beta, for staging), -# and $MCP_SERVE_METRICS (set it to serve the Prometheus exposition at /metrics) +# $MCP_SERVE_METRICS (set it to serve the Prometheus exposition at /metrics), +# and $MCP_MAX_REQUEST_BODY (the cap on an MCP request body in bytes, default 4 MiB) ``` The human-facing pages — the landing page and the `/privacy-policy`, diff --git a/src/auth.rs b/src/auth.rs index fafe0f8..88508b5 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -3073,6 +3073,25 @@ fn token_ttl(default: Duration, grant_expiration_ns: Option, now_ns: u64) - } } +#[cfg(test)] +impl AuthStore { + /// A live bearer token for a made-up session, so a test in the crate root + /// can get past [`require_token`] without an Internet Identity handshake. + pub(crate) async fn test_bearer(&self) -> String { + let token = format!("mcp-token-test-{}", Uuid::new_v4()); + self.tokens.write().await.insert( + token.clone(), + TokenInfo { + principal: "test-principal".into(), + session_id: "test-session".into(), + created: Instant::now(), + ttl: Duration::from_secs(600), + }, + ); + token + } +} + /// Mint + store an access token bound to the session key's principal, its /// lifetime matched to the II grant's expiration (see [`token_ttl`]). async fn issue_token(store: &AuthStore, session_id: &str) -> Response { diff --git a/src/body_limit.rs b/src/body_limit.rs new file mode 100644 index 0000000..6ce607a --- /dev/null +++ b/src/body_limit.rs @@ -0,0 +1,151 @@ +//! A cap on the size of an MCP request body. +//! +//! The streamable-HTTP transport reads each request body whole into memory +//! before parsing it as JSON-RPC, and axum's default extractor limit never +//! reaches it (the transport is a tower service, not an extractor), so without +//! this layer a token holder could send a body of any size and have the process +//! buffer all of it. The layer sits INSIDE the bearer gate: an unauthenticated +//! request is turned away on its headers alone, so only an authenticated +//! caller's body is read, and only up to the cap. + +use axum::{ + body::Body, + extract::{Request, State}, + http::{header, StatusCode}, + middleware::Next, + response::{IntoResponse, Response}, +}; + +/// The default cap on an MCP request body, in bytes. An IC ingress message +/// carries at most 2 MiB, so any one canister call fits well under this even +/// once its argument is spelled out as JSON. +pub const DEFAULT_MAX_REQUEST_BODY: usize = 4 * 1024 * 1024; + +/// Refuse a request whose body is over `max` bytes with `413 Payload Too +/// Large`, before buffering more than that of it: a declared `Content-Length` +/// over the cap is refused unread, and an undeclared (chunked) body is read up +/// to the cap and no further. A body within the cap is handed on whole, which +/// costs nothing extra since the transport buffers it anyway. +pub(crate) async fn limit_request_body( + State(max): State, + request: Request, + next: Next, +) -> Response { + let (parts, body) = request.into_parts(); + let declared = parts + .headers + .get(header::CONTENT_LENGTH) + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.trim().parse::().ok()); + if declared.is_some_and(|len| len > max as u64) { + return too_large(max); + } + let bytes = match axum::body::to_bytes(body, max).await { + Ok(bytes) => bytes, + Err(e) => { + return if e.into_inner().is::() { + too_large(max) + } else { + (StatusCode::BAD_REQUEST, "could not read the request body").into_response() + }; + } + }; + next.run(Request::from_parts(parts, Body::from(bytes))).await +} + +fn too_large(max: usize) -> Response { + (StatusCode::PAYLOAD_TOO_LARGE, format!("request body exceeds the {max}-byte limit")) + .into_response() +} + +#[cfg(test)] +mod tests { + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }; + + use axum::{ + body::{Body, Bytes}, + http::{header, Request, StatusCode}, + middleware, + routing::post, + Router, + }; + use tower::ServiceExt; + + const MAX: usize = 1024; + + /// A service behind the cap that answers with how many body bytes reached it. + fn app() -> Router { + Router::new() + .route("/", post(|body: Bytes| async move { body.len().to_string() })) + .layer(middleware::from_fn_with_state(MAX, super::limit_request_body)) + } + + /// A body streamed in `chunk`-byte pieces with no `Content-Length`, counting + /// how many pieces the server pulled. + fn chunked(total: usize, chunk: usize) -> (Body, Arc) { + let pulled = Arc::new(AtomicUsize::new(0)); + let counter = pulled.clone(); + let pieces = (0..total.div_ceil(chunk)).map(move |i| { + counter.fetch_add(1, Ordering::SeqCst); + Ok::<_, std::convert::Infallible>(Bytes::from(vec![b'x'; chunk.min(total - i * chunk)])) + }); + (Body::from_stream(futures_util::stream::iter(pieces)), pulled) + } + + async fn text(resp: axum::response::Response) -> String { + let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX).await.unwrap(); + String::from_utf8(bytes.to_vec()).unwrap() + } + + #[tokio::test] + async fn a_body_up_to_the_cap_reaches_the_service_whole() { + for len in [0, 1, MAX] { + let resp = app() + .oneshot(Request::post("/").body(Body::from(vec![b'x'; len])).unwrap()) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK, "{len} bytes"); + assert_eq!(text(resp).await, len.to_string(), "{len} bytes"); + } + } + + #[tokio::test] + async fn a_declared_length_over_the_cap_is_refused_unread() { + let (body, pulled) = chunked(MAX + 1, 16); + let resp = app() + .oneshot( + Request::post("/") + .header(header::CONTENT_LENGTH, (MAX + 1).to_string()) + .body(body) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE); + assert_eq!(pulled.load(Ordering::SeqCst), 0, "the body must not be read"); + assert!(text(resp).await.contains(&format!("{MAX}-byte limit"))); + } + + #[tokio::test] + async fn an_undeclared_body_over_the_cap_is_refused_at_the_cap() { + let (body, pulled) = chunked(8 * MAX, 16); + let resp = app().oneshot(Request::post("/").body(body).unwrap()).await.unwrap(); + assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE); + let read = pulled.load(Ordering::SeqCst) * 16; + assert!( + read > MAX && read <= MAX + 16, + "reading must stop at the cap, not drain the body: {read} bytes pulled" + ); + } + + #[tokio::test] + async fn an_undeclared_body_within_the_cap_passes() { + let (body, _) = chunked(MAX, 16); + let resp = app().oneshot(Request::post("/").body(body).unwrap()).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(text(resp).await, MAX.to_string()); + } +} diff --git a/src/lib.rs b/src/lib.rs index 0d72665..90149a4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -74,6 +74,7 @@ //! and `II_URL_PROD` / `II_CANISTER_ID_PROD` (the Internet Identity instances). mod auth; +mod body_limit; mod branding; // Docs live in the module itself (`//!` in src/metrics.rs): an outer doc // comment here would resolve its intra-doc links in *this* scope rather than @@ -108,6 +109,7 @@ use serde_json::json; use tokio_util::sync::CancellationToken; pub use auth::SharedClients; +pub use body_limit::DEFAULT_MAX_REQUEST_BODY; /// The components this server is composed from, re-exported so embedders keep /// importing everything from `imcp2` (the types are `imcp2-core`'s — one /// defining crate, so type identity is preserved). @@ -170,6 +172,9 @@ pub struct McpServer { store: auth::AuthStore, public_url: String, mcp_path: String, + /// The largest MCP request body the endpoint reads, in bytes (see + /// [`Self::with_max_request_body`]). + max_request_body: usize, /// Cancels the streamable-HTTP sessions and the session reaper on /// [`Self::shutdown`]. ct: CancellationToken, @@ -205,10 +210,22 @@ impl McpServer { store, public_url, mcp_path, + max_request_body: DEFAULT_MAX_REQUEST_BODY, ct: CancellationToken::new(), } } + /// Cap the MCP request body at `bytes` instead of the default + /// [`DEFAULT_MAX_REQUEST_BODY`]. The streamable-HTTP transport buffers each + /// request body whole before parsing it, so the cap is what bounds the + /// memory one authenticated request can take; a body over it is refused + /// with `413` before more than the cap is read. The `imcp2` binary sets + /// it from `$MCP_MAX_REQUEST_BODY`. + pub fn with_max_request_body(mut self, bytes: usize) -> Self { + self.max_request_body = bytes; + self + } + /// The operational-files directory this instance uses — where imcp2 creates /// the state a restart must survive (today, the client-registration store). /// Reported straight from the client store, so it is the directory files @@ -256,7 +273,8 @@ impl McpServer { /// fallback, so the bare mount path, its trailing-slash form, and /// unmatched sub-paths all reach it), bearer-token gated, with the CORS /// preflight answered before authentication and `WWW-Authenticate` - /// exposed cross-origin. + /// exposed cross-origin, and its request body capped behind the gate + /// (see [`Self::with_max_request_body`]). /// /// Nest with `nest_service` (which forwards the bare trailing-slash form /// too, unlike `nest`). Nesting at the application root is the one axum @@ -317,8 +335,16 @@ impl McpServer { axum::http::header::WWW_AUTHENTICATE, axum::http::HeaderName::from_static("mcp-session-id"), ]); + // Layers run outermost-last: CORS, then the bearer gate, then the body + // cap — so an unauthenticated request is refused on its headers without + // its body being read, and an authenticated one is read only up to the + // cap before the transport buffers it. let gated_mcp = Router::new() .fallback_service(mcp_service) + .layer(middleware::from_fn_with_state( + self.max_request_body, + body_limit::limit_request_body, + )) .layer(middleware::from_fn_with_state(self.store.clone(), auth::require_token)) .layer(mcp_cors); @@ -633,7 +659,12 @@ mod lib_tests { allowed_hosts_for, normalize_mount_path, normalize_public_url, Agent, IiInstance, McpConfig, McpServer, SharedClients, IC_URL, }; + use axum::{ + body::Body, + http::{Request, StatusCode}, + }; use std::path::PathBuf; + use tower::ServiceExt; /// Build a server whose `clients` load and `state_dir` field both point at /// `dir` (the correct pairing). Construction is pure — no network. @@ -745,4 +776,45 @@ mod lib_tests { assert_eq!(hosts.iter().filter(|h| *h == "localhost").count(), 1); assert!(has(&hosts, "localhost:8000"), "{hosts:?}"); } + + /// The body cap guards the MCP endpoint from behind the bearer gate: an + /// authenticated oversized request gets `413`, an unauthenticated one still + /// gets the `401` challenge with its body left unread. + #[tokio::test] + async fn the_mcp_endpoint_caps_authenticated_request_bodies() { + let dir = std::env::temp_dir().join("imcp2-lib-body-limit-test"); + let server = server_with_dir(dir).with_max_request_body(64); + let bearer = format!("Bearer {}", server.store.test_bearer().await); + let oversized = || Body::from(vec![b'x'; 65]); + + let resp = server + .mcp_router() + .oneshot(Request::post("/").header("authorization", &bearer).body(oversized()).unwrap()) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE); + + let resp = server + .mcp_router() + .oneshot(Request::post("/").body(oversized()).unwrap()) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED, "the gate comes first"); + + // Within the cap and authenticated: the transport answers, not the cap. + let resp = server + .mcp_router() + .oneshot( + Request::post("/") + .header("authorization", &bearer) + .header("content-type", "application/json") + .header("accept", "application/json, text/event-stream") + .body(Body::from(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_ne!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE); + assert_ne!(resp.status(), StatusCode::UNAUTHORIZED); + } } diff --git a/src/main.rs b/src/main.rs index d2418bc..7d6b7b8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -69,6 +69,20 @@ fn require_resource() -> bool { } } +/// The largest MCP request body the server reads, in bytes: `$MCP_MAX_REQUEST_BODY` +/// when set, else the library default ([`imcp2::DEFAULT_MAX_REQUEST_BODY`]). A +/// value that is not a whole number is a configuration error and refuses to +/// start, rather than silently running with the default. +fn max_request_body() -> anyhow::Result { + match std::env::var("MCP_MAX_REQUEST_BODY") { + Ok(v) => v + .trim() + .parse() + .map_err(|e| anyhow::anyhow!("MCP_MAX_REQUEST_BODY must be a byte count: {e}")), + Err(_) => Ok(imcp2::DEFAULT_MAX_REQUEST_BODY), + } +} + /// Whether to also serve the beta Internet Identity instance at `/mcp-beta`. /// Off unless `$MCP_SERVE_BETA` is truthy (`1`/`true`/`yes`/`on`), so a /// production deployment serves only `/mcp` (production II) and the staging @@ -309,6 +323,9 @@ async fn main() -> anyhow::Result<()> { // the operational directory. let clients = SharedClients::load(&state_dir); + // One request-body cap for every instance (see `max_request_body`). + let max_request_body = max_request_body()?; + // Client ID Metadata Documents: on for every instance unless the deploy says no. let cimd_on = cimd_enabled(); if !cimd_on { @@ -327,7 +344,8 @@ async fn main() -> anyhow::Result<()> { state_dir: state_dir.clone(), require_resource: require_resource(), cimd_enabled: cimd_on, - }); + }) + .with_max_request_body(max_request_body); prod.spawn_session_reaper(); // Beta Internet Identity at `/mcp-beta`: opt in via `$MCP_SERVE_BETA`, so @@ -342,7 +360,8 @@ async fn main() -> anyhow::Result<()> { state_dir, require_resource: require_resource(), cimd_enabled: cimd_on, - }); + }) + .with_max_request_body(max_request_body); beta.spawn_session_reaper(); Some(beta) } else {