diff --git a/src/server/adapter.rs b/src/server/adapter.rs index 6953479..48c4126 100644 --- a/src/server/adapter.rs +++ b/src/server/adapter.rs @@ -738,32 +738,43 @@ mod tests { let _ = ProtocolHandler::handle(&adapter, conn, &auth).await; }); - let mut client = client; - client + let (mut reader_client, mut writer_client) = tokio::io::split(client); + writer_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 writer = tokio::spawn(async move { + let chunk = vec![b'a'; 64 * 1024]; + let chunk_header = format!("{:x}\r\n", chunk.len()); + for _ in 0..(MCP_BODY_LIMIT / chunk.len()) + 1 { + if writer_client + .write_all(chunk_header.as_bytes()) + .await + .is_err() + { + return; + } + if writer_client.write_all(&chunk).await.is_err() { + return; + } + if writer_client.write_all(b"\r\n").await.is_err() { + return; + } + } + let _ = writer_client.write_all(b"0\r\n\r\n").await; + }); let mut response = Vec::new(); - let _ = tokio::time::timeout( + let read = tokio::time::timeout( std::time::Duration::from_secs(10), - client.read_to_end(&mut response), + reader_client.read_to_end(&mut response), ) - .await - .expect("read timed out") - .unwrap(); + .await; + writer.abort(); + read.expect("read timed out").unwrap(); let text = String::from_utf8_lossy(&response); assert!( diff --git a/tasks/server/review-001-mcp-body-limit.md b/tasks/server/review-001-mcp-body-limit.md index 219e1c0..0a95224 100644 --- a/tasks/server/review-001-mcp-body-limit.md +++ b/tasks/server/review-001-mcp-body-limit.md @@ -1,7 +1,7 @@ --- id: review-001-mcp-body-limit name: Cap the /mcp nest body size (SRV-03, gated on mcp feature) -status: pending +status: completed depends_on: [] scope: narrow risk: low @@ -17,8 +17,8 @@ nests rmcp's `StreamableHttpService`; the bearer middleware only stashes identity. axum's 2 MiB `DefaultBodyLimit` applies to axum *extractors*, but the nested rmcp service collects the raw body itself (`body.collect().await`, verified against rmcp 1.8.0) with no cap — a -single multi-GB chunked `POST /mcp` is buffered entirely in memory; a few -concurrent requests OOM the process. Gateway routes are correctly capped +single multi-GB chunked `POST /mcp` is buffered entirely in memory; a +few concurrent requests OOM the process. Gateway routes are correctly capped at 2 MiB by the extractors; `/mcp` is the one uncapped surface. Fix: wrap the `/mcp` nest with an explicit `DefaultBodyLimit` (or an @@ -29,10 +29,10 @@ default. ## Acceptance Criteria -- [ ] Explicit body limit applied to the `/mcp` nest; limit value documented in code and ADR-039 or http-server.md if touched -- [ ] Test: oversized `POST /mcp` body → 413 (Content-Length and streaming/chunked variants) -- [ ] Normal-size MCP initialize + tools/call round-trip still passes -- [ ] `cargo test --all-features` passes (feature-gated code) +- [x] Explicit body limit applied to the `/mcp` nest; limit value documented in code and ADR-039 or http-server.md if touched +- [x] Test: oversized `POST /mcp` body → 413 (Content-Length and streaming/chunked variants) +- [x] Normal-size MCP initialize + tools/call round-trip still passes +- [x] `cargo test --all-features` passes (feature-gated code) ## References @@ -41,8 +41,52 @@ default. ## Notes -> Agent fills during implementation. +`DefaultBodyLimit` was rejected on source evidence, not guesswork: its +`Layer` implementation (axum-core 0.5.6 +`extract/default_body_limit.rs`) only inserts an extension into the +request — the limit is enforced inside `FromRequest` extractors. rmcp's +`StreamableHttpService` consumes the raw body via +`server_side_http::expect_json` → `body.collect()` and never checks that +extension, so the layer cannot intercept a raw-body-collecting nested +service. A bare `http-body-util::Limited` wrapper was also insufficient: +its body-read error surfaces through rmcp as `500` (rmcp maps +`expect_json` body errors to `INTERNAL_SERVER_ERROR`), not `413`. + +The implemented fix is a small `from_fn` middleware (`mcp_body_limit`, +feature-gated in `src/server/adapter.rs`) layered directly on the +`/mcp` nest, inner to the bearer-auth layer: + +1. A `Content-Length`-declared body over the cap is rejected with `413` + before any body bytes are read. +2. Otherwise the body is wrapped in a counting stream + (`CountingBody`, 8 MiB budget) that flags exceedance in a request + extension and errors on the first chunk that crosses the cap; the + middleware post-checks the flag and answers `413` regardless of what + the inner service produced. The cut-off also bounds buffering — + rmcp never sees a body larger than the cap plus one chunk. + +Limit choice: **8 MiB** (`MCP_BODY_LIMIT`), headroom over the gateway's +2 MiB whole-body default for legitimate JSON-RPC batch payloads on the +MCP surface; documented on the constant. + +http-server.md §"What" lists the `/mcp` surface but does not document +per-surface body limits anywhere (the gateway's 2 MiB is extractor +behavior, not prose), so no doc addition was needed; the constant's doc +comment is the normative statement. ## Summary -> Filled on completion. \ No newline at end of file +Capped the `/mcp` nest (SRV-03): an 8 MiB `mcp_body_limit` middleware +now wraps the rmcp `StreamableHttpService` nest in `build_router` +(feature `mcp`), inner to the bearer-auth layer, leaving auth ordering +and the reserved-paths probe untouched. `DefaultBodyLimit` does not +apply (extension-only mechanism; rmcp collects the raw body), so the +middleware both caps and status-sources: declared oversizes are rejected +before reading, streaming oversizes are cut mid-read and answered `413`. +Tests: `mcp_rejects_oversized_body_declared_content_length_with_413`, +`mcp_rejects_oversized_chunked_body_with_413` (raw chunked HTTP/1.1 over +the duplex socket, concurrent write/read split), and the normal-size +`mcp_endpoint_serves_four_gateway_tools_bearer_gated` round-trip stays +green. Verified: `cargo test` (219), `cargo test --features mcp --lib` +(257), `cargo test --all-features` (269 + integration suites), clippy +default and `--all-features` `-D warnings`, `cargo fmt --check`. \ No newline at end of file