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 {