From c2f8ac2da773e3a8ff2e7ba29a540e12da25d162 Mon Sep 17 00:00:00 2001 From: Luke Valenta Date: Thu, 30 Jul 2026 12:37:39 -0400 Subject: [PATCH 1/8] mirror_worker: stream and incrementally commit add-entries uploads Replace the whole-body buffer with a streaming decode: gunzip and parse the body as it arrives through a small stream-buffer adapter (retry the parse as chunks land, consume on success), so a large upload no longer holds the entire body in memory. Persist incrementally every `commit_packages` packages (default 32, config capped at 1024) so a long or interrupted upload advances the frontier as it streams rather than only at the end. Each flush commits from the last persisted frontier and is resumable; the running frontier is threaded locally across flushes so the DO is not re-queried between chunks. Tiles are immutable and content-addressed and the DO advance is a monotone compare-and-swap, so a repeated or concurrent flush of the same range is harmless. --- crates/mirror_worker/config.dev.json | 1 + crates/mirror_worker/config.schema.json | 7 + crates/mirror_worker/config/src/lib.rs | 34 +- crates/mirror_worker/src/add_entries.rs | 551 +++++++++++++++------- crates/mirror_worker/src/body.rs | 254 +++++++--- crates/mirror_worker/src/commit.rs | 36 ++ crates/mirror_worker/src/lib.rs | 1 + crates/mirror_worker/src/stream_buffer.rs | 195 ++++++++ 8 files changed, 847 insertions(+), 232 deletions(-) create mode 100644 crates/mirror_worker/src/stream_buffer.rs diff --git a/crates/mirror_worker/config.dev.json b/crates/mirror_worker/config.dev.json index dc91a78a..e41633b3 100644 --- a/crates/mirror_worker/config.dev.json +++ b/crates/mirror_worker/config.dev.json @@ -5,6 +5,7 @@ "submission_prefix": "http://localhost:8787/", "monitoring_prefix": "http://localhost:8787/", "clean_interval_secs": 5, + "commit_packages": 2, "logs": { "oid/1.3.6.1.4.1.32473.2": { "description": "Dev-only MTC CA cosigner. Key name is the CA ID; the mirror serves log numbers 1-6 as origins oid/1.3.6.1.4.1.32473.2.0.. log_public_keys holds a dev-only ML-DSA-44 SPKI.", diff --git a/crates/mirror_worker/config.schema.json b/crates/mirror_worker/config.schema.json index 0c39947b..fd5da031 100644 --- a/crates/mirror_worker/config.schema.json +++ b/crates/mirror_worker/config.schema.json @@ -30,6 +30,13 @@ "default": 3600, "description": "How often (in seconds) the per-origin partial-tile cleaner wakes to clean orphaned partial tiles from object storage. Defaults to 3600 (one hour) when omitted." }, + "commit_packages": { + "type": "integer", + "minimum": 1, + "maximum": 1024, + "default": 32, + "description": "How many entry packages add-entries verifies before flushing them to storage and advancing the persisted-entry frontier. Bounds in-memory buffering and gives durable mid-request progress on large uploads. Defaults to 32 (the recommended per-request package budget) when omitted; capped at 1024 to bound worst-case buffering." + }, "logs": { "type": "object", "description": "CAs this mirror mirrors, keyed by log_key_name: the CA cosigner's note-signature name (the CA ID) on the checkpoints it ingests. Used as a signed-note key name at runtime, so per c2sp.org/signed-note it MUST NOT contain '+', whitespace, or control characters.", diff --git a/crates/mirror_worker/config/src/lib.rs b/crates/mirror_worker/config/src/lib.rs index 5197b506..433f605b 100644 --- a/crates/mirror_worker/config/src/lib.rs +++ b/crates/mirror_worker/config/src/lib.rs @@ -62,6 +62,13 @@ pub struct AppConfig { /// back to a one-hour default (see [`Self::clean_interval_secs`]). /// Consumed by [`mirror_worker`](../mirror_worker/)'s `cleaner_do`. pub clean_interval_secs: Option, + /// How many entry packages the `add-entries` handler verifies before + /// flushing them to storage and advancing the persisted-entry + /// frontier. Bounds in-memory buffering and gives durable mid-request + /// progress on large uploads. `None` falls back to a default of 32 + /// (see [`Self::commit_packages`]). Consumed by + /// [`mirror_worker`](../mirror_worker/)'s `add_entries`. + pub commit_packages: Option, /// CAs this mirror mirrors, keyed by `log_key_name`: the CA /// cosigner's note-signature name (the CA ID) carried by the /// checkpoints it ingests. @@ -118,6 +125,17 @@ impl AppConfig { self.clean_interval_secs.unwrap_or(3600) } + /// How many entry packages `add-entries` commits per flush, falling + /// back to 32 when `commit_packages` is unset. 32 matches the + /// per-request package budget clients are recommended to stay within + /// (tlog-mirror "Implementation Considerations"), so a compliant + /// single-request upload still commits once, while larger uploads + /// flush every 32 packages instead of buffering the whole body. + #[must_use] + pub fn commit_packages(&self) -> u64 { + self.commit_packages.unwrap_or(32) + } + /// Validate the configuration beyond what `serde` and the JSON schema /// can express. /// @@ -141,9 +159,9 @@ impl AppConfig { /// signed-note key name length cap, since each origin is itself /// used as a checkpoint origin. /// - /// Simple single-field bounds (e.g. the log-number ranges) are - /// expressed in `config.schema.json` and enforced by the build - /// script, so they are not re-checked here. + /// Simple single-field bounds (e.g. `commit_packages` and the + /// log-number ranges) are expressed in `config.schema.json` and + /// enforced by the build script, so they are not re-checked here. /// /// `log_key_name` uniqueness across log entries is not checked here; /// it is enforced earlier, during deserialization (see @@ -340,6 +358,7 @@ mod tests { submission_prefix: "https://mirror.example/".to_owned(), monitoring_prefix: None, clean_interval_secs: None, + commit_packages: None, logs: HashMap::from([( "example.com/log1".to_owned(), LogParams { @@ -464,6 +483,14 @@ mod tests { .expect("a valid log-number window is accepted"); } + #[test] + fn commit_packages_defaults_to_32() { + let mut cfg = good_app_config(); + assert_eq!(cfg.commit_packages(), 32); + cfg.commit_packages = Some(8); + assert_eq!(cfg.commit_packages(), 8); + } + #[test] fn validate_rejects_inverted_window() { let cfg = with_log(|log| { @@ -532,6 +559,7 @@ mod tests { submission_prefix: "https://mirror.example/".to_owned(), monitoring_prefix: None, clean_interval_secs: None, + commit_packages: None, logs: HashMap::from([( "a".repeat(250), LogParams { diff --git a/crates/mirror_worker/src/add_entries.rs b/crates/mirror_worker/src/add_entries.rs index 9209fb66..7a3e62e5 100644 --- a/crates/mirror_worker/src/add_entries.rs +++ b/crates/mirror_worker/src/add_entries.rs @@ -4,19 +4,21 @@ //! `POST /add-entries` handler. //! //! Implements the [c2sp.org/tlog-mirror `add-entries`][add-e] endpoint: -//! read the (optionally gzip) request body, verify each [`EntryPackage`] +//! stream the (optionally gzip) request body, verify each [`EntryPackage`] //! against the target pending checkpoint with a subtree consistency proof, -//! persist the verified entries as bundles and hash tiles (see -//! [`crate::commit`]), and advance the persisted-entry frontier. +//! and incrementally persist the verified entries as bundles and hash +//! tiles (see [`crate::commit`]), advancing the persisted-entry frontier. //! -//! The whole body is buffered, then all its packages are committed in a -//! single pass; the mirror checkpoint is cosigned only once the upload is -//! durably committed. A follow-up commit adds incremental streaming so a -//! large upload does not have to buffer the entire body in memory. +//! Packages are committed in chunks of `commit_packages` (config, +//! default 32) rather than buffering the whole body: every +//! `commit_packages` verified packages are flushed to storage and the +//! frontier is advanced, bounding in-memory buffering and giving durable +//! mid-request progress, per the tlog-mirror streaming model. The mirror +//! checkpoint is cosigned only once the whole upload is durably committed. //! //! A complete upload writes the cosigned checkpoint and returns 200 with //! the mirror's [cosignature][cosig] line(s); a client-truncated upload -//! keeps the persisted prefix and returns 202 with the advanced next entry +//! keeps the flushed prefix and returns 202 with the advanced next entry //! so the client can resume (see [Processing][proc]). //! //! [add-e]: https://c2sp.org/tlog-mirror#add-entries @@ -43,7 +45,7 @@ use tlog_mirror::{ #[allow(clippy::wildcard_imports)] use worker::*; -use generic_log_worker::util::now_millis; +use generic_log_worker::{ObjectBackend, util::now_millis}; use crate::{ body, commit, @@ -54,15 +56,16 @@ use crate::{ NextEntry, PendingCheckpoint, state_stub, }, storage::load_origin_bucket, + stream_buffer::StreamBuffer, }; /// Handle `POST /add-entries`. /// -/// See the module-level comment for the full flow: read and verify entry -/// packages from the request body, persist the verified entries and -/// advance the persisted-entry frontier, and either cosign the mirror -/// checkpoint (200) or, for a truncated upload, persist the verified -/// prefix and return 202. +/// See the module-level comment for the full flow: parse and verify entry +/// packages over a streamed (optionally gzip) request body, persist the +/// verified entries and advance the persisted-entry frontier, and either +/// cosign the mirror checkpoint (200) or, for a truncated upload, persist +/// the verified prefix and return 202. #[worker::send] pub(crate) async fn add_entries( State(env): State, @@ -85,19 +88,12 @@ pub(crate) async fn add_entries( // Considerations"). The Workers runtime does not decompress request // bodies, so gzip-encoded bodies are gunzipped here; unknown encodings // are 415'd (see `crate::body`). - // - // The whole (decoded) body is buffered before processing; a follow-up - // commit streams it instead of buffering it all. - let raw = body::read_decoded_body(&parts.headers, body).await?; - let mut cursor = Cursor::new(raw.as_slice()); + let stream = body::decoded_stream(&parts.headers, body)?; + let mut buf = StreamBuffer::new(stream); - let header = match AddEntriesRequestHeader::read_from(&mut cursor) { - Ok(header) => header, - Err(e) => { - log::warn!("add-entries: malformed header: {e:?}"); - return Err(AppError::BadRequest(e.to_string())); - } - }; + // Pull chunks until the header parses, retrying on UnexpectedEof. The + // header size is bounded (~131 KB max), so the loop terminates. + let header = parse_header(&mut buf).await?; let Some(verifiers) = log_verifiers(&header.log_origin) else { return Err(AppError::UnknownLogOrigin); @@ -180,15 +176,7 @@ pub(crate) async fn add_entries( ) .await?; - verify_and_persist( - &env, - &header, - &snapshot, - &target, - &mut cursor, - &first_prefix, - ) - .await + stream_and_commit(&env, &header, &snapshot, &target, &mut buf, &first_prefix).await } /// Return true iff the request's `Content-Type` is @@ -205,67 +193,226 @@ fn content_type_is_octet_stream(headers: &axum::http::HeaderMap) -> bool { == "application/octet-stream" } -/// Read, verify, and persist the entry packages for `[upload_start, -/// upload_end)`, then produce the HTTP response. +/// Stream, verify, and incrementally persist the entry packages for +/// `[upload_start, upload_end)`, then produce the HTTP response. +/// +/// Packages are read and verified in canonical order and buffered until +/// `commit_packages` (config, default 32) have accumulated, then flushed +/// as entry bundles + hash tiles with the persisted-entry frontier +/// advanced in the DO (see [`flush_chunk`]). This bounds in-memory +/// buffering and gives durable mid-request progress rather than deferring +/// every write to the end of the body. Entries below the frontier at +/// request start are already persisted, so they are verified but not +/// re-saved (spec: "skip saving already-written entries"). /// -/// Packages are read and verified in canonical order from the buffered -/// body. Entries below the frontier at request start are already -/// persisted, so they are verified but not re-saved (spec: "skip saving -/// already-written entries"); the remaining entries are committed in a -/// single [`commit::persist_entries`] call and the persisted-entry -/// frontier is advanced once in the DO. A follow-up commit flushes -/// incrementally instead of once at the end. +/// The running frontier `(size, hash)` is threaded locally across flushes +/// from each [`commit::persist_entries`] result, so the DO is not +/// re-queried between chunks. Because tiles are immutable and +/// content-addressed and the DO advance is a monotone compare-and-swap, a +/// repeated or concurrent flush of the same range is harmless. /// /// Response cases (see [Processing][proc]): /// /// * every package received and the recomputed tree matches the target: /// cosign the mirror checkpoint, advance it, and return 200 with the /// cosignature line(s). -/// * client truncation after at least one complete package: persist the -/// verified prefix and return 202 with the advanced next entry. +/// * client truncation after at least one complete package: return 202 +/// with the advanced next entry so the client resumes. /// * truncation before the first complete package, or a malformed body: /// 400. A package that fails subtree-consistency verification: 422. /// /// # Errors /// -/// Returns an error on a storage failure while persisting, or (500) if the -/// recomputed root of a complete upload disagrees with the target -/// checkpoint. +/// Returns an error on a transport failure reading the body stream, a +/// storage failure while flushing, or (500) if the recomputed root of a +/// complete upload disagrees with the target checkpoint. /// /// [proc]: https://c2sp.org/tlog-mirror#processing -#[allow(clippy::too_many_lines)] -async fn verify_and_persist( +async fn stream_and_commit( env: &Env, header: &AddEntriesRequestHeader, snapshot: &MirrorStateSnapshot, target: &PendingCheckpoint, - cursor: &mut Cursor<&[u8]>, + buf: &mut StreamBuffer, first_prefix: &[Vec], -) -> ApiResult { +) -> ApiResult +where + S: futures_util::Stream>> + Unpin, +{ let bucket = load_origin_bucket(env, &header.log_origin)?; + let result = persist_packages( + env, + header, + target, + buf, + first_prefix, + &bucket, + &snapshot.next_entry, + ) + .await?; + + let persisted_new = result.frontier_size > snapshot.next_entry.size; + + if result.truncated { + log::info!( + "add-entries: client-truncated after {} complete packages; persisted through {}", + result.packages_received, + result.frontier_size, + ); + return Ok(mirror_info_202( + env, + snapshot, + &header.log_origin, + result.frontier_size, + )); + } + + // Spec: the mirror updates its checkpoint to `upload_end` only once + // "the next entry will be greater or equal to `upload_end`", i.e. all + // entries up to `upload_end` are durably persisted. A request that + // persists nothing (e.g. an empty body, or one whose packages are all + // already-persisted) must not let us cosign past our frontier: without + // this guard `upload_end` above `next_entry` would sign a checkpoint at + // a size we never wrote tiles for. When the frontier has not reached + // `upload_end`, treat it like a truncated upload and 202 so the client + // resumes from the advertised next entry. + if result.frontier_size < header.upload_end { + log::info!( + "add-entries: frontier {} below upload_end {}; nothing to persist this request, \ + returning 202 to resume", + result.frontier_size, + header.upload_end, + ); + return Ok(mirror_info_202( + env, + snapshot, + &header.log_origin, + result.frontier_size, + )); + } + + // Every canonical package was received. Any bytes still in the stream + // are discarded, not rejected: the spec says "the mirror discards any + // partial bytes after the last successfully authenticated entry + // package", and the only defined 400 is when no package authenticated + // at all (handled in persist_packages). Rejecting here would also be + // dishonest, since the entries were already persisted and the frontier + // advanced. + + // When we persisted new entries, the recomputed tree MUST match the + // pending checkpoint the packages were proven against. A mismatch + // means proof verification and tile computation disagree: an internal + // error, never a client fault. When nothing new was persisted (a + // re-upload of an already-persisted range), the log-signed target is + // trusted directly: pending checkpoints are consistency-chained on the + // add-checkpoint path and tickets only revive our own past pendings, so + // the target is on the same branch as the persisted tree by + // construction. + if persisted_new + && (result.frontier_size != header.upload_end || result.frontier_hash != target.hash) + { + log::error!( + "add-entries: recomputed frontier ({}, {}) != target ({}, {})", + result.frontier_size, + result.frontier_hash, + target.size, + target.hash, + ); + return Err(AppError::InternalServerError( + "recomputed root mismatch".to_owned(), + )); + } + + // Frontier ahead of `upload_end` (a prior upload persisted past it, or a + // racing client): the edge tiles were written at the larger frontier, so + // the narrower "cut" tiles a tree of size `upload_end` needs may be + // missing and a cosignature at `upload_end` would be unverifiable against + // our own tiles. Spec requires the tree at `upload_end` be servable before + // cosigning, so synthesize those cut tiles first. + if result.frontier_size > header.upload_end { + commit::ensure_cut_tiles( + &bucket, + header.upload_end, + result.frontier_size, + result.frontier_hash, + ) + .await?; + } + + cosign_and_serve(env, header, target, snapshot).await +} + +/// The persisted-entry frontier reached by [`persist_packages`], plus +/// whether the client truncated the stream and how many complete packages +/// were verified. +struct StreamResult { + frontier_size: u64, + frontier_hash: Hash, + packages_received: u64, + truncated: bool, +} + +/// Read, verify, and incrementally flush the entry packages for +/// `[upload_start, upload_end)`, resuming from the frontier `start`. +/// +/// Verified entries are buffered until `commit_packages` (config, default +/// 32) packages have accumulated, then flushed to storage with the +/// frontier advanced in the DO (see [`flush_chunk`]); the trailing partial +/// chunk is flushed before returning. Entries below `start.size` are +/// already persisted and skipped (spec: "skip saving already-written +/// entries"). The frontier `(size, hash)` is threaded locally across +/// flushes, so the DO is not re-queried between chunks. +/// +/// # Errors +/// +/// 400 if the body is malformed or truncates before the first complete +/// package, 422 if a package fails subtree-consistency verification, or a +/// transport/storage error from reading the body or flushing a chunk. +async fn persist_packages( + env: &Env, + header: &AddEntriesRequestHeader, + target: &PendingCheckpoint, + buf: &mut StreamBuffer, + first_prefix: &[Vec], + bucket: &O, + start: &NextEntry, +) -> ApiResult +where + S: futures_util::Stream>> + Unpin, + O: ObjectBackend, +{ + // config.schema.json caps commit_packages (max 1024), enforced by the + // build script, so this always fits usize; the fallback is unreachable. + let commit_packages = usize::try_from(crate::CONFIG.commit_packages()).unwrap_or(usize::MAX); + // Entries below the request-start frontier are already persisted; new // persistence begins at this fixed boundary. - let initial_next = snapshot.next_entry.size; + let initial_next = start.size; + + // Running persisted-entry frontier, threaded locally across flushes. + let mut frontier_size = start.size; + let mut frontier_hash = start.hash; - // Newly-received entries collected across all packages for a single - // commit, and the leaf index one past the last collected entry. - let mut entries: Vec> = Vec::new(); - let mut collected_end = initial_next; + // Entries buffered for the next flush, the leaf index one past the last + // buffered entry (the flush target), and how many packages are buffered. + let mut chunk: Vec> = Vec::new(); + let mut chunk_end = frontier_size; + let mut chunk_pkgs = 0usize; let mut packages_received: u64 = 0; let mut truncated = false; for (pkg_start, pkg_end) in package_ranges(header.upload_start, header.upload_end) { let num_entries = pkg_end - pkg_start; - let pkg = match read_package(cursor, num_entries) { - PackageOutcome::Ok(pkg) => pkg, - // The body ended between or partway through a package: client - // truncation. Keep whatever complete packages were verified. - PackageOutcome::Eof => { + let pkg = match parse_next_package(buf, num_entries).await? { + ParseOutcome::Ok(pkg) => pkg, + // Clean and mid-package EOF are both client truncation: keep + // whatever complete packages were already flushed/buffered. + ParseOutcome::CleanEof | ParseOutcome::MidPackageEof => { truncated = true; break; } - PackageOutcome::Err(e) => { + ParseOutcome::Err(e) => { // Spec: once at least one package has been authenticated // and saved the mirror MUST respond 202, not 400. Treat a // malformed later package like truncation so the verified @@ -301,111 +448,60 @@ async fn verify_and_persist( } packages_received += 1; - // Collect only the not-yet-persisted tail of this package. + // Buffer only the not-yet-persisted tail of this package. if pkg_end > initial_next { let skip = usize::try_from(initial_next.saturating_sub(pkg_start)) .map_err(|_| Error::from("skip count overflows usize"))?; - entries.extend(pkg.entries.into_iter().skip(skip)); - collected_end = pkg_end; + chunk.extend(pkg.entries.into_iter().skip(skip)); + chunk_end = pkg_end; + } + chunk_pkgs += 1; + + if chunk_pkgs == commit_packages { + if !chunk.is_empty() { + (frontier_size, frontier_hash) = flush_chunk( + bucket, + env, + &header.log_origin, + frontier_size, + frontier_hash, + chunk_end, + &mut chunk, + ) + .await?; + } + chunk_pkgs = 0; } } // Truncation before the first complete package is a hard 400. if truncated && packages_received == 0 { - log::warn!("add-entries: body truncated before the first complete package"); + log::warn!("add-entries: stream truncated before the first complete package"); return Err(AppError::BadRequest( "no complete entry package received".to_owned(), )); } - // Persist all newly-received entries in a single commit, advancing the - // persisted-entry frontier once. - let mut frontier_size = initial_next; - let mut frontier_hash = snapshot.next_entry.hash; - if !entries.is_empty() { - let root = commit::persist_entries( - &bucket, + // Flush the trailing partial chunk. + if !chunk.is_empty() { + (frontier_size, frontier_hash) = flush_chunk( + bucket, + env, + &header.log_origin, frontier_size, frontier_hash, - collected_end, - &entries, + chunk_end, + &mut chunk, ) .await?; - advance_next_entry(env, &header.log_origin, collected_end, root).await?; - frontier_size = collected_end; - frontier_hash = root; } - let persisted_new = frontier_size > initial_next; - if truncated { - log::info!( - "add-entries: client-truncated after {packages_received} complete packages; \ - persisted through {frontier_size}", - ); - return Ok(mirror_info_202( - env, - snapshot, - &header.log_origin, - frontier_size, - )); - } - - // Frontier below `upload_end`: not every entry up to `upload_end` is - // persisted yet (a truncated body, or an already-persisted request that - // stops short). This is the spec's "not yet received all packages" case, - // so 202 with the advanced frontier for the client to resume from. - if frontier_size < header.upload_end { - log::info!( - "add-entries: frontier {frontier_size} below upload_end {}; returning 202 to resume", - header.upload_end, - ); - return Ok(mirror_info_202( - env, - snapshot, - &header.log_origin, - frontier_size, - )); - } - - // Every canonical package was received. Any bytes past the last one - // are discarded, not rejected: the spec says "the mirror discards any - // partial bytes after the last successfully authenticated entry - // package", and the only defined 400 is when no package authenticated - // at all (handled above). Rejecting here would also be dishonest, since - // the entries were already persisted and the frontier advanced. - - // When we persisted new entries, the recomputed tree MUST match the - // pending checkpoint the packages were proven against. A mismatch means - // proof verification and tile computation disagree: an internal error, - // never a client fault. When nothing new was persisted (a re-upload of - // an already-persisted range), the log-signed target is trusted - // directly: pending checkpoints are consistency-chained on the - // add-checkpoint path and tickets only revive our own past pendings, so - // the target is on the same branch as the persisted tree by - // construction. - if persisted_new && (frontier_size != header.upload_end || frontier_hash != target.hash) { - log::error!( - "add-entries: recomputed frontier ({frontier_size}, {frontier_hash}) != target ({}, {})", - target.size, - target.hash, - ); - return Err(AppError::InternalServerError( - "recomputed root mismatch".to_owned(), - )); - } - - // Frontier ahead of `upload_end` (a prior upload persisted past it, or a - // racing client): the edge tiles were written at the larger frontier, so - // the narrower "cut" tiles a tree of size `upload_end` needs may be - // missing and a cosignature at `upload_end` would be unverifiable against - // our own tiles. Spec requires the tree at `upload_end` be servable before - // we cosign, so synthesize those cut tiles first. - if frontier_size > header.upload_end { - let bucket = load_origin_bucket(env, &header.log_origin)?; - commit::ensure_cut_tiles(&bucket, header.upload_end, frontier_size, frontier_hash).await?; - } - - cosign_and_serve(env, header, target, snapshot).await + Ok(StreamResult { + frontier_size, + frontier_hash, + packages_received, + truncated, + }) } /// The spec's `excess_entries = min(upload_end, next_entry) - @@ -418,31 +514,6 @@ fn excess_entries(upload_start: u64, upload_end: u64, next_entry: u64) -> u64 { next_entry.min(upload_end).saturating_sub(upload_start) } -/// Outcome of reading the next entry package from the buffered body. -enum PackageOutcome { - Ok(EntryPackage), - /// The body ended between packages or partway through one: a client - /// truncation. Complete packages already read are persisted (partial - /// progress); truncation before the first complete package is a 400. - Eof, - Err(ParseError), -} - -/// Read the next entry package from `cursor`, returning -/// [`PackageOutcome::Eof`] when the buffered body is exhausted (cleanly -/// between packages or mid-package) so the caller treats it as a client -/// truncation. -fn read_package(cursor: &mut Cursor<&[u8]>, num_entries: u64) -> PackageOutcome { - if usize::try_from(cursor.position()).unwrap_or(usize::MAX) >= cursor.get_ref().len() { - return PackageOutcome::Eof; - } - match EntryPackage::read_from(cursor, num_entries) { - Ok(pkg) => PackageOutcome::Ok(pkg), - Err(ParseError::Io(ref e)) if e.kind() == ErrorKind::UnexpectedEof => PackageOutcome::Eof, - Err(e) => PackageOutcome::Err(e), - } -} - /// Cosign the target checkpoint with the mirror key, advance the durable /// mirror checkpoint via the DO, and return the 200 response carrying the /// mirror cosignature line(s). @@ -562,6 +633,34 @@ async fn cosign_and_serve( .into_response()) } +/// Persist the buffered `chunk` (the entries `[frontier_size, chunk_end)`) +/// as entry bundles + hash tiles resuming from the current frontier, +/// advance the persisted-entry frontier in the DO, and return the new +/// frontier `(chunk_end, root)`. Clears `chunk`. +/// +/// [`commit::persist_entries`] writes immutable, content-addressed tiles +/// and the DO advance is a monotone compare-and-swap, so a repeated or +/// concurrent flush of the same range is a harmless no-op. +/// +/// # Errors +/// +/// Returns an error on a storage failure or if the DO advance RPC fails. +async fn flush_chunk( + bucket: &O, + env: &Env, + origin: &str, + frontier_size: u64, + frontier_hash: Hash, + chunk_end: u64, + chunk: &mut Vec>, +) -> Result<(u64, Hash)> { + let root = + commit::persist_entries(bucket, frontier_size, frontier_hash, chunk_end, chunk).await?; + advance_next_entry(env, origin, chunk_end, root).await?; + chunk.clear(); + Ok((chunk_end, root)) +} + /// Read the persisted-leaf prefix required to verify a non-256-aligned /// first package: the leaves `[subtree_start, upload_start)` where /// `subtree_start` is `upload_start` rounded down to a 256 boundary. @@ -672,6 +771,97 @@ async fn advance_next_entry(env: &Env, origin: &str, size: u64, hash: Hash) -> R } } +/// Outcome of attempting to read the next entry package from the stream +/// buffer. +/// +/// `CleanEof` (stream ended cleanly between packages) and `MidPackageEof` +/// (stream ended partway through a package) are kept distinct for +/// diagnostics, though the handler treats both as a client truncation: +/// complete packages already received are persisted (partial progress), +/// and a truncation before the first complete package is a 400. +enum ParseOutcome { + Ok(EntryPackage), + CleanEof, + MidPackageEof, + Err(ParseError), +} + +/// Read the `add-entries` request header from `buf`, pulling more +/// chunks from the underlying stream until the header parses or the +/// stream errors. Returns `Ok(Ok(header))` on success or `Ok(Err(resp))` +/// where `resp` is a fully-formed 400 response on malformed input. +/// +/// The header has a bounded maximum size (u16 origin + u64s + u16 +/// ticket + hash + u8 proof-size + 63 hashes <= ~131 KB), so the +/// retry-on-`UnexpectedEof` loop terminates. +async fn parse_header(buf: &mut StreamBuffer) -> ApiResult +where + S: futures_util::Stream>> + Unpin, +{ + loop { + let mut cursor = Cursor::new(buf.buffered()); + match AddEntriesRequestHeader::read_from(&mut cursor) { + Ok(header) => { + let consumed = usize::try_from(cursor.position()).unwrap_or(usize::MAX); + buf.consume(consumed); + return Ok(header); + } + Err(ParseError::Io(ref e)) if e.kind() == ErrorKind::UnexpectedEof => { + // Need more bytes to parse the header. Pull another + // chunk; if the stream is already at EOF, the header + // is fundamentally malformed (truncated before being + // complete). + if !buf.pull_one().await? { + log::warn!( + "add-entries: stream ended before header was complete \ + ({} bytes buffered)", + buf.len() + ); + return Err(AppError::BadRequest( + "malformed (truncated header)".to_owned(), + )); + } + } + Err(e) => { + log::warn!("add-entries: malformed header: {e:?}"); + return Err(AppError::BadRequest(e.to_string())); + } + } + } +} + +/// Read the next entry package from `buf`, pulling more chunks from +/// the underlying stream until the package parses or the stream ends. +/// See [`ParseOutcome`] for the four cases. +async fn parse_next_package(buf: &mut StreamBuffer, num_entries: u64) -> Result +where + S: futures_util::Stream>> + Unpin, +{ + // EOF with an empty buffer: clean truncation between packages. + if buf.is_eof() && buf.len() == 0 { + return Ok(ParseOutcome::CleanEof); + } + loop { + let mut cursor = Cursor::new(buf.buffered()); + match EntryPackage::read_from(&mut cursor, num_entries) { + Ok(pkg) => { + let consumed = usize::try_from(cursor.position()).unwrap_or(usize::MAX); + buf.consume(consumed); + return Ok(ParseOutcome::Ok(pkg)); + } + Err(ParseError::Io(ref e)) if e.kind() == ErrorKind::UnexpectedEof => { + if !buf.pull_one().await? { + if buf.len() == 0 { + return Ok(ParseOutcome::CleanEof); + } + return Ok(ParseOutcome::MidPackageEof); + } + } + Err(e) => return Ok(ParseOutcome::Err(e)), + } + } +} + /// Read the per-origin DO state snapshot. A non-200 status or RPC failure /// is a transport-level error the handler maps to 500. async fn fetch_snapshot(env: &Env, origin: &str) -> Result { @@ -1056,6 +1246,17 @@ mod tests { verify_package(&prefix, &pkg, 256, 512, &cp).expect("non-aligned package verifies"); } + #[test] + fn verify_nonaligned_first_package_max_prefix_ok() { + // Largest prefix a non-aligned first package can carry: upload_start + // = 511 leaves persisted [256, 511) as a 255-leaf prefix, with a + // single uploaded entry [511, 512) closing the subtree. + let (prefix, pkg, cp) = fixture(1000, 256, 511, 512); + assert_eq!(prefix.len(), 255); + assert_eq!(pkg.entries.len(), 1); + verify_package(&prefix, &pkg, 256, 512, &cp).expect("max-prefix package verifies"); + } + #[test] fn verify_first_ever_package_from_zero_ok() { // Subtree rooted at 0 (first bundle), partial last package. diff --git a/crates/mirror_worker/src/body.rs b/crates/mirror_worker/src/body.rs index 298e50e5..eb7f7cab 100644 --- a/crates/mirror_worker/src/body.rs +++ b/crates/mirror_worker/src/body.rs @@ -6,77 +6,187 @@ //! [c2sp.org/tlog-mirror][spec] requires mirrors to accept //! `Content-Encoding: gzip` request bodies (see [Request Body][reqbody]); //! clients MAY send gzip without negotiating first. The Cloudflare Workers -//! runtime does not transparently decompress request bodies (unlike -//! responses, which it (de)compresses based on `Accept-Encoding`), so the -//! mirror must gunzip the body itself. +//! runtime does not transparently decompress request bodies (unlike responses, +//! which it (de)compresses based on `Accept-Encoding`), so the mirror must +//! gunzip the body itself. //! -//! This module reads the whole body into memory and inflates it in one -//! pass. A follow-up commit replaces this with incremental streaming so a -//! large upload does not have to buffer the entire body (Workers isolates -//! have a ~128 MB memory ceiling). +//! `add-entries` bodies are unbounded and Workers isolates have a ~128 MB +//! memory ceiling, so the body must be decompressed *incrementally* rather +//! than slurped-then-inflated. [`gunzip`] wraps the runtime's chunked +//! [`futures_util::Stream`] body in a decoding stream that inflates each +//! compressed chunk as it arrives and yields the plaintext chunks, keeping +//! the same streaming contract the identity path relies on (see +//! [`crate::stream_buffer`]). //! //! [spec]: https://c2sp.org/tlog-mirror#add-entries //! [reqbody]: https://c2sp.org/tlog-mirror#request-body -use std::io::Read as _; +use std::io::Write as _; +use std::pin::Pin; -use flate2::read::GzDecoder; -use futures_util::StreamExt as _; +use flate2::write::GzDecoder; +use futures_util::stream::{Stream, StreamExt as _}; #[allow(clippy::wildcard_imports)] use worker::*; use crate::frontend_worker::{ApiResult, AppError}; -/// Read the whole request body and decode it per `Content-Encoding`. +/// A boxed, `Unpin` body stream, the common type the `add-entries` +/// handler feeds to [`crate::stream_buffer::StreamBuffer`] regardless of +/// whether the request body was identity- or gzip-encoded. +pub(crate) type BodyStream = Pin>>>>; + +/// Open the request body as a decoded chunk stream, honoring +/// `Content-Encoding`. /// -/// `identity` (or an absent header) returns the body unchanged; -/// `gzip`/`x-gzip` is inflated. Any other encoding is unsupported: the -/// mirror can't authenticate a body it can't read, so this returns 415. +/// `identity` (or an absent header) passes the body through unchanged; +/// `gzip`/`x-gzip` is inflated incrementally via [`gunzip`]. Any other +/// encoding is unsupported: the mirror can't authenticate a body it can't +/// read, so this returns 415. /// /// # Errors /// /// Returns [`AppError::UnsupportedMediaType`] for an unrecognized -/// `Content-Encoding`, [`AppError::BadRequest`] for a malformed/truncated -/// gzip body, or a transport error while reading the body stream. -pub(crate) async fn read_decoded_body( +/// `Content-Encoding`. +pub(crate) fn decoded_stream( headers: &axum::http::HeaderMap, body: axum::body::Body, -) -> ApiResult> { +) -> ApiResult { let encoding = headers .get(axum::http::header::CONTENT_ENCODING) .and_then(|v| v.to_str().ok()) .unwrap_or_default() .trim() .to_ascii_lowercase(); + // Adapt axum's body data stream (`Result`) to the + // `Result>` chunk contract the buffer/gunzip pipeline expects. + let raw = body.into_data_stream().map(|r| { + r.map(|b| b.to_vec()) + .map_err(|e| Error::from(e.to_string())) + }); + let stream: BodyStream = match encoding.as_str() { + "" | "identity" => Box::pin(raw), + "gzip" | "x-gzip" => gunzip(raw), + other => { + return Err(AppError::UnsupportedMediaType(format!( + "Unsupported Content-Encoding: {other}" + ))); + } + }; + Ok(stream) +} - let mut raw = Vec::new(); - let mut stream = body.into_data_stream(); - while let Some(chunk) = stream.next().await { - let chunk = chunk.map_err(|e| Error::from(e.to_string()))?; - raw.extend_from_slice(&chunk); - } +/// Incremental gzip inflater: feed compressed bytes with [`Self::push`] +/// and drain the plaintext produced so far; call [`Self::finish`] once the +/// compressed input ends to validate the gzip trailer (CRC-32 + ISIZE). +/// +/// Backed by [`flate2`]'s pure-Rust `rust_backend` (`miniz_oxide`), so it +/// compiles to and runs under WASM. `flate2::write::GzDecoder` handles all +/// gzip framing: the 10-byte header, optional FNAME/FEXTRA/etc. fields +/// (buffered across `push` calls if split across chunks), and the trailer. +struct GzipInflater { + decoder: GzDecoder>, +} - match encoding.as_str() { - "" | "identity" => Ok(raw), - "gzip" | "x-gzip" => { - let mut out = Vec::new(); - GzDecoder::new(raw.as_slice()) - .read_to_end(&mut out) - .map_err(|e| AppError::BadRequest(format!("gzip decode failed: {e}")))?; - Ok(out) +impl GzipInflater { + fn new() -> Self { + Self { + decoder: GzDecoder::new(Vec::new()), } - other => Err(AppError::UnsupportedMediaType(format!( - "Unsupported Content-Encoding: {other}" - ))), + } + + /// Feed one compressed chunk and return the plaintext bytes produced. + /// May return an empty `Vec` if `input` only completed part of the + /// gzip header or a DEFLATE block that hasn't emitted output yet. + fn push(&mut self, input: &[u8]) -> Result> { + self.decoder + .write_all(input) + .map_err(|e| Error::from(format!("gzip decode failed: {e}")))?; + Ok(std::mem::take(self.decoder.get_mut())) + } + + /// Finish decompression, returning any trailing plaintext. Errors if + /// the gzip stream was truncated or its CRC-32/ISIZE trailer does not + /// match the decompressed data. + fn finish(self) -> Result> { + self.decoder + .finish() + .map_err(|e| Error::from(format!("gzip stream incomplete or corrupt: {e}"))) } } +/// Wrap a compressed body `Stream` in a decoding stream that yields the +/// gunzipped plaintext chunks. +/// +/// The returned stream inflates lazily: each poll pulls compressed chunks +/// from `inner` until it can emit at least one plaintext byte, so memory +/// use stays bounded by the chunk size rather than the whole body. A +/// decode error (malformed gzip) or a truncated stream surfaces as a +/// terminal `Err` item, after which the stream ends. +pub(crate) fn gunzip(inner: S) -> BodyStream +where + S: Stream>> + Unpin + 'static, +{ + struct DecodeState { + inner: S, + inflater: Option, + done: bool, + } + + Box::pin(futures_util::stream::unfold( + DecodeState { + inner, + inflater: Some(GzipInflater::new()), + done: false, + }, + |mut st| async move { + if st.done { + return None; + } + loop { + match st.inner.next().await { + Some(Ok(chunk)) => { + let out = match st.inflater.as_mut().expect("inflater present").push(&chunk) + { + Ok(out) => out, + Err(e) => { + st.done = true; + return Some((Err(e), st)); + } + }; + // A chunk may not yet yield any plaintext (partial + // header / block); pull more instead of emitting an + // empty item. + if out.is_empty() { + continue; + } + return Some((Ok(out), st)); + } + Some(Err(e)) => { + st.done = true; + return Some((Err(e), st)); + } + None => { + st.done = true; + let inflater = st.inflater.take().expect("inflater present"); + return match inflater.finish() { + Ok(tail) if !tail.is_empty() => Some((Ok(tail), st)), + Ok(_) => None, + Err(e) => Some((Err(e), st)), + }; + } + } + } + }, + )) +} + #[cfg(test)] mod tests { use super::*; use flate2::Compression; use flate2::write::GzEncoder; - use std::io::Write as _; + use futures_util::stream; /// gzip-compress `data` into a single buffer for test input. fn gzip(data: &[u8]) -> Vec { @@ -85,43 +195,79 @@ mod tests { enc.finish().unwrap() } - /// Inflate `compressed` the same way [`read_decoded_body`] does. - fn gunzip(compressed: &[u8]) -> Result> { + /// Split `bytes` into chunks of `size` and build a stream of them. + fn chunked_stream( + bytes: &[u8], + size: usize, + ) -> impl Stream>> + Unpin + 'static { + let chunks: Vec>> = + bytes.chunks(size.max(1)).map(|c| Ok(c.to_vec())).collect(); + stream::iter(chunks) + } + + async fn collect(mut s: impl Stream>> + Unpin) -> Result> { let mut out = Vec::new(); - GzDecoder::new(compressed) - .read_to_end(&mut out) - .map_err(|e| Error::from(format!("gzip decode failed: {e}")))?; + while let Some(item) = s.next().await { + out.extend_from_slice(&item?); + } Ok(out) } - #[test] - fn gzip_roundtrips() { + #[tokio::test] + async fn roundtrips_across_chunk_sizes() { + // A payload big enough to span multiple DEFLATE flushes, and + // compressible enough to exercise real inflation. let mut plain: Vec = Vec::new(); for i in 0..20_000u32 { plain.extend_from_slice(format!("entry-{i};").as_bytes()); } - let decoded = gunzip(&gzip(&plain)).expect("roundtrip"); - assert_eq!(decoded, plain); + let compressed = gzip(&plain); + // Feeding the compressed stream in a variety of chunk sizes (incl. + // 1-byte chunks that split the header/trailer) must all reconstruct + // the original plaintext. + for size in [1usize, 2, 7, 64, 1024, compressed.len()] { + let decoded = collect(gunzip(chunked_stream(&compressed, size))) + .await + .unwrap_or_else(|e| panic!("chunk size {size} failed: {e}")); + assert_eq!(decoded, plain, "chunk size {size} mismatch"); + } } - #[test] - fn empty_payload_roundtrips() { - let decoded = gunzip(&gzip(b"")).expect("roundtrip"); + #[tokio::test] + async fn empty_payload_roundtrips() { + let compressed = gzip(b""); + let decoded = collect(gunzip(chunked_stream(&compressed, 3))) + .await + .unwrap(); assert!(decoded.is_empty()); } - #[test] - fn truncated_gzip_errors() { + #[tokio::test] + async fn truncated_stream_errors() { let mut compressed = gzip(b"the quick brown fox jumps over the lazy dog"); + // Drop the trailer (and some of the deflate payload) so the stream + // ends mid-member; finish() must report the truncation. compressed.truncate(compressed.len() - 6); - assert!(gunzip(&compressed).is_err(), "truncated gzip must error"); + let err = collect(gunzip(chunked_stream(&compressed, 4))).await; + assert!(err.is_err(), "truncated gzip must surface an error"); } - #[test] - fn corrupt_gzip_errors() { + #[tokio::test] + async fn corrupt_data_errors() { let mut compressed = gzip(b"hello world, this is a test payload for corruption"); + // Corrupt a byte in the middle of the DEFLATE payload. let mid = compressed.len() / 2; compressed[mid] ^= 0xff; - assert!(gunzip(&compressed).is_err(), "corrupt gzip must error"); + let err = collect(gunzip(chunked_stream(&compressed, 5))).await; + assert!(err.is_err(), "corrupt gzip must surface an error"); + } + + #[tokio::test] + async fn upstream_error_propagates() { + let compressed = gzip(b"partial"); + let mut chunks: Vec>> = vec![Ok(compressed[..4].to_vec())]; + chunks.push(Err(Error::from("boom"))); + let err = collect(gunzip(stream::iter(chunks))).await; + assert!(err.is_err(), "upstream stream error must propagate"); } } diff --git a/crates/mirror_worker/src/commit.rs b/crates/mirror_worker/src/commit.rs index f60d1102..e316f251 100644 --- a/crates/mirror_worker/src/commit.rs +++ b/crates/mirror_worker/src/commit.rs @@ -724,6 +724,42 @@ mod tests { ); } + #[tokio::test] + async fn chunked_commit_matches_single_commit() { + // Streaming add-entries flushes in frontier-advancing chunks. Each + // chunk must write the same entry bundles and hash tiles a single + // commit would, so the persisted objects and the final root are + // identical (the chunked run may additionally leave orphaned + // partial tiles behind for the cleaner). + let chunked = MemBackend::default(); + let mut size = 0u64; + let mut hash = EMPTY_HASH; + for end in [256u64, 512, 768, 900] { + hash = persist_entries(&chunked, size, hash, end, &leaves(size..end)) + .await + .unwrap(); + size = end; + } + assert_eq!(hash, reference_root(900)); + + let oneshot = MemBackend::default(); + let root = persist_entries(&oneshot, 0, EMPTY_HASH, 900, &leaves(0..900)) + .await + .unwrap(); + assert_eq!(root, hash); + + // Every object a single commit writes is present, byte-for-byte, + // after the chunked commit. + let chunked = chunked.store.borrow(); + for (key, bytes) in oneshot.store.borrow().iter() { + assert_eq!( + chunked.get(key), + Some(bytes), + "chunked commit missing or differing object {key}" + ); + } + } + #[tokio::test] async fn entry_bundles_roundtrip() { use length_prefixed::ReadLengthPrefixedBytesExt as _; diff --git a/crates/mirror_worker/src/lib.rs b/crates/mirror_worker/src/lib.rs index 4075adc9..47f36603 100644 --- a/crates/mirror_worker/src/lib.rs +++ b/crates/mirror_worker/src/lib.rs @@ -41,6 +41,7 @@ mod commit; mod frontend_worker; mod mirror_state_do; mod storage; +mod stream_buffer; /// The binding name used in `wrangler.jsonc` for the `MirrorState` DO. pub(crate) const MIRROR_STATE_BINDING: &str = "MIRROR_STATE"; diff --git a/crates/mirror_worker/src/stream_buffer.rs b/crates/mirror_worker/src/stream_buffer.rs new file mode 100644 index 00000000..d4a72f87 --- /dev/null +++ b/crates/mirror_worker/src/stream_buffer.rs @@ -0,0 +1,195 @@ +// Copyright (c) 2025-2026 Cloudflare, Inc. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause + +//! Streaming buffer adapter that bridges the worker runtime's async +//! [`futures_util::Stream`] of byte chunks to the synchronous +//! [`std::io::Read`] API expected by [`tlog_mirror::wire`]'s parsers. +//! +//! The adapter holds an internal `Vec` byte buffer and a read +//! offset. Callers pull more bytes from the underlying stream +//! asynchronously (via [`StreamBuffer::pull_one`]), then parse +//! synchronously over a [`Cursor`](std::io::Cursor) wrapping the buffered +//! slice. The natural parse pattern is a retry loop: attempt the parse, +//! and on [`std::io::ErrorKind::UnexpectedEof`] pull another chunk and +//! retry. On success, [`StreamBuffer::consume`] advances the read offset +//! past the parsed bytes so subsequent parses start at the next byte. +//! +//! [`StreamBuffer::consume`] only bumps the read offset (O(1)); the dead +//! prefix is reclaimed lazily by the next [`StreamBuffer::pull_one`], +//! which shifts the live tail to the front before appending. This keeps +//! per-package consumes cheap and bounds the shift to at most once per +//! pull. +//! +//! This pattern lets the `add-entries` handler process each entry +//! package's bytes as soon as enough have arrived, without buffering +//! the entire (potentially 100 MB) request body in memory at once. + +use futures_util::stream::{Stream, StreamExt as _}; + +/// A growable byte buffer fed by an async `Stream, +/// E>>`. See the module-level comment for the intended usage pattern. +pub(crate) struct StreamBuffer { + stream: S, + buf: Vec, + /// Read offset into `buf`: bytes before it have been consumed and + /// are reclaimed on the next [`Self::pull_one`]. The live buffer is + /// `buf[start..]`. + start: usize, + /// Set when the underlying stream has signalled end-of-stream. + /// Subsequent [`Self::pull_one`] calls return `Ok(false)` + /// without polling the (already-finished) stream. + eof: bool, +} + +impl StreamBuffer +where + S: Stream, E>> + Unpin, +{ + /// Construct a new streaming buffer wrapping `stream`. The buffer + /// starts empty; callers must call [`Self::pull_one`] (or its + /// helpers) to populate it before parsing. + pub fn new(stream: S) -> Self { + Self { + stream, + buf: Vec::new(), + start: 0, + eof: false, + } + } + + /// Pull one chunk from the underlying stream and append to the + /// internal buffer. Returns `Ok(true)` if a chunk was appended, + /// `Ok(false)` if the stream ended (clean EOF). Once the stream has + /// ended, all subsequent calls return `Ok(false)` without polling + /// the stream again. + /// + /// # Errors + /// Propagates any error from the underlying stream. + pub async fn pull_one(&mut self) -> std::result::Result { + if self.eof { + return Ok(false); + } + match self.stream.next().await { + Some(Ok(chunk)) => { + // Reclaim the consumed prefix before growing, so the + // buffer tracks only live bytes plus one chunk. + if self.start > 0 { + self.buf.drain(..self.start); + self.start = 0; + } + self.buf.extend_from_slice(&chunk); + Ok(true) + } + Some(Err(e)) => Err(e), + None => { + self.eof = true; + Ok(false) + } + } + } + + /// View the currently-buffered (live) bytes. Used to construct a + /// sync `Cursor` for parsing. + pub fn buffered(&self) -> &[u8] { + &self.buf[self.start..] + } + + /// Number of live bytes currently buffered. + pub fn len(&self) -> usize { + self.buf.len() - self.start + } + + /// `true` if the underlying stream has signalled end-of-stream + /// (regardless of whether bytes remain in the buffer). + pub fn is_eof(&self) -> bool { + self.eof + } + + /// Advance the read offset past the first `n` live bytes. The caller + /// should call this after a successful parse to skip consumed input. + /// The bytes are not freed here; the next [`Self::pull_one`] reclaims + /// them. + /// + /// # Panics + /// Panics if `n > self.len()`. + pub fn consume(&mut self, n: usize) { + assert!( + n <= self.len(), + "StreamBuffer::consume({n}) but only {} bytes buffered", + self.len() + ); + self.start += n; + } +} + +#[cfg(test)] +mod tests { + use super::StreamBuffer; + use futures_util::stream; + + /// Build a `StreamBuffer` over an in-memory iterator-of-Results. + /// Used to unit-test the buffering behaviour without a real worker + /// runtime. + fn from_chunks( + chunks: Vec>, + ) -> StreamBuffer>> + Unpin> { + StreamBuffer::new(stream::iter( + chunks.into_iter().map(Ok::<_, std::io::Error>), + )) + } + + #[tokio::test(flavor = "current_thread")] + async fn pull_one_returns_chunks_in_order() { + let mut buf = from_chunks(vec![b"foo".to_vec(), b"bar".to_vec()]); + assert!(buf.pull_one().await.unwrap()); + assert_eq!(buf.buffered(), b"foo"); + assert!(buf.pull_one().await.unwrap()); + assert_eq!(buf.buffered(), b"foobar"); + // Stream is now empty; further pulls return Ok(false). + assert!(!buf.pull_one().await.unwrap()); + assert!(buf.is_eof()); + assert!(!buf.pull_one().await.unwrap()); + } + + #[tokio::test(flavor = "current_thread")] + async fn consume_advances_past_used_bytes() { + let mut buf = from_chunks(vec![b"abcdef".to_vec()]); + buf.pull_one().await.unwrap(); + buf.consume(2); + assert_eq!(buf.buffered(), b"cdef"); + buf.consume(4); + assert_eq!(buf.buffered(), b""); + assert_eq!(buf.len(), 0); + } + + #[tokio::test(flavor = "current_thread")] + async fn pull_reclaims_consumed_prefix() { + let mut buf = from_chunks(vec![b"abcd".to_vec(), b"efgh".to_vec()]); + buf.pull_one().await.unwrap(); + buf.consume(3); + assert_eq!(buf.buffered(), b"d"); + // The next pull shifts the live tail ("d") to the front, then + // appends the new chunk. + buf.pull_one().await.unwrap(); + assert_eq!(buf.buffered(), b"defgh"); + assert_eq!(buf.len(), 5); + } + + #[tokio::test(flavor = "current_thread")] + #[should_panic(expected = "StreamBuffer::consume")] + async fn consume_past_end_panics() { + let mut buf = from_chunks(vec![b"ab".to_vec()]); + buf.pull_one().await.unwrap(); + buf.consume(99); + } + + #[tokio::test(flavor = "current_thread")] + async fn errors_propagate() { + let chunks: Vec>> = + vec![Ok(b"foo".to_vec()), Err(std::io::Error::other("boom"))]; + let mut buf = StreamBuffer::new(stream::iter(chunks)); + assert!(buf.pull_one().await.unwrap()); + let err = buf.pull_one().await.unwrap_err(); + assert_eq!(err.to_string(), "boom"); + } +} From 7117a7216be1505ce9fcda3a8643e0b4108f7522 Mon Sep 17 00:00:00 2001 From: Luke Valenta Date: Fri, 31 Jul 2026 12:26:29 -0400 Subject: [PATCH 2/8] mirror_worker: map malformed-gzip bodies to 400, not 500 Gzip decode/trailer failures were surfaced as worker::Error, which the handler's blanket From mapped to 500. A malformed or truncated gzip request body is a client fault. Introduce a typed BodyError (Decode vs Transport) on the body stream so decode faults map to 400 while genuine transport failures stay 500; parse_next_package now returns ApiResult so the distinction survives to the handler. --- crates/mirror_worker/src/add_entries.rs | 16 +++--- crates/mirror_worker/src/body.rs | 54 ++++++++++++++++----- crates/mirror_worker/src/frontend_worker.rs | 10 ++++ 3 files changed, 62 insertions(+), 18 deletions(-) diff --git a/crates/mirror_worker/src/add_entries.rs b/crates/mirror_worker/src/add_entries.rs index 7a3e62e5..514a2ebc 100644 --- a/crates/mirror_worker/src/add_entries.rs +++ b/crates/mirror_worker/src/add_entries.rs @@ -48,7 +48,8 @@ use worker::*; use generic_log_worker::{ObjectBackend, util::now_millis}; use crate::{ - body, commit, + body::{self, BodyError}, + commit, frontend_worker::{ApiResult, AppError}, load_mirror_signer, load_ticket_sealer, log_verifiers, mirror_state_do::{ @@ -237,7 +238,7 @@ async fn stream_and_commit( first_prefix: &[Vec], ) -> ApiResult where - S: futures_util::Stream>> + Unpin, + S: futures_util::Stream, BodyError>> + Unpin, { let bucket = load_origin_bucket(env, &header.log_origin)?; @@ -379,7 +380,7 @@ async fn persist_packages( start: &NextEntry, ) -> ApiResult where - S: futures_util::Stream>> + Unpin, + S: futures_util::Stream, BodyError>> + Unpin, O: ObjectBackend, { // config.schema.json caps commit_packages (max 1024), enforced by the @@ -796,7 +797,7 @@ enum ParseOutcome { /// retry-on-`UnexpectedEof` loop terminates. async fn parse_header(buf: &mut StreamBuffer) -> ApiResult where - S: futures_util::Stream>> + Unpin, + S: futures_util::Stream, BodyError>> + Unpin, { loop { let mut cursor = Cursor::new(buf.buffered()); @@ -833,9 +834,12 @@ where /// Read the next entry package from `buf`, pulling more chunks from /// the underlying stream until the package parses or the stream ends. /// See [`ParseOutcome`] for the four cases. -async fn parse_next_package(buf: &mut StreamBuffer, num_entries: u64) -> Result +async fn parse_next_package( + buf: &mut StreamBuffer, + num_entries: u64, +) -> ApiResult where - S: futures_util::Stream>> + Unpin, + S: futures_util::Stream, BodyError>> + Unpin, { // EOF with an empty buffer: clean truncation between packages. if buf.is_eof() && buf.len() == 0 { diff --git a/crates/mirror_worker/src/body.rs b/crates/mirror_worker/src/body.rs index eb7f7cab..4f70d610 100644 --- a/crates/mirror_worker/src/body.rs +++ b/crates/mirror_worker/src/body.rs @@ -31,10 +31,31 @@ use worker::*; use crate::frontend_worker::{ApiResult, AppError}; +/// An error surfaced by a [`BodyStream`], distinguishing a client-side +/// decode fault from a transport failure so the `add-entries` handler can +/// map each to the right HTTP status (see `From for AppError`). +#[derive(Debug)] +pub(crate) enum BodyError { + /// Malformed or truncated gzip body: a client fault, mapped to 400. + Decode(String), + /// Transport failure reading the underlying request body: mapped to + /// 500. + Transport(Error), +} + +impl std::fmt::Display for BodyError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + BodyError::Decode(e) => write!(f, "{e}"), + BodyError::Transport(e) => write!(f, "{e}"), + } + } +} + /// A boxed, `Unpin` body stream, the common type the `add-entries` /// handler feeds to [`crate::stream_buffer::StreamBuffer`] regardless of /// whether the request body was identity- or gzip-encoded. -pub(crate) type BodyStream = Pin>>>>; +pub(crate) type BodyStream = Pin, BodyError>>>>; /// Open the request body as a decoded chunk stream, honoring /// `Content-Encoding`. @@ -62,7 +83,7 @@ pub(crate) fn decoded_stream( // `Result>` chunk contract the buffer/gunzip pipeline expects. let raw = body.into_data_stream().map(|r| { r.map(|b| b.to_vec()) - .map_err(|e| Error::from(e.to_string())) + .map_err(|e| BodyError::Transport(Error::from(e.to_string()))) }); let stream: BodyStream = match encoding.as_str() { "" | "identity" => Box::pin(raw), @@ -125,7 +146,7 @@ impl GzipInflater { /// terminal `Err` item, after which the stream ends. pub(crate) fn gunzip(inner: S) -> BodyStream where - S: Stream>> + Unpin + 'static, + S: Stream, BodyError>> + Unpin + 'static, { struct DecodeState { inner: S, @@ -151,7 +172,7 @@ where Ok(out) => out, Err(e) => { st.done = true; - return Some((Err(e), st)); + return Some((Err(BodyError::Decode(e.to_string())), st)); } }; // A chunk may not yet yield any plaintext (partial @@ -172,7 +193,7 @@ where return match inflater.finish() { Ok(tail) if !tail.is_empty() => Some((Ok(tail), st)), Ok(_) => None, - Err(e) => Some((Err(e), st)), + Err(e) => Some((Err(BodyError::Decode(e.to_string())), st)), }; } } @@ -199,13 +220,15 @@ mod tests { fn chunked_stream( bytes: &[u8], size: usize, - ) -> impl Stream>> + Unpin + 'static { - let chunks: Vec>> = + ) -> impl Stream, BodyError>> + Unpin + 'static { + let chunks: Vec, BodyError>> = bytes.chunks(size.max(1)).map(|c| Ok(c.to_vec())).collect(); stream::iter(chunks) } - async fn collect(mut s: impl Stream>> + Unpin) -> Result> { + async fn collect( + mut s: impl Stream, BodyError>> + Unpin, + ) -> std::result::Result, BodyError> { let mut out = Vec::new(); while let Some(item) = s.next().await { out.extend_from_slice(&item?); @@ -249,7 +272,10 @@ mod tests { // ends mid-member; finish() must report the truncation. compressed.truncate(compressed.len() - 6); let err = collect(gunzip(chunked_stream(&compressed, 4))).await; - assert!(err.is_err(), "truncated gzip must surface an error"); + assert!( + matches!(err, Err(BodyError::Decode(_))), + "truncated gzip must surface a client decode error" + ); } #[tokio::test] @@ -259,14 +285,18 @@ mod tests { let mid = compressed.len() / 2; compressed[mid] ^= 0xff; let err = collect(gunzip(chunked_stream(&compressed, 5))).await; - assert!(err.is_err(), "corrupt gzip must surface an error"); + assert!( + matches!(err, Err(BodyError::Decode(_))), + "corrupt gzip must surface a client decode error" + ); } #[tokio::test] async fn upstream_error_propagates() { let compressed = gzip(b"partial"); - let mut chunks: Vec>> = vec![Ok(compressed[..4].to_vec())]; - chunks.push(Err(Error::from("boom"))); + let mut chunks: Vec, BodyError>> = + vec![Ok(compressed[..4].to_vec())]; + chunks.push(Err(BodyError::Transport(Error::from("boom")))); let err = collect(gunzip(stream::iter(chunks))).await; assert!(err.is_err(), "upstream stream error must propagate"); } diff --git a/crates/mirror_worker/src/frontend_worker.rs b/crates/mirror_worker/src/frontend_worker.rs index bd20f91c..982206c0 100644 --- a/crates/mirror_worker/src/frontend_worker.rs +++ b/crates/mirror_worker/src/frontend_worker.rs @@ -155,6 +155,16 @@ impl From for AppError { } } +impl From for AppError { + fn from(err: crate::body::BodyError) -> Self { + match err { + // Malformed/truncated gzip is a client fault. + crate::body::BodyError::Decode(msg) => Self::BadRequest(msg), + crate::body::BodyError::Transport(e) => Self::InternalServerError(e.to_string()), + } + } +} + impl IntoResponse for AppError { fn into_response(self) -> axum::response::Response { match self { From ba5676822a60465d1023c56791893e4667972f67 Mon Sep 17 00:00:00 2001 From: Luke Valenta Date: Fri, 31 Jul 2026 12:37:59 -0400 Subject: [PATCH 3/8] mirror_worker: address streaming review nits - Count only entry-contributing packages toward a chunk flush, so a run of already-persisted packages no longer triggers empty flushes; drop the now-redundant empty-chunk guard (commit_packages >= 1). - Mark StreamBuffer eof on a stream error so is_eof() stays consistent and a recovering caller won't re-poll the failed stream. - Explain the clean-vs-mid-package truncation branch in parse_next_package, and fix the stale parse_header doc that described a nested-Result return the code no longer uses. --- crates/mirror_worker/src/add_entries.rs | 26 ++++++++++++++++------- crates/mirror_worker/src/stream_buffer.rs | 23 ++++++++++++++------ 2 files changed, 34 insertions(+), 15 deletions(-) diff --git a/crates/mirror_worker/src/add_entries.rs b/crates/mirror_worker/src/add_entries.rs index 514a2ebc..6a31edd3 100644 --- a/crates/mirror_worker/src/add_entries.rs +++ b/crates/mirror_worker/src/add_entries.rs @@ -449,17 +449,21 @@ where } packages_received += 1; - // Buffer only the not-yet-persisted tail of this package. + // Buffer only the not-yet-persisted tail of this package. Packages + // wholly below the request-start frontier are already persisted, so + // they contribute no entries and don't count toward a chunk flush; + // `chunk_pkgs` therefore tracks only packages that added buffered + // entries. if pkg_end > initial_next { let skip = usize::try_from(initial_next.saturating_sub(pkg_start)) .map_err(|_| Error::from("skip count overflows usize"))?; chunk.extend(pkg.entries.into_iter().skip(skip)); chunk_end = pkg_end; - } - chunk_pkgs += 1; + chunk_pkgs += 1; - if chunk_pkgs == commit_packages { - if !chunk.is_empty() { + // `commit_packages >= 1` (config), so a full chunk always holds + // at least one package's entries: no empty-flush guard needed. + if chunk_pkgs == commit_packages { (frontier_size, frontier_hash) = flush_chunk( bucket, env, @@ -470,8 +474,8 @@ where &mut chunk, ) .await?; + chunk_pkgs = 0; } - chunk_pkgs = 0; } } @@ -789,8 +793,9 @@ enum ParseOutcome { /// Read the `add-entries` request header from `buf`, pulling more /// chunks from the underlying stream until the header parses or the -/// stream errors. Returns `Ok(Ok(header))` on success or `Ok(Err(resp))` -/// where `resp` is a fully-formed 400 response on malformed input. +/// stream errors. Returns the parsed header, or +/// [`AppError::BadRequest`] (400) if the input is malformed or truncated +/// before the header is complete. /// /// The header has a bounded maximum size (u16 origin + u64s + u16 /// ticket + hash + u8 proof-size + 63 hashes <= ~131 KB), so the @@ -855,6 +860,11 @@ where } Err(ParseError::Io(ref e)) if e.kind() == ErrorKind::UnexpectedEof => { if !buf.pull_one().await? { + // Stream ended mid-parse. An empty buffer means the + // previous package consumed exactly all buffered bytes + // and this call started a fresh (never-arriving) + // package: a clean between-package truncation. A + // non-empty buffer holds a partial package. if buf.len() == 0 { return Ok(ParseOutcome::CleanEof); } diff --git a/crates/mirror_worker/src/stream_buffer.rs b/crates/mirror_worker/src/stream_buffer.rs index d4a72f87..107d788b 100644 --- a/crates/mirror_worker/src/stream_buffer.rs +++ b/crates/mirror_worker/src/stream_buffer.rs @@ -35,9 +35,9 @@ pub(crate) struct StreamBuffer { /// are reclaimed on the next [`Self::pull_one`]. The live buffer is /// `buf[start..]`. start: usize, - /// Set when the underlying stream has signalled end-of-stream. - /// Subsequent [`Self::pull_one`] calls return `Ok(false)` - /// without polling the (already-finished) stream. + /// Set when the underlying stream has signalled end-of-stream or + /// yielded an error. Subsequent [`Self::pull_one`] calls return + /// `Ok(false)` without polling the (already-finished) stream. eof: bool, } @@ -60,11 +60,13 @@ where /// Pull one chunk from the underlying stream and append to the /// internal buffer. Returns `Ok(true)` if a chunk was appended, /// `Ok(false)` if the stream ended (clean EOF). Once the stream has - /// ended, all subsequent calls return `Ok(false)` without polling - /// the stream again. + /// ended or errored, all subsequent calls return `Ok(false)` without + /// polling the stream again. /// /// # Errors - /// Propagates any error from the underlying stream. + /// Propagates any error from the underlying stream. An error is + /// terminal: [`Self::is_eof`] is set so a caller that recovers from + /// the error does not re-poll the already-failed stream. pub async fn pull_one(&mut self) -> std::result::Result { if self.eof { return Ok(false); @@ -80,7 +82,10 @@ where self.buf.extend_from_slice(&chunk); Ok(true) } - Some(Err(e)) => Err(e), + Some(Err(e)) => { + self.eof = true; + Err(e) + } None => { self.eof = true; Ok(false) @@ -191,5 +196,9 @@ mod tests { assert!(buf.pull_one().await.unwrap()); let err = buf.pull_one().await.unwrap_err(); assert_eq!(err.to_string(), "boom"); + // An error is terminal: eof is set and further pulls are no-ops + // rather than re-polling the already-failed stream. + assert!(buf.is_eof()); + assert!(!buf.pull_one().await.unwrap()); } } From 7a678b8028db7996080f000bd19c01f6c6df9a1a Mon Sep 17 00:00:00 2001 From: Luke Valenta Date: Fri, 21 Aug 2026 11:48:57 -0400 Subject: [PATCH 4/8] mirror_worker: retry header parse on log_origin truncation A valid header split partway through log_origin returned 400 instead of pulling the next chunk. read_from special-cased that short read as a distinct LogOriginTruncated variant, so parse_header's single Io(UnexpectedEof) arm missed it. Drop the variant so every short read in the fixed-layout header is a plain Io(UnexpectedEof) that parse_header treats as "need more bytes", 400ing only once the stream ends. Adds streaming parse_header tests for byte-by-byte delivery, a split inside log_origin, and genuine truncation. --- crates/mirror_worker/src/add_entries.rs | 73 +++++++++++++++++++++++-- crates/tlog_mirror/src/error.rs | 8 --- crates/tlog_mirror/src/wire.rs | 12 +--- 3 files changed, 69 insertions(+), 24 deletions(-) diff --git a/crates/mirror_worker/src/add_entries.rs b/crates/mirror_worker/src/add_entries.rs index 6a31edd3..84001954 100644 --- a/crates/mirror_worker/src/add_entries.rs +++ b/crates/mirror_worker/src/add_entries.rs @@ -812,11 +812,9 @@ where buf.consume(consumed); return Ok(header); } + // Short read: the whole header is not buffered yet. Pull more + // and retry; a real EOF here means a truncated header. Err(ParseError::Io(ref e)) if e.kind() == ErrorKind::UnexpectedEof => { - // Need more bytes to parse the header. Pull another - // chunk; if the stream is already at EOF, the header - // is fundamentally malformed (truncated before being - // complete). if !buf.pull_one().await? { log::warn!( "add-entries: stream ended before header was complete \ @@ -1189,13 +1187,16 @@ impl HashReader for MapReader<'_> { #[cfg(test)] mod tests { use super::{ - CONTENT_TYPE, MapReader, content_type_is_octet_stream, excess_entries, verify_package, + CONTENT_TYPE, MapReader, content_type_is_octet_stream, excess_entries, parse_header, + verify_package, }; + use crate::body::BodyError; use crate::mirror_state_do::PendingCheckpoint; + use crate::stream_buffer::StreamBuffer; use std::collections::HashMap; use tlog_core::{Hash, Subtree, stored_hash_index, stored_hashes, tree_hash}; use tlog_mirror::EntryPackage; - use tlog_mirror::PACKAGE_ALIGNMENT; + use tlog_mirror::{AddEntriesRequestHeader, PACKAGE_ALIGNMENT}; /// Deterministic distinct entry bytes for leaf `i`. fn entry(i: u64) -> Vec { @@ -1382,4 +1383,64 @@ mod tests { headers.insert(CONTENT_TYPE, "application/json".parse().unwrap()); assert!(!content_type_is_octet_stream(&headers)); } + + fn header_bytes() -> Vec { + let header = AddEntriesRequestHeader { + log_origin: "rome.ct.example.com/2026h1".to_owned(), + upload_start: 256, + upload_end: 512, + ticket: b"opaque-ticket".to_vec(), + }; + let mut buf = Vec::new(); + header.write_to(&mut buf).unwrap(); + buf + } + + fn stream_buffer( + chunks: Vec>, + ) -> StreamBuffer, BodyError>> + Unpin> { + StreamBuffer::new(futures_util::stream::iter( + chunks.into_iter().map(Ok::<_, BodyError>), + )) + } + + // A header delivered as single-byte chunks must reassemble. + #[tokio::test(flavor = "current_thread")] + async fn parse_header_reassembles_byte_by_byte() { + let bytes = header_bytes(); + let chunks: Vec> = bytes.iter().map(|b| vec![*b]).collect(); + let mut buf = stream_buffer(chunks); + let Ok(header) = parse_header(&mut buf).await else { + panic!("header reassembles"); + }; + assert_eq!(header.log_origin, "rome.ct.example.com/2026h1"); + assert_eq!(header.upload_start, 256); + assert_eq!(header.upload_end, 512); + assert_eq!(header.ticket, b"opaque-ticket"); + } + + // A header split partway through log_origin must reassemble. + #[tokio::test(flavor = "current_thread")] + async fn parse_header_split_inside_log_origin() { + let bytes = header_bytes(); + let (first, rest) = bytes.split_at(3); + let mut buf = stream_buffer(vec![first.to_vec(), rest.to_vec()]); + let Ok(header) = parse_header(&mut buf).await else { + panic!("header reassembles"); + }; + assert_eq!(header.log_origin, "rome.ct.example.com/2026h1"); + } + + // A stream that ends partway through log_origin is genuinely truncated + // and must be a 400, not an endless pull loop. + #[tokio::test(flavor = "current_thread")] + async fn parse_header_truncated_in_log_origin_is_bad_request() { + let bytes = header_bytes(); + let truncated = bytes[..5].to_vec(); + let mut buf = stream_buffer(vec![truncated]); + assert!(matches!( + parse_header(&mut buf).await, + Err(super::AppError::BadRequest(_)) + )); + } } diff --git a/crates/tlog_mirror/src/error.rs b/crates/tlog_mirror/src/error.rs index 7258fe9c..0959637c 100644 --- a/crates/tlog_mirror/src/error.rs +++ b/crates/tlog_mirror/src/error.rs @@ -12,14 +12,6 @@ pub enum ParseError { #[error("io: {0}")] Io(#[from] std::io::Error), - /// The `log_origin_size` u16 prefix advertised more bytes than were - /// available in the input. - #[error("log_origin truncated: advertised {advertised} bytes")] - LogOriginTruncated { - /// Size advertised by the wire `log_origin_size` u16. - advertised: u16, - }, - /// The `log_origin` bytes were not valid UTF-8. #[error("log_origin is not valid UTF-8")] LogOriginNotUtf8, diff --git a/crates/tlog_mirror/src/wire.rs b/crates/tlog_mirror/src/wire.rs index c591733a..16898b03 100644 --- a/crates/tlog_mirror/src/wire.rs +++ b/crates/tlog_mirror/src/wire.rs @@ -61,15 +61,7 @@ impl AddEntriesRequestHeader { pub fn read_from(mut reader: R) -> Result { let log_origin_size = reader.read_u16::()?; let mut log_origin_bytes = vec![0u8; usize::from(log_origin_size)]; - reader.read_exact(&mut log_origin_bytes).map_err(|e| { - if e.kind() == io::ErrorKind::UnexpectedEof { - ParseError::LogOriginTruncated { - advertised: log_origin_size, - } - } else { - ParseError::Io(e) - } - })?; + reader.read_exact(&mut log_origin_bytes)?; let log_origin = String::from_utf8(log_origin_bytes).map_err(|_| ParseError::LogOriginNotUtf8)?; @@ -481,7 +473,7 @@ mod tests { let err = AddEntriesRequestHeader::read_from(Cursor::new(&buf)).unwrap_err(); assert!(matches!( err, - ParseError::LogOriginTruncated { advertised: 5 } + ParseError::Io(e) if e.kind() == io::ErrorKind::UnexpectedEof )); } From 830070dc8f948f267cd557ba8fb0079c55943a7b Mon Sep 17 00:00:00 2001 From: Luke Valenta Date: Fri, 21 Aug 2026 12:09:17 -0400 Subject: [PATCH 5/8] mirror_worker: bound gunzip output per run against zip bombs GzipInflater::push wrote whole compressed chunks into the decoder and took the entire accumulated plaintext at once, so a small but highly compressible chunk could inflate into a Vec far larger than the isolate's memory ceiling. Feed the decoder in INFLATE_STEP slices, drain the sink per slice, and emit plaintext in <=INFLATE_STEP runs that gunzip queues and yields one per poll. Adds a zip-bomb test. --- crates/mirror_worker/src/body.rs | 117 +++++++++++++++++++++++++------ 1 file changed, 94 insertions(+), 23 deletions(-) diff --git a/crates/mirror_worker/src/body.rs b/crates/mirror_worker/src/body.rs index 4f70d610..22ddb649 100644 --- a/crates/mirror_worker/src/body.rs +++ b/crates/mirror_worker/src/body.rs @@ -97,6 +97,15 @@ pub(crate) fn decoded_stream( Ok(stream) } +/// Largest compressed slice fed to the decoder per step, and the point at +/// which accumulated plaintext is drained. DEFLATE can inflate a small +/// input by a large factor (a hostile "zip bomb" reaches ~1000x), so a +/// single unbounded `write_all` of a whole chunk could balloon the sink +/// `Vec` far past the isolate's memory ceiling. Feeding the decoder in +/// bounded slices and draining between them keeps the plaintext held at +/// once proportional to this window, not to the compression ratio. +const INFLATE_STEP: usize = 64 * 1024; + /// Incremental gzip inflater: feed compressed bytes with [`Self::push`] /// and drain the plaintext produced so far; call [`Self::finish`] once the /// compressed input ends to validate the gzip trailer (CRC-32 + ISIZE). @@ -105,6 +114,11 @@ pub(crate) fn decoded_stream( /// compiles to and runs under WASM. `flate2::write::GzDecoder` handles all /// gzip framing: the 10-byte header, optional FNAME/FEXTRA/etc. fields /// (buffered across `push` calls if split across chunks), and the trailer. +/// +/// [`Self::push`] feeds the decoder at most [`INFLATE_STEP`] compressed +/// bytes at a time and drains after each step, so the sink never holds +/// more than one step's worth of inflated output regardless of how large +/// or compressible the caller's chunk is. struct GzipInflater { decoder: GzDecoder>, } @@ -116,14 +130,28 @@ impl GzipInflater { } } - /// Feed one compressed chunk and return the plaintext bytes produced. - /// May return an empty `Vec` if `input` only completed part of the - /// gzip header or a DEFLATE block that hasn't emitted output yet. - fn push(&mut self, input: &[u8]) -> Result> { - self.decoder - .write_all(input) - .map_err(|e| Error::from(format!("gzip decode failed: {e}")))?; - Ok(std::mem::take(self.decoder.get_mut())) + /// Feed one compressed chunk, invoking `emit` with each bounded + /// plaintext run produced. `emit` may be called zero times (the input + /// only advanced the gzip header or a not-yet-emitting DEFLATE block), + /// once, or many times for a highly compressible chunk. + /// + /// Input is written to the decoder one [`INFLATE_STEP`] slice at a + /// time, draining the sink after each, and every drained run is further + /// split into at most [`INFLATE_STEP`]-byte emissions. So no single + /// emitted run (what flows downstream) exceeds one step regardless of + /// the (attacker-controlled) compression ratio, and the sink is drained + /// per input step instead of accumulating the whole chunk's output. + fn push(&mut self, input: &[u8], mut emit: impl FnMut(Vec)) -> Result<()> { + for slice in input.chunks(INFLATE_STEP) { + self.decoder + .write_all(slice) + .map_err(|e| Error::from(format!("gzip decode failed: {e}")))?; + let produced = std::mem::take(self.decoder.get_mut()); + for run in produced.chunks(INFLATE_STEP) { + emit(run.to_vec()); + } + } + Ok(()) } /// Finish decompression, returning any trailing plaintext. Errors if @@ -140,10 +168,13 @@ impl GzipInflater { /// gunzipped plaintext chunks. /// /// The returned stream inflates lazily: each poll pulls compressed chunks -/// from `inner` until it can emit at least one plaintext byte, so memory +/// from `inner` until it can emit at least one plaintext run, so memory /// use stays bounded by the chunk size rather than the whole body. A -/// decode error (malformed gzip) or a truncated stream surfaces as a -/// terminal `Err` item, after which the stream ends. +/// single compressed chunk that inflates a lot is emitted as several +/// [`INFLATE_STEP`]-bounded runs, drained from `pending` across polls, so +/// a hostile compression ratio cannot force one giant allocation. A decode +/// error (malformed gzip) or a truncated stream surfaces as a terminal +/// `Err` item, after which the stream ends. pub(crate) fn gunzip(inner: S) -> BodyStream where S: Stream, BodyError>> + Unpin + 'static, @@ -151,6 +182,9 @@ where struct DecodeState { inner: S, inflater: Option, + /// Plaintext runs decoded from the last compressed chunk but not + /// yet emitted, drained one per poll (front to back). + pending: std::collections::VecDeque>, done: bool, } @@ -158,30 +192,34 @@ where DecodeState { inner, inflater: Some(GzipInflater::new()), + pending: std::collections::VecDeque::new(), done: false, }, |mut st| async move { + if let Some(run) = st.pending.pop_front() { + return Some((Ok(run), st)); + } if st.done { return None; } loop { match st.inner.next().await { Some(Ok(chunk)) => { - let out = match st.inflater.as_mut().expect("inflater present").push(&chunk) - { - Ok(out) => out, - Err(e) => { - st.done = true; - return Some((Err(BodyError::Decode(e.to_string())), st)); - } - }; + let inflater = st.inflater.as_mut().expect("inflater present"); + let mut runs = std::collections::VecDeque::new(); + if let Err(e) = inflater.push(&chunk, |run| runs.push_back(run)) { + st.done = true; + return Some((Err(BodyError::Decode(e.to_string())), st)); + } // A chunk may not yet yield any plaintext (partial // header / block); pull more instead of emitting an - // empty item. - if out.is_empty() { + // empty item. Otherwise emit the first run now and + // queue the rest for subsequent polls. + let Some(first) = runs.pop_front() else { continue; - } - return Some((Ok(out), st)); + }; + st.pending = runs; + return Some((Ok(first), st)); } Some(Err(e)) => { st.done = true; @@ -291,6 +329,39 @@ mod tests { ); } + // A highly compressible payload whose plaintext far exceeds + // INFLATE_STEP, delivered as a single compressed chunk, must inflate + // to the exact original but be emitted as several bounded runs so no + // single allocation exceeds the step. This is the zip-bomb guard. + #[tokio::test] + async fn large_ratio_chunk_emits_bounded_runs() { + let plain = vec![0u8; INFLATE_STEP * 10 + 123]; + let compressed = gzip(&plain); + assert!( + compressed.len() < INFLATE_STEP, + "test payload should compress to well under one step" + ); + // Feed the whole compressed body as one chunk. + let mut s = gunzip(stream::iter(vec![Ok::<_, BodyError>(compressed)])); + let mut total = 0usize; + let mut runs = 0usize; + while let Some(item) = s.next().await { + let run = item.unwrap(); + assert!( + run.len() <= INFLATE_STEP, + "run of {} bytes exceeds INFLATE_STEP {INFLATE_STEP}", + run.len() + ); + total += run.len(); + runs += 1; + } + assert_eq!(total, plain.len()); + assert!( + runs > 1, + "a >step payload must span multiple runs, got {runs}" + ); + } + #[tokio::test] async fn upstream_error_propagates() { let compressed = gzip(b"partial"); From 201aa0780a45cee0911a656cb60d688f4347dcd0 Mon Sep 17 00:00:00 2001 From: Luke Valenta Date: Fri, 21 Aug 2026 12:15:52 -0400 Subject: [PATCH 6/8] mirror_worker: cap buffered entries by bytes, not just package count commit_packages bounds only the number of buffered packages, but a single package can hold 256 entries of up to 65535 bytes (~16 MiB), so a count-only cap could buffer far past the isolate's memory ceiling. Add max_chunk_bytes (default 16 MiB) and flush early once buffered entry bytes reach it, whichever cap trips first. --- crates/mirror_worker/config.schema.json | 6 +++++ crates/mirror_worker/config/src/lib.rs | 32 +++++++++++++++++++++++++ crates/mirror_worker/src/add_entries.rs | 23 ++++++++++++++---- 3 files changed, 57 insertions(+), 4 deletions(-) diff --git a/crates/mirror_worker/config.schema.json b/crates/mirror_worker/config.schema.json index fd5da031..8eecfe47 100644 --- a/crates/mirror_worker/config.schema.json +++ b/crates/mirror_worker/config.schema.json @@ -37,6 +37,12 @@ "default": 32, "description": "How many entry packages add-entries verifies before flushing them to storage and advancing the persisted-entry frontier. Bounds in-memory buffering and gives durable mid-request progress on large uploads. Defaults to 32 (the recommended per-request package budget) when omitted; capped at 1024 to bound worst-case buffering." }, + "max_chunk_bytes": { + "type": "integer", + "minimum": 1, + "default": 16777216, + "description": "Byte ceiling on the entries buffered between flushes. add-entries flushes early once buffered entry bytes reach this many, even if fewer than commit_packages packages have accumulated. Bounds peak memory independent of package sizes, since a single package can hold up to ~16 MiB. Defaults to 16777216 (16 MiB) when omitted." + }, "logs": { "type": "object", "description": "CAs this mirror mirrors, keyed by log_key_name: the CA cosigner's note-signature name (the CA ID) on the checkpoints it ingests. Used as a signed-note key name at runtime, so per c2sp.org/signed-note it MUST NOT contain '+', whitespace, or control characters.", diff --git a/crates/mirror_worker/config/src/lib.rs b/crates/mirror_worker/config/src/lib.rs index 433f605b..f578ba46 100644 --- a/crates/mirror_worker/config/src/lib.rs +++ b/crates/mirror_worker/config/src/lib.rs @@ -69,6 +69,16 @@ pub struct AppConfig { /// (see [`Self::commit_packages`]). Consumed by /// [`mirror_worker`](../mirror_worker/)'s `add_entries`. pub commit_packages: Option, + /// Byte ceiling on the entries buffered between flushes. `add-entries` + /// flushes early once the buffered entry bytes reach this many, even if + /// fewer than `commit_packages` packages have accumulated. This bounds + /// peak memory independent of package sizes: `commit_packages` alone + /// caps only the package *count*, and a single package can hold up to + /// 256 entries of 65535 bytes (~16 MiB), so a count-only bound can + /// exceed the isolate's memory ceiling. `None` falls back to a default + /// (see [`Self::max_chunk_bytes`]). Consumed by + /// [`mirror_worker`](../mirror_worker/)'s `add_entries`. + pub max_chunk_bytes: Option, /// CAs this mirror mirrors, keyed by `log_key_name`: the CA /// cosigner's note-signature name (the CA ID) carried by the /// checkpoints it ingests. @@ -136,6 +146,18 @@ impl AppConfig { self.commit_packages.unwrap_or(32) } + /// Byte ceiling on buffered entries before `add-entries` flushes early, + /// falling back to 16 MiB when `max_chunk_bytes` is unset. This bounds + /// peak in-memory buffering regardless of how large individual packages + /// are, complementing the `commit_packages` count cap. 16 MiB matches + /// the worst-case size of a single spec-maximal package (256 entries * + /// 65535 bytes), so the default never flushes mid-package for a + /// compliant upload yet still caps a pathological one. + #[must_use] + pub fn max_chunk_bytes(&self) -> u64 { + self.max_chunk_bytes.unwrap_or(16 * 1024 * 1024) + } + /// Validate the configuration beyond what `serde` and the JSON schema /// can express. /// @@ -359,6 +381,7 @@ mod tests { monitoring_prefix: None, clean_interval_secs: None, commit_packages: None, + max_chunk_bytes: None, logs: HashMap::from([( "example.com/log1".to_owned(), LogParams { @@ -491,6 +514,14 @@ mod tests { assert_eq!(cfg.commit_packages(), 8); } + #[test] + fn max_chunk_bytes_defaults_to_16_mib() { + let mut cfg = good_app_config(); + assert_eq!(cfg.max_chunk_bytes(), 16 * 1024 * 1024); + cfg.max_chunk_bytes = Some(1024); + assert_eq!(cfg.max_chunk_bytes(), 1024); + } + #[test] fn validate_rejects_inverted_window() { let cfg = with_log(|log| { @@ -560,6 +591,7 @@ mod tests { monitoring_prefix: None, clean_interval_secs: None, commit_packages: None, + max_chunk_bytes: None, logs: HashMap::from([( "a".repeat(250), LogParams { diff --git a/crates/mirror_worker/src/add_entries.rs b/crates/mirror_worker/src/add_entries.rs index 84001954..ca116b48 100644 --- a/crates/mirror_worker/src/add_entries.rs +++ b/crates/mirror_worker/src/add_entries.rs @@ -386,6 +386,11 @@ where // config.schema.json caps commit_packages (max 1024), enforced by the // build script, so this always fits usize; the fallback is unreachable. let commit_packages = usize::try_from(crate::CONFIG.commit_packages()).unwrap_or(usize::MAX); + // Byte ceiling on buffered entries; flush early when reached so peak + // memory is bounded regardless of package sizes. Saturating to + // usize::MAX on a 32-bit target just means "never trip the byte cap", + // leaving the package-count cap in force. + let max_chunk_bytes = usize::try_from(crate::CONFIG.max_chunk_bytes()).unwrap_or(usize::MAX); // Entries below the request-start frontier are already persisted; new // persistence begins at this fixed boundary. @@ -400,6 +405,7 @@ where let mut chunk: Vec> = Vec::new(); let mut chunk_end = frontier_size; let mut chunk_pkgs = 0usize; + let mut chunk_bytes = 0usize; let mut packages_received: u64 = 0; let mut truncated = false; @@ -457,13 +463,21 @@ where if pkg_end > initial_next { let skip = usize::try_from(initial_next.saturating_sub(pkg_start)) .map_err(|_| Error::from("skip count overflows usize"))?; - chunk.extend(pkg.entries.into_iter().skip(skip)); + let tail = pkg.entries.into_iter().skip(skip); + for entry in tail { + chunk_bytes = chunk_bytes.saturating_add(entry.len()); + chunk.push(entry); + } chunk_end = pkg_end; chunk_pkgs += 1; - // `commit_packages >= 1` (config), so a full chunk always holds - // at least one package's entries: no empty-flush guard needed. - if chunk_pkgs == commit_packages { + // Flush when either cap trips: `commit_packages` bounds the + // package count, `max_chunk_bytes` bounds peak memory when + // individual packages are large. `commit_packages >= 1` + // (config), so a full chunk always holds at least one package's + // entries; the byte cap only fires after entries were buffered, + // so neither branch can flush an empty chunk. + if chunk_pkgs == commit_packages || chunk_bytes >= max_chunk_bytes { (frontier_size, frontier_hash) = flush_chunk( bucket, env, @@ -475,6 +489,7 @@ where ) .await?; chunk_pkgs = 0; + chunk_bytes = 0; } } } From d9d73e0a866d5a2bd8391134de9c4dd72e75aad3 Mon Sep 17 00:00:00 2001 From: Luke Valenta Date: Fri, 21 Aug 2026 12:28:47 -0400 Subject: [PATCH 7/8] mirror_worker: parse entry packages incrementally parse_next_package rebuilt a Cursor at byte zero and re-ran EntryPackage::read_from on every pull, re-parsing and re-allocating all already-read entries; a large package arriving in small chunks was O(n^2) in copies. Parse entry-by-entry instead, consuming each entry and the proof from the StreamBuffer as soon as it is fully buffered, so a short read resumes from the next unread unit. Adds tests for byte-by-byte reassembly, back-to-back packages, both truncation classes, and the num_hashes limit. --- crates/mirror_worker/src/add_entries.rs | 240 +++++++++++++++++++++--- 1 file changed, 214 insertions(+), 26 deletions(-) diff --git a/crates/mirror_worker/src/add_entries.rs b/crates/mirror_worker/src/add_entries.rs index ca116b48..58bfd788 100644 --- a/crates/mirror_worker/src/add_entries.rs +++ b/crates/mirror_worker/src/add_entries.rs @@ -849,9 +849,53 @@ where } } -/// Read the next entry package from `buf`, pulling more chunks from -/// the underlying stream until the package parses or the stream ends. -/// See [`ParseOutcome`] for the four cases. +/// Pull chunks until at least `n` bytes are buffered. Returns `Ok(true)` +/// once `n` bytes are available, or `Ok(false)` if the stream ended first. +/// Nothing is consumed; the caller parses the now-buffered bytes. +async fn fill_at_least(buf: &mut StreamBuffer, n: usize) -> ApiResult +where + S: futures_util::Stream, BodyError>> + Unpin, +{ + while buf.len() < n { + if !buf.pull_one().await? { + return Ok(false); + } + } + Ok(true) +} + +/// Read one length-prefixed entry (`u16 len || len bytes`) from the front +/// of `buf`, consuming exactly the bytes read. Returns `Ok(None)` if the +/// stream ends before a complete entry is buffered (truncation). Unlike a +/// whole-package reparse, each call consumes what it reads, so pulling more +/// bytes for a later entry never re-copies the entries already taken. +async fn read_one_entry(buf: &mut StreamBuffer) -> ApiResult>> +where + S: futures_util::Stream, BodyError>> + Unpin, +{ + if !fill_at_least(buf, 2).await? { + return Ok(None); + } + let bytes = buf.buffered(); + let len = usize::from(u16::from_be_bytes([bytes[0], bytes[1]])); + let total = 2 + len; + if !fill_at_least(buf, total).await? { + return Ok(None); + } + let entry = buf.buffered()[2..total].to_vec(); + buf.consume(total); + Ok(Some(entry)) +} + +/// Read the next entry package from `buf`, pulling more chunks from the +/// underlying stream until the package parses or the stream ends. See +/// [`ParseOutcome`] for the four cases. +/// +/// The package is parsed incrementally, consuming each entry and the proof +/// from `buf` as soon as it is fully buffered. A short read pulls one more +/// chunk and resumes from the next unread unit, so a large package +/// arriving in small chunks is parsed and copied once, not re-parsed from +/// byte zero on every chunk. async fn parse_next_package( buf: &mut StreamBuffer, num_entries: u64, @@ -859,34 +903,62 @@ async fn parse_next_package( where S: futures_util::Stream, BodyError>> + Unpin, { + // Reject oversized counts before allocating, matching + // EntryPackage::read_from so the two parse paths agree on limits. + if num_entries > PACKAGE_ALIGNMENT { + return Ok(ParseOutcome::Err(ParseError::TooManyEntries(num_entries))); + } // EOF with an empty buffer: clean truncation between packages. if buf.is_eof() && buf.len() == 0 { return Ok(ParseOutcome::CleanEof); } - loop { - let mut cursor = Cursor::new(buf.buffered()); - match EntryPackage::read_from(&mut cursor, num_entries) { - Ok(pkg) => { - let consumed = usize::try_from(cursor.position()).unwrap_or(usize::MAX); - buf.consume(consumed); - return Ok(ParseOutcome::Ok(pkg)); - } - Err(ParseError::Io(ref e)) if e.kind() == ErrorKind::UnexpectedEof => { - if !buf.pull_one().await? { - // Stream ended mid-parse. An empty buffer means the - // previous package consumed exactly all buffered bytes - // and this call started a fresh (never-arriving) - // package: a clean between-package truncation. A - // non-empty buffer holds a partial package. - if buf.len() == 0 { - return Ok(ParseOutcome::CleanEof); - } - return Ok(ParseOutcome::MidPackageEof); - } + + let num_entries = usize::try_from(num_entries).unwrap_or(usize::MAX); + let mut entries = Vec::with_capacity(num_entries); + for _ in 0..num_entries { + match read_one_entry(buf).await? { + Some(entry) => entries.push(entry), + None => { + // Stream ended before this entry completed. Nothing buffered + // and no partial bytes means a clean between-package + // truncation; otherwise a partial package was left behind. + return Ok(package_eof(entries.is_empty(), buf.len() > 0)); } - Err(e) => return Ok(ParseOutcome::Err(e)), } } + + // Proof: `u8 num_hashes || num_hashes * HASH_SIZE bytes`. + if !fill_at_least(buf, 1).await? { + return Ok(package_eof(entries.is_empty(), buf.len() > 0)); + } + let num_hashes = buf.buffered()[0]; + if num_hashes > tlog_mirror::MAX_HASHES_PER_PROOF { + return Ok(ParseOutcome::Err(ParseError::TooManyHashes(num_hashes))); + } + let proof_bytes = 1 + usize::from(num_hashes) * tlog_core::HASH_SIZE; + if !fill_at_least(buf, proof_bytes).await? { + return Ok(package_eof(entries.is_empty(), buf.len() > 0)); + } + let mut proof = Vec::with_capacity(usize::from(num_hashes)); + for i in 0..usize::from(num_hashes) { + let off = 1 + i * tlog_core::HASH_SIZE; + let mut hash = [0u8; tlog_core::HASH_SIZE]; + hash.copy_from_slice(&buf.buffered()[off..off + tlog_core::HASH_SIZE]); + proof.push(Hash(hash)); + } + buf.consume(proof_bytes); + Ok(ParseOutcome::Ok(EntryPackage { entries, proof })) +} + +/// Classify a stream end reached partway through [`parse_next_package`]: +/// a clean between-package truncation when nothing of this package had +/// been read, otherwise a mid-package truncation. +fn package_eof(no_entries_yet: bool, saw_partial: bool) -> ParseOutcome { + if no_entries_yet && !saw_partial { + ParseOutcome::CleanEof + } else { + ParseOutcome::MidPackageEof + } } /// Read the per-origin DO state snapshot. A non-200 status or RPC failure @@ -1202,8 +1274,8 @@ impl HashReader for MapReader<'_> { #[cfg(test)] mod tests { use super::{ - CONTENT_TYPE, MapReader, content_type_is_octet_stream, excess_entries, parse_header, - verify_package, + CONTENT_TYPE, MapReader, ParseOutcome, content_type_is_octet_stream, excess_entries, + parse_header, parse_next_package, verify_package, }; use crate::body::BodyError; use crate::mirror_state_do::PendingCheckpoint; @@ -1458,4 +1530,120 @@ mod tests { Err(super::AppError::BadRequest(_)) )); } + + fn sample_package() -> EntryPackage { + EntryPackage { + entries: vec![ + b"first-entry".to_vec(), + Vec::new(), + b"a-much-longer-third-entry-with-more-bytes".to_vec(), + vec![0xab; 300], + ], + proof: vec![ + Hash([0x11; tlog_core::HASH_SIZE]), + Hash([0x22; tlog_core::HASH_SIZE]), + ], + } + } + + fn package_bytes(pkg: &EntryPackage) -> Vec { + let mut buf = Vec::new(); + pkg.write_to(&mut buf).unwrap(); + buf + } + + // parse_next_package returns ApiResult, whose Err (AppError) is not + // Debug, so unwrap the transport layer by hand for the tests. + fn ok_outcome(res: super::ApiResult) -> ParseOutcome { + let Ok(outcome) = res else { + panic!("unexpected transport error from parse_next_package"); + }; + outcome + } + + // The incremental parser must reconstruct a package identical to a + // one-shot read regardless of how the wire bytes are chunked, including + // single-byte chunks that split every length prefix, entry, and proof + // hash. This is the regression for re-parsing from byte zero. + #[tokio::test(flavor = "current_thread")] + async fn parse_next_package_reassembles_across_chunk_sizes() { + let pkg = sample_package(); + let bytes = package_bytes(&pkg); + let num_entries = pkg.entries.len() as u64; + for size in [1usize, 2, 3, 7, 33, bytes.len()] { + let chunks: Vec> = bytes.chunks(size).map(<[u8]>::to_vec).collect(); + let mut buf = stream_buffer(chunks); + let ParseOutcome::Ok(parsed) = + ok_outcome(parse_next_package(&mut buf, num_entries).await) + else { + panic!("chunk size {size} should parse Ok"); + }; + assert_eq!(parsed.entries, pkg.entries, "chunk size {size} entries"); + assert_eq!(parsed.proof, pkg.proof, "chunk size {size} proof"); + } + } + + // Two packages back to back: the first parse must consume exactly its + // bytes, leaving the second intact for the next call. + #[tokio::test(flavor = "current_thread")] + async fn parse_next_package_leaves_following_package_intact() { + let pkg = sample_package(); + let mut bytes = package_bytes(&pkg); + bytes.extend(package_bytes(&pkg)); + let num_entries = pkg.entries.len() as u64; + // 5-byte chunks so package boundaries fall mid-chunk. + let chunks: Vec> = bytes.chunks(5).map(<[u8]>::to_vec).collect(); + let mut buf = stream_buffer(chunks); + for which in ["first", "second"] { + let ParseOutcome::Ok(parsed) = + ok_outcome(parse_next_package(&mut buf, num_entries).await) + else { + panic!("{which} package should parse Ok"); + }; + assert_eq!(parsed.entries, pkg.entries, "{which} entries"); + assert_eq!(parsed.proof, pkg.proof, "{which} proof"); + } + } + + // Stream ending exactly on a package boundary is a clean truncation. + #[tokio::test(flavor = "current_thread")] + async fn parse_next_package_clean_eof_between_packages() { + let pkg = sample_package(); + let mut buf = stream_buffer(vec![package_bytes(&pkg)]); + let num_entries = pkg.entries.len() as u64; + assert!(matches!( + ok_outcome(parse_next_package(&mut buf, num_entries).await), + ParseOutcome::Ok(_) + )); + // Buffer now empty and stream ended: next call is a clean EOF. + assert!(matches!( + ok_outcome(parse_next_package(&mut buf, num_entries).await), + ParseOutcome::CleanEof + )); + } + + // Stream ending partway through a package is a mid-package truncation. + #[tokio::test(flavor = "current_thread")] + async fn parse_next_package_mid_package_eof() { + let pkg = sample_package(); + let bytes = package_bytes(&pkg); + let mut buf = stream_buffer(vec![bytes[..10].to_vec()]); + let num_entries = pkg.entries.len() as u64; + assert!(matches!( + ok_outcome(parse_next_package(&mut buf, num_entries).await), + ParseOutcome::MidPackageEof + )); + } + + // An oversized num_hashes must be an Err, matching read_from's limit. + #[tokio::test(flavor = "current_thread")] + async fn parse_next_package_rejects_too_many_hashes() { + // One zero-length entry, then num_hashes = 64 (> spec max 63). + let bytes = vec![0x00, 0x00, 64]; + let mut buf = stream_buffer(vec![bytes]); + assert!(matches!( + ok_outcome(parse_next_package(&mut buf, 1).await), + ParseOutcome::Err(super::ParseError::TooManyHashes(64)) + )); + } } From 53a322ffe5673e31b54220f06eaa8bda416b686b Mon Sep 17 00:00:00 2001 From: Luke Valenta Date: Mon, 24 Aug 2026 11:26:26 -0400 Subject: [PATCH 8/8] mirror_worker: bound gunzip input per step INFLATE_STEP capped the compressed slice handed to write_all, but write_all inflates its whole slice into the decoder sink before returning, so a 64 KiB slice could produce ~64 MiB of plaintext at DEFLATE's ~1032x. Shrink the input step to 4 KiB, capping transient plaintext at ~4 MiB, and stop feeding once a step produces output, parking the rest of the chunk for the next poll. Downstream runs stay capped by a separate EMIT_STEP. --- crates/mirror_worker/src/body.rs | 207 ++++++++++++++++++++----------- 1 file changed, 138 insertions(+), 69 deletions(-) diff --git a/crates/mirror_worker/src/body.rs b/crates/mirror_worker/src/body.rs index 22ddb649..73975bcb 100644 --- a/crates/mirror_worker/src/body.rs +++ b/crates/mirror_worker/src/body.rs @@ -97,28 +97,33 @@ pub(crate) fn decoded_stream( Ok(stream) } -/// Largest compressed slice fed to the decoder per step, and the point at -/// which accumulated plaintext is drained. DEFLATE can inflate a small -/// input by a large factor (a hostile "zip bomb" reaches ~1000x), so a -/// single unbounded `write_all` of a whole chunk could balloon the sink -/// `Vec` far past the isolate's memory ceiling. Feeding the decoder in -/// bounded slices and draining between them keeps the plaintext held at -/// once proportional to this window, not to the compression ratio. -const INFLATE_STEP: usize = 64 * 1024; +/// Compressed bytes fed to the decoder per [`GzipInflater::step`] call. +/// `write_all` inflates its whole slice into the decoder's sink before +/// returning, and DEFLATE reaches ~1032x, so a step holds up to +/// `INFLATE_INPUT_STEP * 1032` bytes of plaintext transiently (~4 MiB +/// here). A small input step keeps that under the isolate's ~128 MB +/// ceiling even for a zip bomb; a larger chunk or the whole body is never +/// handed to `write_all` at once. +const INFLATE_INPUT_STEP: usize = 4 * 1024; -/// Incremental gzip inflater: feed compressed bytes with [`Self::push`] -/// and drain the plaintext produced so far; call [`Self::finish`] once the -/// compressed input ends to validate the gzip trailer (CRC-32 + ISIZE). +/// Largest plaintext run emitted downstream. The plaintext drained after a +/// step is split into runs no larger than this. +const EMIT_STEP: usize = 64 * 1024; + +/// Incremental gzip inflater: feed a bounded compressed slice with +/// [`Self::step`], draining the plaintext it produced; call +/// [`Self::finish`] once the compressed input ends to validate the gzip +/// trailer (CRC-32 + ISIZE). /// /// Backed by [`flate2`]'s pure-Rust `rust_backend` (`miniz_oxide`), so it /// compiles to and runs under WASM. `flate2::write::GzDecoder` handles all /// gzip framing: the 10-byte header, optional FNAME/FEXTRA/etc. fields -/// (buffered across `push` calls if split across chunks), and the trailer. +/// (buffered across calls if split across steps), and the trailer. /// -/// [`Self::push`] feeds the decoder at most [`INFLATE_STEP`] compressed -/// bytes at a time and drains after each step, so the sink never holds -/// more than one step's worth of inflated output regardless of how large -/// or compressible the caller's chunk is. +/// The caller (see [`gunzip`]) feeds at most [`INFLATE_INPUT_STEP`] bytes +/// per [`Self::step`] and stops once a step yields plaintext, leaving the +/// rest of the compressed chunk unfed. At most one step is inflated ahead +/// of what has been emitted. struct GzipInflater { decoder: GzDecoder>, } @@ -130,26 +135,24 @@ impl GzipInflater { } } - /// Feed one compressed chunk, invoking `emit` with each bounded - /// plaintext run produced. `emit` may be called zero times (the input - /// only advanced the gzip header or a not-yet-emitting DEFLATE block), - /// once, or many times for a highly compressible chunk. + /// Feed one compressed slice (at most [`INFLATE_INPUT_STEP`] bytes), + /// invoking `emit` with each [`EMIT_STEP`]-bounded plaintext run the + /// step produced. `emit` may be called zero times (the slice only + /// advanced the gzip header or a not-yet-emitting DEFLATE block), once, + /// or several times for a highly compressible slice. /// - /// Input is written to the decoder one [`INFLATE_STEP`] slice at a - /// time, draining the sink after each, and every drained run is further - /// split into at most [`INFLATE_STEP`]-byte emissions. So no single - /// emitted run (what flows downstream) exceeds one step regardless of - /// the (attacker-controlled) compression ratio, and the sink is drained - /// per input step instead of accumulating the whole chunk's output. - fn push(&mut self, input: &[u8], mut emit: impl FnMut(Vec)) -> Result<()> { - for slice in input.chunks(INFLATE_STEP) { - self.decoder - .write_all(slice) - .map_err(|e| Error::from(format!("gzip decode failed: {e}")))?; - let produced = std::mem::take(self.decoder.get_mut()); - for run in produced.chunks(INFLATE_STEP) { - emit(run.to_vec()); - } + /// # Panics + /// + /// Debug-asserts `slice.len() <= INFLATE_INPUT_STEP`; a larger slice + /// breaks the transient-plaintext bound. + fn step(&mut self, slice: &[u8], mut emit: impl FnMut(Vec)) -> Result<()> { + debug_assert!(slice.len() <= INFLATE_INPUT_STEP, "input slice too large"); + self.decoder + .write_all(slice) + .map_err(|e| Error::from(format!("gzip decode failed: {e}")))?; + let produced = std::mem::take(self.decoder.get_mut()); + for run in produced.chunks(EMIT_STEP) { + emit(run.to_vec()); } Ok(()) } @@ -167,14 +170,15 @@ impl GzipInflater { /// Wrap a compressed body `Stream` in a decoding stream that yields the /// gunzipped plaintext chunks. /// -/// The returned stream inflates lazily: each poll pulls compressed chunks -/// from `inner` until it can emit at least one plaintext run, so memory -/// use stays bounded by the chunk size rather than the whole body. A -/// single compressed chunk that inflates a lot is emitted as several -/// [`INFLATE_STEP`]-bounded runs, drained from `pending` across polls, so -/// a hostile compression ratio cannot force one giant allocation. A decode -/// error (malformed gzip) or a truncated stream surfaces as a terminal -/// `Err` item, after which the stream ends. +/// Each poll feeds the decoder at most [`INFLATE_INPUT_STEP`] compressed +/// bytes and stops once a step yields plaintext, leaving the rest of the +/// current chunk unfed until the next poll (tracked by `current`/`pos`). +/// At most one input step is inflated ahead of what has been emitted, so +/// the transient plaintext is bounded by `INFLATE_INPUT_STEP * ratio` +/// rather than the whole chunk's inflation. A step that produces more than +/// [`EMIT_STEP`] of plaintext yields several bounded runs drained from +/// `pending`. A decode error (malformed gzip) or a truncated stream +/// surfaces as a terminal `Err` item, after which the stream ends. pub(crate) fn gunzip(inner: S) -> BodyStream where S: Stream, BodyError>> + Unpin + 'static, @@ -182,8 +186,13 @@ where struct DecodeState { inner: S, inflater: Option, - /// Plaintext runs decoded from the last compressed chunk but not - /// yet emitted, drained one per poll (front to back). + /// Compressed chunk currently being fed, and the offset of the + /// next unfed byte. Bytes at `pos..` are not fed until the current + /// runs drain, capping how far ahead the decoder inflates. + current: Vec, + pos: usize, + /// Plaintext runs from the last step not yet emitted, drained one + /// per poll (front to back). pending: std::collections::VecDeque>, done: bool, } @@ -192,6 +201,8 @@ where DecodeState { inner, inflater: Some(GzipInflater::new()), + current: Vec::new(), + pos: 0, pending: std::collections::VecDeque::new(), done: false, }, @@ -203,24 +214,30 @@ where return None; } loop { - match st.inner.next().await { - Some(Ok(chunk)) => { - let inflater = st.inflater.as_mut().expect("inflater present"); - let mut runs = std::collections::VecDeque::new(); - if let Err(e) = inflater.push(&chunk, |run| runs.push_back(run)) { - st.done = true; - return Some((Err(BodyError::Decode(e.to_string())), st)); - } - // A chunk may not yet yield any plaintext (partial - // header / block); pull more instead of emitting an - // empty item. Otherwise emit the first run now and - // queue the rest for subsequent polls. - let Some(first) = runs.pop_front() else { - continue; - }; + // Feed the unfed tail of the current chunk one bounded + // step at a time, stopping as soon as a step emits. + while st.pos < st.current.len() { + let end = (st.pos + INFLATE_INPUT_STEP).min(st.current.len()); + let slice = st.current[st.pos..end].to_vec(); + st.pos = end; + let inflater = st.inflater.as_mut().expect("inflater present"); + let mut runs = std::collections::VecDeque::new(); + if let Err(e) = inflater.step(&slice, |run| runs.push_back(run)) { + st.done = true; + return Some((Err(BodyError::Decode(e.to_string())), st)); + } + if let Some(first) = runs.pop_front() { st.pending = runs; return Some((Ok(first), st)); } + } + + // Current chunk fully fed; pull the next compressed chunk. + match st.inner.next().await { + Some(Ok(chunk)) => { + st.current = chunk; + st.pos = 0; + } Some(Err(e)) => { st.done = true; return Some((Err(e), st)); @@ -329,17 +346,16 @@ mod tests { ); } - // A highly compressible payload whose plaintext far exceeds - // INFLATE_STEP, delivered as a single compressed chunk, must inflate - // to the exact original but be emitted as several bounded runs so no - // single allocation exceeds the step. This is the zip-bomb guard. + // A compressible payload whose plaintext exceeds EMIT_STEP, delivered + // as one compressed chunk, reconstructs and is emitted as + // EMIT_STEP-bounded runs rather than a single unbounded one. #[tokio::test] async fn large_ratio_chunk_emits_bounded_runs() { - let plain = vec![0u8; INFLATE_STEP * 10 + 123]; + let plain = vec![0u8; EMIT_STEP * 10 + 123]; let compressed = gzip(&plain); assert!( - compressed.len() < INFLATE_STEP, - "test payload should compress to well under one step" + compressed.len() < EMIT_STEP, + "test payload should compress to well under one emit step" ); // Feed the whole compressed body as one chunk. let mut s = gunzip(stream::iter(vec![Ok::<_, BodyError>(compressed)])); @@ -348,8 +364,8 @@ mod tests { while let Some(item) = s.next().await { let run = item.unwrap(); assert!( - run.len() <= INFLATE_STEP, - "run of {} bytes exceeds INFLATE_STEP {INFLATE_STEP}", + run.len() <= EMIT_STEP, + "run of {} bytes exceeds EMIT_STEP {EMIT_STEP}", run.len() ); total += run.len(); @@ -362,6 +378,59 @@ mod tests { ); } + // Feeding INFLATE_INPUT_STEP-bounded slices (as gunzip does) round-trips + // a body whose compressed form spans several steps. + #[test] + fn inflater_step_bounds_input() { + // A poorly-compressible body whose compressed form spans several + // input steps. + let mut plain = Vec::new(); + for i in 0u32..0x1_0000 { + plain.extend_from_slice(&i.wrapping_mul(2_654_435_761).to_le_bytes()); + } + let compressed = gzip(&plain); + assert!( + compressed.len() > INFLATE_INPUT_STEP, + "test needs a body spanning multiple input steps" + ); + + let mut inflater = GzipInflater::new(); + let mut total = 0usize; + // Feed as gunzip does: INFLATE_INPUT_STEP-bounded slices. + for slice in compressed.chunks(INFLATE_INPUT_STEP) { + inflater + .step(slice, |run| { + assert!(run.len() <= EMIT_STEP, "run exceeds EMIT_STEP"); + total += run.len(); + }) + .unwrap(); + } + total += inflater.finish().unwrap().len(); + assert_eq!(total, plain.len()); + } + + // One compressed chunk whose inflation exceeds a step: the first poll + // returns an EMIT_STEP-bounded run, leaving the rest unfed, and the + // full stream reconstructs the original. + #[tokio::test] + async fn large_ratio_chunk_emits_before_full_feed() { + let plain = vec![0u8; EMIT_STEP * 8]; + let compressed = gzip(&plain); + let mut s = gunzip(stream::iter(vec![Ok::<_, BodyError>(compressed)])); + let first = s.next().await.expect("a run").expect("no error"); + assert!( + first.len() <= EMIT_STEP, + "first emitted run must be bounded, got {}", + first.len() + ); + // The remaining runs complete the original. + let mut total = first.len(); + while let Some(item) = s.next().await { + total += item.unwrap().len(); + } + assert_eq!(total, plain.len()); + } + #[tokio::test] async fn upstream_error_propagates() { let compressed = gzip(b"partial");