diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b25cf93ee..dd730ed34 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -117,6 +117,10 @@ jobs: run: cargo clippy -p structured-zstd --features hash,std,dict-builder -- -D warnings - name: Clippy (testing facade) run: cargo clippy -p structured-zstd --features hash,std,dict-builder,bench-internals -- -D warnings + - name: Clippy (caller-workspace, all targets) + # The feature swaps the entropy tables' build scratch for inline + # storage and adds the workspace decoder; lint both with their tests. + run: cargo clippy -p structured-zstd --all-targets --features hash,std,dict-builder,lsm,caller-workspace -- -D warnings - name: Clippy (ffi-bench targets) # The libzstd-linking benches/tests/examples live in `ffi-bench`. # Lint them with the features their `required-features` gate on so the @@ -130,8 +134,8 @@ jobs: env: RUSTDOCFLAGS: -D warnings run: | - cargo doc -p structured-zstd --no-deps --features std,hash,ldm,dict-builder,lsm - cargo doc -p structured-zstd --no-deps --document-private-items --features std,hash,ldm,dict-builder,lsm + cargo doc -p structured-zstd --no-deps --features std,hash,ldm,dict-builder,lsm,caller-workspace + cargo doc -p structured-zstd --no-deps --document-private-items --features std,hash,ldm,dict-builder,lsm,caller-workspace cargo doc -p structured-zstd-c -p structured-zstd-wasm --no-deps --document-private-items - name: Gate — compare_ffi/memory must not pull bench-internals # Rationale: both `compare_ffi` (timing) and `compare_ffi_memory` @@ -385,9 +389,16 @@ jobs: if: needs.changes.outputs.code == 'true' working-directory: zstd run: cargo nextest run --profile ci -p structured-zstd --features hash,std,dict-builder + - name: Test (caller-workspace) + # The whole suite again with the feature on: the shared decode core + # then runs on inline table storage, and the workspace decoder's own + # tests and its zero-allocation test join in. + if: needs.changes.outputs.code == 'true' + working-directory: zstd + run: cargo nextest run --profile ci -p structured-zstd --features hash,std,dict-builder,caller-workspace - name: Doc tests if: needs.changes.outputs.code == 'true' - run: cargo test --doc -p structured-zstd --features hash,std,dict-builder + run: cargo test --doc -p structured-zstd --features hash,std,dict-builder,caller-workspace - name: FFI parity + cross-validation tests if: needs.changes.outputs.code == 'true' # The libzstd-linking parity / cross-validation / corpus tests live in @@ -478,6 +489,9 @@ jobs: - name: Clippy (embedded_minimal — kernel-scalar + std) working-directory: zstd run: cargo clippy -p structured-zstd --no-default-features --features kernel-scalar,std,hash -- -D warnings + - name: Clippy (no_std + caller-workspace) + working-directory: zstd + run: cargo clippy -p structured-zstd --no-default-features --features caller-workspace,hash -- -D warnings # The clippy steps above compile for the host, which has SSE2 and a # float unit. A kernel target has neither (x86_64-unknown-none is # -sse,-sse2,...,+soft-float), so a vector type that @@ -493,6 +507,7 @@ jobs: for profile in dev release; do cargo build -p structured-zstd --lib --profile "$profile" --target x86_64-unknown-none --no-default-features --features hash cargo build -p structured-zstd --lib --profile "$profile" --target x86_64-unknown-none --no-default-features --features kernel-avx2,hash + cargo build -p structured-zstd --lib --profile "$profile" --target x86_64-unknown-none --no-default-features --features caller-workspace,hash done wasm: @@ -527,8 +542,24 @@ jobs: # where a missing cargo-feature gate would go unnoticed, so keep it built. - name: Clippy (wasm32 simd128 feature without +simd128) run: cargo clippy -p structured-zstd -p structured-zstd-wasm --target wasm32-unknown-unknown --no-default-features --features kernel-simd128 -- -D warnings + - name: Clippy (wasm32 caller-workspace) + run: cargo clippy -p structured-zstd --target wasm32-unknown-unknown --no-default-features --features caller-workspace,hash,kernel-simd128 --config 'target.wasm32-unknown-unknown.rustflags=["-C","target-feature=+simd128"]' -- -D warnings + # The caller-workspace decoder is built for browser bundles, so its tests + # and its zero-allocation test run as wasm too, under wasmtime (the WASI + # target supplies the test harness's std). + - name: Install wasmtime + uses: taiki-e/install-action@1168f6810a3cb30a3514bf5756a3b7f0422c6df8 # v2.87.28 + with: + tool: wasmtime,cargo-nextest + - name: Workspace decoder tests (wasm32-wasip1) + env: + CARGO_TARGET_WASM32_WASIP1_RUNNER: wasmtime + run: | + rustup target add wasm32-wasip1 + cargo nextest run --profile ci -p structured-zstd --target wasm32-wasip1 --no-default-features --features caller-workspace,hash,std --lib -E 'test(/decoding::workspace::tests/)' + cargo nextest run --profile ci -p structured-zstd --target wasm32-wasip1 --no-default-features --features caller-workspace,hash,std --test workspace_no_alloc - name: Install wasm-pack - uses: taiki-e/install-action@v2 + uses: taiki-e/install-action@1168f6810a3cb30a3514bf5756a3b7f0422c6df8 # v2.87.28 with: tool: wasm-pack - uses: actions/setup-node@v7 @@ -594,7 +625,7 @@ jobs: with: prefix-key: wasm-bench - name: Install wasm-pack - uses: taiki-e/install-action@v2 + uses: taiki-e/install-action@1168f6810a3cb30a3514bf5756a3b7f0422c6df8 # v2.87.28 with: tool: wasm-pack - uses: actions/setup-node@v7 @@ -763,7 +794,7 @@ jobs: prefix-key: fuzz workspaces: zstd/fuzz - name: Install cargo-fuzz - uses: taiki-e/install-action@v2 + uses: taiki-e/install-action@1168f6810a3cb30a3514bf5756a3b7f0422c6df8 # v2.87.28 with: tool: cargo-fuzz - name: Replay regression corpus @@ -793,6 +824,11 @@ jobs: cargo fuzz run --target "$FUZZ_TARGET_TRIPLE" "$target" -- -max_total_time=90 -timeout=30 echo "::endgroup::" done + - name: Fuzz decode_workspace (90s, caller-workspace build) + # Its own step: the target needs the library's caller-workspace + # feature, which the targets above must not get. + working-directory: zstd/fuzz + run: cargo fuzz run --target "$FUZZ_TARGET_TRIPLE" --features workspace decode_workspace -- -max_total_time=90 -timeout=30 - name: Upload new crash artifacts on failure if: failure() uses: actions/upload-artifact@v7 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index d75414e72..c99e28461 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -48,7 +48,7 @@ jobs: with: targets: wasm32-unknown-unknown - name: Install wasm-pack - uses: taiki-e/install-action@v2 + uses: taiki-e/install-action@1168f6810a3cb30a3514bf5756a3b7f0422c6df8 # v2.87.28 with: tool: wasm-pack - uses: actions/setup-node@v7 diff --git a/README.md b/README.md index 4e2a26272..bd4332f38 100644 --- a/README.md +++ b/README.md @@ -231,6 +231,19 @@ reads the first frame header and reports any problem with the input, and the source untouched, so a container format can decide not to decompress an entry and copy its raw bytes instead. +### Decoding in caller-owned memory + +Behind the `caller-workspace` feature, `decoding::workspace::WorkspaceDecoder` +decodes a stream chunk by chunk inside one byte slice the caller supplies, +and never allocates. `workspace_bytes(max_window)` says how large the slice +must be, so a consumer that accounts its own memory (a browser bundle +verifying a download, say) admits it before decoding starts. Each step is +bounded by the window a frame may declare, the total output, and a work +budget charged before the work it pays for; a refusal is reported apart from +malformed or truncated input. It works in `no_std` builds and in WebAssembly, +and decodes through the same code as the other decoders. Frames that need a +dictionary are refused. + ### Dictionaries ```rust,no_run @@ -294,6 +307,7 @@ in pure Rust: | `kernel-scalar` | ✅ | Marker for the always-compiled scalar fallback | | `dict-builder` | ❌ | Pure-Rust COVER / FastCOVER / legacy dictionary training | | `lsm` | ❌ | [Storage-format extensions](#storage-format-extensions) | +| `caller-workspace` | ❌ | [Decoding in caller-owned memory](#decoding-in-caller-owned-memory) | Each flag gates its tier wherever that tier exists. `kernel-sse`, `kernel-bmi2`, `kernel-avx2`, `kernel-neon` and `kernel-simd128` cover both diff --git a/zstd/Cargo.toml b/zstd/Cargo.toml index d519405e8..32bc82d02 100644 --- a/zstd/Cargo.toml +++ b/zstd/Cargo.toml @@ -69,7 +69,7 @@ required-features = ["std", "dict-builder", "ldm"] # `#[cfg_attr(docsrs, doc(cfg(...)))]` annotations that render feature # badges on each item. [package.metadata.docs.rs] -features = ["std", "hash", "ldm", "dict-builder", "lsm"] +features = ["std", "hash", "ldm", "dict-builder", "lsm", "caller-workspace"] rustdoc-args = ["--cfg", "docsrs"] [dependencies] @@ -182,6 +182,11 @@ ldm = ["hash"] # targets with atomic pointer support this feature is a no-op. critical-section = ["dep:critical-section"] std = [] +# Allocation-free streaming decoder over caller-owned workspace memory +# (`decoding::workspace`). Also stores the entropy tables' build scratch +# inline instead of in `Vec`s, which the workspace decoder needs; with the +# feature off the other decoders build exactly as before. +caller-workspace = [] # --------------------------------------------------------------------------- # Internal features. NOT part of the public API: they widen the crate's @@ -229,6 +234,23 @@ lsm = [] # stable interface of this crate. rustc-dep-of-std = ["dep:compiler_builtins", "dep:core", "dep:alloc"] +[[test]] +name = "workspace_no_alloc" +path = "tests/workspace_no_alloc.rs" +required-features = ["caller-workspace", "std"] + +# Pure-Rust examples of the caller-workspace decoder; unlike the examples the +# `ffi-bench` crate builds, they link no C code. +[[example]] +name = "workspace_stream_verify" +path = "examples/workspace_stream_verify.rs" +required-features = ["caller-workspace", "std"] + +[[example]] +name = "workspace_decode_loop" +path = "examples/workspace_decode_loop.rs" +required-features = ["caller-workspace", "std"] + # All benches, integration tests, and FFI diagnostic examples are targets of # the non-published `ffi-bench` crate (they link the C `zstd` bindings); their # source files remain under `benches/`, `tests/`, and `examples/` here and are diff --git a/zstd/examples/workspace_decode_loop.rs b/zstd/examples/workspace_decode_loop.rs new file mode 100644 index 000000000..bd66e97cf --- /dev/null +++ b/zstd/examples/workspace_decode_loop.rs @@ -0,0 +1,135 @@ +//! Decode loop comparing the caller-workspace decoder with the existing ones +//! on the same build, for `perf stat` A/B runs. +//! +//! ```text +//! workspace_decode_loop [file|synthetic] [window_log] +//! ``` +//! +//! `mode` is `workspace` (the step decoder, 64 KiB input and output chunks), +//! `stream` (`StreamingDecoder` read through a 64 KiB buffer, the existing +//! streaming path) or `all` (`FrameDecoder::decode_all` into one buffer). The +//! input is compressed once with this crate's encoder at `level`; the default +//! file is the in-tree `decodecorpus_files/z000033`, and `synthetic` is random +//! bytes, which compress to Raw blocks. `window_log` caps the encoder's window +//! (default: the level's own) so small and large windows can be compared. + +use std::hint::black_box; +use std::io::Read; + +use structured_zstd::decoding::workspace::{ + DecodeLimits, StepState, WorkBudget, WorkspaceDecoder, workspace_bytes, +}; +use structured_zstd::decoding::{FrameDecoder, StreamingDecoder, read_frame_header_info}; +use structured_zstd::encoding::{ + CompressionLevel, CompressionParameters, compress_with_parameters, +}; + +const CHUNK: usize = 64 * 1024; + +fn source(arg: Option<&str>) -> Vec { + match arg { + Some(s) if s.eq_ignore_ascii_case("synthetic") => { + let mut x = 0x517c_c1b7_2722_0a95u64; + (0..(1 << 20)) + .map(|_| { + x = x + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + (x >> 56) as u8 + }) + .collect() + } + Some(path) => std::fs::read(path).expect("read input file"), + None => { + let path = concat!(env!("CARGO_MANIFEST_DIR"), "/decodecorpus_files/z000033"); + std::fs::read(path).expect("read decodecorpus_files/z000033") + } + } +} + +fn compress(data: &[u8], level: i32, window_log: Option) -> Vec { + let level = CompressionLevel::Level(level); + let mut parameters = CompressionParameters::builder(level); + if let Some(log) = window_log { + parameters = parameters.window_log(log); + } + compress_with_parameters(data, ¶meters.build().expect("parameters")) +} + +fn main() { + let args: Vec = std::env::args().collect(); + let level: i32 = args.get(1).and_then(|s| s.parse().ok()).unwrap_or(3); + let iters: u32 = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(1_000); + let mode = args.get(3).map_or("workspace", String::as_str); + let data = source(args.get(4).map(String::as_str)); + let window_log: Option = args.get(5).and_then(|s| s.parse().ok()); + + let compressed = compress(&data, level, window_log); + let window = read_frame_header_info(&compressed, false) + .expect("frame header") + .window_size + .max(1024); + eprintln!( + "{} bytes -> {} at level {level}, window {window}, mode {mode}", + data.len(), + compressed.len() + ); + + let mut total = 0u64; + let start = std::time::Instant::now(); + match mode { + "workspace" => { + let limits = DecodeLimits::new(window); + let mut workspace = vec![0u8; workspace_bytes(window).expect("window")]; + let mut decoder = WorkspaceDecoder::new(&mut workspace, limits).expect("workspace"); + let mut out = vec![0u8; CHUNK]; + for _ in 0..iters { + decoder.reset(); + let mut work = WorkBudget::new(u64::MAX); + let mut input = black_box(compressed.as_slice()); + loop { + let chunk = &input[..input.len().min(CHUNK)]; + let last = chunk.len() == input.len(); + let progress = decoder + .step(chunk, last, &mut out, &mut work) + .expect("workspace decode"); + total += black_box(&out[..progress.produced]).len() as u64; + input = &input[progress.consumed..]; + if progress.state == StepState::Done { + break; + } + } + } + } + "stream" => { + let mut out = vec![0u8; CHUNK]; + for _ in 0..iters { + let mut decoder = StreamingDecoder::new(black_box(compressed.as_slice())); + loop { + let n = decoder.read(&mut out).expect("stream decode"); + if n == 0 { + break; + } + total += black_box(&out[..n]).len() as u64; + } + } + } + "all" => { + let mut target = vec![0u8; data.len() + 64]; + let mut decoder = FrameDecoder::new(); + for _ in 0..iters { + total += decoder + .decode_all(black_box(&compressed), &mut target) + .expect("decode_all") as u64; + } + } + other => panic!("unknown mode {other}; use workspace, stream or all"), + } + let elapsed = start.elapsed(); + assert_eq!(total, data.len() as u64 * u64::from(iters)); + eprintln!( + "{iters} iters, {:.0} ns/iter, {:.1} MB/s", + elapsed.as_nanos() as f64 / f64::from(iters), + total as f64 / elapsed.as_secs_f64() / 1e6 + ); +} diff --git a/zstd/examples/workspace_stream_verify.rs b/zstd/examples/workspace_stream_verify.rs new file mode 100644 index 000000000..d1030095d --- /dev/null +++ b/zstd/examples/workspace_stream_verify.rs @@ -0,0 +1,139 @@ +//! Verify a compressed blob in a bounded amount of memory before using it. +//! +//! The workspace is admitted up front from the limits and allocated fallibly; +//! the blob is decoded in fixed chunks through a reusable output buffer, and +//! every produced chunk is counted and hashed instead of collected. The blob is +//! accepted only when the stream ends cleanly, the size is exact and the digest +//! matches; any error rejects it with nothing partial kept. +//! +//! A real consumer would use its own digest (SHA-256, say); FNV-1a keeps this +//! example free of dependencies. +//! +//! ```text +//! cargo run --example workspace_stream_verify --features caller-workspace +//! ``` + +use structured_zstd::decoding::workspace::{ + DecodeLimits, StepState, WorkBudget, WorkspaceDecodeError, WorkspaceDecoder, workspace_bytes, +}; +use structured_zstd::encoding::{CompressionLevel, compress_slice_to_vec}; + +/// What the manifest says about the blob. +struct Manifest { + size: u64, + digest: u64, +} + +#[derive(Debug)] +enum Rejected { + /// The workspace could not be admitted. + Memory, + /// The decoder refused the stream (limit, format, truncation, ...). + Decode(WorkspaceDecodeError), + /// The stream decoded, but not to the manifest's size or digest. + Mismatch, +} + +impl core::fmt::Display for Rejected { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + Rejected::Memory => write!(f, "workspace not admitted"), + Rejected::Decode(e) => write!(f, "decode refused: {e}"), + Rejected::Mismatch => write!(f, "size or digest differs from the manifest"), + } + } +} + +/// How a verification ended, for printing. +fn outcome(result: Result<(), Rejected>) -> String { + match result { + Ok(()) => "accepted".to_string(), + Err(e) => format!("rejected ({e})"), + } +} + +struct Fnv1a(u64); + +impl Fnv1a { + fn new() -> Self { + Self(0xcbf2_9ce4_8422_2325) + } + fn update(&mut self, bytes: &[u8]) { + for &b in bytes { + self.0 = (self.0 ^ u64::from(b)).wrapping_mul(0x0100_0000_01b3); + } + } +} + +fn verify(compressed: &[u8], manifest: &Manifest, max_window: u64) -> Result<(), Rejected> { + // 1. Size the workspace from the limits and admit it before decoding. + let limits = DecodeLimits::new(max_window).with_max_output(manifest.size); + let needed = workspace_bytes(limits.max_window_size).map_err(|_| Rejected::Memory)?; + let mut workspace = Vec::new(); + workspace + .try_reserve_exact(needed) + .map_err(|_| Rejected::Memory)?; + workspace.resize(needed, 0); + let mut decoder = + WorkspaceDecoder::new(&mut workspace, limits).map_err(|_| Rejected::Memory)?; + + // 2. Feed chunks, hash and count each produced slice. + let mut work = WorkBudget::new(16 * (compressed.len() as u64 + manifest.size) + (1 << 20)); + let mut out = [0u8; 16 * 1024]; + let mut hasher = Fnv1a::new(); + let mut produced = 0u64; + let mut input = compressed; + loop { + let chunk = &input[..input.len().min(64 * 1024)]; + let last = chunk.len() == input.len(); + let progress = decoder + .step(chunk, last, &mut out, &mut work) + .map_err(Rejected::Decode)?; + hasher.update(&out[..progress.produced]); + produced += progress.produced as u64; + input = &input[progress.consumed..]; + if progress.state == StepState::Done { + break; + } + } + + // 3. Exact size and digest, then commit. + if produced != manifest.size || hasher.0 != manifest.digest { + return Err(Rejected::Mismatch); + } + Ok(()) +} + +fn main() { + let content: Vec = (0..2_000_000u32) + .map(|i| (i % 251) as u8 ^ (i >> 11) as u8) + .collect(); + let compressed = compress_slice_to_vec(&content, CompressionLevel::Default); + let mut hasher = Fnv1a::new(); + hasher.update(&content); + let manifest = Manifest { + size: content.len() as u64, + digest: hasher.0, + }; + + verify(&compressed, &manifest, 8 << 20).expect("the blob verifies"); + println!( + "verified {} bytes from {} compressed in a {}-byte workspace", + manifest.size, + compressed.len(), + workspace_bytes(8 << 20).unwrap() + ); + + let mut corrupt = compressed.clone(); + let middle = corrupt.len() / 2; + corrupt[middle] ^= 0x40; + println!( + "corrupted blob: {}", + outcome(verify(&corrupt, &manifest, 8 << 20)) + ); + let truncated = &compressed[..compressed.len() - 10]; + println!( + "truncated blob: {}", + outcome(verify(truncated, &manifest, 8 << 20)) + ); +} diff --git a/zstd/fuzz/Cargo.toml b/zstd/fuzz/Cargo.toml index da1d9a13c..91e50ae50 100644 --- a/zstd/fuzz/Cargo.toml +++ b/zstd/fuzz/Cargo.toml @@ -18,6 +18,11 @@ features = ["fuzz-exports"] [dependencies.libfuzzer-sys] git = "https://github.com/rust-fuzz/libfuzzer-sys.git" +# Builds the library with its caller-workspace storage; kept off for the other +# targets so they keep fuzzing the default build. +[features] +workspace = ["structured-zstd/caller-workspace"] + # Prevent this from interfering with workspaces [workspace] members = ["."] @@ -45,3 +50,8 @@ path = "fuzz_targets/huff0.rs" [[bin]] name = "fse" path = "fuzz_targets/fse.rs" + +[[bin]] +name = "decode_workspace" +path = "fuzz_targets/decode_workspace.rs" +required-features = ["workspace"] diff --git a/zstd/fuzz/fuzz_targets/decode_workspace.rs b/zstd/fuzz/fuzz_targets/decode_workspace.rs new file mode 100644 index 000000000..984e13456 --- /dev/null +++ b/zstd/fuzz/fuzz_targets/decode_workspace.rs @@ -0,0 +1,65 @@ +#![no_main] +#[macro_use] +extern crate libfuzzer_sys; + +use structured_zstd::decoding::workspace::{ + DecodeLimits, StepState, WorkBudget, WorkspaceDecodeError, WorkspaceDecoder, workspace_bytes, +}; +use structured_zstd::decoding::{ContentChecksum, FrameDecoder}; + +const MAX_OUTPUT: usize = 1 << 20; + +// The first four bytes choose the window limit, the input and output chunk +// sizes and the work budget; the rest is the stream. Whatever they are, the +// decoder must not panic or stall, and where it decodes or rejects for a +// format reason the one-shot decoder must agree. +fuzz_target!(|data: &[u8]| { + let Some((&[w, i, o, b], stream)) = data.split_first_chunk::<4>() else { + return; + }; + let window = 1u64 << (10 + u32::from(w % 11)); + let in_chunk = usize::from(i).max(1); + let out_chunk = usize::from(o).max(1); + let budget = if b == 0 { u64::MAX } else { u64::from(b) << 12 }; + + let limits = DecodeLimits::new(window).with_max_output(MAX_OUTPUT as u64); + let mut workspace = vec![0u8; workspace_bytes(window).unwrap()]; + let mut decoder = WorkspaceDecoder::new(&mut workspace, limits).unwrap(); + let mut work = WorkBudget::new(budget); + let mut out = vec![0u8; out_chunk]; + let mut decoded = Vec::new(); + let mut pos = 0; + let result = loop { + let end = (pos + in_chunk).min(stream.len()); + match decoder.step(&stream[pos..end], end == stream.len(), &mut out, &mut work) { + Ok(progress) => { + assert!( + progress.consumed > 0 + || progress.produced > 0 + || progress.state == StepState::Done + ); + decoded.extend_from_slice(&out[..progress.produced]); + pos += progress.consumed; + if progress.state == StepState::Done { + break Ok(()); + } + } + Err(e) => break Err(e), + } + }; + if matches!( + result, + Err(WorkspaceDecodeError::Limit(_) | WorkspaceDecodeError::Unsupported(_)) + ) { + return; + } + + let mut reference = FrameDecoder::new(); + reference.set_content_checksum(ContentChecksum::Verify); + let mut target = vec![0u8; MAX_OUTPUT]; + let expected = reference + .decode_all(stream, &mut target) + .ok() + .map(|n| &target[..n]); + assert_eq!(result.is_ok().then_some(decoded.as_slice()), expected); +}); diff --git a/zstd/src/common/inline_vec.rs b/zstd/src/common/inline_vec.rs new file mode 100644 index 000000000..9a4dc9a8d --- /dev/null +++ b/zstd/src/common/inline_vec.rs @@ -0,0 +1,150 @@ +//! Fixed-capacity, allocation-free stand-in for the `Vec` scratch the entropy +//! table builders keep, selected by the `caller-workspace` feature so the +//! tables can live in caller-owned memory. + +use core::mem::MaybeUninit; +use core::ops::{Deref, DerefMut}; + +/// A `Vec`-shaped buffer of at most `N` elements stored inline. +/// +/// Every user is bounded by the format below `N` and checks its input against +/// that bound before appending (the Huffman weight count, the FSE alphabet), so +/// growing past `N` is a broken invariant: the bounds-checked write panics +/// rather than leaving the buffer. +pub(crate) struct InlineVec { + len: usize, + buf: [MaybeUninit; N], +} + +impl InlineVec { + pub(crate) const fn new() -> Self { + Self { + len: 0, + buf: [const { MaybeUninit::uninit() }; N], + } + } + + /// `Vec::with_capacity` signature; the capacity is fixed at `N`. + pub(crate) const fn with_capacity(_capacity: usize) -> Self { + Self::new() + } + + /// Initialise an empty buffer in place without building it on the stack. + /// + /// # Safety + /// `this` must be valid for writes and aligned for `Self`. + pub(crate) unsafe fn init_in_place(this: *mut Self) { + // SAFETY: the caller guarantees `this` is writable and aligned; the + // element array is `MaybeUninit` and needs no initialisation. + unsafe { core::ptr::addr_of_mut!((*this).len).write(0) }; + } + + #[inline(always)] + pub(crate) fn len(&self) -> usize { + self.len + } + + #[inline(always)] + pub(crate) fn is_empty(&self) -> bool { + self.len == 0 + } + + #[inline(always)] + pub(crate) fn clear(&mut self) { + self.len = 0; + } + + /// No-op: the storage is fixed. Callers size their requests below `N`. + #[inline(always)] + pub(crate) fn reserve(&mut self, additional: usize) { + debug_assert!(self.len + additional <= N); + } + + #[inline(always)] + pub(crate) fn push(&mut self, value: T) { + self.buf[self.len] = MaybeUninit::new(value); + self.len += 1; + } + + pub(crate) fn resize(&mut self, new_len: usize, value: T) { + if new_len > self.len { + for slot in &mut self.buf[self.len..new_len] { + *slot = MaybeUninit::new(value); + } + } + self.len = new_len; + } + + pub(crate) fn extend_from_slice(&mut self, values: &[T]) { + let end = self.len + values.len(); + for (slot, &value) in self.buf[self.len..end].iter_mut().zip(values) { + *slot = MaybeUninit::new(value); + } + self.len = end; + } + + /// The writable tail past the live prefix. + #[inline(always)] + pub(crate) fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit] { + &mut self.buf[self.len..] + } + + /// # Safety + /// `new_len <= N`, and every element below `new_len` is initialised. + #[inline(always)] + pub(crate) unsafe fn set_len(&mut self, new_len: usize) { + assert!(new_len <= N, "InlineVec::set_len past its capacity"); + self.len = new_len; + } + + #[inline(always)] + pub(crate) fn as_slice(&self) -> &[T] { + self + } +} + +impl Deref for InlineVec { + type Target = [T]; + + #[inline(always)] + fn deref(&self) -> &[T] { + debug_assert!(self.len <= N); + // SAFETY: every element below `len <= N` was written by `push`, + // `resize`, `extend_from_slice` or a `set_len` whose caller initialised + // it, each of which bounds `len` by `N`; `MaybeUninit` has `T`'s + // layout. + unsafe { core::slice::from_raw_parts(self.buf.as_ptr().cast::(), self.len) } + } +} + +impl DerefMut for InlineVec { + #[inline(always)] + fn deref_mut(&mut self) -> &mut [T] { + debug_assert!(self.len <= N); + // SAFETY: as for `deref`. + unsafe { core::slice::from_raw_parts_mut(self.buf.as_mut_ptr().cast::(), self.len) } + } +} + +impl Clone for InlineVec { + fn clone(&self) -> Self { + let mut out = Self::new(); + out.extend_from_slice(self); + out + } +} + +impl Default for InlineVec { + fn default() -> Self { + Self::new() + } +} + +impl core::fmt::Debug for InlineVec { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_list().entries(self.iter()).finish() + } +} + +#[cfg(test)] +mod tests; diff --git a/zstd/src/common/inline_vec/tests.rs b/zstd/src/common/inline_vec/tests.rs new file mode 100644 index 000000000..274ed500e --- /dev/null +++ b/zstd/src/common/inline_vec/tests.rs @@ -0,0 +1,50 @@ +use super::InlineVec; + +/// Appends within capacity read back exactly, like a `Vec`. +#[test] +fn push_resize_extend_within_capacity_read_back() { + let mut v: InlineVec = InlineVec::new(); + v.push(1); + v.extend_from_slice(&[2, 3]); + v.resize(5, 9); + assert_eq!(v.as_slice(), &[1, 2, 3, 9, 9]); + v.resize(2, 0); + assert_eq!(v.as_slice(), &[1, 2]); + v.resize(8, 4); + assert_eq!(v.len(), 8); + v.clear(); + assert!(v.is_empty()); +} + +/// A clone holds the same elements. +#[test] +fn clone_copies_the_elements() { + let mut v: InlineVec = InlineVec::new(); + v.extend_from_slice(&[5, 6, 7]); + assert_eq!(&*v.clone(), &[5, 6, 7]); +} + +/// Appending past the capacity is a broken invariant and must stop at the +/// buffer's edge rather than write beyond it. +#[test] +#[should_panic] +fn push_past_capacity_panics() { + let mut v: InlineVec = InlineVec::new(); + v.extend_from_slice(&[1, 2]); + v.push(3); +} + +#[test] +#[should_panic] +fn resize_past_capacity_panics() { + let mut v: InlineVec = InlineVec::new(); + v.resize(4, 0); +} + +#[test] +#[should_panic(expected = "past its capacity")] +fn set_len_past_capacity_panics() { + let mut v: InlineVec = InlineVec::new(); + // SAFETY: the call is expected to panic before exposing anything. + unsafe { v.set_len(3) }; +} diff --git a/zstd/src/common/mod.rs b/zstd/src/common/mod.rs index cc38b0924..07d0a725c 100644 --- a/zstd/src/common/mod.rs +++ b/zstd/src/common/mod.rs @@ -43,3 +43,6 @@ pub const MAX_DECODER_WINDOW_SIZE: u64 = 1 << 31; /// `1 << 30` on 32-bit ones, where a larger window could not be addressed. #[cfg(not(target_pointer_width = "64"))] pub const MAX_DECODER_WINDOW_SIZE: u64 = 1 << 30; + +#[cfg(feature = "caller-workspace")] +pub(crate) mod inline_vec; diff --git a/zstd/src/decoding/block_decoder.rs b/zstd/src/decoding/block_decoder.rs index fa6807b75..049c19d97 100644 --- a/zstd/src/decoding/block_decoder.rs +++ b/zstd/src/decoding/block_decoder.rs @@ -63,6 +63,23 @@ pub(crate) fn block_fits_the_maximum( Ok(()) } +/// A compressed block's content is held to the frame's block maximum too +/// (RFC 8878 3.1.1.2.4), as upstream's `ZSTD_decompressBlock_internal` refuses +/// `srcSize > blockSizeMax`. Checked before the block body runs, whatever the +/// content would decode to. +#[inline] +fn compressed_content_fits( + header: &BlockHeader, + window_size: usize, +) -> Result<(), DecompressBlockError> { + let size = header.content_size as usize; + let maximum = block_maximum(window_size); + if size > maximum { + return Err(DecompressBlockError::ContentPastBlockMaximum { size, maximum }); + } + Ok(()) +} + /// A block with no sequences: its literals ARE its output. Out of line, so the /// per-block body keeps the shape the sequence executor is laid out around. /// @@ -233,6 +250,7 @@ impl BlockDecoder { panic!("Reserved-type block should be rejected during header parsing"); } BlockType::Compressed => { + compressed_content_fits(header, workspace.split().buffer.window_size)?; let n = header.content_size as usize; if source.len() < n { return Err(DecodeBlockContentError::ReadError { @@ -311,6 +329,7 @@ impl BlockDecoder { } BlockType::Compressed => { + compressed_content_fits(header, workspace.split().buffer.window_size)?; self.decompress_block(header, workspace, dict, source)?; self.internal_state = DecoderState::ReadyToDecodeNextHeader; @@ -332,6 +351,7 @@ impl BlockDecoder { // decode path (`decode_all`) skips this copy by passing // a borrowed slice of the input straight to // `decompress_block_inplace_with_parts`. + use crate::decoding::scratch::ByteScratch; let parts = workspace.split(); parts .block_content_buffer @@ -404,14 +424,18 @@ impl BlockDecoder { /// lifetime-widening trick is gone — disjoint-field borrows /// type-check naturally. #[allow(clippy::too_many_arguments)] - fn decompress_block_inplace_with_parts<'d, B: super::buffer_backend::BufferBackend>( + fn decompress_block_inplace_with_parts< + 'd, + B: super::buffer_backend::BufferBackend, + L: crate::decoding::scratch::ByteScratch, + >( &mut self, header: &BlockHeader, huf: &mut crate::decoding::scratch::HuffmanScratch, fse: &'d mut crate::decoding::scratch::FSEScratch, buffer: &mut crate::decoding::decode_buffer::DecodeBuffer, offset_hist: &mut [u32; 3], - literals_buffer: &mut alloc::vec::Vec, + literals_buffer: &mut L, raw: &[u8], dict: Option<&'d crate::decoding::dictionary::Dictionary>, ) -> Result<(), DecompressBlockError> { diff --git a/zstd/src/decoding/errors.rs b/zstd/src/decoding/errors.rs index 54208fb33..ce7da7ef3 100644 --- a/zstd/src/decoding/errors.rs +++ b/zstd/src/decoding/errors.rs @@ -282,6 +282,13 @@ pub enum DecompressBlockError { size: usize, maximum: usize, }, + /// A compressed block's content is `size` bytes, past the block maximum + /// (RFC 8878 3.1.1.2.4, `Block_Maximum_Size`) of `maximum`, the smaller of + /// the frame's window and 128 KiB, whatever it would decode to. + ContentPastBlockMaximum { + size: usize, + maximum: usize, + }, /// A block with no sequences whose literals do not fit a fixed-capacity /// backend: `requested` bytes at `tail` against `capacity`. The block is /// within the block maximum, so this says the caller's slice is short, and @@ -334,6 +341,10 @@ impl core::fmt::Display for DecompressBlockError { f, "Block expands to {size} bytes, past this frame's maximum of {maximum}" ), + DecompressBlockError::ContentPastBlockMaximum { size, maximum } => write!( + f, + "Compressed block of {size} bytes, past this frame's maximum of {maximum}" + ), DecompressBlockError::LiteralsOutputOverflow { tail, requested, diff --git a/zstd/src/decoding/frame_decoder/tests.rs b/zstd/src/decoding/frame_decoder/tests.rs index 4fcf8d3f0..ced345ea7 100644 --- a/zstd/src/decoding/frame_decoder/tests.rs +++ b/zstd/src/decoding/frame_decoder/tests.rs @@ -4025,3 +4025,57 @@ fn a_frame_larger_than_the_callers_buffer_is_refused_up_front() { Err(FrameDecoderError::TargetTooSmall) )); } + +/// A compressed block whose content is larger than the frame's block maximum +/// is malformed (RFC 8878 3.1.1.2.4), as upstream's `ZSTD_decompressBlock_internal` +/// refuses `srcSize > blockSizeMax`, even when it would produce no output. Here +/// a single-segment frame declaring 0 bytes (window and block maximum 0) +/// carries a 2-byte compressed block holding an empty literals section and no +/// sequences. +#[test] +fn a_compressed_block_past_the_block_maximum_is_refused() { + let frame = [ + 0x28, 0xB5, 0x2F, 0xFD, // magic + 0x20, 0x00, // single segment, Frame_Content_Size 0 + 0x15, 0x00, 0x00, // last, Compressed, Block_Size 2 + 0x00, 0x00, // raw literals of 0 bytes, 0 sequences + ]; + // Pin the cause, not only the refusal: any other error here would mean the + // block was refused for a reason this check does not own. + let expect_past_maximum = |err: &FrameDecoderError| { + let source = match err { + FrameDecoderError::FailedToReadBlockBody(source) => source, + #[cfg(feature = "lsm")] + FrameDecoderError::FailedToReadBlockBodyAt { source, .. } => source, + other => panic!("expected a block-body refusal, got {other:?}"), + }; + assert!( + matches!( + source, + DecodeBlockContentError::DecompressBlockError( + DecompressBlockError::ContentPastBlockMaximum { + size: 2, + maximum: 0 + } + ), + ), + "expected the block-maximum refusal, got {err:?}", + ); + }; + let mut out = [0u8; 16]; + let err = FrameDecoder::new() + .decode_all(&frame, &mut out) + .expect_err("decode_all must refuse the block"); + expect_past_maximum(&err); + let mut decoded = Vec::new(); + let mut streaming = crate::decoding::StreamingDecoder::new(&frame[..]); + let err = crate::io::Read::read_to_end(&mut streaming, &mut decoded) + .expect_err("the streaming decoder must refuse the block"); + match err + .get_ref() + .and_then(|e| e.downcast_ref::()) + { + Some(err) => expect_past_maximum(err), + None => panic!("expected a FrameDecoderError, got {err:?}"), + } +} diff --git a/zstd/src/decoding/literals_section_decoder.rs b/zstd/src/decoding/literals_section_decoder.rs index 0e692e46c..1c6fdd0c1 100644 --- a/zstd/src/decoding/literals_section_decoder.rs +++ b/zstd/src/decoding/literals_section_decoder.rs @@ -3,7 +3,7 @@ use super::super::blocks::literals_section::{LiteralsSection, LiteralsSectionType}; use super::buffer_backend::WILDCOPY_OVERLENGTH; -use super::scratch::HuffmanScratch; +use super::scratch::{ByteScratch, HuffmanScratch}; use crate::bit_io::BitReaderReversed; #[cfg(all(target_arch = "x86_64", feature = "kernel-avx2"))] use crate::cpu_kernel::Avx2Kernel; @@ -28,6 +28,7 @@ use crate::cpu_kernel::{CpuKernel, CpuKernelTag, ScalarKernel}; use crate::decoding::dictionary::Dictionary; use crate::decoding::errors::DecompressLiteralsError; use crate::huff0::HuffmanDecoder; +#[cfg(test)] use alloc::vec::Vec; /// Decode and decompress the provided literals section into `target`, returning the number of bytes read. @@ -107,14 +108,14 @@ pub struct LiteralsView<'a> { // The block's literals inputs, each a borrow or a scalar the caller already // holds; a struct here would marshal a record per block to satisfy the lint. #[allow(clippy::too_many_arguments)] -pub fn decode_literals_zerocopy<'a>( +pub fn decode_literals_zerocopy<'a, L: ByteScratch>( section: &LiteralsSection, scratch: &mut HuffmanScratch, dict: Option<&Dictionary>, source: &'a [u8], payload_len: usize, needs_slack: bool, - target: &'a mut Vec, + target: &'a mut L, kernel: CpuKernelTag, ) -> Result, DecompressLiteralsError> { // Snapshot `target.len()` before any decode work — the returned @@ -188,8 +189,8 @@ fn bytes_used_raw(section: &LiteralsSection) -> u32 { /// view over both. The slack is written, not merely reserved: the copiers read /// it, and reading uninitialised memory is undefined behaviour here even though /// the bytes are discarded. -fn pad_with_slack( - target: &mut Vec, +fn pad_with_slack( + target: &mut L, base: usize, n: usize, bytes_used: u32, @@ -199,7 +200,7 @@ fn pad_with_slack( target.resize(base + n + WILDCOPY_OVERLENGTH, 0); } Ok(LiteralsView { - data: &target[base..], + data: &target.as_slice()[base..], len: n, bytes_used, }) @@ -209,12 +210,12 @@ fn pad_with_slack( /// This function is used when the literals section is `Compressed` or `Treeless` /// /// Returns the number of bytes read. -fn decompress_literals( +fn decompress_literals( section: &LiteralsSection, scratch: &mut HuffmanScratch, dict: Option<&Dictionary>, source: &[u8], - target: &mut Vec, + target: &mut L, kernel: CpuKernelTag, ) -> Result { // `kernel` was resolved once, before the first block, by whoever owns the @@ -257,7 +258,7 @@ fn decompress_literals( // past. #[cfg(all(target_arch = "aarch64", feature = "kernel-neon"))] CpuKernelTag::Neon => { - decompress_literals_impl::(section, scratch, dict, source, target) + decompress_literals_impl::(section, scratch, dict, source, target) } #[cfg(all( target_arch = "aarch64", @@ -265,22 +266,22 @@ fn decompress_literals( any(feature = "std", target_feature = "sve"), ))] CpuKernelTag::Sve => { - decompress_literals_impl::(section, scratch, dict, source, target) + decompress_literals_impl::(section, scratch, dict, source, target) } - _ => decompress_literals_impl::(section, scratch, dict, source, target), + _ => decompress_literals_impl::(section, scratch, dict, source, target), } } #[cfg(all(target_arch = "x86_64", feature = "kernel-avx2"))] #[target_feature(enable = "bmi2,avx2")] -unsafe fn decompress_literals_avx2( +unsafe fn decompress_literals_avx2( section: &LiteralsSection, scratch: &mut HuffmanScratch, dict: Option<&Dictionary>, source: &[u8], - target: &mut Vec, + target: &mut L, ) -> Result { - decompress_literals_impl::(section, scratch, dict, source, target) + decompress_literals_impl::(section, scratch, dict, source, target) } #[cfg(all( @@ -288,34 +289,34 @@ unsafe fn decompress_literals_avx2( feature = "kernel-bmi2" ))] #[target_feature(enable = "bmi2")] -unsafe fn decompress_literals_bmi2( +unsafe fn decompress_literals_bmi2( section: &LiteralsSection, scratch: &mut HuffmanScratch, dict: Option<&Dictionary>, source: &[u8], - target: &mut Vec, + target: &mut L, ) -> Result { - decompress_literals_impl::(section, scratch, dict, source, target) + decompress_literals_impl::(section, scratch, dict, source, target) } #[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))] #[target_feature(enable = "avx512vbmi2,avx512f,avx512vl,avx512bw,bmi2,avx2")] -unsafe fn decompress_literals_vbmi2( +unsafe fn decompress_literals_vbmi2( section: &LiteralsSection, scratch: &mut HuffmanScratch, dict: Option<&Dictionary>, source: &[u8], - target: &mut Vec, + target: &mut L, ) -> Result { - decompress_literals_impl::(section, scratch, dict, source, target) + decompress_literals_impl::(section, scratch, dict, source, target) } -fn decompress_literals_impl( +fn decompress_literals_impl( section: &LiteralsSection, scratch: &mut HuffmanScratch, dict: Option<&Dictionary>, source: &[u8], - target: &mut Vec, + target: &mut L, ) -> Result { use DecompressLiteralsError as err; diff --git a/zstd/src/decoding/mod.rs b/zstd/src/decoding/mod.rs index e1c652cca..67f27cd31 100644 --- a/zstd/src/decoding/mod.rs +++ b/zstd/src/decoding/mod.rs @@ -1,13 +1,16 @@ //! RFC 8878 Zstandard decoder. //! -//! Three entry points are exposed, each with progressively lower-level -//! control: +//! The entry points: //! //! * [`StreamingDecoder`] — implements [`crate::io::Read`] over a compressed //! byte stream, transparently parsing the frame header and concatenated //! frames. The typical choice for application code. //! * [`FrameDecoder`] — single-frame interface; use when the caller manages //! the input buffer manually (zero-copy slices, network framing, etc). +//! * `workspace::WorkspaceDecoder` (feature `caller-workspace`): streaming +//! decoder over a byte slice the caller sizes and owns, with no allocation +//! and limits on window, output and work. For consumers that account their +//! memory themselves, such as browser bundles verifying a download. //! * [`DictionaryHandle`] — pre-parsed dictionary handle. Parse the //! dictionary bytes once with [`DictionaryHandle::decode_dict`] and reuse //! the handle across every subsequent decode; saves the per-frame @@ -32,6 +35,9 @@ pub mod errors; mod frame_decoder; mod streaming_decoder; +#[cfg(feature = "caller-workspace")] +#[cfg_attr(docsrs, doc(cfg(feature = "caller-workspace")))] +pub mod workspace; pub use dictionary::{Dictionary, DictionaryHandle, MAGIC_NUM as DICTIONARY_MAGIC}; pub use frame_decoder::{BlockDecodingStrategy, ContentChecksum, FrameDecoder}; diff --git a/zstd/src/decoding/ringbuffer.rs b/zstd/src/decoding/ringbuffer.rs index 0bb3e7dc5..9df73302e 100644 --- a/zstd/src/decoding/ringbuffer.rs +++ b/zstd/src/decoding/ringbuffer.rs @@ -47,6 +47,10 @@ pub struct RingBuffer { /// not doubled for one block. A need beyond it still grows as before. /// `usize::MAX` (the default) leaves the doubling alone. growth_limit: usize, + /// `buf` is caller-owned memory ([`Self::from_borrowed`]): never grown, + /// never freed. + #[cfg(feature = "caller-workspace")] + borrowed: bool, } // SAFETY: RingBuffer does not hold any thread specific values -> it can be sent to another thread -> RingBuffer is Send @@ -66,6 +70,31 @@ impl RingBuffer { tail: 0, max_capacity: usize::MAX, growth_limit: usize::MAX, + #[cfg(feature = "caller-workspace")] + borrowed: false, + } + } + + /// A ring over caller-owned memory of `cap + WILDCOPY_OVERLENGTH` bytes, + /// `cap` of them indexed by the ring and the rest the copiers' overshoot + /// slack. It never grows: the caller sizes `cap` for the most the decode + /// can hold at once. + /// + /// # Safety + /// `buf` is valid for reads and writes of `cap + WILDCOPY_OVERLENGTH` + /// initialised bytes for as long as the ring is used, and nothing else + /// accesses that memory meanwhile. `cap >= 1`. + #[cfg(feature = "caller-workspace")] + pub(crate) unsafe fn from_borrowed(buf: NonNull, cap: usize) -> Self { + debug_assert!(cap >= 1); + RingBuffer { + buf, + cap, + head: 0, + tail: 0, + max_capacity: usize::MAX, + growth_limit: usize::MAX, + borrowed: true, } } @@ -252,6 +281,15 @@ impl RingBuffer { #[inline(never)] #[cold] fn reserve_amortized(&mut self, amount: usize) { + // A borrowed ring is sized for the window plus two block maxima (see + // `decoding::workspace`), which covers every write a block can issue + // before the per-block ceiling refuses it, so growth is unreachable. + // Growing would free memory this ring does not own. + #[cfg(feature = "caller-workspace")] + assert!( + !self.borrowed, + "a caller-owned ring cannot grow; its capacity bounds every block" + ); // SAFETY: if we were successfully able to construct this layout when we allocated then it's also valid do so now let current_layout = unsafe { Layout::array::(self.cap + WILDCOPY_OVERLENGTH).unwrap_unchecked() }; @@ -1304,6 +1342,10 @@ impl Drop for RingBuffer { if self.cap == 0 { return; } + #[cfg(feature = "caller-workspace")] + if self.borrowed { + return; + } // SAFETY: if we were successfully able to construct this layout when we allocated then it's also valid do so now. // Layout matches `reserve_amortized` which inflates by WILDCOPY_OVERLENGTH. diff --git a/zstd/src/decoding/scratch.rs b/zstd/src/decoding/scratch.rs index d50dd26d9..f721bcaae 100644 --- a/zstd/src/decoding/scratch.rs +++ b/zstd/src/decoding/scratch.rs @@ -22,7 +22,13 @@ use crate::blocks::sequence_section::{ /// #132. pub struct DecoderScratch { /// The decoder used for Huffman blocks. + #[cfg(not(feature = "caller-workspace"))] pub huf: HuffmanScratch, + /// The decoder used for Huffman blocks. Boxed, like `fse`, because + /// `caller-workspace` stores the decode table inline (~6 KiB), which would + /// otherwise grow every owned decoder's stack frame. + #[cfg(feature = "caller-workspace")] + pub huf: Box, /// The decoder used for FSE blocks. /// /// Boxed so `DecoderScratch` -- and the `FrameDecoder` that holds one @@ -52,13 +58,83 @@ pub struct DecoderScratch { /// The lifetime `'a` is the shorter-of (a) the underlying owner's /// lifetime and (b) the active borrow. The backend type `B` flows /// through to [`super::decode_buffer::DecodeBuffer`]. -pub struct WorkspaceRef<'a, B: BufferBackend> { +pub struct WorkspaceRef<'a, B: BufferBackend, L: ByteScratch = Vec> { pub huf: &'a mut HuffmanScratch, pub fse: &'a mut FSEScratch, pub buffer: &'a mut DecodeBuffer, pub offset_hist: &'a mut [u32; 3], - pub literals_buffer: &'a mut Vec, - pub block_content_buffer: &'a mut Vec, + pub literals_buffer: &'a mut L, + pub block_content_buffer: &'a mut L, +} + +/// Per-block byte scratch the literals and block-content stages write into: +/// the `Vec` the owned decoders keep, or (under `caller-workspace`) a +/// fixed region of caller memory. The methods are the `Vec` subset those +/// stages use, so the `Vec` monomorph compiles to the same code as before. +pub(crate) trait ByteScratch { + fn len(&self) -> usize; + fn clear(&mut self); + fn truncate(&mut self, len: usize); + /// Room for `additional` more bytes past `len`. + fn reserve(&mut self, additional: usize); + fn push(&mut self, byte: u8); + fn resize(&mut self, new_len: usize, value: u8); + fn extend_from_slice(&mut self, bytes: &[u8]); + fn as_mut_ptr(&mut self) -> *mut u8; + /// # Safety + /// Every byte below `new_len` is initialised and lies within the reserved + /// capacity. + unsafe fn set_len(&mut self, new_len: usize); + fn as_slice(&self) -> &[u8]; + fn as_mut_slice(&mut self) -> &mut [u8]; +} + +impl ByteScratch for Vec { + #[inline(always)] + fn len(&self) -> usize { + Vec::len(self) + } + #[inline(always)] + fn clear(&mut self) { + Vec::clear(self) + } + #[inline(always)] + fn truncate(&mut self, len: usize) { + Vec::truncate(self, len) + } + #[inline(always)] + fn reserve(&mut self, additional: usize) { + Vec::reserve(self, additional) + } + #[inline(always)] + fn push(&mut self, byte: u8) { + Vec::push(self, byte) + } + #[inline(always)] + fn resize(&mut self, new_len: usize, value: u8) { + Vec::resize(self, new_len, value) + } + #[inline(always)] + fn extend_from_slice(&mut self, bytes: &[u8]) { + Vec::extend_from_slice(self, bytes) + } + #[inline(always)] + fn as_mut_ptr(&mut self) -> *mut u8 { + Vec::as_mut_ptr(self) + } + #[inline(always)] + unsafe fn set_len(&mut self, new_len: usize) { + // SAFETY: forwarded to the caller of this method. + unsafe { Vec::set_len(self, new_len) } + } + #[inline(always)] + fn as_slice(&self) -> &[u8] { + Vec::as_slice(self) + } + #[inline(always)] + fn as_mut_slice(&mut self) -> &mut [u8] { + Vec::as_mut_slice(self) + } } /// Polymorphic accessor for the decoder's per-call scratch state. @@ -76,11 +152,13 @@ pub struct WorkspaceRef<'a, B: BufferBackend> { /// `&mut huf` and `&mut literals_buffer` simultaneously. pub(crate) trait Workspace { type Backend: BufferBackend; - fn split(&mut self) -> WorkspaceRef<'_, Self::Backend>; + type Bytes: ByteScratch; + fn split(&mut self) -> WorkspaceRef<'_, Self::Backend, Self::Bytes>; } impl Workspace for DecoderScratch { type Backend = B; + type Bytes = Vec; fn split(&mut self) -> WorkspaceRef<'_, B> { WorkspaceRef { huf: &mut self.huf, @@ -124,6 +202,7 @@ pub struct DirectScratch<'o, 'p> { impl<'o, 'p> Workspace for DirectScratch<'o, 'p> { type Backend = super::user_slice_buf::UserSliceBackend<'o>; + type Bytes = Vec; fn split(&mut self) -> WorkspaceRef<'_, Self::Backend> { // Reborrow the `&'p mut` fields to `&'_ mut` so the returned // WorkspaceRef's lifetime is tied to the `&mut self` of this @@ -144,10 +223,13 @@ impl<'o, 'p> Workspace for DirectScratch<'o, 'p> { impl DecoderScratch { pub fn new(window_size: usize) -> DecoderScratch { DecoderScratch { + #[cfg(not(feature = "caller-workspace"))] huf: HuffmanScratch { table: HuffmanTable::new(), table_source: TableSource::Local, }, + #[cfg(feature = "caller-workspace")] + huf: Box::new(HuffmanScratch::new()), fse: Box::new(FSEScratch { offsets: AlignedFSETable::new(MAX_OFFSET_CODE), literal_lengths: AlignedFSETable::new(MAX_LITERAL_LENGTH_CODE), @@ -175,10 +257,17 @@ impl DecoderScratch { /// than inline in this struct: charge `size_of` for the boxed value on top /// of the build-scratch vectors `FSEScratch::heap_bytes` sums. pub fn workspace_bytes(&self) -> usize { + // Under `caller-workspace` the Huffman scratch is boxed too. + let boxed_huf = if cfg!(feature = "caller-workspace") { + core::mem::size_of::() + } else { + 0 + }; self.buffer.capacity() + self.literals_buffer.capacity() + self.block_content_buffer.capacity() + self.huf.heap_bytes() + + boxed_huf + core::mem::size_of::() + self.fse.heap_bytes() } @@ -234,54 +323,7 @@ impl DecoderScratch { } self.buffer.reset(window_size); - - // Mirror upstream zstd's `ZSTD_decompressBegin`: it never clears the - // entropy tables per frame (marks them invalid by flag, rebuilds lazily). - // Resetting an already-empty table is a no-op on the observable state - // (`accuracy_log`/`max_num_bits` stay 0, the buffers stay empty), so a - // RAW / RLE / zero-sequence frame — which never builds a table — skips - // the clears (and the Huffman `[0; 13]` rank-count memset) entirely. - // Byte-identical: the table is cleared exactly when it was built. - if self.fse.literal_lengths.is_populated() { - self.fse.literal_lengths.reset(); - } - if self.fse.match_lengths.is_populated() { - self.fse.match_lengths.reset(); - } - if self.fse.offsets.is_populated() { - self.fse.offsets.reset(); - } - // Reset the cached pipeline-gate signal alongside the FSE - // table reset — otherwise scratch reuse across frames could - // engage the long pipeline on a new frame's Repeat-mode - // header based on the previous frame's offset distribution - // (or vice versa: skip the pipeline when the new frame - // actually has long offsets). - self.fse.offsets_long_share = 0; - // Revert any dictionary copy-on-write attachment: a scratch - // reused from a dict-attached frame must not read the previous - // dictionary's tables on the next (possibly dict-less) frame. - self.fse.detach_dict(); - // Pair the one-shot cold-dict flag with `reset`: a scratch - // reused from a dictionary-attached frame whose blocks never - // entered sequence decoding (raw-/RLE-only blocks, zero-seq - // compressed blocks) would otherwise carry the flag into the - // next frame and mis-apply the cold-dict gate there. Cleared - // alongside `offsets_long_share` so the no-dict path keeps - // the documented "no behaviour change" property. - self.fse.ddict_is_cold = false; - - // Same lazy-entropy skip as the FSE tables above: a Huffman table that - // was never built this round (`max_num_bits == 0`) is already in its - // reset state, so re-running `reset` (clearing 6 buffers + the - // `[0; 13]` rank-count memset) is wasted no-op work on RAW / RLE / raw- - // literal frames. Byte-identical. - if self.huf.table.max_num_bits != 0 { - self.huf.table.reset(); - } - // Mirror the FSE detach: a reused workspace must not read a - // previous frame's dictionary Huffman table. - self.huf.detach_dict(); + reset_entropy(&mut self.huf, &mut self.fse); } pub fn init_from_dict(&mut self, dict: &DictionaryHandle) { @@ -306,6 +348,58 @@ impl DecoderScratch { } } +/// Per-frame entropy reset shared by every decoder that owns a set of tables. +#[inline(always)] +pub(crate) fn reset_entropy(huf: &mut HuffmanScratch, fse: &mut FSEScratch) { + // Mirror upstream zstd's `ZSTD_decompressBegin`: it never clears the + // entropy tables per frame (marks them invalid by flag, rebuilds lazily). + // Resetting an already-empty table is a no-op on the observable state + // (`accuracy_log`/`max_num_bits` stay 0, the buffers stay empty), so a + // RAW / RLE / zero-sequence frame, which never builds a table, skips + // the clears (and the Huffman `[0; 13]` rank-count memset) entirely. + // Byte-identical: the table is cleared exactly when it was built. + if fse.literal_lengths.is_populated() { + fse.literal_lengths.reset(); + } + if fse.match_lengths.is_populated() { + fse.match_lengths.reset(); + } + if fse.offsets.is_populated() { + fse.offsets.reset(); + } + // Reset the cached pipeline-gate signal alongside the FSE + // table reset; otherwise scratch reuse across frames could + // engage the long pipeline on a new frame's Repeat-mode + // header based on the previous frame's offset distribution + // (or vice versa: skip the pipeline when the new frame + // actually has long offsets). + fse.offsets_long_share = 0; + // Revert any dictionary copy-on-write attachment: a scratch + // reused from a dict-attached frame must not read the previous + // dictionary's tables on the next (possibly dict-less) frame. + fse.detach_dict(); + // Pair the one-shot cold-dict flag with `reset`: a scratch + // reused from a dictionary-attached frame whose blocks never + // entered sequence decoding (raw-/RLE-only blocks, zero-seq + // compressed blocks) would otherwise carry the flag into the + // next frame and mis-apply the cold-dict gate there. Cleared + // alongside `offsets_long_share` so the no-dict path keeps + // the documented "no behaviour change" property. + fse.ddict_is_cold = false; + + // Same lazy-entropy skip as the FSE tables above: a Huffman table that + // was never built this round (`max_num_bits == 0`) is already in its + // reset state, so re-running `reset` (clearing 6 buffers + the + // `[0; 13]` rank-count memset) is wasted no-op work on RAW / RLE / raw- + // literal frames. Byte-identical. + if huf.table.max_num_bits != 0 { + huf.table.reset(); + } + // Mirror the FSE detach: a reused workspace must not read a + // previous frame's dictionary Huffman table. + huf.detach_dict(); +} + #[derive(Clone)] pub struct HuffmanScratch { pub table: HuffmanTable, @@ -330,6 +424,20 @@ impl HuffmanScratch { } } + /// Initialise an empty scratch in place, in caller-owned memory. + /// + /// # Safety + /// `this` must be valid for writes and aligned for `Self`. + #[cfg(feature = "caller-workspace")] + pub(crate) unsafe fn init_in_place(this: *mut Self) { + use core::ptr::addr_of_mut; + // SAFETY: forwarded from the caller; every field is written. + unsafe { + HuffmanTable::init_in_place(addr_of_mut!((*this).table)); + addr_of_mut!((*this).table_source).write(TableSource::Local); + } + } + /// Heap bytes owned by this scratch: the locally-built Huffman table. /// A `Dict`-sourced table is read through a shared, ref-counted handle /// (not owned here), so it is excluded, mirroring upstream not charging @@ -500,6 +608,34 @@ impl FSEScratch { } } + /// Initialise an empty scratch in place, in caller-owned memory, without + /// materialising the ~12 KiB of tables on the stack. + /// + /// # Safety + /// `this` must be valid for writes and aligned for `Self`. + #[cfg(feature = "caller-workspace")] + pub(crate) unsafe fn init_in_place(this: *mut Self) { + use crate::fse::SeqFSETable; + use core::ptr::addr_of_mut; + // SAFETY: forwarded from the caller; every field is written. + unsafe { + SeqFSETable::init_in_place(addr_of_mut!((*this).offsets.0), MAX_OFFSET_CODE); + SeqFSETable::init_in_place( + addr_of_mut!((*this).literal_lengths.0), + MAX_LITERAL_LENGTH_CODE, + ); + SeqFSETable::init_in_place( + addr_of_mut!((*this).match_lengths.0), + MAX_MATCH_LENGTH_CODE, + ); + addr_of_mut!((*this).offsets_long_share).write(0); + addr_of_mut!((*this).ddict_is_cold).write(false); + addr_of_mut!((*this).ll_source).write(SeqTableSource::Local); + addr_of_mut!((*this).of_source).write(SeqTableSource::Local); + addr_of_mut!((*this).ml_source).write(SeqTableSource::Local); + } + } + /// Snapshot the live (COW-resolved) sequence tables into `self` as /// owned `Local` copies. Used by the LSM resume snapshot/restore /// path (`FrameDecoder::export_entropy` / `restore_entropy`): the diff --git a/zstd/src/decoding/workspace.rs b/zstd/src/decoding/workspace.rs new file mode 100644 index 000000000..056c6b417 --- /dev/null +++ b/zstd/src/decoding/workspace.rs @@ -0,0 +1,1301 @@ +//! Streaming decoder that runs entirely in caller-owned memory. +//! +//! [`WorkspaceDecoder`] decodes a stream of zstd frames chunk by chunk without +//! touching the heap: the window, the entropy tables, the literals and the +//! staging for blocks split across input chunks all live in one byte slice the +//! caller supplies. [`workspace_bytes`] says how large that slice must be for a +//! given window limit, so the caller can admit the memory in its own accounting +//! before decoding starts, and keep the charge for as long as the slice lives. +//! +//! Every step is bounded three ways, all chosen by the caller: the window a +//! frame may declare and the total output ([`DecodeLimits`]), and the work a +//! step may do ([`WorkBudget`]). A refusal is reported apart from malformed +//! input, an unsupported feature, and input that ends mid-frame, and it ends +//! the stream: every later step fails with [`WorkspaceDecodeError::Failed`] +//! until [`WorkspaceDecoder::reset`]. +//! +//! Frames that name a dictionary are refused with +//! [`Unsupported::Dictionary`]; this decoder keeps no dictionary registry. +//! +//! The decoding itself is the same code the other decoders run, so output is +//! identical to theirs. +//! +//! # Examples +//! +//! ``` +//! use structured_zstd::decoding::workspace::{ +//! workspace_bytes, DecodeLimits, StepState, WorkBudget, WorkspaceDecoder, +//! }; +//! use structured_zstd::encoding::{compress_slice_to_vec, CompressionLevel}; +//! +//! let data = b"the same words, the same words, the same words".repeat(100); +//! let compressed = compress_slice_to_vec(&data, CompressionLevel::Default); +//! +//! let limits = DecodeLimits::new(1 << 20).with_max_output(1 << 20); +//! let mut workspace = vec![0u8; workspace_bytes(limits.max_window_size).unwrap()]; +//! let mut decoder = WorkspaceDecoder::new(&mut workspace, limits).unwrap(); +//! let mut work = WorkBudget::new(1 << 24); +//! +//! let mut decoded = Vec::new(); +//! let mut out = [0u8; 512]; +//! let mut input = &compressed[..]; +//! loop { +//! // Feed at most 100 bytes at a time; say when nothing more will come. +//! let chunk = &input[..input.len().min(100)]; +//! let last = chunk.len() == input.len(); +//! let progress = decoder.step(chunk, last, &mut out, &mut work).unwrap(); +//! decoded.extend_from_slice(&out[..progress.produced]); +//! input = &input[progress.consumed..]; +//! if progress.state == StepState::Done { +//! break; +//! } +//! } +//! assert_eq!(decoded, data); +//! ``` + +use super::block_decoder::{self, BlockDecoder}; +use super::buffer_backend::{BufferBackend, WILDCOPY_OVERLENGTH}; +use super::decode_buffer::DecodeBuffer; +use super::errors::{ + BlockHeaderReadError, DecodeBlockContentError, FrameHeaderError, ReadFrameHeaderError, +}; +use super::frame; +use super::ringbuffer::RingBuffer; +use super::scratch::{ByteScratch, FSEScratch, HuffmanScratch, Workspace, WorkspaceRef}; +use super::user_slice_buf::UserSliceBackend; +use crate::blocks::block::{BlockHeader, BlockType}; +use crate::common::{MAGIC_NUM, MAX_DECODER_WINDOW_SIZE, MIN_WINDOW_SIZE}; +use core::marker::PhantomData; +use core::ptr::NonNull; + +/// Work units charged before a compressed block decodes, on top of its input +/// bytes and its output bound: the most its entropy-table builds can cost +/// (one Huffman table of `1 << 11` entries and three sequence tables of +/// `1 << 9`, rounded up). +pub const COMPRESSED_BLOCK_UNITS: u64 = 4096; + +/// Largest frame header: magic, descriptor, window byte, 4-byte dictionary id +/// and 8-byte content size (RFC 8878 3.1.1.1). +const MAX_FRAME_HEADER: usize = 18; + +/// Alignment the entropy tables need at the start of the workspace. +const TABLE_ALIGN: usize = { + let fse = core::mem::align_of::(); + let huf = core::mem::align_of::(); + if fse > huf { fse } else { huf } +}; + +/// Bytes of workspace a [`WorkspaceDecoder`] needs to decode frames whose +/// window is at most `max_window_size`. +/// +/// The figure covers every byte the decoder writes: the window with room for +/// the block being decoded, the entropy tables, the literals, the staging for a +/// block split across input chunks, and alignment slack, so any sub-slice of +/// that length works whatever its address. +/// +/// # Errors +/// +/// [`WorkspaceConfigError::WindowLimitOutOfRange`] when `max_window_size` lies +/// outside 1 KiB ..= [`MAX_DECODER_WINDOW_SIZE`], +/// and [`WorkspaceConfigError::SizeOverflow`] when the size does not fit in +/// `usize` on this target. +/// +/// # Examples +/// +/// ``` +/// use structured_zstd::decoding::workspace::workspace_bytes; +/// let small = workspace_bytes(64 * 1024).unwrap(); +/// let large = workspace_bytes(8 << 20).unwrap(); +/// assert!(small < large); +/// assert!(workspace_bytes(100).is_err()); +/// ``` +pub fn workspace_bytes(max_window_size: u64) -> Result { + Layout::for_window(max_window_size).map(|layout| layout.total) +} + +/// Where each part of the workspace sits, relative to its aligned start. +struct Layout { + huf_offset: usize, + tables_len: usize, + ring_cap: usize, + literals_len: usize, + staging_len: usize, + total: usize, +} + +impl Layout { + fn for_window(max_window_size: u64) -> Result { + use WorkspaceConfigError::SizeOverflow; + if !(MIN_WINDOW_SIZE..=MAX_DECODER_WINDOW_SIZE).contains(&max_window_size) { + return Err(WorkspaceConfigError::WindowLimitOutOfRange { + requested: max_window_size, + }); + } + let window = usize::try_from(max_window_size).map_err(|_| SizeOverflow)?; + let block_max = block_decoder::block_maximum(window); + let huf_offset = core::mem::size_of::() + .next_multiple_of(core::mem::align_of::()); + let tables_len = huf_offset + core::mem::size_of::(); + // The window, plus two block maxima: a block's matches stop at one + // block past the window (the per-block ceiling), and literals copied + // after its last match can add up to one more before the block-size + // check rejects it. One more slot is the ring's empty/full sentinel. + let ring_cap = window + .checked_add(2 * block_max) + .and_then(|n| n.checked_add(1)) + .ok_or(SizeOverflow)?; + let literals_len = block_max + WILDCOPY_OVERLENGTH; + let staging_len = block_max; + let total = [ + tables_len, + ring_cap, + WILDCOPY_OVERLENGTH, + literals_len, + staging_len, + ] + .iter() + .try_fold(TABLE_ALIGN - 1, |sum, &part| sum.checked_add(part)) + .ok_or(SizeOverflow)?; + Ok(Layout { + huf_offset, + tables_len, + ring_cap, + literals_len, + staging_len, + total, + }) + } +} + +/// Limits a [`WorkspaceDecoder`] holds every frame of the stream to. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DecodeLimits { + /// Largest window a frame may declare. Also sizes the workspace; see + /// [`workspace_bytes`]. + pub max_window_size: u64, + /// Most decompressed bytes the whole stream may produce. + pub max_output: u64, +} + +impl DecodeLimits { + /// Limits with the given window ceiling and no output limit. + pub const fn new(max_window_size: u64) -> Self { + Self { + max_window_size, + max_output: u64::MAX, + } + } + + /// The same limits with the total output capped at `max_output` bytes. + pub const fn with_max_output(self, max_output: u64) -> Self { + Self { max_output, ..self } + } +} + +/// Work a caller allows its decode steps, in units charged before the work +/// they pay for runs. +/// +/// One unit is charged per input byte consumed (frame and block headers, +/// block contents, skippable frames, checksums), and one per byte a block may +/// produce, before it is produced: a Raw block's bytes as each part of it is +/// copied, an RLE block's stated size, and a compressed block's frame block +/// maximum plus [`COMPRESSED_BLOCK_UNITS`] for its table builds. A charge the +/// budget cannot cover refuses the step before that work starts. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct WorkBudget { + remaining: u64, +} + +impl WorkBudget { + /// A budget of `units`. + pub const fn new(units: u64) -> Self { + Self { remaining: units } + } + + /// Units still available. + pub const fn remaining(&self) -> u64 { + self.remaining + } + + /// Add `units` to the budget, stopping at `u64::MAX`, which is already + /// more work than any stream can ask for. + pub fn add(&mut self, units: u64) { + self.remaining = self.remaining.saturating_add(units); + } + + fn charge(&mut self, units: u64) -> Result<(), WorkspaceDecodeError> { + match self.remaining.checked_sub(units) { + Some(left) => { + self.remaining = left; + Ok(()) + } + None => Err(WorkspaceDecodeError::Limit(LimitExceeded::Work { + needed: units, + remaining: self.remaining, + })), + } + } +} + +/// What one [`WorkspaceDecoder::step`] did. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Progress { + /// Input bytes consumed; the next step starts after them. + pub consumed: usize, + /// Bytes written to the front of the output slice. + pub produced: usize, + /// What the decoder needs next. + pub state: StepState, +} + +/// Why a [`WorkspaceDecoder::step`] returned. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum StepState { + /// Every input byte was consumed; call again with more, or with the + /// final-input flag set once there is no more. + NeedInput, + /// The output slice is full; call again with room. + NeedOutput, + /// The input was final and every frame in it is decoded and delivered. + Done, +} + +/// A workspace size or limit [`workspace_bytes`] or [`WorkspaceDecoder::new`] +/// refused. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum WorkspaceConfigError { + /// The window limit lies outside 1 KiB ..= + /// [`MAX_DECODER_WINDOW_SIZE`]. + WindowLimitOutOfRange { + /// The limit asked for. + requested: u64, + }, + /// The workspace size does not fit in `usize` on this target. + SizeOverflow, + /// The workspace slice is shorter than [`workspace_bytes`] asks for. + WorkspaceTooSmall { + /// Bytes the limits require. + needed: usize, + /// Bytes supplied. + provided: usize, + }, +} + +impl core::fmt::Display for WorkspaceConfigError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + Self::WindowLimitOutOfRange { requested } => write!( + f, + "window limit {requested} lies outside {MIN_WINDOW_SIZE}..={MAX_DECODER_WINDOW_SIZE}" + ), + Self::SizeOverflow => write!(f, "workspace size does not fit in usize"), + Self::WorkspaceTooSmall { needed, provided } => write!( + f, + "workspace of {provided} bytes is smaller than the {needed} the limits require" + ), + } + } +} + +#[cfg(feature = "std")] +impl std::error::Error for WorkspaceConfigError {} + +/// Why a [`WorkspaceDecoder::step`] failed. Every variant ends the stream: +/// later steps return [`Self::Failed`] until [`WorkspaceDecoder::reset`]. +#[derive(Debug)] +#[non_exhaustive] +pub enum WorkspaceDecodeError { + /// A caller-chosen limit refused the stream. + Limit(LimitExceeded), + /// The input is not a valid zstd stream. + Malformed(MalformedInput), + /// The input is valid but needs something this decoder does not do. + Unsupported(Unsupported), + /// The input was marked final in the middle of a frame. + Truncated, + /// An earlier step failed; call [`WorkspaceDecoder::reset`] first. + Failed, +} + +/// A limit from [`DecodeLimits`] or [`WorkBudget`] that refused the stream. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum LimitExceeded { + /// A frame declared a window above [`DecodeLimits::max_window_size`]. + Window { + /// The window the frame declared. + requested: u64, + /// The limit. + limit: u64, + }, + /// The stream would produce more than [`DecodeLimits::max_output`]. + Output { + /// The limit. + limit: u64, + }, + /// The next unit of work costs more than the budget has left. + Work { + /// Units the refused work costs. + needed: u64, + /// Units the budget had. + remaining: u64, + }, +} + +/// Something this decoder does not do, met in otherwise valid input. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum Unsupported { + /// The frame needs dictionary `id`, and this decoder takes none. + Dictionary { + /// The dictionary id the frame header names. + id: u32, + }, +} + +/// How the input breaks the zstd format. +#[derive(Debug)] +#[non_exhaustive] +pub enum MalformedInput { + /// A frame header, or the magic number before it, is invalid. + FrameHeader(ReadFrameHeaderError), + /// The frame header declares an invalid window. + Window(FrameHeaderError), + /// A block header is invalid. + BlockHeader(BlockHeaderReadError), + /// A block states a size above its frame's block maximum + /// (RFC 8878 3.1.1.2.4). + BlockTooLarge { + /// The size the block header states. + size: usize, + /// The frame's block maximum. + maximum: usize, + }, + /// A block's content does not decode. + Block(DecodeBlockContentError), + /// The frame's blocks do not add up to its declared content size. + ContentSize { + /// The size the frame header declares. + declared: u64, + /// The size its blocks produce (so far, when they overrun it). + produced: u64, + }, + /// The content checksum does not match the decoded output. + Checksum { + /// The checksum the frame carries. + expected: u32, + /// The low 32 bits of the XXH64 of the decoded frame. + calculated: u32, + }, +} + +impl core::fmt::Display for WorkspaceDecodeError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + Self::Limit(limit) => match limit { + LimitExceeded::Window { requested, limit } => write!( + f, + "frame window of {requested} bytes exceeds the {limit}-byte limit" + ), + LimitExceeded::Output { limit } => { + write!(f, "output exceeds the {limit}-byte limit") + } + LimitExceeded::Work { needed, remaining } => write!( + f, + "work of {needed} units exceeds the {remaining} remaining" + ), + }, + Self::Malformed(malformed) => match malformed { + MalformedInput::FrameHeader(e) => write!(f, "invalid frame header: {e}"), + MalformedInput::Window(e) => write!(f, "invalid frame window: {e}"), + MalformedInput::BlockHeader(e) => write!(f, "invalid block header: {e}"), + MalformedInput::BlockTooLarge { size, maximum } => write!( + f, + "block of {size} bytes exceeds the frame's block maximum of {maximum}" + ), + MalformedInput::Block(e) => write!(f, "invalid block: {e}"), + MalformedInput::ContentSize { declared, produced } => write!( + f, + "frame declares {declared} bytes but its blocks produce {produced}" + ), + MalformedInput::Checksum { + expected, + calculated, + } => write!( + f, + "content checksum {calculated:#010x} does not match the stored {expected:#010x}" + ), + }, + Self::Unsupported(Unsupported::Dictionary { id }) => { + write!( + f, + "frame needs dictionary {id}, which this decoder does not take" + ) + } + Self::Truncated => write!(f, "input ended in the middle of a frame"), + Self::Failed => write!(f, "an earlier step failed; reset the decoder"), + } + } +} + +#[cfg(feature = "std")] +impl std::error::Error for WorkspaceDecodeError {} + +/// A fixed region of the workspace standing in for the `Vec` the owned +/// decoders keep for literals. +/// +/// A push past the end is dropped but still counted, so the literals decoder's +/// own count check rejects the block as it would with a growable buffer. Every +/// other write is sized by the caller below the region's length; the asserts +/// are what keeps a broken size argument from writing outside it. +struct SliceScratch<'w> { + buf: &'w mut [u8], + len: usize, +} + +impl ByteScratch for SliceScratch<'_> { + #[inline(always)] + fn len(&self) -> usize { + self.len + } + #[inline(always)] + fn clear(&mut self) { + self.len = 0; + } + #[inline(always)] + fn truncate(&mut self, len: usize) { + self.len = self.len.min(len); + } + #[inline(always)] + fn reserve(&mut self, additional: usize) { + assert!( + self.len + .checked_add(additional) + .is_some_and(|end| end <= self.buf.len()), + "literals exceed the workspace region sized for a block" + ); + } + #[inline(always)] + fn push(&mut self, byte: u8) { + if let Some(slot) = self.buf.get_mut(self.len) { + *slot = byte; + } + self.len += 1; + } + fn resize(&mut self, new_len: usize, value: u8) { + if new_len > self.len { + self.buf[self.len..new_len].fill(value); + } + self.len = new_len; + } + fn extend_from_slice(&mut self, bytes: &[u8]) { + let end = self.len + bytes.len(); + self.buf[self.len..end].copy_from_slice(bytes); + self.len = end; + } + #[inline(always)] + fn as_mut_ptr(&mut self) -> *mut u8 { + self.buf.as_mut_ptr() + } + #[inline(always)] + unsafe fn set_len(&mut self, new_len: usize) { + assert!(new_len <= self.buf.len()); + self.len = new_len; + } + #[inline(always)] + fn as_slice(&self) -> &[u8] { + &self.buf[..self.len.min(self.buf.len())] + } + #[inline(always)] + fn as_mut_slice(&mut self) -> &mut [u8] { + let len = self.len.min(self.buf.len()); + &mut self.buf[..len] + } +} + +/// The current frame's window over the workspace's window region, chosen per +/// frame as `FrameDecoder` chooses its backend: a single-segment frame's whole +/// content is its window, so it decodes linearly into the region, the shape +/// `decode_all` writes into a caller's slice; any other frame decodes into a +/// ring. +enum Window<'w> { + /// No frame has started since construction or reset. + Idle, + Ring(DecodeBuffer), + Linear(DecodeBuffer>), +} + +/// Run `$body` with `$buffer` bound to the active window's buffer, or yield +/// `$idle` when no frame has started. +macro_rules! with_window { + ($window:expr, $buffer:ident => $body:expr, idle => $idle:expr) => { + match $window { + Window::Ring($buffer) => $body, + Window::Linear($buffer) => $body, + Window::Idle => $idle, + } + }; +} + +/// Output bytes `buffer` has produced this frame. The linear backend's inline +/// sequence path skips the running counter, so its write cursor counts there; +/// the ring keeps the counter, its cursor wrapping. +fn produced_so_far(buffer: &DecodeBuffer) -> u64 { + if B::INLINE_EXEC_MAINTAINS_OUTPUT_COUNTER { + buffer.total_output() + } else { + buffer.buffer_ref().tail() as u64 + } +} + +/// The state the shared block decoder works on for one block: the tables and +/// literals in the workspace, and the active window's buffer. +struct BlockScratch<'a, 'w, B: BufferBackend> { + huf: &'a mut HuffmanScratch, + fse: &'a mut FSEScratch, + buffer: &'a mut DecodeBuffer, + offset_hist: &'a mut [u32; 3], + literals: &'a mut SliceScratch<'w>, + /// The Read-sourced block path's buffer; this decoder stages blocks + /// itself and hands them over as slices, so it stays empty. + block_content: &'a mut SliceScratch<'w>, +} + +impl<'w, B: BufferBackend> Workspace for BlockScratch<'_, 'w, B> { + type Backend = B; + type Bytes = SliceScratch<'w>; + fn split(&mut self) -> WorkspaceRef<'_, B, SliceScratch<'w>> { + WorkspaceRef { + huf: &mut *self.huf, + fse: &mut *self.fse, + buffer: &mut *self.buffer, + offset_hist: &mut *self.offset_hist, + literals_buffer: &mut *self.literals, + block_content_buffer: &mut *self.block_content, + } + } +} + +/// Decode one block whose whole content is `content`, returning the bytes it +/// produced. +fn run_block( + decoder: &mut BlockDecoder, + scratch: &mut BlockScratch<'_, '_, B>, + header: &BlockHeader, + mut content: &[u8], +) -> Result { + let before = produced_so_far(scratch.buffer); + decoder.decode_block_content_from_slice(header, scratch, None, &mut content)?; + Ok(produced_so_far(scratch.buffer) - before) +} + +/// What the decoder reads next. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +enum Phase { + /// The 4-byte magic number of the next frame. + Magic, + /// The descriptor byte, which gives the header's length. + Descriptor, + /// The rest of the frame header. + FrameHeader, + /// The 4-byte length of a skippable frame. + SkippableSize, + /// The payload of a skippable frame. + Skipping { remaining: u64 }, + /// A 3-byte block header. + BlockHeader, + /// A block's content. + BlockBody(PendingBlock), + /// The 4-byte content checksum. + Checksum, +} + +/// A block whose header was read, carried while its content arrives. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +struct PendingBlock { + last_block: bool, + block_type: BlockType, + decompressed_size: u32, + content_size: u32, +} + +impl PendingBlock { + fn header(self) -> BlockHeader { + BlockHeader { + last_block: self.last_block, + block_type: self.block_type, + decompressed_size: self.decompressed_size, + content_size: self.content_size, + } + } +} + +/// The frame being decoded. +#[derive(Clone, Copy)] +struct FrameState { + declared_size: Option, + has_checksum: bool, + block_maximum: usize, + produced: u64, +} + +/// Allocation-free streaming zstd decoder over a caller-owned workspace. +/// +/// See the [module documentation](self) for the contract and an example. +pub struct WorkspaceDecoder<'w> { + huf: &'w mut HuffmanScratch, + fse: &'w mut FSEScratch, + offset_hist: [u32; 3], + literals: SliceScratch<'w>, + block_content: SliceScratch<'w>, + window: Window<'w>, + /// The workspace region every frame's window is laid over, held as a + /// pointer because frames alternate between a ring (which addresses it by + /// pointer) and a linear buffer (which borrows it), never both at once. + region: NonNull, + /// `region`'s length: the ring's capacity plus the copiers' slack. + region_len: usize, + ring_cap: usize, + _region: PhantomData<&'w mut [u8]>, + staging: &'w mut [u8], + staged: usize, + block_decoder: BlockDecoder, + kernel: crate::cpu_kernel::CpuKernelTag, + limits: DecodeLimits, + phase: Phase, + header: [u8; MAX_FRAME_HEADER], + header_len: usize, + header_target: usize, + frame: FrameState, + /// Every block of the current frame is decoded, so all of the window is + /// output to deliver rather than history to keep. + blocks_done: bool, + stream_produced: u64, + failed: bool, +} + +// SAFETY: `region` points into the `&'w mut [u8]` workspace this decoder +// borrows exclusively for `'w` (`_region` records that borrow), and every other +// field is `Send`. Moving the decoder moves that exclusive borrow with it, as +// moving the `&mut [u8]` itself would. +unsafe impl Send for WorkspaceDecoder<'_> {} + +impl<'w> WorkspaceDecoder<'w> { + /// Build a decoder in `workspace`, which must hold at least + /// [`workspace_bytes`]`(limits.max_window_size)` bytes. Nothing is + /// allocated, now or later; the workspace stays borrowed for the decoder's + /// life. + /// + /// # Errors + /// + /// [`WorkspaceConfigError`] when the window limit is out of range or the + /// workspace is too small for it. + pub fn new( + workspace: &'w mut [u8], + limits: DecodeLimits, + ) -> Result { + let layout = Layout::for_window(limits.max_window_size)?; + let provided = workspace.len(); + if provided < layout.total { + return Err(WorkspaceConfigError::WorkspaceTooSmall { + needed: layout.total, + provided, + }); + } + // `layout.total` holds `TABLE_ALIGN - 1` bytes of slack, and the + // offset a byte pointer needs to reach an alignment is below it. + let skip = workspace.as_ptr().align_offset(TABLE_ALIGN); + debug_assert!(skip < TABLE_ALIGN); + let rest = &mut workspace[skip..]; + let (tables, rest) = rest.split_at_mut(layout.tables_len); + let (ring, rest) = rest.split_at_mut(layout.ring_cap + WILDCOPY_OVERLENGTH); + let (literals, rest) = rest.split_at_mut(layout.literals_len); + let staging = &mut rest[..layout.staging_len]; + + let tables = tables.as_mut_ptr(); + let fse_ptr = tables.cast::(); + // SAFETY: `tables` starts at a `TABLE_ALIGN` boundary and spans + // `tables_len` bytes: the FSE scratch at its start and the Huffman + // scratch at `huf_offset`, a multiple of its alignment, both in bounds. + // Each is fully initialised in place before a reference to it exists, + // and the bytes are reached through nothing else while `workspace` + // stays borrowed for `'w`. Neither type owns heap memory under this + // feature, so never running their destructors leaks nothing. + let (fse, huf) = unsafe { + let huf_ptr = tables.add(layout.huf_offset).cast::(); + FSEScratch::init_in_place(fse_ptr); + HuffmanScratch::init_in_place(huf_ptr); + (&mut *fse_ptr, &mut *huf_ptr) + }; + let region_len = ring.len(); + let region = NonNull::from(ring).cast::(); + + let kernel = crate::cpu_kernel::detect_cpu_kernel(); + Ok(WorkspaceDecoder { + huf, + fse, + offset_hist: [1, 4, 8], + literals: SliceScratch { + buf: literals, + len: 0, + }, + block_content: SliceScratch { + buf: &mut [], + len: 0, + }, + window: Window::Idle, + region, + region_len, + ring_cap: layout.ring_cap, + _region: PhantomData, + staging, + staged: 0, + block_decoder: block_decoder::with_kernel(kernel), + kernel, + limits, + phase: Phase::Magic, + header: [0; MAX_FRAME_HEADER], + header_len: 0, + header_target: 4, + frame: FrameState { + declared_size: None, + has_checksum: false, + block_maximum: 0, + produced: 0, + }, + blocks_done: false, + stream_produced: 0, + failed: false, + }) + } + + /// Start a new stream: forget the current one, including a failure, and + /// expect a frame's magic number next. The workspace is reused as it is. + pub fn reset(&mut self) { + self.window = Window::Idle; + self.literals.clear(); + self.staged = 0; + self.block_decoder = block_decoder::with_kernel(self.kernel); + self.phase = Phase::Magic; + self.header_len = 0; + self.header_target = 4; + self.blocks_done = false; + self.stream_produced = 0; + self.failed = false; + } + + /// Decompressed bytes the stream has produced so far, delivered or not. + pub fn total_output(&self) -> u64 { + self.stream_produced + } + + /// Decode from `input` into `output` until one of them runs out or the + /// stream ends. + /// + /// `input_is_final` says no input follows this slice. Without it, running + /// out of input asks for more ([`StepState::NeedInput`]); with it, input + /// that ends between frames finishes the stream ([`StepState::Done`]) and + /// input that ends inside one is [`WorkspaceDecodeError::Truncated`]. + /// Unconsumed input must be passed again, first, on the next call. + /// + /// Concatenated frames decode one after another, and skippable frames are + /// consumed without output. + /// + /// # Errors + /// + /// See [`WorkspaceDecodeError`]. Any error ends the stream; the output + /// already written by this call is then not counted in a [`Progress`]. + pub fn step( + &mut self, + input: &[u8], + input_is_final: bool, + output: &mut [u8], + work: &mut WorkBudget, + ) -> Result { + if self.failed { + return Err(WorkspaceDecodeError::Failed); + } + let result = self.step_inner(input, input_is_final, output, work); + if result.is_err() { + self.failed = true; + } + result + } + + fn step_inner( + &mut self, + input: &[u8], + input_is_final: bool, + output: &mut [u8], + work: &mut WorkBudget, + ) -> Result { + let mut consumed = 0; + let mut produced = 0; + loop { + produced += self.deliver(&mut output[produced..]); + if self.pending() != 0 { + return Ok(Progress { + consumed, + produced, + state: StepState::NeedOutput, + }); + } + let rest = &input[consumed..]; + match self.phase { + Phase::Skipping { remaining } => { + if remaining == 0 { + self.expect_frame(); + continue; + } + if rest.is_empty() { + return self.out_of_input(input_is_final, consumed, produced); + } + let take = rest + .len() + .min(usize::try_from(remaining).unwrap_or(usize::MAX)); + work.charge(take as u64)?; + consumed += take; + self.phase = Phase::Skipping { + remaining: remaining - take as u64, + }; + } + Phase::BlockBody(block) if block.block_type == BlockType::Raw => { + // A Raw block's content is its output: it goes straight + // into the window as it arrives, never through staging. + // `staged` counts the bytes written so far. + let remaining = block.content_size as usize - self.staged; + if remaining == 0 { + self.finish_block(block)?; + continue; + } + if rest.is_empty() { + return self.out_of_input(input_is_final, consumed, produced); + } + let take = rest.len().min(remaining); + // One unit for the input byte, one for the output byte. + work.charge(2 * take as u64)?; + self.push_raw(&rest[..take])?; + self.staged += take; + consumed += take; + } + Phase::BlockBody(block) => { + let size = block.content_size as usize; + if self.staged == 0 && rest.len() >= size { + work.charge(size as u64)?; + consumed += size; + self.decode_block(block, &rest[..size], work)?; + continue; + } + if rest.is_empty() { + return self.out_of_input(input_is_final, consumed, produced); + } + let take = rest.len().min(size - self.staged); + work.charge(take as u64)?; + self.staging[self.staged..self.staged + take].copy_from_slice(&rest[..take]); + self.staged += take; + consumed += take; + if self.staged == size { + self.staged = 0; + // The staging region is a separate field from the + // state the block decoder writes, so it lends the + // content in place. + let staging = core::mem::take(&mut self.staging); + let decoded = self.decode_block(block, &staging[..size], work); + self.staging = staging; + decoded?; + } + } + _ => { + if self.header_len < self.header_target { + if rest.is_empty() { + return self.out_of_input(input_is_final, consumed, produced); + } + let take = rest.len().min(self.header_target - self.header_len); + work.charge(take as u64)?; + self.header[self.header_len..self.header_len + take] + .copy_from_slice(&rest[..take]); + self.header_len += take; + consumed += take; + } + if self.header_len == self.header_target { + self.header_complete()?; + } + } + } + } + } + + /// Input ran out. Only between frames does final input end the stream. + fn out_of_input( + &self, + input_is_final: bool, + consumed: usize, + produced: usize, + ) -> Result { + if !input_is_final { + return Ok(Progress { + consumed, + produced, + state: StepState::NeedInput, + }); + } + if self.phase == Phase::Magic && self.header_len == 0 { + return Ok(Progress { + consumed, + produced, + state: StepState::Done, + }); + } + Err(WorkspaceDecodeError::Truncated) + } + + /// Bytes ready to deliver: past the window while the frame decodes, all + /// of it once its last block is in. + fn pending(&self) -> usize { + let blocks_done = self.blocks_done; + with_window!(&self.window, buffer => if blocks_done { + buffer.can_drain() + } else { + buffer.can_drain_to_window_size().unwrap_or(0) + }, idle => 0) + } + + fn deliver(&mut self, output: &mut [u8]) -> usize { + if output.is_empty() { + return 0; + } + let blocks_done = self.blocks_done; + let drained = with_window!(&mut self.window, buffer => if blocks_done { + buffer.read_all(output) + } else { + buffer.read_reporting_pending(output).map(|(read, _)| read) + }, idle => Ok(0)); + // Draining into a slice cannot fail: the copy reports no error. + drained.unwrap_or(0) + } + + fn expect_frame(&mut self) { + self.phase = Phase::Magic; + self.header_len = 0; + self.header_target = 4; + } + + fn expect_header(&mut self, phase: Phase, target: usize) { + self.phase = phase; + self.header_target = target; + } + + /// The bytes the current phase collects are all in `header`. + fn header_complete(&mut self) -> Result<(), WorkspaceDecodeError> { + use WorkspaceDecodeError::Malformed; + match self.phase { + Phase::Magic => { + let magic = u32::from_le_bytes([ + self.header[0], + self.header[1], + self.header[2], + self.header[3], + ]); + if (0x184D2A50..=0x184D2A5F).contains(&magic) { + self.expect_header(Phase::SkippableSize, 8); + } else if magic == MAGIC_NUM { + self.expect_header(Phase::Descriptor, 5); + } else { + return Err(Malformed(MalformedInput::FrameHeader( + ReadFrameHeaderError::BadMagicNumber(magic), + ))); + } + } + Phase::SkippableSize => { + let length = u32::from_le_bytes([ + self.header[4], + self.header[5], + self.header[6], + self.header[7], + ]); + self.phase = Phase::Skipping { + remaining: u64::from(length), + }; + } + Phase::Descriptor => { + let length = frame::measure_header_from_slice(&self.header[..5]) + .map_err(|e| Malformed(MalformedInput::FrameHeader(e)))?; + self.expect_header(Phase::FrameHeader, length); + } + Phase::FrameHeader => self.start_frame()?, + Phase::BlockHeader => self.start_block()?, + Phase::Checksum => { + self.check_checksum()?; + self.expect_frame(); + } + Phase::Skipping { .. } | Phase::BlockBody(_) => { + unreachable!("phases without a header buffer are handled by the step loop") + } + } + Ok(()) + } + + fn start_frame(&mut self) -> Result<(), WorkspaceDecodeError> { + use WorkspaceDecodeError::{Limit, Malformed}; + let mut source = &self.header[..self.header_target]; + let (header, _) = frame::read_frame_header_from_slice(&mut source, false) + .map_err(|e| Malformed(MalformedInput::FrameHeader(e)))?; + if let Some(id) = header.dictionary_id() { + return Err(WorkspaceDecodeError::Unsupported(Unsupported::Dictionary { + id, + })); + } + let window = header + .window_size() + .map_err(|e| Malformed(MalformedInput::Window(e)))?; + if window > self.limits.max_window_size { + return Err(Limit(LimitExceeded::Window { + requested: window, + limit: self.limits.max_window_size, + })); + } + let declared_size = header.fcs_declared().then(|| header.frame_content_size()); + if let Some(declared) = declared_size + && self + .stream_produced + .checked_add(declared) + .is_none_or(|total| total > self.limits.max_output) + { + return Err(Limit(LimitExceeded::Output { + limit: self.limits.max_output, + })); + } + // Fits `usize`: it is at most the limit, which `Layout` converted. + let window = window as usize; + let has_checksum = header.descriptor.content_checksum_flag(); + + // End the previous frame's view of the region before laying out the + // next one, so the two never coexist. + self.window = Window::Idle; + self.window = if header.descriptor.single_segment_flag() { + // SAFETY: `region` is `region_len` initialised bytes borrowed for + // `'w`, and no other view of it is alive: the previous window was + // dropped just above. + let region = + unsafe { core::slice::from_raw_parts_mut(self.region.as_ptr(), self.region_len) }; + Window::Linear(DecodeBuffer::from_backend( + UserSliceBackend::from_slice(region), + window, + )) + } else { + // SAFETY: as above; the ring addresses `ring_cap` bytes of the + // region plus the copiers' slack, `region_len` in all, and + // `ring_cap >= 1`. + let ring = unsafe { RingBuffer::from_borrowed(self.region, self.ring_cap) }; + Window::Ring(DecodeBuffer::from_backend(ring, window)) + }; + with_window!(&mut self.window, buffer => { + buffer.set_declared_content(declared_size); + #[cfg(feature = "hash")] + buffer.set_compute_hash(has_checksum); + }, idle => {}); + self.offset_hist = [1, 4, 8]; + self.literals.clear(); + super::scratch::reset_entropy(self.huf, self.fse); + + self.frame = FrameState { + declared_size, + has_checksum, + block_maximum: block_decoder::block_maximum(window), + produced: 0, + }; + self.blocks_done = false; + self.header_len = 0; + self.expect_header(Phase::BlockHeader, 3); + Ok(()) + } + + fn start_block(&mut self) -> Result<(), WorkspaceDecodeError> { + use WorkspaceDecodeError::Malformed; + let (header, _) = self + .block_decoder + .read_block_header(&self.header[..3]) + .map_err(|e| Malformed(MalformedInput::BlockHeader(e)))?; + // RFC 8878 3.1.1.2.4: Block_Size is at most Block_Maximum_Size for + // every block type; for RLE it is the regenerated size. + let size = match header.block_type { + BlockType::RLE => header.decompressed_size, + _ => header.content_size, + } as usize; + let maximum = self.frame.block_maximum; + if size > maximum { + return Err(Malformed(MalformedInput::BlockTooLarge { size, maximum })); + } + self.header_len = 0; + self.staged = 0; + self.phase = Phase::BlockBody(PendingBlock { + last_block: header.last_block, + block_type: header.block_type, + decompressed_size: header.decompressed_size, + content_size: header.content_size, + }); + Ok(()) + } + + /// Decode one block whose whole content is `content`, charging its work + /// first and holding its output to the limits. + fn decode_block( + &mut self, + block: PendingBlock, + content: &[u8], + work: &mut WorkBudget, + ) -> Result<(), WorkspaceDecodeError> { + use WorkspaceDecodeError::{Limit, Malformed}; + let output_bound = match block.block_type { + BlockType::Compressed => self.frame.block_maximum as u64 + COMPRESSED_BLOCK_UNITS, + _ => u64::from(block.decompressed_size), + }; + work.charge(output_bound)?; + let output_limit = Limit(LimitExceeded::Output { + limit: self.limits.max_output, + }); + if block.block_type != BlockType::Compressed + && self + .stream_produced + .checked_add(u64::from(block.decompressed_size)) + .is_none_or(|total| total > self.limits.max_output) + { + return Err(output_limit); + } + + let header = block.header(); + let produced = with_window!(&mut self.window, buffer => run_block( + &mut self.block_decoder, + &mut BlockScratch { + huf: &mut *self.huf, + fse: &mut *self.fse, + buffer, + offset_hist: &mut self.offset_hist, + literals: &mut self.literals, + block_content: &mut self.block_content, + }, + &header, + content, + ), idle => unreachable!("a block is decoded only inside a started frame")) + .map_err(|e| Malformed(MalformedInput::Block(e)))?; + + self.stream_produced += produced; + self.frame.produced += produced; + if self.stream_produced > self.limits.max_output { + return Err(output_limit); + } + if let Some(declared) = self.frame.declared_size + && self.frame.produced > declared + { + return Err(Malformed(MalformedInput::ContentSize { + declared, + produced: self.frame.produced, + })); + } + self.finish_block(block) + } + + /// Append the next part of a Raw block to the window, after holding it to + /// the output limit and to the frame's declared size. + fn push_raw(&mut self, bytes: &[u8]) -> Result<(), WorkspaceDecodeError> { + use crate::cpu_kernel::{BaselineKernel, CpuKernelTag, ScalarKernel}; + use WorkspaceDecodeError::{Limit, Malformed}; + let len = bytes.len() as u64; + let stream_total = self.stream_produced + len; + if stream_total > self.limits.max_output { + return Err(Limit(LimitExceeded::Output { + limit: self.limits.max_output, + })); + } + let frame_total = self.frame.produced + len; + if let Some(declared) = self.frame.declared_size + && frame_total > declared + { + return Err(Malformed(MalformedInput::ContentSize { + declared, + produced: frame_total, + })); + } + // One copy per piece, under the kernel a Raw block takes elsewhere: + // the scalar tier stays portable. + let scalar = self.kernel == CpuKernelTag::Scalar; + let pushed = with_window!(&mut self.window, buffer => if scalar { + buffer.try_push::(bytes) + } else { + buffer.try_push::(bytes) + }, idle => unreachable!("a block is decoded only inside a started frame")); + // The block was held to the block maximum and the window region holds + // a window and two block maxima, so only a frame that outruns its own + // window can reach here. + pushed.map_err(|_| { + Malformed(MalformedInput::Block( + DecodeBlockContentError::BackendOverflow { + step: BlockType::Raw, + }, + )) + })?; + self.stream_produced = stream_total; + self.frame.produced = frame_total; + Ok(()) + } + + /// The block's content is all decoded: check a last block against the + /// declared size and read what follows it. + fn finish_block(&mut self, block: PendingBlock) -> Result<(), WorkspaceDecodeError> { + if let Some(declared) = self.frame.declared_size + && block.last_block + && self.frame.produced != declared + { + return Err(WorkspaceDecodeError::Malformed( + MalformedInput::ContentSize { + declared, + produced: self.frame.produced, + }, + )); + } + self.staged = 0; + self.header_len = 0; + if !block.last_block { + self.expect_header(Phase::BlockHeader, 3); + } else { + self.blocks_done = true; + if self.frame.has_checksum { + self.expect_header(Phase::Checksum, 4); + } else { + self.expect_frame(); + } + } + Ok(()) + } + + /// Compare the stored checksum with the decoded frame's. The step loop + /// reaches this only once the whole frame was delivered, which is when the + /// running digest covers all of it. Without the `hash` feature the + /// checksum is read and not checked. + fn check_checksum(&mut self) -> Result<(), WorkspaceDecodeError> { + #[cfg(feature = "hash")] + { + use core::hash::Hasher; + let expected = u32::from_le_bytes([ + self.header[0], + self.header[1], + self.header[2], + self.header[3], + ]); + let calculated = + with_window!(&self.window, buffer => buffer.hash.finish(), idle => 0) as u32; + if expected != calculated { + return Err(WorkspaceDecodeError::Malformed(MalformedInput::Checksum { + expected, + calculated, + })); + } + } + Ok(()) + } +} + +#[cfg(test)] +mod tests; diff --git a/zstd/src/decoding/workspace/tests.rs b/zstd/src/decoding/workspace/tests.rs new file mode 100644 index 000000000..58be93615 --- /dev/null +++ b/zstd/src/decoding/workspace/tests.rs @@ -0,0 +1,566 @@ +use super::*; +use crate::decoding::FrameDecoder; +use crate::decoding::read_frame_header_info; +use crate::encoding::{CompressionLevel, compress_slice_to_vec}; +use alloc::vec; +use alloc::vec::Vec; + +const UNLIMITED: u64 = u64::MAX; + +fn text(len: usize) -> Vec { + let words = [ + "alpha ", + "beta ", + "gamma ", + "delta ", + "manifest ", + "blob ", + "chunk ", + ]; + let mut out = Vec::with_capacity(len); + let mut i = 0usize; + while out.len() < len { + out.extend_from_slice(words[(i * 7 + i / 13) % words.len()].as_bytes()); + out.extend_from_slice(&(i as u32).to_le_bytes()[..(i % 3)]); + i += 1; + } + out.truncate(len); + out +} + +fn noise(len: usize) -> Vec { + let mut x = 0x9E37_79B9_7F4A_7C15u64; + (0..len) + .map(|_| { + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + x as u8 + }) + .collect() +} + +/// Feed `compressed` `in_chunk` bytes at a time with `out_chunk` bytes of +/// output room per step, asserting every step makes progress. +fn run( + decoder: &mut WorkspaceDecoder<'_>, + compressed: &[u8], + in_chunk: usize, + out_chunk: usize, + work: &mut WorkBudget, +) -> Result, WorkspaceDecodeError> { + let mut decoded = Vec::new(); + let mut out = vec![0u8; out_chunk]; + let mut pos = 0; + loop { + let end = (pos + in_chunk).min(compressed.len()); + let last = end == compressed.len(); + let progress = decoder.step(&compressed[pos..end], last, &mut out, work)?; + decoded.extend_from_slice(&out[..progress.produced]); + match progress.state { + StepState::Done => return Ok(decoded), + StepState::NeedInput => assert_eq!(progress.consumed, end - pos), + StepState::NeedOutput => assert_eq!(progress.produced, out_chunk), + } + assert!( + progress.consumed > 0 || progress.produced > 0, + "a step that is not done must make progress" + ); + pos += progress.consumed; + } +} + +fn decode_chunked( + compressed: &[u8], + limits: DecodeLimits, + in_chunk: usize, + out_chunk: usize, + work: &mut WorkBudget, +) -> Result, WorkspaceDecodeError> { + let mut workspace = vec![0u8; workspace_bytes(limits.max_window_size).unwrap()]; + let mut decoder = WorkspaceDecoder::new(&mut workspace, limits).unwrap(); + run(&mut decoder, compressed, in_chunk, out_chunk, work) +} + +fn decode_whole(compressed: &[u8], limits: DecodeLimits) -> Result, WorkspaceDecodeError> { + decode_chunked( + compressed, + limits, + compressed.len().max(1), + 1 << 20, + &mut WorkBudget::new(UNLIMITED), + ) +} + +fn reference_decode(compressed: &[u8], capacity: usize) -> Option> { + let mut decoder = FrameDecoder::new(); + decoder.set_content_checksum(crate::decoding::ContentChecksum::Verify); + let mut target = vec![0u8; capacity]; + let n = decoder.decode_all(compressed, &mut target).ok()?; + target.truncate(n); + Some(target) +} + +/// A frame without content size or dictionary, with window `1 << (10 + exp)` +/// and the given `(block_type, payload, regenerated_size)` blocks. +fn handmade_frame(exp: u8, checksum: Option, blocks: &[(u32, &[u8], u32)]) -> Vec { + let mut frame = Vec::new(); + frame.extend_from_slice(&MAGIC_NUM.to_le_bytes()); + frame.push(if checksum.is_some() { 1 << 2 } else { 0 }); + frame.push(exp << 3); + for (i, &(block_type, payload, size)) in blocks.iter().enumerate() { + let last = u32::from(i + 1 == blocks.len()); + let header = last | (block_type << 1) | (size << 3); + frame.extend_from_slice(&header.to_le_bytes()[..3]); + frame.extend_from_slice(payload); + } + if let Some(sum) = checksum { + frame.extend_from_slice(&sum.to_le_bytes()); + } + frame +} + +fn skippable_frame(payload: &[u8]) -> Vec { + let mut frame = Vec::new(); + frame.extend_from_slice(&0x184D_2A53u32.to_le_bytes()); + frame.extend_from_slice(&(payload.len() as u32).to_le_bytes()); + frame.extend_from_slice(payload); + frame +} + +/// Every chunking of input and output reproduces the data and agrees with the +/// one-shot decoder, on compressible, incompressible and run-length input. +#[test] +fn chunked_decode_matches_frame_decoder() { + let mut mixed = text(150_000); + mixed.extend(core::iter::repeat_n(0xA5, 140_000)); + mixed.extend(noise(150_000)); + let fixtures = [text(200_000), noise(150_000), vec![0u8; 300_000], mixed]; + let levels = [ + CompressionLevel::Uncompressed, + CompressionLevel::Fastest, + CompressionLevel::Default, + CompressionLevel::Best, + ]; + let limits = DecodeLimits::new(1 << 20); + for data in &fixtures { + for level in levels { + let compressed = compress_slice_to_vec(data, level); + let reference = reference_decode(&compressed, data.len()).unwrap(); + assert_eq!(&reference, data); + for (in_chunk, out_chunk) in [(7, 13), (100, 4096), (compressed.len(), data.len())] { + let decoded = decode_chunked( + &compressed, + limits, + in_chunk, + out_chunk, + &mut WorkBudget::new(UNLIMITED), + ) + .unwrap(); + assert!( + decoded == *data, + "level {level:?} chunks {in_chunk}/{out_chunk}" + ); + } + } + } +} + +/// Compress with the window capped at `1 << window_log`, so content larger +/// than that becomes a multi-segment frame decoded through the ring. +fn compress_windowed(data: &[u8], level: i32, window_log: u32) -> Vec { + use crate::encoding::{CompressionParameters, compress_with_parameters}; + let parameters = CompressionParameters::builder(CompressionLevel::Level(level)) + .window_log(window_log) + .build() + .unwrap(); + compress_with_parameters(data, ¶meters) +} + +/// Multi-segment frames, whose ring wraps many times over a small window, and +/// single-segment frames, which decode linearly, both reproduce the data at +/// every chunking. +#[test] +fn ring_and_linear_windows_decode() { + let mut data = text(120_000); + data.extend(noise(30_000)); + data.extend(text(50_000)); + for (window_log, level) in [(10, 3), (12, 7), (17, 19)] { + let compressed = compress_windowed(&data, level, window_log); + let info = read_frame_header_info(&compressed, false).unwrap(); + assert!(info.window_size < data.len() as u64, "multi-segment frame"); + for (in_chunk, out_chunk) in [(1, 7), (999, 1_000), (compressed.len(), 1 << 16)] { + let decoded = decode_chunked( + &compressed, + DecodeLimits::new(info.window_size.max(1024)), + in_chunk, + out_chunk, + &mut WorkBudget::new(UNLIMITED), + ) + .unwrap(); + assert!( + decoded == data, + "window_log {window_log} chunks {in_chunk}/{out_chunk}" + ); + } + } + let single = compress_slice_to_vec(&data[..60_000], CompressionLevel::Default); + let info = read_frame_header_info(&single, false).unwrap(); + assert_eq!(info.window_size, 60_000, "single-segment frame"); + let decoded = decode_chunked( + &single, + DecodeLimits::new(60_000), + 333, + 77, + &mut WorkBudget::new(UNLIMITED), + ) + .unwrap(); + assert_eq!(decoded, data[..60_000]); +} + +/// Byte-at-a-time input and output, the most a header or block can be split. +#[test] +fn single_byte_steps_decode() { + let data = text(20_000); + let compressed = compress_slice_to_vec(&data, CompressionLevel::Default); + let decoded = decode_chunked( + &compressed, + DecodeLimits::new(1 << 16), + 1, + 1, + &mut WorkBudget::new(UNLIMITED), + ) + .unwrap(); + assert_eq!(decoded, data); +} + +/// Concatenated frames decode in order, skippable frames add no output, and a +/// skippable frame split across steps is still skipped whole. +#[test] +fn concatenated_and_skippable_frames() { + let a = text(10_000); + let b = noise(5_000); + let mut stream = skippable_frame(b"leading metadata"); + stream.extend(compress_slice_to_vec(&a, CompressionLevel::Fastest)); + stream.extend(skippable_frame(&[])); + stream.extend(skippable_frame(&noise(300))); + stream.extend(compress_slice_to_vec(&b, CompressionLevel::Default)); + let mut expected = a.clone(); + expected.extend_from_slice(&b); + for chunk in [1, 5, stream.len()] { + let decoded = decode_chunked( + &stream, + DecodeLimits::new(1 << 16), + chunk, + 977, + &mut WorkBudget::new(UNLIMITED), + ) + .unwrap(); + assert_eq!(decoded, expected); + } +} + +/// No input at all is an empty stream; an empty frame decodes to nothing. +#[test] +fn empty_stream_and_empty_frame() { + let limits = DecodeLimits::new(1 << 16); + assert_eq!(decode_whole(&[], limits).unwrap(), Vec::::new()); + let empty = compress_slice_to_vec(&[], CompressionLevel::Default); + assert_eq!(decode_whole(&empty, limits).unwrap(), Vec::::new()); +} + +/// Input that stops inside a frame is `Truncated` when final and a request +/// for more input when not, at every cut point. +#[test] +fn every_cut_is_truncation_when_final() { + let data = text(3_000); + let compressed = compress_slice_to_vec(&data, CompressionLevel::Default); + let limits = DecodeLimits::new(1 << 16); + let mut workspace = vec![0u8; workspace_bytes(limits.max_window_size).unwrap()]; + let mut decoder = WorkspaceDecoder::new(&mut workspace, limits).unwrap(); + let mut out = vec![0u8; data.len()]; + for cut in 1..compressed.len() { + decoder.reset(); + let mut work = WorkBudget::new(UNLIMITED); + let progress = decoder + .step(&compressed[..cut], false, &mut out, &mut work) + .unwrap(); + assert_eq!(progress.state, StepState::NeedInput, "cut {cut}"); + assert_eq!(progress.consumed, cut); + assert!(matches!( + decoder.step(&[], true, &mut out, &mut work), + Err(WorkspaceDecodeError::Truncated) + )); + } +} + +/// The exact workspace size works at every address; one byte less is refused +/// before anything is decoded, as is a window limit out of range. +#[test] +fn workspace_size_boundaries() { + let limit = 64 * 1024; + let needed = workspace_bytes(limit).unwrap(); + let data = text(50_000); + let compressed = compress_slice_to_vec(&data, CompressionLevel::Default); + let mut backing = vec![0u8; needed + TABLE_ALIGN]; + for offset in 0..TABLE_ALIGN { + let workspace = &mut backing[offset..offset + needed]; + let mut decoder = WorkspaceDecoder::new(workspace, DecodeLimits::new(limit)).unwrap(); + let decoded = run( + &mut decoder, + &compressed, + 333, + 4096, + &mut WorkBudget::new(UNLIMITED), + ) + .unwrap(); + assert_eq!(decoded, data, "offset {offset}"); + } + assert!(matches!( + WorkspaceDecoder::new(&mut backing[..needed - 1], DecodeLimits::new(limit)), + Err(WorkspaceConfigError::WorkspaceTooSmall { needed: n, provided }) + if n == needed && provided == needed - 1 + )); + for bad in [ + 0, + MIN_WINDOW_SIZE - 1, + MAX_DECODER_WINDOW_SIZE + 1, + u64::MAX, + ] { + assert_eq!( + workspace_bytes(bad), + Err(WorkspaceConfigError::WindowLimitOutOfRange { requested: bad }) + ); + } +} + +/// A frame's window is admitted up to the limit and refused one byte past it. +#[test] +fn window_limit_at_and_over() { + let payload = text(1_000); + let frame = handmade_frame(6, None, &[(0, &payload, 1_000)]); + let window = 1u64 << 16; + assert_eq!( + decode_whole(&frame, DecodeLimits::new(window)).unwrap(), + payload + ); + assert!(matches!( + decode_whole(&frame, DecodeLimits::new(window - 1)), + Err(WorkspaceDecodeError::Limit(LimitExceeded::Window { requested, limit })) + if requested == window && limit == window - 1 + )); +} + +/// Output is admitted up to the limit and refused past it, for blocks that +/// state their size (before they run) and for a frame that declares one (at +/// its header). +#[test] +fn output_limit_at_and_over() { + let payload = text(2_000); + let frame = handmade_frame(6, None, &[(0, &payload, 2_000), (1, &[7], 1_000)]); + let mut expected = payload.clone(); + expected.extend(core::iter::repeat_n(7, 1_000)); + let limits = DecodeLimits::new(1 << 16); + assert_eq!( + decode_whole(&frame, limits.with_max_output(3_000)).unwrap(), + expected + ); + assert!(matches!( + decode_whole(&frame, limits.with_max_output(2_999)), + Err(WorkspaceDecodeError::Limit(LimitExceeded::Output { + limit: 2_999 + })) + )); + + let data = text(40_000); + let compressed = compress_slice_to_vec(&data, CompressionLevel::Default); + assert_eq!( + decode_whole(&compressed, limits.with_max_output(40_000)).unwrap(), + data + ); + assert!(matches!( + decode_whole(&compressed, limits.with_max_output(39_999)), + Err(WorkspaceDecodeError::Limit(LimitExceeded::Output { .. })) + )); +} + +/// The exact work a stream costs succeeds and one unit less is refused; and +/// at any budget, no byte is produced that was not paid for first. +#[test] +fn work_budget_at_and_over() { + let mut data = text(200_000); + data.extend(noise(60_000)); + let compressed = compress_slice_to_vec(&data, CompressionLevel::Default); + let limits = DecodeLimits::new(1 << 20); + let mut work = WorkBudget::new(UNLIMITED); + decode_chunked(&compressed, limits, 4096, 4096, &mut work).unwrap(); + let cost = UNLIMITED - work.remaining(); + assert!(cost >= compressed.len() as u64 + data.len() as u64); + + let mut exact = WorkBudget::new(cost); + assert_eq!( + decode_chunked(&compressed, limits, 4096, 4096, &mut exact).unwrap(), + data + ); + assert_eq!(exact.remaining(), 0); + assert!(matches!( + decode_chunked( + &compressed, + limits, + 4096, + 4096, + &mut WorkBudget::new(cost - 1) + ), + Err(WorkspaceDecodeError::Limit(LimitExceeded::Work { .. })) + )); + + let mut workspace = vec![0u8; workspace_bytes(limits.max_window_size).unwrap()]; + let mut decoder = WorkspaceDecoder::new(&mut workspace, limits).unwrap(); + for budget in (0..cost).step_by(9_973) { + decoder.reset(); + let result = run( + &mut decoder, + &compressed, + 4096, + 4096, + &mut WorkBudget::new(budget), + ); + assert!(matches!( + result, + Err(WorkspaceDecodeError::Limit(LimitExceeded::Work { .. })) + )); + assert!(decoder.total_output() <= budget, "budget {budget}"); + } +} + +/// A block larger than its frame's block maximum is malformed, at the edge. +#[test] +fn block_size_at_and_over_the_maximum() { + let fits = noise(1_024); + let frame = handmade_frame(0, None, &[(0, &fits, 1_024)]); + let limits = DecodeLimits::new(1 << 16); + assert_eq!(decode_whole(&frame, limits).unwrap(), fits); + let over = noise(1_025); + let frame = handmade_frame(0, None, &[(0, &over, 1_025)]); + assert!(matches!( + decode_whole(&frame, limits), + Err(WorkspaceDecodeError::Malformed( + MalformedInput::BlockTooLarge { + size: 1_025, + maximum: 1_024 + } + )) + )); +} + +/// The content checksum is verified, and a wrong one is reported as such. +#[cfg(feature = "hash")] +#[test] +fn checksum_verified() { + use core::hash::Hasher; + let payload = text(5_000); + let mut hasher = twox_hash::XxHash64::with_seed(0); + hasher.write(&payload); + let sum = hasher.finish() as u32; + let limits = DecodeLimits::new(1 << 16); + let good = handmade_frame(6, Some(sum), &[(0, &payload, 5_000)]); + for chunk in [1, 3, good.len()] { + let decoded = decode_chunked(&good, limits, chunk, 64, &mut WorkBudget::new(UNLIMITED)); + assert_eq!(decoded.unwrap(), payload); + } + let bad = handmade_frame(6, Some(sum ^ 1), &[(0, &payload, 5_000)]); + assert!(matches!( + decode_whole(&bad, limits), + Err(WorkspaceDecodeError::Malformed(MalformedInput::Checksum { expected, calculated })) + if expected == sum ^ 1 && calculated == sum + )); +} + +/// A frame naming a dictionary is refused as unsupported, not as malformed. +#[test] +fn dictionary_frame_is_unsupported() { + let frame = include_bytes!("../../../dict_tests/files/ModemManager.service.zst"); + let id = read_frame_header_info(frame, false) + .unwrap() + .dictionary_id + .unwrap(); + assert!(matches!( + decode_whole(frame, DecodeLimits::new(1 << 20)), + Err(WorkspaceDecodeError::Unsupported(Unsupported::Dictionary { id: got })) if got == id + )); +} + +/// Bytes after the last frame that are not a frame are malformed; the error +/// holds until a reset, after which the decoder works again. +#[test] +fn failure_is_terminal_until_reset() { + let data = text(4_000); + let mut stream = compress_slice_to_vec(&data, CompressionLevel::Default); + stream.extend_from_slice(b"junk"); + let limits = DecodeLimits::new(1 << 16); + let mut workspace = vec![0u8; workspace_bytes(limits.max_window_size).unwrap()]; + let mut decoder = WorkspaceDecoder::new(&mut workspace, limits).unwrap(); + let mut work = WorkBudget::new(UNLIMITED); + assert!(matches!( + run(&mut decoder, &stream, 64, 64, &mut work), + Err(WorkspaceDecodeError::Malformed( + MalformedInput::FrameHeader(ReadFrameHeaderError::BadMagicNumber(_)) + )) + )); + let mut out = [0u8; 16]; + assert!(matches!( + decoder.step(&stream, true, &mut out, &mut work), + Err(WorkspaceDecodeError::Failed) + )); + decoder.reset(); + let valid = &stream[..stream.len() - 4]; + assert_eq!(run(&mut decoder, valid, 64, 64, &mut work).unwrap(), data); +} + +/// The decoder can move to another thread with its workspace, as the +/// `&mut [u8]` it borrows can; the raw window pointer must not take that away. +#[test] +fn decoder_is_send() { + fn assert_send() {} + assert_send::>(); +} + +/// Corrupting any byte of a frame is accepted or refused exactly as the +/// one-shot decoder decides, with the same output when accepted, and never +/// panics or stalls. +#[test] +fn corrupted_frames_agree_with_frame_decoder() { + let data = text(6_000); + let limits = DecodeLimits::new(1 << 20); + // A single-segment frame (linear window) and a multi-segment one whose + // 1 KiB window makes the ring wrap. + let frames = [ + compress_slice_to_vec(&data, CompressionLevel::Default), + compress_windowed(&data, 3, 10), + ]; + for compressed in &frames { + for index in 0..compressed.len() { + for flip in [0x01u8, 0x40, 0xFF] { + let mut corrupt = compressed.clone(); + corrupt[index] ^= flip; + let ours = + decode_chunked(&corrupt, limits, 17, 251, &mut WorkBudget::new(UNLIMITED)); + if matches!( + ours, + Err(WorkspaceDecodeError::Limit(_) | WorkspaceDecodeError::Unsupported(_)) + ) { + // A limit or a dictionary id the corruption introduced: the + // one-shot decoder has neither, so there is nothing to compare. + continue; + } + let theirs = reference_decode(&corrupt, 1 << 20); + assert_eq!( + ours.as_ref().ok(), + theirs.as_ref(), + "byte {index} ^ {flip:#04x}: ours {:?}", + ours.as_ref().err() + ); + } + } + } +} diff --git a/zstd/src/fse/fse_decoder.rs b/zstd/src/fse/fse_decoder.rs index a8e56adfd..8f70a2b06 100644 --- a/zstd/src/fse/fse_decoder.rs +++ b/zstd/src/fse/fse_decoder.rs @@ -3,6 +3,22 @@ use crate::cpu_kernel::CpuKernel; use crate::decoding::errors::{FSEDecoderError, FSETableError}; use alloc::vec::Vec; +/// Build-time probability list: at most 256 symbols (`u8` alphabet). +#[cfg(not(feature = "caller-workspace"))] +type ProbabilityBuf = Vec; +/// Build-time probability list: at most 256 symbols (`u8` alphabet), inline so +/// a table can live in caller-owned memory. +#[cfg(feature = "caller-workspace")] +type ProbabilityBuf = crate::common::inline_vec::InlineVec; + +/// Symbol-spread scratch: one byte per state of the largest table (`1 << 9`). +#[cfg(not(feature = "caller-workspace"))] +type SpreadBuf = Vec; +/// Symbol-spread scratch: one byte per state of the largest table (`1 << 9`), +/// inline so a table can live in caller-owned memory. +#[cfg(feature = "caller-workspace")] +type SpreadBuf = crate::common::inline_vec::InlineVec; + // Visibility = `pub` so the `pub type FSEDecoder` / `pub type // SeqFSEDecoder` aliases below don't expose a more-private struct // (compiler `private_interfaces` warning). External reachability is @@ -263,7 +279,7 @@ pub struct FSETableImpl { /// Number of live entries in `decode` (`1 << accuracy_log`). decode_len: usize, /// Reused scratch buffer for symbol spreading to avoid per-build allocations. - symbol_spread_buffer: Vec, + symbol_spread_buffer: SpreadBuf, /// The size of the table is stored in logarithm base 2 format, /// with the **size of the table** being equal to `(1 << accuracy_log)`. /// This value is used so that the decoder knows how many bits to read from the bitstream. @@ -278,7 +294,11 @@ pub struct FSETableImpl { /// /// If a symbol probability is set to `-1`, it means that the probability of a symbol /// occurring in the data is less than one. + #[cfg(not(feature = "caller-workspace"))] pub symbol_probabilities: Vec, //used while building the decode Vector + /// See the `Vec` form of this field; stored inline under `caller-workspace`. + #[cfg(feature = "caller-workspace")] + pub(crate) symbol_probabilities: ProbabilityBuf, } /// Type alias preserved for HUF-weight-stream callers and existing @@ -313,8 +333,8 @@ impl FSETableImpl { // i32 here was a heap allocation paid per fresh table (3 seq + the // HUF weight table per `DecoderScratch`), part of the per-frame fixed // cost that dominates tiny-frame decode. - symbol_probabilities: Vec::new(), - symbol_spread_buffer: Vec::new(), + symbol_probabilities: ProbabilityBuf::new(), + symbol_spread_buffer: SpreadBuf::new(), // Uninitialised — no `CAP`-entry memset; the build fills the live // span before any read (see the field doc). decode: [const { core::mem::MaybeUninit::uninit() }; CAP], @@ -326,11 +346,38 @@ impl FSETableImpl { /// Heap bytes owned by this table. The `decode` table is a fixed inline /// array (counted by `size_of`, not here); only the build-scratch vectors /// are heap-allocated. + #[cfg(not(feature = "caller-workspace"))] pub fn heap_bytes(&self) -> usize { self.symbol_spread_buffer.capacity() + self.symbol_probabilities.capacity() * core::mem::size_of::() } + /// Heap bytes owned by this table: none, the build scratch is inline. + #[cfg(feature = "caller-workspace")] + pub fn heap_bytes(&self) -> usize { + 0 + } + + /// Initialise an empty table in place, without building it on the stack: + /// the decode array and the build scratch are left uninitialised, as `new` + /// leaves them. + /// + /// # Safety + /// `this` must be valid for writes and aligned for `Self`. + #[cfg(feature = "caller-workspace")] + pub(crate) unsafe fn init_in_place(this: *mut Self, max_symbol: u8) { + use core::ptr::addr_of_mut; + // SAFETY: the caller guarantees `this` is writable and aligned; every + // field that is not `MaybeUninit` storage is written here. + unsafe { + addr_of_mut!((*this).max_symbol).write(max_symbol); + addr_of_mut!((*this).decode_len).write(0); + addr_of_mut!((*this).accuracy_log).write(0); + SpreadBuf::init_in_place(addr_of_mut!((*this).symbol_spread_buffer)); + ProbabilityBuf::init_in_place(addr_of_mut!((*this).symbol_probabilities)); + } + } + /// Live decode entries (`decode[..decode_len]`). #[inline(always)] pub fn decode(&self) -> &[E] { @@ -499,6 +546,11 @@ impl FSETableImpl { max: ENTRY_MAX_ACCURACY_LOG, }); } + // Refused before anything is stored: the probability store is bounded + // by the alphabet, and this also bounds the sum below. + if probs.len() > self.max_symbol as usize + 1 { + return Err(FSETableError::TooManySymbols { got: probs.len() }); + } // Probability sum check: `build_decoding_table` assumes the // sum of positive probabilities plus the count of `-1` // entries (each contributing one slot at the top of the @@ -559,7 +611,7 @@ impl FSETableImpl { return Err(FSETableError::ProbabilityCounterMismatch { got: probability_sum, expected_sum: table_size, - symbol_probabilities: probs.to_vec(), + symbol_probabilities: mismatch_diagnostic(probs), }); } self.symbol_probabilities.clear(); @@ -575,14 +627,12 @@ impl FSETableImpl { /// allocation for the symbol counter (stack-allocated since the /// max symbol count is bounded by `u8::MAX + 1 = 256`). /// - /// Wraps `build_decoding_table_inner` so the `symbol_spread_buffer` - /// scratch is unconditionally restored to `self` on every exit - /// path (success OR error) — otherwise an early `Err` from the - /// inner pass would drop the taken buffer and force a fresh - /// allocation on the next build. + /// Wraps `build_decoding_table_inner`, which borrows the table's fields + /// one by one so the spread scratch is lent in place rather than moved + /// out and back. /// - /// On `Err` the table is also fully `reset()` after the buffer - /// restore: the inner pass mutates `self.decode` (partial push) + /// On `Err` the table is fully `reset()`: the inner pass mutates + /// `self.decode` (partial push) /// while `self.accuracy_log` / `self.symbol_probabilities` were /// already set by the caller (`build_from_probabilities`); leaving /// that inconsistency in place would let a subsequent `init_state` @@ -595,26 +645,38 @@ impl FSETableImpl { } fn build_decoding_table_meta(&mut self, meta: SeqMeta<'_>) -> Result<(), FSETableError> { - let mut spread = core::mem::take(&mut self.symbol_spread_buffer); - let result = self.build_decoding_table_inner(&mut spread, meta); - self.symbol_spread_buffer = spread; + let result = Self::build_decoding_table_inner( + &mut self.decode, + &mut self.decode_len, + &mut self.symbol_spread_buffer, + &self.symbol_probabilities, + self.max_symbol, + self.accuracy_log, + meta, + ); if result.is_err() { self.reset(); } result } + // The table's fields, borrowed one by one; see `build_decoding_table_meta`. + #[allow(clippy::too_many_arguments)] fn build_decoding_table_inner( - &mut self, - spread: &mut Vec, + decode: &mut [core::mem::MaybeUninit; CAP], + decode_len: &mut usize, + spread: &mut SpreadBuf, + symbol_probabilities: &ProbabilityBuf, + max_symbol: u8, + accuracy_log: u8, meta: SeqMeta<'_>, ) -> Result<(), FSETableError> { - let nb_symbols = self.symbol_probabilities.len(); - if nb_symbols > self.max_symbol as usize + 1 { + let nb_symbols = symbol_probabilities.len(); + if nb_symbols > max_symbol as usize + 1 { return Err(FSETableError::TooManySymbols { got: nb_symbols }); } - let table_size = 1 << self.accuracy_log; + let table_size = 1 << accuracy_log; // The decode array is `[E; CAP]` and the fast-spread scratch is a fixed // `FSE_FAST_SPREAD_BUF`; a `table_size` above CAP would overrun both. // The wire path never exceeds CAP (sequence acc_log <= 9, HUF <= 6), but @@ -623,7 +685,7 @@ impl FSETableImpl { // oversized table as a typed error instead of panic-overrunning. if table_size > CAP { return Err(FSETableError::AccLogTooBig { - got: self.accuracy_log, + got: accuracy_log, max: CAP.trailing_zeros() as u8, }); } @@ -661,7 +723,7 @@ impl FSETableImpl { // a tighter scalar loop without the Iterator::next state // machine. The enumerate+take iterator chain was visible as // ~1.8% combined self-time on the decode flamegraph. - let probs = self.symbol_probabilities.as_slice(); + let probs = symbol_probabilities.as_slice(); let mut negative_idx = table_size; for symbol in 0..nb_symbols { let prob = probs[symbol]; @@ -752,7 +814,6 @@ impl FSETableImpl { // so `nextState == 1`, `high_bit == 1`, `nbBits == // accuracy_log`, `newState == 0` — exactly the -1 entry // shape. - let accuracy_log = self.accuracy_log; let table_size_u32 = table_size as u32; // Write the `table_size` live entries directly into the inline, // value-initialised `decode` array (no `MaybeUninit` / `set_len` @@ -767,7 +828,7 @@ impl FSETableImpl { table_size <= CAP, "FSE table_size {table_size} exceeds monomorphized CAP {CAP}", ); - self.decode_len = 0; + *decode_len = 0; // Slice index instead of `spread.iter().take(table_size)`: // if `spread.len() < table_size` (a future refactor breaking @@ -818,7 +879,7 @@ impl FSETableImpl { // is dropped here harmlessly (`MaybeUninit` has no // Drop). return Err(FSETableError::TableInvariantViolation { - prob: self.symbol_probabilities[symbol as usize], + prob: symbol_probabilities[symbol as usize], symbol, num_bits: nb, accuracy_log, @@ -840,7 +901,7 @@ impl FSETableImpl { // instead of re-walking the finished table. `from_raw` // zero-inits the meta fields, so the no-meta arms match // the post-pass results exactly. - self.decode[state_idx] = core::mem::MaybeUninit::new(match meta { + decode[state_idx] = core::mem::MaybeUninit::new(match meta { SeqMeta::None => entry, SeqMeta::Packed(packed) => { let m = packed.get(symbol as usize).copied().unwrap_or(0); @@ -858,7 +919,7 @@ impl FSETableImpl { // Commit the live span. The loop wrote every `state_idx ∈ // [0, table_size)`; an Err above returns with `decode_len` still 0. - self.decode_len = table_size; + *decode_len = table_size; Ok(()) } @@ -948,6 +1009,15 @@ impl FSETableImpl { let prob = (value as i32) - 1; + // Stop as soon as the list outgrows the alphabet, as upstream's + // `FSE_readNCount` stops at `charnum > maxSV1`, instead of letting + // a malformed run of zero probabilities grow it without bound. + let max_symbols = self.max_symbol as usize + 1; + if self.symbol_probabilities.len() >= max_symbols { + return Err(FSETableError::TooManySymbols { + got: self.symbol_probabilities.len() + 1, + }); + } self.symbol_probabilities.push(prob); if prob != 0 { if prob > 0 { @@ -962,8 +1032,11 @@ impl FSETableImpl { loop { let skip_amount = read_bits!(2) as usize; - self.symbol_probabilities - .resize(self.symbol_probabilities.len() + skip_amount, 0); + let new_len = self.symbol_probabilities.len() + skip_amount; + if new_len > max_symbols { + return Err(FSETableError::TooManySymbols { got: new_len }); + } + self.symbol_probabilities.resize(new_len, 0); if skip_amount != 3 { break; } @@ -975,14 +1048,11 @@ impl FSETableImpl { return Err(FSETableError::ProbabilityCounterMismatch { got: probability_counter, expected_sum: probability_sum, - symbol_probabilities: self.symbol_probabilities.clone(), - }); - } - if self.symbol_probabilities.len() > self.max_symbol as usize + 1 { - return Err(FSETableError::TooManySymbols { - got: self.symbol_probabilities.len(), + symbol_probabilities: mismatch_diagnostic(&self.symbol_probabilities), }); } + // Held inside the loop, before every append. + debug_assert!(self.symbol_probabilities.len() <= self.max_symbol as usize + 1); let bytes_read = if bit_pos.is_multiple_of(8) { bit_pos / 8 @@ -1232,6 +1302,19 @@ fn highest_bit_set(x: u32) -> u32 { u32::BITS - x.leading_zeros() } +/// The probabilities a sum-mismatch error carries for diagnosis. +#[cfg(not(feature = "caller-workspace"))] +fn mismatch_diagnostic(probs: &[i32]) -> Vec { + probs.to_vec() +} + +/// Empty under `caller-workspace`: copying the list would allocate on +/// malformed input, which that decoder promises never to do. +#[cfg(feature = "caller-workspace")] +fn mismatch_diagnostic(_probs: &[i32]) -> Vec { + Vec::new() +} + //utility functions for building the decoding table from probabilities /// Calculate the position of the next entry of the table given the current /// position and size of the table. diff --git a/zstd/src/fse/fse_decoder/tests.rs b/zstd/src/fse/fse_decoder/tests.rs index b93671a58..69929e405 100644 --- a/zstd/src/fse/fse_decoder/tests.rs +++ b/zstd/src/fse/fse_decoder/tests.rs @@ -105,6 +105,37 @@ fn build_from_probabilities_accepts_exact_sum() { assert!(result.is_ok(), "expected Ok for exact sum, got {result:?}"); } +/// More probabilities than the table's alphabet holds are refused before they +/// are stored: the store is bounded by the alphabet under `caller-workspace`, +/// so copying them first would panic instead of returning the typed error. +#[test] +fn build_from_probabilities_rejects_more_symbols_than_the_alphabet() { + let mut t = FSETable::new(255); + let mut probs = [0i32; 300]; + probs[0] = 16; + let result = t.build_from_probabilities(4, &probs); + assert!( + matches!(result, Err(FSETableError::TooManySymbols { got: 300 })), + "expected TooManySymbols {{ got: 300 }}, got {result:?}", + ); +} + +/// A table description whose zero-probability run outgrows the alphabet is +/// refused as soon as it does, as upstream's `FSE_readNCount` stops at +/// `charnum > maxSV1`, rather than collecting zeros until the input runs out. +/// Accuracy log 5, one zero probability, then 2-bit "three more zeros" fields +/// for the rest of the input. +#[test] +fn zero_run_past_the_alphabet_is_refused_early() { + let mut source = alloc::vec![0x10u8, 0xFE]; + source.extend(core::iter::repeat_n(0xFF, 200)); + let mut table = SeqFSETable::new(35); + match table.build_decoder(&source, 9) { + Err(FSETableError::TooManySymbols { got }) => assert!(got > 36 && got <= 39), + other => panic!("expected TooManySymbols, got {other:?}"), + } +} + /// The decoder accepts tables with `accuracy_log` up to /// `ENTRY_MAX_ACCURACY_LOG = 16`, but the encoder table builder indexes a /// fixed 512-entry (`1 << 9`) scratch. A decoder table with `accuracy_log > @@ -117,7 +148,8 @@ fn to_encoder_table_rejects_accuracy_log_above_nine() { // A table with accuracy_log 10 (1 << 10 = 1024 entries) and a plausible // probability distribution summing to 1024. t.accuracy_log = 10; - t.symbol_probabilities = alloc::vec![512, 512]; + t.symbol_probabilities.clear(); + t.symbol_probabilities.extend_from_slice(&[512, 512]); assert!( t.to_encoder_table().is_none(), "accuracy_log > 9 must not be reusable as an encoder table" diff --git a/zstd/src/huff0/huff0_decoder.rs b/zstd/src/huff0/huff0_decoder.rs index 6f8356500..f8421cc73 100644 --- a/zstd/src/huff0/huff0_decoder.rs +++ b/zstd/src/huff0/huff0_decoder.rs @@ -8,6 +8,22 @@ use alloc::vec::Vec; /// The Zstandard specification limits the maximum length of a code to 11 bits. pub(crate) const MAX_MAX_NUM_BITS: u8 = 11; +/// Decode lookup table, `1 << max_num_bits` packed entries. +#[cfg(not(feature = "caller-workspace"))] +type DecodeTableBuf = Vec; +/// Decode lookup table, at most `1 << 11` packed entries, inline so the table +/// can live in caller-owned memory. +#[cfg(feature = "caller-workspace")] +type DecodeTableBuf = crate::common::inline_vec::InlineVec; + +/// Per-symbol weights or code lengths: at most 256 symbols. +#[cfg(not(feature = "caller-workspace"))] +type SymbolBuf = Vec; +/// Per-symbol weights or code lengths: at most 256 symbols, inline so the +/// table can live in caller-owned memory. +#[cfg(feature = "caller-workspace")] +type SymbolBuf = crate::common::inline_vec::InlineVec; + pub struct HuffmanDecoder<'table> { table: &'table HuffmanTable, /// State is used to index into the table. @@ -139,19 +155,19 @@ pub struct HuffmanTable { /// in the low byte lets the 4-stream decoder shift directly by the entry /// (modulo 64), without extracting it on the dependent bitstream path. /// Entries remain two bytes: at most 4 KiB for the 11-bit table. - pub(crate) packed_decode: Vec, + pub(crate) packed_decode: DecodeTableBuf, /// The weight of a symbol is the number of occurrences in a table. /// This value is used in constructing a binary tree referred to as /// a Huffman tree. Once this tree is constructed, it can be used to build the /// lookup table - weights: Vec, + weights: SymbolBuf, /// The maximum size in bits a prefix code in the encoded data can be. /// This value is used so that the decoder knows how many bits /// to read from the bitstream before checking the table. This /// value must be 11 or lower. pub max_num_bits: u8, state_mask: u64, - bits: Vec, + bits: SymbolBuf, bit_ranks: [u32; (MAX_MAX_NUM_BITS as usize) + 1], /// Running `sum(1 << (w - 1))` accumulated WHILE weights decode /// (upstream zstd `HUF_readStats` keeps the same running stats), replacing a @@ -171,6 +187,7 @@ pub struct HuffmanTable { impl HuffmanTable { /// Heap bytes owned by this table (decode lookup table plus the /// weight-decoding scratch vectors and the weight-stream FSE table). + #[cfg(not(feature = "caller-workspace"))] pub fn heap_bytes(&self) -> usize { self.packed_decode.capacity() * core::mem::size_of::() + self.weights.capacity() @@ -178,15 +195,44 @@ impl HuffmanTable { + self.fse_table.heap_bytes() } + /// Heap bytes owned by this table: none, every buffer is inline. + #[cfg(feature = "caller-workspace")] + pub fn heap_bytes(&self) -> usize { + 0 + } + + /// Initialise an empty table in place, without building it on the stack. + /// + /// # Safety + /// `this` must be valid for writes and aligned for `Self`. + #[cfg(feature = "caller-workspace")] + pub(crate) unsafe fn init_in_place(this: *mut Self) { + use core::ptr::addr_of_mut; + // SAFETY: the caller guarantees `this` is writable and aligned; every + // field that is not `MaybeUninit` storage is written here. + unsafe { + DecodeTableBuf::init_in_place(addr_of_mut!((*this).packed_decode)); + SymbolBuf::init_in_place(addr_of_mut!((*this).weights)); + addr_of_mut!((*this).max_num_bits).write(0); + addr_of_mut!((*this).state_mask).write(0); + SymbolBuf::init_in_place(addr_of_mut!((*this).bits)); + addr_of_mut!((*this).bit_ranks).write([0; (MAX_MAX_NUM_BITS as usize) + 1]); + addr_of_mut!((*this).weight_sum).write(0); + addr_of_mut!((*this).weight_rank_count).write([0; (MAX_MAX_NUM_BITS as usize) + 1]); + addr_of_mut!((*this).last_weight).write(0); + FSETable::init_in_place(addr_of_mut!((*this).fse_table), 255); + } + } + /// Create a new, empty table. pub fn new() -> HuffmanTable { HuffmanTable { - packed_decode: Vec::new(), + packed_decode: DecodeTableBuf::new(), - weights: Vec::with_capacity(256), + weights: SymbolBuf::with_capacity(256), max_num_bits: 0, state_mask: 0, - bits: Vec::with_capacity(256), + bits: SymbolBuf::with_capacity(256), bit_ranks: [0; (MAX_MAX_NUM_BITS as usize) + 1], weight_sum: 0, weight_rank_count: [0; (MAX_MAX_NUM_BITS as usize) + 1], @@ -456,7 +502,7 @@ impl HuffmanTable { // the per-index parity branch defeated unrolling and read // every byte twice. The running stats accumulate alongside // (upstream zstd `HUF_readStats` shape). - let mut push_weight = |w: u8, weights: &mut Vec| -> Result<(), err> { + let mut push_weight = |w: u8, weights: &mut SymbolBuf| -> Result<(), err> { if w > MAX_MAX_NUM_BITS { return Err(err::WeightBiggerThanMaxNumBits { got: w }); } diff --git a/zstd/src/huff0/huff0_decoder/tests.rs b/zstd/src/huff0/huff0_decoder/tests.rs index c39fbd0bd..46a522c50 100644 --- a/zstd/src/huff0/huff0_decoder/tests.rs +++ b/zstd/src/huff0/huff0_decoder/tests.rs @@ -1,21 +1,21 @@ use super::*; -use alloc::vec; fn test_table() -> HuffmanTable { // Packed `num_bits | (symbol << 8)` per state index (upstream zstd `HUF_DEltX1`). - let packed_decode = vec![ + let mut packed_decode = DecodeTableBuf::new(); + packed_decode.extend_from_slice(&[ 1u16 | (u16::from(b'A') << 8), 2u16 | (u16::from(b'B') << 8), 1u16 | (u16::from(b'C') << 8), 2u16 | (u16::from(b'D') << 8), - ]; + ]); HuffmanTable { packed_decode, - weights: Vec::new(), + weights: SymbolBuf::new(), max_num_bits: 2, state_mask: 0b11, - bits: Vec::new(), + bits: SymbolBuf::new(), bit_ranks: [0; (MAX_MAX_NUM_BITS as usize) + 1], weight_sum: 0, weight_rank_count: [0; (MAX_MAX_NUM_BITS as usize) + 1], diff --git a/zstd/src/lib.rs b/zstd/src/lib.rs index d4bdd1aba..9b555adba 100644 --- a/zstd/src/lib.rs +++ b/zstd/src/lib.rs @@ -10,6 +10,9 @@ //! compression levels ([`encoding::CompressionLevel`]). //! * [`dictionary`] (feature `dict-builder`) — COVER / FastCOVER training //! plus raw-to-finalized dictionary helpers. +//! * `decoding::workspace` (feature `caller-workspace`): streaming decoder +//! that runs entirely in caller-owned memory, never allocating, with +//! caller-set limits on window, output and work. //! //! No FFI, no cmake, no system zstd. `no_std` builds are supported by //! disabling the default `std` feature. diff --git a/zstd/tests/workspace_no_alloc.rs b/zstd/tests/workspace_no_alloc.rs new file mode 100644 index 000000000..5c2ccc2f8 --- /dev/null +++ b/zstd/tests/workspace_no_alloc.rs @@ -0,0 +1,285 @@ +//! The caller-workspace decoder never allocates: from construction through +//! every block type, table changes, limits, malformed and truncated input, +//! reset and drop. A counting global allocator, armed on this thread only +//! around the decoder's use, proves it. + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::cell::Cell; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use structured_zstd::decoding::workspace::{ + DecodeLimits, LimitExceeded, StepState, WorkBudget, WorkspaceDecodeError, WorkspaceDecoder, + workspace_bytes, +}; +use structured_zstd::encoding::{ + CompressionLevel, CompressionParameters, compress_slice_to_vec, compress_with_parameters, +}; + +struct CountingAllocator; + +static ALLOCATIONS: AtomicUsize = AtomicUsize::new(0); + +thread_local! { + static ARMED: Cell = const { Cell::new(false) }; +} + +fn note_allocation() { + if ARMED.with(Cell::get) { + ALLOCATIONS.fetch_add(1, Ordering::Relaxed); + } +} + +// SAFETY: forwards every call to `System` unchanged; it only counts. +unsafe impl GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + note_allocation(); + unsafe { System.alloc(layout) } + } + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + note_allocation(); + unsafe { System.alloc_zeroed(layout) } + } + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + note_allocation(); + unsafe { System.realloc(ptr, layout, new_size) } + } + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { System.dealloc(ptr, layout) } + } +} + +#[global_allocator] +static GLOBAL: CountingAllocator = CountingAllocator; + +/// Run `f` and require that it allocated nothing on this thread. +fn without_allocating(what: &str, f: impl FnOnce() -> R) -> R { + ALLOCATIONS.store(0, Ordering::Relaxed); + ARMED.with(|armed| armed.set(true)); + let result = f(); + ARMED.with(|armed| armed.set(false)); + let count = ALLOCATIONS.load(Ordering::Relaxed); + assert_eq!(count, 0, "{what}: {count} allocations"); + result +} + +fn text(len: usize) -> Vec { + let words = ["alpha ", "beta ", "gamma ", "delta ", "manifest ", "chunk "]; + let mut out = Vec::with_capacity(len); + let mut i = 0usize; + while out.len() < len { + out.extend_from_slice(words[(i * 7 + i / 13) % words.len()].as_bytes()); + out.extend_from_slice(&(i as u32).to_le_bytes()[..i % 3]); + i += 1; + } + out.truncate(len); + out +} + +fn noise(len: usize) -> Vec { + let mut x = 0x2545_F491_4F6C_DD1Du64; + (0..len) + .map(|_| { + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + x as u8 + }) + .collect() +} + +/// Decode `stream` in `chunk`-byte steps through a fixed output buffer, +/// comparing each delivered piece against `expected` in place. +fn decode_and_compare( + decoder: &mut WorkspaceDecoder<'_>, + stream: &[u8], + expected: &[u8], + chunk: usize, + work: &mut WorkBudget, +) -> Result<(), WorkspaceDecodeError> { + let mut out = [0u8; 1500]; + let mut pos = 0; + let mut matched = 0; + loop { + let end = (pos + chunk).min(stream.len()); + let progress = decoder.step(&stream[pos..end], end == stream.len(), &mut out, work)?; + let piece = &out[..progress.produced]; + assert_eq!(piece, &expected[matched..matched + piece.len()]); + matched += piece.len(); + pos += progress.consumed; + if progress.state == StepState::Done { + assert_eq!(matched, expected.len()); + return Ok(()); + } + } +} + +/// The guard sees an allocation, so a clean run of the test below means +/// something. +#[test] +#[should_panic(expected = "1 allocations")] +fn guard_counts_allocations() { + without_allocating("a Vec", || std::hint::black_box(vec![1u8; 16])); +} + +#[test] +fn workspace_decoder_never_allocates() { + let mut mixed = text(200_000); + mixed.extend(noise(100_000)); + mixed.extend(std::iter::repeat_n(0x5A, 150_000)); + mixed.extend(text(60_000)); + let mut streams = Vec::new(); + for level in [ + CompressionLevel::Uncompressed, + CompressionLevel::Fastest, + CompressionLevel::Default, + CompressionLevel::Best, + CompressionLevel::Level(19), + ] { + streams.push((compress_slice_to_vec(&mixed, level), mixed.clone())); + } + // Multi-segment with a 1 KiB window: decoded through a ring that wraps. + let ring_parameters = CompressionParameters::builder(CompressionLevel::Level(3)) + .window_log(10) + .build() + .unwrap(); + streams.push(( + compress_with_parameters(&mixed, &ring_parameters), + mixed.clone(), + )); + let small = text(9_000); + let mut concatenated = compress_slice_to_vec(&small, CompressionLevel::Default); + concatenated.extend_from_slice(&0x184D_2A51u32.to_le_bytes()); + concatenated.extend_from_slice(&3u32.to_le_bytes()); + concatenated.extend_from_slice(b"pad"); + concatenated.extend(compress_slice_to_vec( + &noise(4_000), + CompressionLevel::Fastest, + )); + let mut concatenated_expected = small.clone(); + concatenated_expected.extend(noise(4_000)); + streams.push((concatenated, concatenated_expected)); + streams.push(( + compress_slice_to_vec(&[], CompressionLevel::Default), + Vec::new(), + )); + + let corrupt: Vec> = { + let frame = compress_slice_to_vec(&small, CompressionLevel::Default); + (0..frame.len()) + .flat_map(|i| { + [0x01u8, 0x80, 0xFF].map(|flip| { + let mut c = frame.clone(); + c[i] ^= flip; + c + }) + }) + .collect() + }; + let truncated = compress_slice_to_vec(&small, CompressionLevel::Best); + let dictionary_frame = include_bytes!("../dict_tests/files/ModemManager.service.zst"); + + let limits = DecodeLimits::new(1 << 20); + let mut workspace = vec![0u8; workspace_bytes(limits.max_window_size).unwrap()]; + let mut small_workspace = vec![0u8; workspace_bytes(4 << 10).unwrap()]; + + without_allocating("construction, every stream, failures, reset, drop", || { + let mut decoder = WorkspaceDecoder::new(&mut workspace, limits).unwrap(); + for chunk in [1, 977, usize::MAX] { + for (stream, expected) in &streams { + decoder.reset(); + decode_and_compare( + &mut decoder, + stream, + expected, + chunk.min(stream.len().max(1)), + &mut WorkBudget::new(u64::MAX), + ) + .unwrap(); + } + } + + for c in &corrupt { + decoder.reset(); + let mut out = [0u8; 700]; + let mut work = WorkBudget::new(u64::MAX); + let mut pos = 0; + while let Ok(progress) = decoder.step(&c[pos..], true, &mut out, &mut work) { + pos += progress.consumed; + if progress.state == StepState::Done { + break; + } + } + } + + for cut in 1..truncated.len() { + decoder.reset(); + let mut out = [0u8; 700]; + let mut work = WorkBudget::new(u64::MAX); + let mut pos = 0; + let error = loop { + match decoder.step(&truncated[pos..cut], true, &mut out, &mut work) { + Ok(progress) => pos += progress.consumed, + Err(e) => break e, + } + }; + assert!(matches!(error, WorkspaceDecodeError::Truncated)); + } + + decoder.reset(); + let mut out = [0u8; 64]; + assert!(matches!( + decoder.step( + dictionary_frame, + true, + &mut out, + &mut WorkBudget::new(u64::MAX) + ), + Err(WorkspaceDecodeError::Unsupported(_)) + )); + + let (stream, expected) = &streams[2]; + decoder.reset(); + assert!(matches!( + decode_and_compare( + &mut decoder, + stream, + expected, + 4096, + &mut WorkBudget::new(50_000) + ), + Err(WorkspaceDecodeError::Limit(LimitExceeded::Work { .. })) + )); + drop(decoder); + + let mut capped = WorkspaceDecoder::new( + &mut workspace, + limits.with_max_output(expected.len() as u64 - 1), + ) + .unwrap(); + assert!(matches!( + decode_and_compare( + &mut capped, + stream, + expected, + 4096, + &mut WorkBudget::new(u64::MAX) + ), + Err(WorkspaceDecodeError::Limit(LimitExceeded::Output { .. })) + )); + drop(capped); + + let mut narrow = + WorkspaceDecoder::new(&mut small_workspace, DecodeLimits::new(4 << 10)).unwrap(); + assert!(matches!( + decode_and_compare( + &mut narrow, + stream, + expected, + 4096, + &mut WorkBudget::new(u64::MAX) + ), + Err(WorkspaceDecodeError::Limit(LimitExceeded::Window { .. })) + )); + drop(narrow); + }); +}