fix(gateway): cap /publish body buffering + explicit body-limit layer (GW-15)

Review 002 GW-15 [major]: /publish lost both claimed memory bounds in
the GW-06 streaming rewrite. Unauthenticated POST /publish with chunked
'a'-forever (no newline) grew the heap with the upload until OOM, and
each poll re-scanned the whole buffer (O(n^2) on top).

- BufferedLines: cap the unterminated tail against
  MAX_PUBLISH_LINE_BYTES immediately after every chunk read (checked
  before yielding, even with no \n seen) and re-check on the
  trailing-EOF mem::take path; a breach aborts the whole reader
  (pending lines included) with the same terminal INVALID_INPUT
  LineCap error. Complete lines keep the baseline at-cap semantics.

- Whole gateway router: explicit request-body-limit layer
  (GATEWAY_BODY_LIMIT = 2 MiB + 64 KiB framing headroom). A raw-Body
  handler never consults axum's DefaultBodyLimit (that is an extension
  extractors read), so /publish had no whole-body cap at all. The layer
  pre-rejects oversized declared Content-Length and wraps chunked
  uploads in a counting stream; both answer plain-text 413. It
  deliberately sits above the per-line cap so a single over-cap line
  still surfaces the semantic line-cap error. Upstream body read
  failures are not flagged as limit-exceeded (disconnects are not 413).

- New wire tests: streamed never-newline over-cap rejections (both the
  streamed multi-chunk and trailing-EOF shapes), 413 on chunked
  over-limit uploads, 413 on oversized declared Content-Length, cap
  breach batched with complete lines, at-cap line still round-trips,
  declared-length over-limit pre-rejection.

Module status mapping note (GW-16 tracks the drift): hand-rolled
pre-dispatch rejections use 400/INVALID_INPUT while mid-stream chunk
errors map 422 through gateway::error; normalizing is GW-16, not GW-15.

