From 767890bb7b9fee8d455ea7dc24d2767729f0b02e Mon Sep 17 00:00:00 2001 From: Steffen Ridderbusch Date: Wed, 21 Oct 2026 20:06:18 +0100 Subject: [PATCH 1/5] feat(dwctl): accept gzip, zstd and br request bodies and compressed batch files (DW-985) Compressed bodies used to fail with misleading errors: a gzipped chat request was reported as missing its model, and a .jsonl.gz batch file as invalid UTF-8 on line 1. Batch files repeat the same prompt on every line and shrink 10-30x compressed, so accepting them saves customers most of their upload time. The inference body limit now decodes gzip, zstd and br before any other layer reads the body, applies the limit to the decoded size and refuses other codings with a 415 that lists the accepted ones. File uploads recognise gzip and zstd by their magic numbers and decode them as they stream, so the size limit, parser and stored file see decoded content. --- Cargo.lock | 79 ++++++ config.yaml | 5 +- docs/src/reference/configuration.md | 2 +- dwctl/Cargo.toml | 4 +- dwctl/src/api/handlers/files.rs | 119 ++++++++- dwctl/src/inference/body_limit.rs | 397 +++++++++++++++++++++++++++- dwctl/src/lib.rs | 1 + dwctl/src/test/strict_mode.rs | 78 ++++++ dwctl/src/upload_compression.rs | 255 ++++++++++++++++++ 9 files changed, 923 insertions(+), 17 deletions(-) create mode 100644 dwctl/src/upload_compression.rs diff --git a/Cargo.lock b/Cargo.lock index 25cda4954..f2efaf1e5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -210,6 +210,18 @@ dependencies = [ "serde_json", ] +[[package]] +name = "async-compression" +version = "0.4.51" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d5755640120b24961f8231e05892e4a4f7e0d04e8aeca161508c9c397946970" +dependencies = [ + "compression-codecs", + "compression-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "async-lock" version = "3.4.2" @@ -1372,6 +1384,25 @@ dependencies = [ "tokio-util", ] +[[package]] +name = "compression-codecs" +version = "0.4.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cae36683101b2899f48bd31f8db9839437356188e3b1ff23096e73dab09ba8fe" +dependencies = [ + "compression-core", + "flate2", + "memchr", + "zstd 0.14.3", + "zstd-safe 8.2.0", +] + +[[package]] +name = "compression-core" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ebdda89edc534bf8d5dd2857562633b6495557e5846552ade994b99ed13c36a" + [[package]] name = "concurrent-queue" version = "2.5.0" @@ -1935,6 +1966,7 @@ dependencies = [ "anyhow", "arc-swap", "argon2", + "async-compression", "async-stream", "async-stripe", "async-stripe-billing", @@ -2039,6 +2071,7 @@ dependencies = [ "utoipa-scalar", "uuid", "wiremock", + "zstd 0.13.3", ] [[package]] @@ -8512,3 +8545,49 @@ name = "zmij" version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" + +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe 7.3.0", +] + +[[package]] +name = "zstd" +version = "0.14.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db3ae9bf90af6e40197bef767eeefc46ad6397420fe873830bd85e5fbbfb3c70" +dependencies = [ + "zstd-safe 8.2.0", +] + +[[package]] +name = "zstd-safe" +version = "7.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64d80649ab6db9d9f6f9c80a40becd948eda4714a0a5ac8c4d157a32231c7882" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-safe" +version = "8.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ded520f16b29e1a9b771c722221b2b0a0398c2316d9bb3c42b458a5414d7434b" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.1.2+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "481ada9c543b2323cbb7bbdbb5249177fbf0188ed959ed2c15e4b22421c84ed3" +dependencies = [ + "cc", + "pkg-config", +] diff --git a/config.yaml b/config.yaml index 529e9e74c..c65263a63 100644 --- a/config.yaml +++ b/config.yaml @@ -505,7 +505,8 @@ batches: # Controls file size, request count, and concurrency limits to prevent resource exhaustion # limits: # files: -# # Maximum file size in bytes (default: 100MB) +# # Maximum file size in bytes (default: 100MB). Gzip and zstd files are +# # decoded as they upload, and this limit applies to the decoded size. # # Set to 0 for unlimited (not recommended for production) # max_file_size: 104857600 # # Maximum number of requests (JSONL lines) per file @@ -526,6 +527,8 @@ batches: # # in batch JSONL files and to realtime request bodies on /ai/v1/*. # # Inference bodies are bounded before parsing or request logging; # # oversized bodies return HTTP 413. Multipart uploads use max_file_size. +# # Realtime bodies may be sent with Content-Encoding gzip, zstd or br; +# # the limit applies to the decoded size, and other codings get HTTP 415. # # Set to 0 for unlimited (not recommended for production) # # Default: 10MB (10485760) # # Changing this limit requires a restart to rebuild all routing layers. diff --git a/docs/src/reference/configuration.md b/docs/src/reference/configuration.md index ed521742f..ad74effe4 100644 --- a/docs/src/reference/configuration.md +++ b/docs/src/reference/configuration.md @@ -364,7 +364,7 @@ batches: |-------|------|---------|-------------| | `enabled` | boolean | `true` | Enable `/ai/v1/files` and `/ai/v1/batches` endpoints. | | `allowed_completion_windows` | list | `["24h"]` | SLA options users can select. | -| `files.max_file_size` | integer | `104857600` | Maximum upload size in bytes. | +| `files.max_file_size` | integer | `104857600` | Maximum upload size in bytes. A gzip or zstd file is decoded as it uploads, and the limit applies to its decoded size. | | `files.default_expiry_seconds` | integer | `86400` | Default file retention. | ## Background Services diff --git a/dwctl/Cargo.toml b/dwctl/Cargo.toml index 4fb557d47..f14c5b854 100644 --- a/dwctl/Cargo.toml +++ b/dwctl/Cargo.toml @@ -28,7 +28,7 @@ fusillade = { path = "../fusillade" } fusillade-arsenal = { path = "../fusillade-arsenal" } tokio = { version = "1.0", features = ["full"] } tokio-stream = { version = "0.1", features = ["sync"] } -tokio-util = { version = "0.7", features = ["rt"] } +tokio-util = { version = "0.7", features = ["rt", "io"] } notify = "8.2.0" sqlx = { version = "0.8", features = [ "runtime-tokio-rustls", @@ -109,6 +109,8 @@ outlet = "0.11.0" outlet-postgres = "0.8.0" brotli = "8.0" flate2 = "1" +zstd = "0.13" +async-compression = { version = "0.4", features = ["tokio", "gzip", "zstd"] } lettre = { version = "0.11", features = [ "smtp-transport", "file-transport", diff --git a/dwctl/src/api/handlers/files.rs b/dwctl/src/api/handlers/files.rs index 0fbe2960c..1f57f9da7 100644 --- a/dwctl/src/api/handlers/files.rs +++ b/dwctl/src/api/handlers/files.rs @@ -30,6 +30,7 @@ use crate::errors::{Error, Result}; use crate::image_normalizer::{ImageInput, ImageNormalizer, Mode as ImageNormalizerMode, walker as image_walker}; use crate::reasoning::ModelReasoningPolicy; use crate::types::Resource; +use crate::upload_compression::{self, ChunkError, Detected}; use axum::{ Json, body::Body, @@ -484,6 +485,10 @@ enum FileUploadError { InvalidJson { line: u64, error: String }, /// Invalid UTF-8 encoding in the file InvalidUtf8 { line: u64, byte_offset: i64, error: String }, + /// The file starts as gzip or zstd but isn't valid data for it (corrupt or truncated) + InvalidCompression { format: &'static str, message: String }, + /// The file is compressed in a format we recognise but don't decode + UnsupportedCompression { format: &'static str }, /// No file field in multipart upload NoFile, /// File contains no valid request templates @@ -574,6 +579,14 @@ impl FileUploadError { line, byte_offset, error ), }, + FileUploadError::InvalidCompression { format, message } => Error::BadRequest { + message: format!("File is not valid {format} data: {message}"), + }, + FileUploadError::UnsupportedCompression { format } => Error::BadRequest { + message: format!( + "File is compressed with {format}, which is not supported. Upload it uncompressed, or compressed with gzip or zstd." + ), + }, FileUploadError::NoFile => Error::BadRequest { message: "No file field found in multipart upload".to_string(), }, @@ -826,8 +839,15 @@ fn create_file_stream( } } "file" => { - // Extract filename from the field - metadata.filename = field.file_name().map(|s| s.to_string()); + let filename = field.file_name().map(|s| s.to_string()); + + // A gzip or zstd file is decoded as it streams in: the size + // limit, the parser and the stored file all see decoded content. + let (format, mut chunks) = upload_compression::open(Box::pin(field)).await; + if let Detected::Unsupported(format) = format { + abort!(FileUploadError::UnsupportedCompression { format }); + } + metadata.filename = filename.map(|name| format.decoded_file_name(&name)); // Send metadata before processing file content if tx.send(fusillade::FileStreamItem::Metadata(metadata.clone())).await.is_err() { @@ -835,11 +855,9 @@ fn create_file_stream( } // Now stream and parse the file content - let mut field = field; - loop { - match field.chunk().await { - Ok(Some(chunk)) => { + match chunks.next().await { + Some(Ok(chunk)) => { let chunk_size = chunk.len() as u64; total_size += chunk_size; @@ -1001,11 +1019,17 @@ fn create_file_stream( } } } - Ok(None) => { + None => { // Normal end of stream break; } - Err(e) => { + Some(Err(ChunkError::Decode(e))) => { + abort!(FileUploadError::InvalidCompression { + format: format.name(), + message: e.to_string(), + }); + } + Some(Err(ChunkError::Upload(e))) => { // Log Display and Debug representations tracing::warn!( error_display = %e, @@ -1137,7 +1161,7 @@ fn create_file_stream( Each line must be a valid JSON object containing `custom_id`, `method`, `url`, and `body` fields. The `model` field in the body must reference a model your API key has access to.", request_body( content_type = "multipart/form-data", - description = "Multipart form with `file` (the JSONL file) and `purpose` (must be `batch`)." + description = "Multipart form with `file` (the JSONL file, optionally gzip- or zstd-compressed) and `purpose` (must be `batch`)." ), responses( (status = 201, description = "File uploaded and validated successfully.", body = FileResponse), @@ -2413,6 +2437,83 @@ mod tests { } } + #[dwctl_test_macros::test] + #[test_log::test] + async fn test_upload_compressed_file_stores_decoded_content(pool: PgPool) { + use std::io::Write; + + let (app, _bg_services) = create_test_app(pool.clone(), false).await; + let user = create_test_user_with_roles(&pool, vec![Role::StandardUser, Role::BatchAPIUser]).await; + let group = create_test_group(&pool).await; + add_user_to_group(&pool, user.id, group.id).await; + let deployment = create_test_deployment(&pool, user.id, "gpt-4-model", "gpt-4").await; + add_deployment_to_group(&pool, deployment.id, group.id, user.id).await; + + let jsonl_content = r#"{"custom_id":"request-1","method":"POST","url":"/v1/chat/completions","body":{"model":"gpt-4","messages":[{"role":"user","content":"Hello 1"}]}} +{"custom_id":"request-2","method":"POST","url":"/v1/chat/completions","body":{"model":"gpt-4","messages":[{"role":"user","content":"Hello 2"}]}} +"#; + let mut gzip = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + gzip.write_all(jsonl_content.as_bytes()).unwrap(); + let gzip = gzip.finish().unwrap(); + let zstd = zstd::encode_all(jsonl_content.as_bytes(), 3).unwrap(); + + for (compressed, name) in [(gzip, "test-batch.jsonl.gz"), (zstd, "test-batch.jsonl.zst")] { + let upload_response = app + .post("/ai/v1/files") + .add_header(&add_auth_headers(&user)[0].0, &add_auth_headers(&user)[0].1) + .add_header(&add_auth_headers(&user)[1].0, &add_auth_headers(&user)[1].1) + .multipart( + axum_test::multipart::MultipartForm::new() + .add_text("purpose", "batch") + .add_part("file", axum_test::multipart::Part::bytes(compressed).file_name(name)), + ) + .await; + + upload_response.assert_status(axum::http::StatusCode::CREATED); + let file: FileResponse = upload_response.json(); + assert_eq!(file.filename, "test-batch.jsonl", "{name}"); + assert_eq!(file.bytes, jsonl_content.len() as i64, "{name}: size is the decoded size"); + + let downloaded = app + .get(&format!("/ai/v1/files/{}/content", file.id)) + .add_header(&add_auth_headers(&user)[0].0, &add_auth_headers(&user)[0].1) + .add_header(&add_auth_headers(&user)[1].0, &add_auth_headers(&user)[1].1) + .await + .text(); + let ids: Vec = downloaded + .lines() + .map(|line| serde_json::from_str::(line).unwrap()["custom_id"].to_string()) + .collect(); + assert_eq!(ids, ["\"request-1\"", "\"request-2\""], "{name}"); + } + } + + #[dwctl_test_macros::test] + #[test_log::test] + async fn test_upload_broken_or_unsupported_compression_names_the_problem(pool: PgPool) { + let (app, _bg_services) = create_test_app(pool.clone(), false).await; + let user = create_test_user_with_roles(&pool, vec![Role::StandardUser, Role::BatchAPIUser]).await; + + let truncated_gzip = b"\x1f\x8b\x08\x00\x00\x00\x00\x00\x00\x03\xab\x56".to_vec(); + let xz = b"\xfd7zXZ\x00\x00\x04\xe6\xd6\xb4\x46".to_vec(); + for (content, expected) in [(truncated_gzip, "not valid gzip data"), (xz, "compressed with xz")] { + let upload_response = app + .post("/ai/v1/files") + .add_header(&add_auth_headers(&user)[0].0, &add_auth_headers(&user)[0].1) + .add_header(&add_auth_headers(&user)[1].0, &add_auth_headers(&user)[1].1) + .multipart( + axum_test::multipart::MultipartForm::new() + .add_text("purpose", "batch") + .add_part("file", axum_test::multipart::Part::bytes(content).file_name("batch.jsonl")), + ) + .await; + + upload_response.assert_status(axum::http::StatusCode::BAD_REQUEST); + let body = upload_response.text(); + assert!(body.contains(expected), "{body}"); + } + } + #[dwctl_test_macros::test] #[test_log::test] async fn test_upload_missing_model_field(pool: PgPool) { diff --git a/dwctl/src/inference/body_limit.rs b/dwctl/src/inference/body_limit.rs index 33374e936..b3986b375 100644 --- a/dwctl/src/inference/body_limit.rs +++ b/dwctl/src/inference/body_limit.rs @@ -1,27 +1,76 @@ -//! Bound inference bodies before middleware parses, copies, or logs them. +//! Bound inference bodies before middleware parses, copies, or logs them, and +//! decode compressed ones so every layer behind this one sees plain JSON. + +use std::io::Read; use axum::{ Json, body::{Body, Bytes, to_bytes}, extract::State, - http::{HeaderMap, Request, StatusCode, header::CONTENT_LENGTH}, + http::{ + HeaderMap, HeaderValue, Request, StatusCode, + header::{ACCEPT_ENCODING, CONTENT_ENCODING, CONTENT_LENGTH}, + }, middleware::Next, response::{IntoResponse, Response}, }; use super::rejections; +/// The content codings an inference body may arrive in, as advertised on a 415. +const ACCEPTED_ENCODINGS: &str = "gzip, zstd, br"; + +/// RFC 9659: encoders of the `zstd` content coding must not need a window +/// above 8 MiB, so a frame declaring a larger one is refused rather than +/// allocated. +const ZSTD_WINDOW_LOG_MAX: u32 = 23; + /// Install outside all inference middleware, including request logging and /// daemon dispatch bypasses. File uploads use a separate router and limit. /// A zero limit retains the configured unlimited behavior. +/// +/// A body sent with `Content-Encoding: gzip`, `zstd` or `br` is decoded here, +/// and the limit applies to the decoded size: that is what every inner layer +/// buffers. Any other coding is refused with a 415 before the body is read. pub async fn limit_inference_body(State(max_bytes): State, request: Request, next: Next) -> Response { - if max_bytes == 0 { + let encoding = match content_encoding(request.headers()) { + Ok(encoding) => encoding, + Err(UnsupportedEncoding) => { + let (parts, _) = request.into_parts(); + let response = unsupported_encoding(parts.uri.path()); + return rejections::refuse(&parts, response, "unsupported_content_encoding", None, None); + } + }; + if max_bytes == 0 && encoding.is_none() { return next.run(request).await; } - let (parts, body) = request.into_parts(); + let (mut parts, body) = request.into_parts(); match read_inference_body(&parts.headers, body, max_bytes).await { - Ok(bytes) => next.run(Request::from_parts(parts, Body::from(bytes))).await, + Ok(bytes) => { + let Some(encoding) = encoding else { + return next.run(Request::from_parts(parts, Body::from(bytes))).await; + }; + let limit = if max_bytes == 0 { usize::MAX } else { max_bytes }; + match decode(encoding, bytes, limit).await { + Ok(decoded) => { + metrics::counter!("dwctl_inference_bodies_decoded_total", "encoding" => encoding.name()).increment(1); + parts.headers.remove(CONTENT_ENCODING); + parts.headers.insert(CONTENT_LENGTH, HeaderValue::from(decoded.len())); + next.run(Request::from_parts(parts, Body::from(decoded))).await + } + Err(DecodeError::TooLarge) => { + let message = format!("Decoded request body exceeds the maximum size of {max_bytes} bytes"); + let response = oversized_body(message, parts.uri.path()); + rejections::refuse(&parts, response, "request_too_large", None, None) + } + Err(DecodeError::Invalid) => { + let message = format!("Request body is not valid {} data", encoding.name()); + let response = invalid_request(StatusCode::BAD_REQUEST, parts.uri.path(), "body_decode_failed", &message); + rejections::refuse(&parts, response, "body_decode_failed", None, None) + } + } + } Err(crate::errors::Error::PayloadTooLarge { message }) => { let response = oversized_body(message, parts.uri.path()); rejections::refuse(&parts, response, "request_too_large", None, None) @@ -87,6 +136,120 @@ pub(crate) fn oversized_body(message: String, path: &str) -> Response { (StatusCode::PAYLOAD_TOO_LARGE, Json(body)).into_response() } +/// An `invalid_request_error` in the error shape of the API `path` belongs to. +fn invalid_request(status: StatusCode, path: &str, code: &str, message: &str) -> Response { + let body = if path.ends_with("/messages") { + serde_json::json!({"type": "error", "error": { + "message": message, "type": "invalid_request_error" + }}) + } else { + serde_json::json!({"error": { + "message": message, "type": "invalid_request_error", "code": code + }}) + }; + (status, Json(body)).into_response() +} + +/// RFC 9110 ยง15.5.16: a 415 for a content coding names the codings accepted. +fn unsupported_encoding(path: &str) -> Response { + let message = format!("Unsupported Content-Encoding. Send the body uncompressed or encoded with one of: {ACCEPTED_ENCODINGS}"); + let mut response = invalid_request(StatusCode::UNSUPPORTED_MEDIA_TYPE, path, "unsupported_content_encoding", &message); + response + .headers_mut() + .insert(ACCEPT_ENCODING, HeaderValue::from_static(ACCEPTED_ENCODINGS)); + response +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Encoding { + Gzip, + Zstd, + Brotli, +} + +impl Encoding { + fn name(self) -> &'static str { + match self { + Encoding::Gzip => "gzip", + Encoding::Zstd => "zstd", + Encoding::Brotli => "br", + } + } +} + +#[derive(Debug, PartialEq, Eq)] +struct UnsupportedEncoding; + +/// The body's content coding, or `None` when it has none. Layered codings +/// (`gzip, br`) are valid HTTP but no client sends them, so they are refused +/// like an unknown coding rather than decoded in turn. +fn content_encoding(headers: &HeaderMap) -> Result, UnsupportedEncoding> { + let mut found = None; + for value in headers.get_all(CONTENT_ENCODING) { + let value = value.to_str().map_err(|_| UnsupportedEncoding)?; + for coding in value.split(',').map(str::trim).filter(|coding| !coding.is_empty()) { + let encoding = match coding.to_ascii_lowercase().as_str() { + "identity" => continue, + "gzip" | "x-gzip" => Encoding::Gzip, + "zstd" => Encoding::Zstd, + "br" => Encoding::Brotli, + _ => return Err(UnsupportedEncoding), + }; + if found.replace(encoding).is_some() { + return Err(UnsupportedEncoding); + } + } + } + Ok(found) +} + +#[derive(Debug, PartialEq, Eq)] +enum DecodeError { + TooLarge, + Invalid, +} + +/// Decodes `bytes`, stopping as soon as the output passes `max_bytes`, so a +/// small body that inflates without bound costs at most the limit. Runs on a +/// blocking thread: inflating a full-size body takes tens of milliseconds. +async fn decode(encoding: Encoding, bytes: Bytes, max_bytes: usize) -> Result { + tokio::task::spawn_blocking(move || decode_blocking(encoding, &bytes, max_bytes)) + .await + .unwrap_or(Err(DecodeError::Invalid)) +} + +fn decode_blocking(encoding: Encoding, bytes: &[u8], max_bytes: usize) -> Result { + let reader: Box = match encoding { + // HTTP gzip bodies may hold several concatenated members. + Encoding::Gzip => Box::new(flate2::read::MultiGzDecoder::new(bytes)), + Encoding::Zstd => { + // Most compressors record the decoded size in the frame header: + // answer an oversized one as such rather than failing the window + // check below with a less helpful error. + if let Ok(Some(size)) = zstd::zstd_safe::get_frame_content_size(bytes) + && size > max_bytes as u64 + { + return Err(DecodeError::TooLarge); + } + let mut decoder = zstd::stream::read::Decoder::with_buffer(bytes).map_err(|_| DecodeError::Invalid)?; + // The decoder allocates the window a frame declares up front, so an + // unbounded window would let a tiny body claim gigabytes. + decoder.window_log_max(ZSTD_WINDOW_LOG_MAX).map_err(|_| DecodeError::Invalid)?; + Box::new(decoder) + } + Encoding::Brotli => Box::new(brotli::Decompressor::new(bytes, 4096)), + }; + let mut decoded = Vec::new(); + reader + .take(max_bytes.saturating_add(1) as u64) + .read_to_end(&mut decoded) + .map_err(|_| DecodeError::Invalid)?; + if decoded.len() > max_bytes { + return Err(DecodeError::TooLarge); + } + Ok(Bytes::from(decoded)) +} + #[cfg(test)] mod tests { use super::*; @@ -196,6 +359,230 @@ mod tests { } } + fn compress(encoding: Encoding, data: &[u8]) -> Vec { + use std::io::Write; + match encoding { + Encoding::Gzip => { + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + encoder.write_all(data).unwrap(); + encoder.finish().unwrap() + } + Encoding::Zstd => zstd::encode_all(data, 3).unwrap(), + Encoding::Brotli => { + let mut out = Vec::new(); + { + let mut writer = brotli::CompressorWriter::new(&mut out, 4096, 5, 22); + writer.write_all(data).unwrap(); + } + out + } + } + } + + /// Echoes the body, failing unless the headers describe it as sent onwards. + fn decoding_app(limit: usize) -> Router { + Router::new() + .route( + "/chat/completions", + post(|headers: HeaderMap, body: Bytes| async move { + assert!(headers.get(CONTENT_ENCODING).is_none()); + assert_eq!(headers[CONTENT_LENGTH], body.len().to_string()); + body + }), + ) + .layer(middleware::from_fn_with_state(limit, limit_inference_body)) + .layer(rejections::test_support::layer("m")) + } + + async fn error_of(response: Response) -> serde_json::Value { + serde_json::from_slice(&axum::body::to_bytes(response.into_body(), 10_000).await.unwrap()).unwrap() + } + + #[tokio::test] + async fn decodes_each_accepted_encoding_for_the_layers_behind_it() { + // Compressible, so every encoding is smaller than the limit it decodes to. + let json = format!( + r#"{{"model":"m","messages":[{{"role":"user","content":"{}"}}]}}"#, + "hello ".repeat(50) + ); + let json = json.as_bytes(); + for (encoding, header) in [ + (Encoding::Gzip, "gzip"), + (Encoding::Gzip, "x-gzip"), + (Encoding::Zstd, "zstd"), + (Encoding::Brotli, "br"), + (Encoding::Zstd, "ZSTD"), + ] { + for limit in [json.len(), 0] { + let compressed = compress(encoding, json); + let response = decoding_app(limit) + .oneshot( + Request::post("/chat/completions") + .header(CONTENT_ENCODING, header) + .header(CONTENT_LENGTH, compressed.len()) + .body(Body::from(compressed)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK, "{header} with limit {limit}"); + assert_eq!(axum::body::to_bytes(response.into_body(), 10_000).await.unwrap(), &json[..]); + } + } + } + + #[tokio::test] + async fn refuses_other_encodings_with_415_before_reading_the_body() { + for header in ["deflate", "lz4", "compress", "gzip, br", "snappy"] { + let polls = Arc::new(AtomicUsize::new(0)); + let observed = polls.clone(); + let body = Body::from_stream(futures::stream::poll_fn(move |_| { + observed.fetch_add(1, Ordering::SeqCst); + std::task::Poll::Ready(None::>) + })); + let response = decoding_app(1024) + .oneshot( + Request::post("/chat/completions") + .header(CONTENT_ENCODING, header) + .body(body) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::UNSUPPORTED_MEDIA_TYPE, "{header}"); + assert_eq!(response.headers()[ACCEPT_ENCODING], "gzip, zstd, br"); + assert_eq!(error_of(response).await["error"]["code"], "unsupported_content_encoding"); + assert_eq!(polls.load(Ordering::SeqCst), 0, "{header}"); + } + } + + #[tokio::test] + async fn answers_messages_requests_in_the_anthropic_error_shape() { + let app = Router::new() + .route("/messages", post(|body: Bytes| async move { body })) + .layer(middleware::from_fn_with_state(1024usize, limit_inference_body)); + let response = app + .oneshot( + Request::post("/messages") + .header(CONTENT_ENCODING, "lz4") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::UNSUPPORTED_MEDIA_TYPE); + let body = error_of(response).await; + assert_eq!(body["type"], "error"); + assert_eq!(body["error"]["type"], "invalid_request_error"); + } + + #[tokio::test] + async fn applies_the_limit_to_the_decoded_size() { + // Megabytes of zeros compress to a few kilobytes: a small upload must + // not be able to make inner layers buffer past the limit. + for encoding in [Encoding::Gzip, Encoding::Zstd, Encoding::Brotli] { + let bomb = compress(encoding, &vec![0u8; 4 << 20]); + assert!(bomb.len() < 64 << 10); + let response = decoding_app(1 << 20) + .oneshot( + Request::post("/chat/completions") + .header(CONTENT_ENCODING, encoding.name()) + .body(Body::from(bomb)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE, "{encoding:?}"); + assert_eq!(error_of(response).await["error"]["code"], "request_too_large"); + } + } + + #[test] + fn refuses_zstd_windows_above_the_http_limit() { + use std::io::Write; + // A streaming compressor that doesn't know the input size declares its + // full window, however small the body: the decoder must not allocate + // more than RFC 9659's 8 MiB for it. + let framed = |window_log| { + let mut encoder = zstd::stream::write::Encoder::new(Vec::new(), 3).unwrap(); + encoder.window_log(window_log).unwrap(); + encoder.write_all(b"{}").unwrap(); + encoder.finish().unwrap() + }; + assert_eq!(decode_blocking(Encoding::Zstd, &framed(23), 1024).unwrap(), &b"{}"[..]); + assert_eq!(decode_blocking(Encoding::Zstd, &framed(24), 1024), Err(DecodeError::Invalid)); + } + + #[tokio::test] + async fn rejects_bodies_that_are_not_valid_for_their_encoding() { + for encoding in [Encoding::Gzip, Encoding::Zstd, Encoding::Brotli] { + let response = decoding_app(1024) + .oneshot( + Request::post("/chat/completions") + .header(CONTENT_ENCODING, encoding.name()) + .body(Body::from(r#"{"model":"m"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{encoding:?}"); + assert_eq!(error_of(response).await["error"]["code"], "body_decode_failed"); + } + } + + #[test] + fn counts_refused_encodings() { + let (lines, _) = rejections::test_support::rejections(async { + decoding_app(1024) + .oneshot( + Request::post("/chat/completions") + .header(CONTENT_ENCODING, "lz4") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + decoding_app(1024) + .oneshot( + Request::post("/chat/completions") + .header(CONTENT_ENCODING, "gzip") + .body(Body::from("not gzip")) + .unwrap(), + ) + .await + .unwrap(); + }); + assert!( + lines + .iter() + .any(|l| l.contains("status=\"415\"") && l.contains("code=\"unsupported_content_encoding\"")), + "{lines:?}" + ); + assert!( + lines + .iter() + .any(|l| l.contains("status=\"400\"") && l.contains("code=\"body_decode_failed\"")), + "{lines:?}" + ); + } + + #[test] + fn reads_the_content_encoding_header() { + let headers = |values: &[&str]| { + let mut headers = HeaderMap::new(); + for value in values { + headers.append(CONTENT_ENCODING, value.parse().unwrap()); + } + headers + }; + assert_eq!(content_encoding(&headers(&[])), Ok(None)); + assert_eq!(content_encoding(&headers(&["identity"])), Ok(None)); + assert_eq!(content_encoding(&headers(&[" GZip "])), Ok(Some(Encoding::Gzip))); + assert_eq!(content_encoding(&headers(&["identity, br"])), Ok(Some(Encoding::Brotli))); + assert_eq!(content_encoding(&headers(&["gzip", "gzip"])), Err(UnsupportedEncoding)); + assert_eq!(content_encoding(&headers(&["deflate"])), Err(UnsupportedEncoding)); + } + #[tokio::test] async fn distinguishes_transport_errors_without_exposing_their_contents() { let body = Body::from_stream(futures::stream::once(async { diff --git a/dwctl/src/lib.rs b/dwctl/src/lib.rs index 8d1d99a45..d707f1e20 100644 --- a/dwctl/src/lib.rs +++ b/dwctl/src/lib.rs @@ -186,6 +186,7 @@ pub mod task_retention; pub mod tasks; pub mod telemetry; mod types; +mod upload_compression; pub mod webhooks; // Test modules diff --git a/dwctl/src/test/strict_mode.rs b/dwctl/src/test/strict_mode.rs index 72276e0d6..be15329c3 100644 --- a/dwctl/src/test/strict_mode.rs +++ b/dwctl/src/test/strict_mode.rs @@ -53,6 +53,84 @@ async fn test_inference_body_limit_precedes_processing_in_all_modes(pool: PgPool } } +/// A compressed realtime body is decoded before the inference stack reads it: +/// the upstream receives plain JSON without `Content-Encoding`, and an +/// encoding we don't decode is refused before it gets that far. +async fn assert_compressed_bodies_reach_upstream_decoded(pool: PgPool, strict_mode: bool) { + use std::io::Write; + + let mock = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("POST")) + .and(wiremock::matchers::path("/chat/completions")) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "chatcmpl-compressed", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + }))) + .mount(&mock) + .await; + let mut config = create_test_config(); + config.onwards.strict_mode = strict_mode; + config.cache.enabled = false; + let (server, key, _services) = super::responses::setup_ai_test_with_config(pool, &mock, config).await; + + let json = serde_json::to_vec(&serde_json::json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "compressed hello"}] + })) + .unwrap(); + let mut gzip = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + gzip.write_all(&json).unwrap(); + for (encoding, body) in [("gzip", gzip.finish().unwrap()), ("zstd", zstd::encode_all(&json[..], 3).unwrap())] { + let response = server + .post("/ai/v1/chat/completions") + .add_header("Authorization", format!("Bearer {key}")) + .add_header("Content-Type", "application/json") + .add_header("Content-Encoding", encoding) + .bytes(body.into()) + .await; + response.assert_status_ok(); + } + + let received = mock.received_requests().await.unwrap(); + assert_eq!(received.len(), 2, "strict_mode={strict_mode}"); + for request in received { + assert!(request.headers.get("content-encoding").is_none(), "strict_mode={strict_mode}"); + let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(body["messages"][0]["content"], "compressed hello", "strict_mode={strict_mode}"); + } + + let response = server + .post("/ai/v1/chat/completions") + .add_header("Authorization", format!("Bearer {key}")) + .add_header("Content-Encoding", "lz4") + .bytes(json.into()) + .await; + response.assert_status(axum::http::StatusCode::UNSUPPORTED_MEDIA_TYPE); + assert_eq!( + response.json::()["error"]["code"], + "unsupported_content_encoding" + ); + assert_eq!( + mock.received_requests().await.unwrap().len(), + 2, + "refused before reaching the upstream" + ); +} + +#[dwctl_test_macros::test] +async fn test_compressed_bodies_reach_upstream_decoded_in_strict_mode(pool: PgPool) { + assert_compressed_bodies_reach_upstream_decoded(pool, true).await; +} + +#[dwctl_test_macros::test] +async fn test_compressed_bodies_reach_upstream_decoded_in_passthrough_mode(pool: PgPool) { + assert_compressed_bodies_reach_upstream_decoded(pool, false).await; +} + /// Test that strict mode rejects unknown endpoints with 404 #[dwctl_test_macros::test] #[test_log::test] diff --git a/dwctl/src/upload_compression.rs b/dwctl/src/upload_compression.rs new file mode 100644 index 000000000..6130da9b0 --- /dev/null +++ b/dwctl/src/upload_compression.rs @@ -0,0 +1,255 @@ +//! Compressed batch files. +//! +//! A batch file is mostly the same system prompt and schema repeated on every +//! line, so it compresses far better than a single request. A file part that +//! starts with a gzip or zstd magic number is decoded as it streams in, so the +//! JSONL parser, the file size limit and the stored file all see the decoded +//! content. Formats we can recognise but don't decode are named in the error, +//! rather than failing as invalid UTF-8 on line 1. +//! +//! Detection is by content, not by file name or part `Content-Type`: no JSONL +//! file can start with any of these signatures, while names and part types +//! are often missing or wrong. + +use std::{io, pin::Pin}; + +use async_compression::tokio::bufread::{GzipDecoder, ZstdDecoder}; +use bytes::Bytes; +use futures::{Stream, StreamExt, TryStreamExt, stream}; +use tokio_util::io::{ReaderStream, StreamReader}; + +/// Long enough for the longest signature below (xz). +const SNIFF_LEN: usize = 6; + +/// Size of the decoded chunks handed to the JSONL parser. +const DECODED_CHUNK: usize = 64 * 1024; + +/// The largest zstd window a file may need: zstd's own default ceiling, so any +/// file the `zstd` CLI decompresses without extra flags uploads too. The +/// decoder allocates the declared window up front; uploads are few and +/// concurrency-limited, unlike realtime bodies, which allow 8 MiB. +const ZSTD_WINDOW_LOG_MAX: u32 = 27; + +/// What a file part turned out to contain. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Detected { + Plain, + Gzip, + Zstd, + /// A compression format we recognise but don't decode, by name. + Unsupported(&'static str), +} + +impl Detected { + pub(crate) fn name(self) -> &'static str { + match self { + Detected::Plain => "uncompressed", + Detected::Gzip => "gzip", + Detected::Zstd => "zstd", + Detected::Unsupported(name) => name, + } + } + + /// The file name the decoded file is stored under: `batch.jsonl.gz` is + /// stored as `batch.jsonl`, since downloads return decoded content. + pub(crate) fn decoded_file_name(self, name: &str) -> String { + let suffixes: &[&str] = match self { + Detected::Gzip => &[".gz", ".gzip"], + Detected::Zstd => &[".zst", ".zstd"], + Detected::Plain | Detected::Unsupported(_) => &[], + }; + suffixes + .iter() + .find_map(|suffix| name.strip_suffix(suffix)) + .unwrap_or(name) + .to_string() + } +} + +/// Identifies the format from the first bytes of a file. +pub(crate) fn detect(prefix: &[u8]) -> Detected { + const SIGNATURES: &[(&[u8], Detected)] = &[ + (&[0x1f, 0x8b], Detected::Gzip), + (&[0x28, 0xb5, 0x2f, 0xfd], Detected::Zstd), + (&[0xfd, b'7', b'z', b'X', b'Z', 0x00], Detected::Unsupported("xz")), + (b"BZh", Detected::Unsupported("bzip2")), + (b"PK\x03\x04", Detected::Unsupported("zip")), + (&[0x04, 0x22, 0x4d, 0x18], Detected::Unsupported("lz4")), + ]; + SIGNATURES + .iter() + .find(|(signature, _)| prefix.starts_with(signature)) + .map_or(Detected::Plain, |(_, detected)| *detected) +} + +/// A chunk of decoded file content could not be produced. +#[derive(Debug)] +pub(crate) enum ChunkError { + /// Reading the upload itself failed: passed through untouched, so callers + /// can still recognise a body length limit. + Upload(E), + /// The content is not valid data for the format it starts as, including + /// a stream that ends mid-frame. + Decode(io::Error), +} + +pub(crate) type Chunks<'a, E> = Pin>> + Send + 'a>>; + +/// Reads just enough of `upload` to identify its format, then returns the +/// format and its content as decoded chunks. An unsupported format's chunks are +/// returned as they are, so the caller decides how to refuse it. A read error +/// while identifying is not lost: it comes out of the returned stream. The +/// caller enforces the decoded size. +pub(crate) async fn open<'a, S, E>(mut upload: S) -> (Detected, Chunks<'a, E>) +where + S: Stream> + Send + Unpin + 'a, + E: std::error::Error + Send + Sync + 'static, +{ + let mut head = Vec::new(); + let mut prefix = Vec::with_capacity(SNIFF_LEN); + while prefix.len() < SNIFF_LEN { + match upload.next().await { + Some(Ok(chunk)) => { + prefix.extend_from_slice(&chunk[..chunk.len().min(SNIFF_LEN - prefix.len())]); + head.push(Ok(chunk)); + } + Some(Err(error)) => { + head.push(Err(error)); + break; + } + None => break, + } + } + let detected = detect(&prefix); + let upload = stream::iter(head).chain(upload); + + let decoded: Chunks<'a, E> = match detected { + Detected::Plain | Detected::Unsupported(_) => Box::pin(upload.map_err(ChunkError::Upload)), + Detected::Gzip | Detected::Zstd => { + let reader = StreamReader::new(upload.map_err(io::Error::other)); + let decoder: Pin> = if detected == Detected::Gzip { + let mut decoder = GzipDecoder::new(reader); + decoder.multiple_members(true); + Box::pin(decoder) + } else { + Box::pin(ZstdDecoder::with_params( + reader, + &[async_compression::zstd::DParameter::window_log_max(ZSTD_WINDOW_LOG_MAX)], + )) + }; + Box::pin(ReaderStream::with_capacity(decoder, DECODED_CHUNK).map_err(upload_or_decode_error)) + } + }; + (detected, decoded) +} + +/// Splits the decoder's I/O errors back into the upload's own errors, which +/// the reader wrapped, and errors in the compressed data itself. +fn upload_or_decode_error(error: io::Error) -> ChunkError { + if error.get_ref().is_some_and(|inner| inner.is::()) { + let inner = error.into_inner().expect("checked above"); + return ChunkError::Upload(*inner.downcast::().expect("checked above")); + } + ChunkError::Decode(error) +} + +#[cfg(test)] +mod tests { + use std::io::Write; + + use super::*; + + const JSONL: &[u8] = b"{\"custom_id\":\"a\"}\n{\"custom_id\":\"b\"}\n"; + + fn gzip(data: &[u8]) -> Vec { + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + encoder.write_all(data).unwrap(); + encoder.finish().unwrap() + } + + /// Feeds `data` in `size`-byte chunks, so the signature can straddle them. + fn upload(data: Vec, size: usize) -> impl Stream> + Send + Unpin { + let chunks: Vec<_> = data.chunks(size).map(|chunk| Ok(Bytes::copy_from_slice(chunk))).collect(); + stream::iter(chunks) + } + + async fn read_all(data: Vec, size: usize) -> (Detected, Result, ChunkError>) { + let (detected, mut chunks) = open(upload(data, size)).await; + let mut out = Vec::new(); + while let Some(chunk) = chunks.next().await { + match chunk { + Ok(chunk) => out.extend_from_slice(&chunk), + Err(error) => return (detected, Err(error)), + } + } + (detected, Ok(out)) + } + + #[tokio::test] + async fn decodes_gzip_and_zstd_files_whatever_the_chunking() { + let mut two_members = gzip(&JSONL[..18]); + two_members.extend(gzip(&JSONL[18..])); + for (compressed, expected) in [ + (gzip(JSONL), Detected::Gzip), + (two_members, Detected::Gzip), + (zstd::encode_all(JSONL, 3).unwrap(), Detected::Zstd), + (JSONL.to_vec(), Detected::Plain), + ] { + for size in [1, 3, 1024] { + let (detected, decoded) = read_all(compressed.clone(), size).await; + assert_eq!(detected, expected); + assert_eq!(decoded.unwrap(), JSONL, "{expected:?} in {size}-byte chunks"); + } + } + } + + #[tokio::test] + async fn reports_truncated_and_corrupt_files_as_decode_errors() { + let gz = gzip(JSONL); + let zst = zstd::encode_all(JSONL, 3).unwrap(); + let mut corrupt = gz.clone(); + let middle = corrupt.len() / 2; + corrupt[middle] ^= 0xff; + for broken in [gz[..gz.len() - 4].to_vec(), zst[..zst.len() - 3].to_vec(), corrupt] { + let (_, decoded) = read_all(broken, 7).await; + assert!(matches!(decoded, Err(ChunkError::Decode(_))), "{decoded:?}"); + } + } + + #[tokio::test] + async fn passes_upload_errors_through_untouched() { + for data in [gzip(JSONL), JSONL.to_vec()] { + let failing = upload(data[..4].to_vec(), 2).chain(stream::iter([Err(io::Error::other("reset"))])); + let (_, mut chunks) = open(failing).await; + let mut saw = None; + while let Some(chunk) = chunks.next().await { + if let Err(error) = chunk { + saw = Some(error); + break; + } + } + match saw { + Some(ChunkError::Upload(error)) => assert_eq!(error.to_string(), "reset"), + other => panic!("expected the upload error, got {other:?}"), + } + } + } + + #[test] + fn names_formats_it_does_not_decode() { + assert_eq!(detect(b"\xfd7zXZ\x00rest"), Detected::Unsupported("xz")); + assert_eq!(detect(b"BZh91AY"), Detected::Unsupported("bzip2")); + assert_eq!(detect(b"PK\x03\x04"), Detected::Unsupported("zip")); + assert_eq!(detect(b"\x04\x22\x4d\x18"), Detected::Unsupported("lz4")); + assert_eq!(detect(b"{\"custom_id\""), Detected::Plain); + assert_eq!(detect(b""), Detected::Plain); + } + + #[test] + fn stores_decoded_files_without_the_compression_suffix() { + assert_eq!(Detected::Gzip.decoded_file_name("batch.jsonl.gz"), "batch.jsonl"); + assert_eq!(Detected::Zstd.decoded_file_name("batch.jsonl.zst"), "batch.jsonl"); + assert_eq!(Detected::Gzip.decoded_file_name("batch.jsonl"), "batch.jsonl"); + assert_eq!(Detected::Plain.decoded_file_name("batch.jsonl.gz"), "batch.jsonl.gz"); + } +} From d21bb3e92a67fb8448e318a8c8ae44ae3599cb59 Mon Sep 17 00:00:00 2001 From: Steffen Ridderbusch Date: Wed, 21 Oct 2026 20:24:49 +0100 Subject: [PATCH 2/5] fix(dwctl): bound brotli windows and read every zstd frame of a file (DW-985) brotli::Decompressor accepts the large-window extension, whose window of up to 1 GiB is allocated before any output, so a tiny body could exhaust memory past the output bound. Decode with the strict state, which keeps RFC 7932's 16 MiB ceiling. async-compression's zstd decoder stopped after the first frame, so a file of concatenated frames was stored truncated. Read them all, and recognise a stream that opens with a skippable frame as zstd. --- config.yaml | 3 +- docs/src/reference/configuration.md | 2 +- dwctl/src/inference/body_limit.rs | 75 ++++++++++++++++++++++++++++- dwctl/src/upload_compression.rs | 32 ++++++++++-- 4 files changed, 104 insertions(+), 8 deletions(-) diff --git a/config.yaml b/config.yaml index c65263a63..ec4f0cb6d 100644 --- a/config.yaml +++ b/config.yaml @@ -506,7 +506,8 @@ batches: # limits: # files: # # Maximum file size in bytes (default: 100MB). Gzip and zstd files are -# # decoded as they upload, and this limit applies to the decoded size. +# # decoded as they upload, and this limit applies to the decoded size +# # (the compressed upload may not exceed it either). # # Set to 0 for unlimited (not recommended for production) # max_file_size: 104857600 # # Maximum number of requests (JSONL lines) per file diff --git a/docs/src/reference/configuration.md b/docs/src/reference/configuration.md index ad74effe4..e5feecdbe 100644 --- a/docs/src/reference/configuration.md +++ b/docs/src/reference/configuration.md @@ -364,7 +364,7 @@ batches: |-------|------|---------|-------------| | `enabled` | boolean | `true` | Enable `/ai/v1/files` and `/ai/v1/batches` endpoints. | | `allowed_completion_windows` | list | `["24h"]` | SLA options users can select. | -| `files.max_file_size` | integer | `104857600` | Maximum upload size in bytes. A gzip or zstd file is decoded as it uploads, and the limit applies to its decoded size. | +| `files.max_file_size` | integer | `104857600` | Maximum upload size in bytes. A gzip or zstd file is decoded as it uploads, and the limit applies to its decoded size; the compressed upload itself may not exceed the limit either. | | `files.default_expiry_seconds` | integer | `86400` | Default file retention. | ## Background Services diff --git a/dwctl/src/inference/body_limit.rs b/dwctl/src/inference/body_limit.rs index b3986b375..42255ec5e 100644 --- a/dwctl/src/inference/body_limit.rs +++ b/dwctl/src/inference/body_limit.rs @@ -237,7 +237,7 @@ fn decode_blocking(encoding: Encoding, bytes: &[u8], max_bytes: usize) -> Result decoder.window_log_max(ZSTD_WINDOW_LOG_MAX).map_err(|_| DecodeError::Invalid)?; Box::new(decoder) } - Encoding::Brotli => Box::new(brotli::Decompressor::new(bytes, 4096)), + Encoding::Brotli => return decode_brotli(bytes, max_bytes), }; let mut decoded = Vec::new(); reader @@ -250,6 +250,47 @@ fn decode_blocking(encoding: Encoding, bytes: &[u8], max_bytes: usize) -> Result Ok(Bytes::from(decoded)) } +/// Brotli as RFC 7932 defines it, with windows up to 16 MiB. `brotli::Decompressor` +/// also accepts the large-window extension, whose windows reach 1 GiB and are +/// allocated before any output, so the output bound alone can't stop them. +/// The strict decoder state refuses that extension; it is only reachable +/// through the streaming API. +fn decode_brotli(bytes: &[u8], max_bytes: usize) -> Result { + use brotli::{BrotliDecompressStream, BrotliResult, BrotliState, HeapAlloc, HuffmanCode}; + + let mut state = BrotliState::new_strict( + HeapAlloc::::new(0), + HeapAlloc::::new(0), + HeapAlloc::::new(HuffmanCode::default()), + ); + let mut buffer = vec![0u8; 64 * 1024]; + let mut decoded = Vec::new(); + let (mut available_in, mut input_offset, mut total_out) = (bytes.len(), 0, 0); + loop { + let (mut available_out, mut output_offset) = (buffer.len(), 0); + let result = BrotliDecompressStream( + &mut available_in, + &mut input_offset, + bytes, + &mut available_out, + &mut output_offset, + &mut buffer, + &mut total_out, + &mut state, + ); + decoded.extend_from_slice(&buffer[..output_offset]); + if decoded.len() > max_bytes { + return Err(DecodeError::TooLarge); + } + match result { + BrotliResult::NeedsMoreOutput => continue, + // Bytes after the end of the stream aren't brotli data either. + BrotliResult::ResultSuccess if available_in == 0 => return Ok(Bytes::from(decoded)), + _ => return Err(DecodeError::Invalid), + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -513,6 +554,38 @@ mod tests { assert_eq!(decode_blocking(Encoding::Zstd, &framed(24), 1024), Err(DecodeError::Invalid)); } + #[test] + fn refuses_brotli_large_window_streams() { + use brotli::enc::BrotliEncoderParams; + let params = BrotliEncoderParams { + large_window: true, + lgwin: 30, + ..Default::default() + }; + let mut large = Vec::new(); + brotli::BrotliCompress(&mut &b"{}"[..], &mut large, ¶ms).unwrap(); + // A valid large-window stream, as the permissive decoder shows... + let mut permissive = Vec::new(); + brotli::Decompressor::new(&large[..], 4096).read_to_end(&mut permissive).unwrap(); + assert_eq!(permissive, b"{}"); + // ...which the strict one refuses before allocating its window. + assert_eq!(decode_blocking(Encoding::Brotli, &large, 1024), Err(DecodeError::Invalid)); + let standard = compress(Encoding::Brotli, b"{}"); + assert_eq!(decode_blocking(Encoding::Brotli, &standard, 1024).unwrap(), &b"{}"[..]); + } + + #[test] + fn decodes_every_concatenated_frame() { + for encoding in [Encoding::Gzip, Encoding::Zstd] { + let mut frames = compress(encoding, b"{\"a\":"); + frames.extend(compress(encoding, b"1}")); + assert_eq!(decode_blocking(encoding, &frames, 1024).unwrap(), &b"{\"a\":1}"[..], "{encoding:?}"); + } + let mut trailing = compress(Encoding::Brotli, b"{}"); + trailing.extend_from_slice(b"junk"); + assert_eq!(decode_blocking(Encoding::Brotli, &trailing, 1024), Err(DecodeError::Invalid)); + } + #[tokio::test] async fn rejects_bodies_that_are_not_valid_for_their_encoding() { for encoding in [Encoding::Gzip, Encoding::Zstd, Encoding::Brotli] { diff --git a/dwctl/src/upload_compression.rs b/dwctl/src/upload_compression.rs index 6130da9b0..d96a988fb 100644 --- a/dwctl/src/upload_compression.rs +++ b/dwctl/src/upload_compression.rs @@ -68,6 +68,11 @@ impl Detected { /// Identifies the format from the first bytes of a file. pub(crate) fn detect(prefix: &[u8]) -> Detected { + // A zstd stream may open with a skippable frame (magic 0x184D2A50 to + // 0x184D2A5F, little-endian) before its data frame; the decoder skips it. + if let [0x50..=0x5f, 0x2a, 0x4d, 0x18, ..] = prefix { + return Detected::Zstd; + } const SIGNATURES: &[(&[u8], Detected)] = &[ (&[0x1f, 0x8b], Detected::Gzip), (&[0x28, 0xb5, 0x2f, 0xfd], Detected::Zstd), @@ -132,10 +137,11 @@ where decoder.multiple_members(true); Box::pin(decoder) } else { - Box::pin(ZstdDecoder::with_params( - reader, - &[async_compression::zstd::DParameter::window_log_max(ZSTD_WINDOW_LOG_MAX)], - )) + let mut decoder = + ZstdDecoder::with_params(reader, &[async_compression::zstd::DParameter::window_log_max(ZSTD_WINDOW_LOG_MAX)]); + // Like gzip members, concatenated zstd frames are one file. + decoder.multiple_members(true); + Box::pin(decoder) }; Box::pin(ReaderStream::with_capacity(decoder, DECODED_CHUNK).map_err(upload_or_decode_error)) } @@ -189,9 +195,16 @@ mod tests { async fn decodes_gzip_and_zstd_files_whatever_the_chunking() { let mut two_members = gzip(&JSONL[..18]); two_members.extend(gzip(&JSONL[18..])); + let mut two_frames = zstd::encode_all(&JSONL[..18], 3).unwrap(); + two_frames.extend(zstd::encode_all(&JSONL[18..], 3).unwrap()); + // A skippable frame (magic, 4-byte length, payload) ahead of the data. + let mut skippable_first = vec![0x50, 0x2a, 0x4d, 0x18, 3, 0, 0, 0, b'x', b'y', b'z']; + skippable_first.extend(zstd::encode_all(JSONL, 3).unwrap()); for (compressed, expected) in [ (gzip(JSONL), Detected::Gzip), (two_members, Detected::Gzip), + (two_frames, Detected::Zstd), + (skippable_first, Detected::Zstd), (zstd::encode_all(JSONL, 3).unwrap(), Detected::Zstd), (JSONL.to_vec(), Detected::Plain), ] { @@ -210,7 +223,15 @@ mod tests { let mut corrupt = gz.clone(); let middle = corrupt.len() / 2; corrupt[middle] ^= 0xff; - for broken in [gz[..gz.len() - 4].to_vec(), zst[..zst.len() - 3].to_vec(), corrupt] { + // A later frame is checked like the first, not ignored. + let mut later_frame_truncated = zst.clone(); + later_frame_truncated.extend_from_slice(&zst[..zst.len() - 3]); + for broken in [ + gz[..gz.len() - 4].to_vec(), + zst[..zst.len() - 3].to_vec(), + corrupt, + later_frame_truncated, + ] { let (_, decoded) = read_all(broken, 7).await; assert!(matches!(decoded, Err(ChunkError::Decode(_))), "{decoded:?}"); } @@ -242,6 +263,7 @@ mod tests { assert_eq!(detect(b"PK\x03\x04"), Detected::Unsupported("zip")); assert_eq!(detect(b"\x04\x22\x4d\x18"), Detected::Unsupported("lz4")); assert_eq!(detect(b"{\"custom_id\""), Detected::Plain); + assert_eq!(detect(&[0x5a, 0x2a, 0x4d, 0x18, 0, 0]), Detected::Zstd); assert_eq!(detect(b""), Detected::Plain); } From da882800bde1a382e96cf339de893a085cfba617 Mon Sep 17 00:00:00 2001 From: Steffen Ridderbusch Date: Wed, 21 Oct 2026 20:52:16 +0100 Subject: [PATCH 3/5] fix(dwctl): stop brotli decoding one byte past the body limit (DW-985) Every other codec reads at most max_bytes + 1 decoded bytes; brotli could buffer a further 64 KiB chunk before checking the limit. --- dwctl/src/inference/body_limit.rs | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/dwctl/src/inference/body_limit.rs b/dwctl/src/inference/body_limit.rs index 42255ec5e..506385800 100644 --- a/dwctl/src/inference/body_limit.rs +++ b/dwctl/src/inference/body_limit.rs @@ -267,14 +267,17 @@ fn decode_brotli(bytes: &[u8], max_bytes: usize) -> Result { let mut decoded = Vec::new(); let (mut available_in, mut input_offset, mut total_out) = (bytes.len(), 0, 0); loop { - let (mut available_out, mut output_offset) = (buffer.len(), 0); + // Never decode more than one byte past the limit, as the other + // codecs' `take(max_bytes + 1)` does. + let room = (max_bytes - decoded.len()).saturating_add(1).min(buffer.len()); + let (mut available_out, mut output_offset) = (room, 0); let result = BrotliDecompressStream( &mut available_in, &mut input_offset, bytes, &mut available_out, &mut output_offset, - &mut buffer, + &mut buffer[..room], &mut total_out, &mut state, ); @@ -574,6 +577,20 @@ mod tests { assert_eq!(decode_blocking(Encoding::Brotli, &standard, 1024).unwrap(), &b"{}"[..]); } + #[test] + fn decodes_up_to_the_limit_and_not_a_byte_further() { + let body = vec![b'x'; 200_000]; + for encoding in [Encoding::Gzip, Encoding::Zstd, Encoding::Brotli] { + let compressed = compress(encoding, &body); + assert_eq!(decode_blocking(encoding, &compressed, body.len()).unwrap().len(), body.len()); + assert_eq!( + decode_blocking(encoding, &compressed, body.len() - 1), + Err(DecodeError::TooLarge), + "{encoding:?}" + ); + } + } + #[test] fn decodes_every_concatenated_frame() { for encoding in [Encoding::Gzip, Encoding::Zstd] { From 00a72e371cd3b6c307a695953a1de4756b776ae5 Mon Sep 17 00:00:00 2001 From: Steffen Ridderbusch Date: Wed, 21 Oct 2026 22:24:47 +0100 Subject: [PATCH 4/5] fix(dwctl): keep mislabelled JSON bodies working and bound concurrent decodes (DW-985) Before compression was supported the Content-Encoding header was ignored, so a plain JSON body under a wrong or meaningless coding (a proxy's stray gzip, a client's utf-8) worked. Refusing it now would break those clients for nothing: pass it on as plain and count it instead. Decoding runs before authentication, and a few kilobytes of compressed body can expand to the body limit plus a decoder window. Bound decodes to one per CPU; requests over the bound wait rather than fail, so legitimate traffic is only delayed and uncompressed bodies never queue. --- dwctl/src/inference/body_limit.rs | 169 +++++++++++++++++++++--------- dwctl/src/test/strict_mode.rs | 18 +++- 2 files changed, 133 insertions(+), 54 deletions(-) diff --git a/dwctl/src/inference/body_limit.rs b/dwctl/src/inference/body_limit.rs index 506385800..013d08f90 100644 --- a/dwctl/src/inference/body_limit.rs +++ b/dwctl/src/inference/body_limit.rs @@ -1,7 +1,7 @@ //! Bound inference bodies before middleware parses, copies, or logs them, and //! decode compressed ones so every layer behind this one sees plain JSON. -use std::io::Read; +use std::{io::Read, sync::LazyLock}; use axum::{ Json, @@ -15,6 +15,8 @@ use axum::{ response::{IntoResponse, Response}, }; +use tokio::sync::Semaphore; + use super::rejections; /// The content codings an inference body may arrive in, as advertised on a 415. @@ -25,69 +27,96 @@ const ACCEPTED_ENCODINGS: &str = "gzip, zstd, br"; /// allocated. const ZSTD_WINDOW_LOG_MAX: u32 = 23; +/// Decodes running at once, one per CPU. Decoding runs before authentication, +/// and a few kilobytes of compressed body can expand to the full body limit +/// plus a decoder window of up to 16 MiB. Without a bound, a burst of small +/// anonymous requests could hold hundreds of those at once. Requests over the +/// bound wait for a slot rather than being refused, so legitimate traffic is +/// only ever delayed; uncompressed bodies never take one. +static DECODE_SLOTS: LazyLock = + LazyLock::new(|| Semaphore::new(std::thread::available_parallelism().map_or(1, |cpus| cpus.get()))); + /// Install outside all inference middleware, including request logging and /// daemon dispatch bypasses. File uploads use a separate router and limit. /// A zero limit retains the configured unlimited behavior. /// /// A body sent with `Content-Encoding: gzip`, `zstd` or `br` is decoded here, /// and the limit applies to the decoded size: that is what every inner layer -/// buffers. Any other coding is refused with a 415 before the body is read. +/// buffers. Any other coding is refused with a 415. +/// +/// Before compression was supported the header was ignored, so a plain JSON +/// body under a wrong or meaningless `Content-Encoding` worked. Such a body +/// still does: it is passed on as it is and counted, rather than refused. pub async fn limit_inference_body(State(max_bytes): State, request: Request, next: Next) -> Response { - let encoding = match content_encoding(request.headers()) { - Ok(encoding) => encoding, - Err(UnsupportedEncoding) => { - let (parts, _) = request.into_parts(); - let response = unsupported_encoding(parts.uri.path()); - return rejections::refuse(&parts, response, "unsupported_content_encoding", None, None); - } - }; - if max_bytes == 0 && encoding.is_none() { + let declared = content_encoding(request.headers()); + if max_bytes == 0 && declared == Ok(None) { return next.run(request).await; } let (mut parts, body) = request.into_parts(); - match read_inference_body(&parts.headers, body, max_bytes).await { - Ok(bytes) => { - let Some(encoding) = encoding else { - return next.run(Request::from_parts(parts, Body::from(bytes))).await; - }; + let bytes = match read_inference_body(&parts.headers, body, max_bytes).await { + Ok(bytes) => bytes, + Err(crate::errors::Error::PayloadTooLarge { message }) => { + let response = oversized_body(message, parts.uri.path()); + return rejections::refuse(&parts, response, "request_too_large", None, None); + } + Err(_) => { + let response = ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({"error": { + "message": "Failed to read request body", + "type": "invalid_request_error", + "code": "body_read_failed" + }})), + ) + .into_response(); + return rejections::refuse(&parts, response, "body_read_failed", None, None); + } + }; + + let body = match declared { + Ok(None) => return next.run(Request::from_parts(parts, Body::from(bytes))).await, + Ok(Some(encoding)) => { let limit = if max_bytes == 0 { usize::MAX } else { max_bytes }; - match decode(encoding, bytes, limit).await { + match decode(&DECODE_SLOTS, encoding, bytes.clone(), limit).await { Ok(decoded) => { metrics::counter!("dwctl_inference_bodies_decoded_total", "encoding" => encoding.name()).increment(1); - parts.headers.remove(CONTENT_ENCODING); - parts.headers.insert(CONTENT_LENGTH, HeaderValue::from(decoded.len())); - next.run(Request::from_parts(parts, Body::from(decoded))).await + decoded } Err(DecodeError::TooLarge) => { let message = format!("Decoded request body exceeds the maximum size of {max_bytes} bytes"); let response = oversized_body(message, parts.uri.path()); - rejections::refuse(&parts, response, "request_too_large", None, None) + return rejections::refuse(&parts, response, "request_too_large", None, None); } + Err(DecodeError::Invalid) if is_json(&bytes) => mislabelled(encoding.name(), bytes), Err(DecodeError::Invalid) => { let message = format!("Request body is not valid {} data", encoding.name()); let response = invalid_request(StatusCode::BAD_REQUEST, parts.uri.path(), "body_decode_failed", &message); - rejections::refuse(&parts, response, "body_decode_failed", None, None) + return rejections::refuse(&parts, response, "body_decode_failed", None, None); } } } - Err(crate::errors::Error::PayloadTooLarge { message }) => { - let response = oversized_body(message, parts.uri.path()); - rejections::refuse(&parts, response, "request_too_large", None, None) - } - Err(_) => { - let response = ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({"error": { - "message": "Failed to read request body", - "type": "invalid_request_error", - "code": "body_read_failed" - }})), - ) - .into_response(); - rejections::refuse(&parts, response, "body_read_failed", None, None) + Err(UnsupportedEncoding) if is_json(&bytes) => mislabelled("other", bytes), + Err(UnsupportedEncoding) => { + let response = unsupported_encoding(parts.uri.path()); + return rejections::refuse(&parts, response, "unsupported_content_encoding", None, None); } - } + }; + // Inner layers and the upstream see the body as it now is: plain. + parts.headers.remove(CONTENT_ENCODING); + parts.headers.insert(CONTENT_LENGTH, HeaderValue::from(body.len())); + next.run(Request::from_parts(parts, Body::from(body))).await +} + +/// Whether `bytes` is a complete JSON document, as every inference body is. +fn is_json(bytes: &[u8]) -> bool { + serde_json::from_slice::(bytes).is_ok() +} + +/// A plain JSON body under a `Content-Encoding` it isn't in: passed on as it is. +fn mislabelled(declared: &'static str, bytes: Bytes) -> Bytes { + metrics::counter!("dwctl_inference_bodies_mislabelled_total", "declared" => declared).increment(1); + bytes } /// Shared with the authenticated playground proxy, whose model parser runs @@ -209,10 +238,14 @@ enum DecodeError { Invalid, } -/// Decodes `bytes`, stopping as soon as the output passes `max_bytes`, so a -/// small body that inflates without bound costs at most the limit. Runs on a -/// blocking thread: inflating a full-size body takes tens of milliseconds. -async fn decode(encoding: Encoding, bytes: Bytes, max_bytes: usize) -> Result { +/// Decodes `bytes` once one of `slots` is free, stopping as soon as the output +/// passes `max_bytes`, so a small body that inflates without bound costs at +/// most the limit. Runs on a blocking thread: inflating a full-size body takes +/// tens of milliseconds. +async fn decode(slots: &Semaphore, encoding: Encoding, bytes: Bytes, max_bytes: usize) -> Result { + let Ok(_slot) = slots.acquire().await else { + return Err(DecodeError::Invalid); + }; tokio::task::spawn_blocking(move || decode_blocking(encoding, &bytes, max_bytes)) .await .unwrap_or(Err(DecodeError::Invalid)) @@ -476,19 +509,13 @@ mod tests { } #[tokio::test] - async fn refuses_other_encodings_with_415_before_reading_the_body() { + async fn refuses_other_encodings_with_415() { for header in ["deflate", "lz4", "compress", "gzip, br", "snappy"] { - let polls = Arc::new(AtomicUsize::new(0)); - let observed = polls.clone(); - let body = Body::from_stream(futures::stream::poll_fn(move |_| { - observed.fetch_add(1, Ordering::SeqCst); - std::task::Poll::Ready(None::>) - })); let response = decoding_app(1024) .oneshot( Request::post("/chat/completions") .header(CONTENT_ENCODING, header) - .body(body) + .body(Body::from(&b"\x04\x22\x4d\x18 compressed bytes"[..])) .unwrap(), ) .await @@ -496,10 +523,48 @@ mod tests { assert_eq!(response.status(), StatusCode::UNSUPPORTED_MEDIA_TYPE, "{header}"); assert_eq!(response.headers()[ACCEPT_ENCODING], "gzip, zstd, br"); assert_eq!(error_of(response).await["error"]["code"], "unsupported_content_encoding"); - assert_eq!(polls.load(Ordering::SeqCst), 0, "{header}"); } } + #[tokio::test] + async fn passes_plain_json_under_a_wrong_content_encoding_through() { + // These worked before compression was supported, when the header was + // ignored; they must keep working, as plain bodies. + let json = br#"{"model":"m","messages":[]}"#; + for header in ["utf-8", "gzip", "zstd", "br", "lz4", "gzip, br"] { + let response = decoding_app(1024) + .oneshot( + Request::post("/chat/completions") + .header(CONTENT_ENCODING, header) + .body(Body::from(&json[..])) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK, "{header}"); + assert_eq!( + axum::body::to_bytes(response.into_body(), 1024).await.unwrap(), + &json[..], + "{header}" + ); + } + } + + #[tokio::test] + async fn decodes_wait_for_a_free_slot() { + let slots = Semaphore::new(1); + let compressed = Bytes::from(compress(Encoding::Gzip, b"{}")); + let busy = slots.acquire().await.unwrap(); + let mut waiting = Box::pin(decode(&slots, Encoding::Gzip, compressed, 1024)); + assert!( + futures::poll!(&mut waiting).is_pending(), + "a decode must wait while every slot is taken" + ); + drop(busy); + assert_eq!(waiting.await.unwrap(), &b"{}"[..]); + assert_eq!(slots.available_permits(), 1, "the slot is returned after decoding"); + } + #[tokio::test] async fn answers_messages_requests_in_the_anthropic_error_shape() { let app = Router::new() @@ -610,7 +675,7 @@ mod tests { .oneshot( Request::post("/chat/completions") .header(CONTENT_ENCODING, encoding.name()) - .body(Body::from(r#"{"model":"m"}"#)) + .body(Body::from("neither compressed nor JSON")) .unwrap(), ) .await diff --git a/dwctl/src/test/strict_mode.rs b/dwctl/src/test/strict_mode.rs index be15329c3..c456a3929 100644 --- a/dwctl/src/test/strict_mode.rs +++ b/dwctl/src/test/strict_mode.rs @@ -103,12 +103,26 @@ async fn assert_compressed_bodies_reach_upstream_decoded(pool: PgPool, strict_mo assert_eq!(body["messages"][0]["content"], "compressed hello", "strict_mode={strict_mode}"); } + // Plain JSON under a coding it isn't in worked before compression was + // supported, and reaches the upstream plain. let response = server .post("/ai/v1/chat/completions") .add_header("Authorization", format!("Bearer {key}")) - .add_header("Content-Encoding", "lz4") + .add_header("Content-Type", "application/json") + .add_header("Content-Encoding", "utf-8") .bytes(json.into()) .await; + response.assert_status_ok(); + let received = mock.received_requests().await.unwrap(); + assert_eq!(received.len(), 3, "strict_mode={strict_mode}"); + assert!(received[2].headers.get("content-encoding").is_none(), "strict_mode={strict_mode}"); + + let response = server + .post("/ai/v1/chat/completions") + .add_header("Authorization", format!("Bearer {key}")) + .add_header("Content-Encoding", "lz4") + .bytes((&b"\x04\x22\x4d\x18 compressed bytes"[..]).into()) + .await; response.assert_status(axum::http::StatusCode::UNSUPPORTED_MEDIA_TYPE); assert_eq!( response.json::()["error"]["code"], @@ -116,7 +130,7 @@ async fn assert_compressed_bodies_reach_upstream_decoded(pool: PgPool, strict_mo ); assert_eq!( mock.received_requests().await.unwrap().len(), - 2, + 3, "refused before reaching the upstream" ); } From e6be2db3fd8bd89a73c4cd9f9467013b0db3b114 Mon Sep 17 00:00:00 2001 From: Steffen Ridderbusch Date: Thu, 22 Oct 2026 11:57:55 +0100 Subject: [PATCH 5/5] fix(dwctl): hold a decode slot until the decode ends, not the request (DW-985) A request dropped mid-decode (the client went away) released its slot while the blocking decode ran on, so disconnecting clients could exceed the one-per-CPU bound. Move the permit into the blocking task. --- dwctl/src/inference/body_limit.rs | 45 ++++++++++++++++++++++++------- 1 file changed, 36 insertions(+), 9 deletions(-) diff --git a/dwctl/src/inference/body_limit.rs b/dwctl/src/inference/body_limit.rs index 013d08f90..94b72f366 100644 --- a/dwctl/src/inference/body_limit.rs +++ b/dwctl/src/inference/body_limit.rs @@ -1,7 +1,10 @@ //! Bound inference bodies before middleware parses, copies, or logs them, and //! decode compressed ones so every layer behind this one sees plain JSON. -use std::{io::Read, sync::LazyLock}; +use std::{ + io::Read, + sync::{Arc, LazyLock}, +}; use axum::{ Json, @@ -33,8 +36,8 @@ const ZSTD_WINDOW_LOG_MAX: u32 = 23; /// anonymous requests could hold hundreds of those at once. Requests over the /// bound wait for a slot rather than being refused, so legitimate traffic is /// only ever delayed; uncompressed bodies never take one. -static DECODE_SLOTS: LazyLock = - LazyLock::new(|| Semaphore::new(std::thread::available_parallelism().map_or(1, |cpus| cpus.get()))); +static DECODE_SLOTS: LazyLock> = + LazyLock::new(|| Arc::new(Semaphore::new(std::thread::available_parallelism().map_or(1, |cpus| cpus.get())))); /// Install outside all inference middleware, including request logging and /// daemon dispatch bypasses. File uploads use a separate router and limit. @@ -242,15 +245,26 @@ enum DecodeError { /// passes `max_bytes`, so a small body that inflates without bound costs at /// most the limit. Runs on a blocking thread: inflating a full-size body takes /// tens of milliseconds. -async fn decode(slots: &Semaphore, encoding: Encoding, bytes: Bytes, max_bytes: usize) -> Result { - let Ok(_slot) = slots.acquire().await else { - return Err(DecodeError::Invalid); - }; - tokio::task::spawn_blocking(move || decode_blocking(encoding, &bytes, max_bytes)) +async fn decode(slots: &Arc, encoding: Encoding, bytes: Bytes, max_bytes: usize) -> Result { + in_slot(slots, move || decode_blocking(encoding, &bytes, max_bytes)) .await .unwrap_or(Err(DecodeError::Invalid)) } +/// Runs `work` on a blocking thread once one of `slots` is free. The slot +/// moves into the thread with the work: a request dropped mid-decode (the +/// client went away) stops waiting for the result, but the work runs on and +/// must keep counting against the bound until it ends. +async fn in_slot(slots: &Arc, work: impl FnOnce() -> T + Send + 'static) -> Option { + let slot = slots.clone().acquire_owned().await.ok()?; + tokio::task::spawn_blocking(move || { + let _slot = slot; + work() + }) + .await + .ok() +} + fn decode_blocking(encoding: Encoding, bytes: &[u8], max_bytes: usize) -> Result { let reader: Box = match encoding { // HTTP gzip bodies may hold several concatenated members. @@ -552,7 +566,7 @@ mod tests { #[tokio::test] async fn decodes_wait_for_a_free_slot() { - let slots = Semaphore::new(1); + let slots = Arc::new(Semaphore::new(1)); let compressed = Bytes::from(compress(Encoding::Gzip, b"{}")); let busy = slots.acquire().await.unwrap(); let mut waiting = Box::pin(decode(&slots, Encoding::Gzip, compressed, 1024)); @@ -565,6 +579,19 @@ mod tests { assert_eq!(slots.available_permits(), 1, "the slot is returned after decoding"); } + #[tokio::test] + async fn a_dropped_request_keeps_its_slot_until_the_decode_ends() { + let slots = Arc::new(Semaphore::new(1)); + let (finish, finished) = std::sync::mpsc::channel::<()>(); + let mut running = Box::pin(in_slot(&slots, move || finished.recv().unwrap())); + assert!(futures::poll!(&mut running).is_pending(), "the work is running"); + drop(running); + assert_eq!(slots.available_permits(), 0, "the abandoned work still holds its slot"); + finish.send(()).unwrap(); + // Returned once the work ends: this waits for exactly that. + drop(slots.acquire().await.unwrap()); + } + #[tokio::test] async fn answers_messages_requests_in_the_anthropic_error_shape() { let app = Router::new()