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..ec4f0cb6d 100644 --- a/config.yaml +++ b/config.yaml @@ -505,7 +505,9 @@ 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 +# # (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 @@ -526,6 +528,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..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. | +| `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/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..94b72f366 100644 --- a/dwctl/src/inference/body_limit.rs +++ b/dwctl/src/inference/body_limit.rs @@ -1,30 +1,67 @@ -//! 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, + sync::{Arc, LazyLock}, +}; 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 tokio::sync::Semaphore; + 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; + +/// 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(|| 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. /// 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 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 { - if max_bytes == 0 { + let declared = content_encoding(request.headers()); + if max_bytes == 0 && declared == Ok(None) { return next.run(request).await; } - let (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, + let (mut parts, body) = request.into_parts(); + 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()); - rejections::refuse(&parts, response, "request_too_large", None, None) + return rejections::refuse(&parts, response, "request_too_large", None, None); } Err(_) => { let response = ( @@ -36,9 +73,53 @@ pub async fn limit_inference_body(State(max_bytes): State, request: Reque }})), ) .into_response(); - rejections::refuse(&parts, response, "body_read_failed", None, None) + 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(&DECODE_SLOTS, encoding, bytes.clone(), limit).await { + Ok(decoded) => { + metrics::counter!("dwctl_inference_bodies_decoded_total", "encoding" => encoding.name()).increment(1); + 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()); + 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); + return rejections::refuse(&parts, response, "body_decode_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 @@ -87,6 +168,179 @@ 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` 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: &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. + 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 => return decode_brotli(bytes, max_bytes), + }; + 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)) +} + +/// 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 { + // 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[..room], + &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::*; @@ -196,6 +450,321 @@ 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() { + for header in ["deflate", "lz4", "compress", "gzip, br", "snappy"] { + let response = decoding_app(1024) + .oneshot( + Request::post("/chat/completions") + .header(CONTENT_ENCODING, header) + .body(Body::from(&b"\x04\x22\x4d\x18 compressed bytes"[..])) + .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"); + } + } + + #[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 = 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)); + 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 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() + .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)); + } + + #[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_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] { + 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] { + let response = decoding_app(1024) + .oneshot( + Request::post("/chat/completions") + .header(CONTENT_ENCODING, encoding.name()) + .body(Body::from("neither compressed nor JSON")) + .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..c456a3929 100644 --- a/dwctl/src/test/strict_mode.rs +++ b/dwctl/src/test/strict_mode.rs @@ -53,6 +53,98 @@ 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}"); + } + + // 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-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"], + "unsupported_content_encoding" + ); + assert_eq!( + mock.received_requests().await.unwrap().len(), + 3, + "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..d96a988fb --- /dev/null +++ b/dwctl/src/upload_compression.rs @@ -0,0 +1,277 @@ +//! 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 { + // 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), + (&[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 { + 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)) + } + }; + (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..])); + 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), + ] { + 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; + // 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:?}"); + } + } + + #[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(&[0x5a, 0x2a, 0x4d, 0x18, 0, 0]), Detected::Zstd); + 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"); + } +}