Verification: cargo test 305+5 pass, cargo test --all-features pass,
cargo clippy --all-targets -- -D warnings pass (--all-features too),
cargo fmt --check pass.
This commit is contained in:
2026-08-30 20:29:35 +00:00
parent e2c255d40c
commit 05ebbef43e
+351 -15
View File
@@ -7,6 +7,17 @@
//! `gateway::error`. There is no per-operation `POST /{service}/{op}`
//! direct-call surface (ADR-047). `/publish` lives in this module too
//! (the review-001 "separate module" note is stale).
//!
//! The whole router carries an explicit request-body limit (2 MiB plus
//! 64 KiB framing headroom — GW-15): a raw-`Body` handler such as
//! `/publish`'s never consults axum's `DefaultBodyLimit` (that limit is
//! an extension extractors read), so the gateway installs its own cap
//! instead.
//!
//! Module status mapping note (GW-16 tracks the drift): the
//! hand-rolled pre-dispatch rejections below use 400/`INVALID_INPUT`
//! while mid-stream chunk errors map 422 through `gateway::error`;
//! normalizing them is GW-16, not GW-15.
use std::collections::VecDeque;
use std::convert::Infallible;
@@ -42,6 +53,25 @@ const SERVICES_SCHEMA: &str = "services/schema";
const MAX_BATCH_OPERATIONS: usize = 100;
const MAX_PUBLISH_LINE_BYTES: usize = 2 * 1024 * 1024;
/// The explicit request-body limit for the whole gateway router
/// (GW-15): the 2 MiB convention plus a 64 KiB framing headroom, so a
/// request can carry an at-cap line (2 MiB) plus its NDJSON first-line
/// header framing without the layer pre-empting the route's own
/// per-line cap error — the semantic 400/`INVALID_INPUT` the cap
/// exists to emit. axum's `DefaultBodyLimit` is a request extension
/// consulted by `FromRequest` extractors; the raw-`Body` hand-rolled
/// `/publish` path never consults it, so without this layer that route
/// has no whole-body cap at all. The layer pre-rejects an oversized
/// declared `Content-Length` and wraps the body in a counting stream
/// that errors when chunked (undeclared-length) uploads exceed the
/// limit; the middleware answers `413 Payload Too Large` in both
/// cases. This is a backstop behind the per-line cap, not a second
/// semantic: the per-line cap fires first on any single over-cap
/// line.
const GATEWAY_BODY_LIMIT: usize = MAX_PUBLISH_LINE_BYTES + 64 * 1024;
const GATEWAY_BODY_LIMIT_EXCEEDED: &str = "gateway request body limit exceeded";
/// SSE keep-alive interval on `/subscribe` (GW-13). Shared comment
/// frames (axum's KeepAlive::default) plus a `retry:` field on the
/// stream's first event reconnect the client on drops; 15 s sits under
@@ -96,6 +126,7 @@ pub(crate) fn gateway_router() -> Router<RouterState> {
.route("/batch", post(batch_handler))
.route("/subscribe", post(subscribe_handler))
.route("/publish", post(publish_handler))
.layer(axum::middleware::from_fn(gateway_body_limit))
}
/// The `/call` and `/subscribe` request body: the target operation
@@ -387,11 +418,96 @@ fn invalid_input_response(message: &str) -> Response {
.into_response()
}
/// The explicit gateway request-body limit (GW-15). See
/// [`GATEWAY_BODY_LIMIT`] for why this is a layer and not axum's
/// `DefaultBodyLimit` extension.
async fn gateway_body_limit(req: axum::extract::Request, next: axum::middleware::Next) -> Response {
let (parts, body) = req.into_parts();
if let Some(len) = parts
.headers
.get(axum::http::header::CONTENT_LENGTH)
.and_then(|v| v.to_str().ok())
.and_then(|v| v.parse::<usize>().ok())
{
if len > GATEWAY_BODY_LIMIT {
return gateway_body_limit_exceeded();
}
}
let exceeded = Arc::new(std::sync::atomic::AtomicBool::new(false));
let limited_req = axum::extract::Request::from_parts(
parts,
axum::body::Body::from_stream(LimitedBody {
inner: body.into_data_stream(),
remaining: GATEWAY_BODY_LIMIT,
exceeded: Arc::clone(&exceeded),
}),
);
let response = next.run(limited_req).await;
if exceeded.load(std::sync::atomic::Ordering::Relaxed) {
return gateway_body_limit_exceeded();
}
response
}
fn gateway_body_limit_exceeded() -> Response {
(
StatusCode::PAYLOAD_TOO_LARGE,
"gateway request body limit exceeded",
)
.into_response()
}
struct LimitedBody {
inner: axum::body::BodyDataStream,
remaining: usize,
exceeded: Arc<std::sync::atomic::AtomicBool>,
}
impl futures::Stream for LimitedBody {
type Item = Result<Bytes, std::io::Error>;
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
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(
GATEWAY_BODY_LIMIT_EXCEEDED,
))));
}
this.remaining -= len;
std::task::Poll::Ready(Some(Ok(data)))
}
std::task::Poll::Ready(Some(Err(_))) => std::task::Poll::Ready(Some(Err(
std::io::Error::other(GATEWAY_BODY_LIMIT_EXCEEDED),
))),
std::task::Poll::Pending => std::task::Poll::Pending,
std::task::Poll::Ready(None) => std::task::Poll::Ready(None),
}
}
}
/// A newline-framed reader over the request body's byte stream. Blank
/// lines are skipped; a trailing unterminated line is yielded as a
/// final line. Lines are capped at `MAX_PUBLISH_LINE_BYTES` (in
/// addition to axum's own 2 MiB default body limit); a byte-level read
/// failure surfaces as a terminal error.
/// final line. Each chunk read extends the buffer, drains every
/// complete `\n`-terminated line into the pending queue (capped at
/// `MAX_PUBLISH_LINE_BYTES`), then caps the unterminated tail **even
/// when no `\n` was seen** (GW-15): the tail can therefore never grow
/// past the cap across chunks, and the trailing-EOF `mem::take` path
/// re-checks the cap before yielding. A cap breach or a byte-level
/// read failure is terminal and aborts the whole reader (pending
/// lines included — a breach poisons the body framing).
struct BufferedLines {
bytes: ByteStream,
buffer: Vec<u8>,
@@ -409,6 +525,13 @@ impl BufferedLines {
}
}
fn abort_on_cap(&mut self) -> LineError {
self.done = true;
self.buffer.clear();
self.pending.clear();
LineError::LineCap
}
async fn next_line(&mut self) -> Result<Option<Vec<u8>>, LineError> {
loop {
if let Some(line) = self.pending.pop_front() {
@@ -422,10 +545,26 @@ impl BufferedLines {
self.buffer.clear();
return Ok(None);
}
if self.buffer.len() > MAX_PUBLISH_LINE_BYTES {
return Err(self.abort_on_cap());
}
return Ok(Some(std::mem::take(&mut self.buffer)));
}
match self.bytes.next().await {
Some(Ok(bytes)) => self.buffer.extend_from_slice(&bytes),
Some(Ok(bytes)) => {
self.buffer.extend_from_slice(&bytes);
while let Some(pos) = self.buffer.iter().position(|b| *b == b'\n') {
let line: Vec<u8> = self.buffer.drain(..=pos).collect();
let line = &line[..line.len() - 1];
if line.len() > MAX_PUBLISH_LINE_BYTES {
return Err(self.abort_on_cap());
}
self.pending.push_back(line.to_vec());
}
if self.buffer.len() > MAX_PUBLISH_LINE_BYTES {
return Err(self.abort_on_cap());
}
}
Some(Err(_)) => {
self.done = true;
return Err(LineError::Read);
@@ -435,17 +574,6 @@ impl BufferedLines {
continue;
}
}
while let Some(pos) = self.buffer.iter().position(|b| *b == b'\n') {
let line: Vec<u8> = self.buffer.drain(..=pos).collect();
let line = &line[..line.len() - 1];
if line.len() > MAX_PUBLISH_LINE_BYTES {
self.done = true;
self.buffer.clear();
self.pending.clear();
return Err(LineError::LineCap);
}
self.pending.push_back(line.to_vec());
}
}
}
}
@@ -2346,6 +2474,214 @@ mod tests {
);
}
fn registry_with_cap_witness_sink() -> Arc<OperationRegistry> {
let mut registry = OperationRegistry::new();
registry
.register(HandlerRegistration::new(
OperationSpec::new(
"ingest/big",
OperationType::Pub,
Visibility::External,
json!({}),
json!({}),
vec![],
AccessControl::default(),
None,
),
HandlerKind::Sink(make_sink_handler(|_input, ctx, mut chunks| async move {
use futures::StreamExt;
let mut last_error = None;
while let Some(chunk) = chunks.next().await {
if let Err(e) = chunk {
last_error = Some(e);
break;
}
}
let error = last_error.expect("a capped line must produce an error item");
ResponseEnvelope::error(ctx.request_id, error)
})),
OperationProvenance::Local,
None,
None,
Capabilities::new(),
))
.unwrap();
Arc::new(registry)
}
#[tokio::test]
async fn publish_streamed_never_newline_body_over_cap_is_rejected_pre_extend() {
let router = build_router(registry_with_cap_witness_sink(), unused_provider());
let first_line = serde_json::to_vec(&json!({
"operation": "ingest/big",
"chunk": { "n": 1 }
}))
.unwrap();
let mut body = first_line;
body.push(b'\n');
body.extend_from_slice(&vec![b'a'; MAX_PUBLISH_LINE_BYTES + 1]);
let chunks: Vec<Bytes> = body.chunks(64 * 1024).map(Bytes::copy_from_slice).collect();
let req = Request::builder()
.method("POST")
.uri("/publish")
.header("content-type", "application/x-ndjson")
.header("transfer-encoding", "chunked")
.body(Body::from_stream(futures::stream::iter(
chunks.into_iter().map(Ok::<_, std::convert::Infallible>),
)))
.unwrap();
let (status, resp) = send(router, req).await;
assert_eq!(
status,
StatusCode::UNPROCESSABLE_ENTITY,
"a streamed never-newline body over the cap must surface the line cap before \
any further chunk: {resp}"
);
assert_eq!(resp.get("code"), Some(&json!("INVALID_INPUT")));
assert!(
resp.get("message")
.and_then(|m| m.as_str())
.map(|m| m.contains("publish line exceeds the per-line cap"))
.unwrap_or(false),
"expected the pre-extend line-cap message, got: {resp}"
);
assert_eq!(resp.get("retryable"), Some(&json!(false)));
}
#[tokio::test]
async fn publish_eof_without_newline_over_cap_is_rejected() {
let router = build_router(registry_with_cap_witness_sink(), unused_provider());
let first_line = serde_json::to_vec(&json!({
"operation": "ingest/big",
"chunk": { "n": 1 }
}))
.unwrap();
let mut body = first_line;
body.push(b'\n');
body.extend_from_slice(&vec![b'a'; MAX_PUBLISH_LINE_BYTES + 1]);
let (status, resp) = send(router, raw_request("POST", "/publish", body)).await;
assert_eq!(
status,
StatusCode::BAD_REQUEST,
"EOF with an over-cap unterminated tail must not yield the buffer: {resp}"
);
assert_eq!(resp.get("code"), Some(&json!("INVALID_INPUT")));
assert!(
resp.get("message")
.and_then(|m| m.as_str())
.map(|m| m.contains("publish line exceeds the per-line cap"))
.unwrap_or(false),
"expected the line-cap message, got: {resp}"
);
}
#[tokio::test]
async fn publish_body_over_gateway_limit_returns_413() {
let router = build_router(publish_registry(), unused_provider());
let mut body = serde_json::to_vec(&json!({
"operation": "ingest/push",
"chunk": { "n": 0 }
}))
.unwrap();
body.push(b'\n');
let filler = vec![b'a'; 4096 - 10];
for _ in 0..600 {
body.extend_from_slice(b"{\"n\":\"");
body.extend_from_slice(&filler);
body.extend_from_slice(b"\"}\n");
}
let chunks: Vec<Bytes> = body.chunks(64 * 1024).map(Bytes::copy_from_slice).collect();
let req = Request::builder()
.method("POST")
.uri("/publish")
.header("content-type", "application/x-ndjson")
.body(Body::from_stream(futures::stream::iter(
chunks.into_iter().map(Ok::<_, std::convert::Infallible>),
)))
.unwrap();
let resp = router.oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::PAYLOAD_TOO_LARGE,
"a chunked upload of many under-cap lines over the body limit must be cut off with 413"
);
let bytes = resp.into_body().collect().await.unwrap().to_bytes();
assert_eq!(
&bytes[..],
GATEWAY_BODY_LIMIT_EXCEEDED.as_bytes(),
"the layer answers plain text 413, not an error envelope"
);
}
#[tokio::test]
async fn publish_declared_content_length_over_gateway_limit_returns_413() {
let router = build_router(publish_registry(), unused_provider());
let req = Request::builder()
.method("POST")
.uri("/publish")
.header("content-type", "application/x-ndjson")
.header("content-length", (GATEWAY_BODY_LIMIT + 1).to_string())
.body(Body::from(vec![b'a'; 64 * 1024]))
.unwrap();
let (status, resp) = send(router, req).await;
assert_eq!(
status,
StatusCode::PAYLOAD_TOO_LARGE,
"a declared content-length over the limit must be pre-rejected with 413"
);
assert_eq!(
resp,
Value::Null,
"the layer answers plain text 413, not an error envelope: {resp}"
);
}
#[tokio::test]
async fn publish_line_cap_breach_when_batched_with_complete_lines_is_still_rejected() {
let router = build_router(registry_with_cap_witness_sink(), unused_provider());
let mut body = ndjson(&[
json!({ "operation": "ingest/big", "chunk": { "n": 1 } }),
json!({ "n": 2 }),
]);
body.extend_from_slice(&vec![b'a'; MAX_PUBLISH_LINE_BYTES + 1]);
body.push(b'\n');
let (status, resp) = send(router, raw_request("POST", "/publish", body)).await;
assert_eq!(
status,
StatusCode::BAD_REQUEST,
"an over-cap line batched with complete lines must still abort the request with the \
cap error (today's status; GW-16 normalizes): {resp}"
);
assert_eq!(resp.get("code"), Some(&json!("INVALID_INPUT")));
}
#[tokio::test]
async fn publish_body_at_line_cap_within_limit_still_round_trips() {
let router = build_router(publish_registry(), unused_provider());
let first_line = serde_json::to_vec(&json!({
"operation": "ingest/push",
"chunk": { "n": 1 }
}))
.unwrap();
let mut body = first_line;
body.push(b'\n');
body.extend_from_slice(b"\"");
body.extend_from_slice(&vec![b'a'; MAX_PUBLISH_LINE_BYTES - 2]);
body.extend_from_slice(b"\"");
body.push(b'\n');
let (status, resp) = send(router, raw_request("POST", "/publish", body)).await;
assert_eq!(
status,
StatusCode::OK,
"a body under the limit with at-cap lines must pass the layer: {resp}"
);
assert_eq!(
resp.get("result"),
Some(&json!("ok")),
"a second line exactly at the per-line cap must be accepted (the cap is >, not >=): {resp}"
);
}
#[tokio::test]
async fn publish_client_disconnect_before_dispatch_signals_error_item() {
let router = build_router(publish_registry(), unused_provider());