trg/

inflate.rs

//! Choosing between streaming and whole-buffer inflate.
//!
//! Whole-buffer (libdeflate) allocates the full decompressed size; streaming
//! (flate2 on zlib-ng) stays around 10 MB. Buffering was once measured ~14x
//! cheaper on CPU, against the gzip binary and before zlib-ng. Re-measured on
//! 2026-10-23 (bench/inflate_modes.py) it never used less total CPU than
//! streaming, so `auto` streams and `--inflate=buffer` is an explicit opt-in.
//! The buffered path reads the gzip ISIZE trailer as a hint, but correctness
//! comes from it growing and finally falling back, not from trusting the hint.
//!
//! `--inflate-budget` is a **hard ceiling in every mode**, not only `auto`:
//! explicit `--inflate=buffer` picks the strategy, never a licence to exceed
//! the budget. Two terms have to fit inside it — the whole compressed file
//! (`inflate_buffered` must `fs::read` it regardless of the eventual
//! decompressed size) and the decompressed output — so both the initial
//! capacity and the retry ceiling are derived from the budget, and a
//! compressed file bigger than the budget is refused before it is ever read
//! into memory.

use std::io::{self, Read, Seek, SeekFrom};
use std::path::Path;

use crate::cli::InflateMode;

/// Below this, buffering buys nothing over streaming, so it is the floor for
/// both the initial capacity and the effective ceiling — even under a
/// stingy or zero `--inflate-budget`.
const MIN_CAPACITY: u64 = 1 << 16; // 64 KiB

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Strategy {
    Stream,
    Buffer { capacity: usize, ceiling: usize },
}

/// Does this file begin with the gzip magic? The buffered path is gzip-only.
fn is_gzip(path: &Path) -> bool {
    use std::io::Read;
    let Ok(mut f) = std::fs::File::open(path) else { return false };
    let mut m = [0u8; 2];
    f.read_exact(&mut m).is_ok() && m == [0x1f, 0x8b]
}

/// Last four bytes of a gzip stream: uncompressed size modulo 2^32.
pub fn isize_hint(path: &Path) -> Option<u64> {
    let mut f = std::fs::File::open(path).ok()?;
    let len = f.metadata().ok()?.len();
    if len < 18 {
        return None; // shorter than the smallest possible gzip member
    }
    f.seek(SeekFrom::End(-4)).ok()?;
    let mut b = [0u8; 4];
    f.read_exact(&mut b).ok()?;
    Some(u32::from_le_bytes(b) as u64)
}

/// `budget`, floored to `MIN_CAPACITY`. Applied uniformly to the compressed-
/// size guard and to the derived capacity/ceiling, so a zero or tiny budget
/// still leaves room for genuinely small archives to use the fast path
/// instead of being refused outright.
fn effective_ceiling(budget: u64) -> u64 {
    budget.max(MIN_CAPACITY)
}

pub fn choose(mode: InflateMode, path: &Path, budget: u64) -> Strategy {
    match mode {
        InflateMode::Stream => Strategy::Stream,
        // The whole-buffer path is libdeflate, which is gzip-only, and
        // `isize_hint` reads gzip's 4-byte trailer. Any other compressed format
        // streams — which is the safe direction: it costs a slower path, never
        // a wrong answer.
        InflateMode::Buffer if !is_gzip(path) => Strategy::Stream,
        InflateMode::Buffer => buffer_within_budget(path, budget, isize_hint(path)),
        // Streaming, always. `auto` used to buffer archives whose trailer
        // fitted the budget, from a measurement that predates zlib-ng; with
        // zlib-ng streaming, the buffered path never used less total CPU on
        // any machine measured on 2026-10-23, and on archives inflating to
        // 42 MB it used 8% (x86-64) to 59% (arm64) more, at 3.6x the memory.
        // The mode stays a separate variant so the default can change again
        // on new evidence without changing the CLI.
        InflateMode::Auto => Strategy::Stream,
    }
}

/// Build a `Buffer` strategy whose capacity and ceiling are both derived from
/// `budget`, falling back to `Stream` when the compressed file itself would
/// already blow the budget — `inflate_buffered` reads that file whole
/// (`fs::read`) no matter how the decompressed side turns out, so it is the
/// other unbounded term and must be checked before anything is read into
/// memory.
fn buffer_within_budget(path: &Path, budget: u64, hint: Option<u64>) -> Strategy {
    let ceiling = effective_ceiling(budget);
    let compressed = std::fs::metadata(path).map(|m| m.len()).unwrap_or(u64::MAX);
    if compressed > ceiling {
        return Strategy::Stream;
    }
    let capacity = hint.unwrap_or(1 << 20).min(budget).max(MIN_CAPACITY);
    Strategy::Buffer { capacity: capacity as usize, ceiling: ceiling as usize }
}

/// Whole-buffer inflate. Doubles the output buffer and retries while the
/// decompressor reports insufficient space, up to `ceiling` bytes — an
/// absolute byte count derived from `--inflate-budget`, not a multiplier of
/// `capacity`. A lying ISIZE trailer can make this retry a few times, but it
/// can never make it allocate past the budget the operator set: past
/// `ceiling` this returns `Err` instead of growing further, and the caller
/// falls back to streaming.
pub fn inflate_buffered(path: &Path, capacity: usize, ceiling: usize) -> io::Result<Vec<u8>> {
    let raw = std::fs::read(path)?;
    inflate_bytes(&raw, path, capacity, ceiling)
}

/// The retry loop of [`inflate_buffered`], separated from the read so that a
/// caller which has already read the compressed bytes — `archive.rs` does,
/// through a file handle it has run `resource::prepare` on — does not have to
/// open and read the file a second time. `path` is used only for messages.
pub fn inflate_bytes(
    raw: &[u8],
    path: &Path,
    capacity: usize,
    ceiling: usize,
) -> io::Result<Vec<u8>> {
    let mut cap = capacity.max(MIN_CAPACITY as usize);
    let ceiling = ceiling.max(cap);

    // The trailer is the last member's size mod 2^32, so the output is at
    // least that long. Over the ceiling the attempt is certain to fail, and
    // making it anyway costs a full decompression before the stream fallback
    // decompresses everything again.
    let trailer = raw.last_chunk::<4>().map_or(0, |b| u32::from_le_bytes(*b) as usize);
    if trailer > ceiling {
        return Err(io::Error::other(format!(
            "the {ceiling}-byte inflate budget would be exceeded: \
             the archive inflates to at least {trailer} bytes"
        )));
    }

    loop {
        let mut d = libdeflater::Decompressor::new();
        let mut out = vec![0u8; cap];
        match d.gzip_decompress(raw, &mut out) {
            Ok(n) => {
                out.truncate(n);
                return Ok(out);
            }
            Err(libdeflater::DecompressionError::InsufficientSpace) if cap < ceiling => {
                cap = cap.saturating_mul(4).min(ceiling);
            }
            Err(libdeflater::DecompressionError::InsufficientSpace) => {
                // Not a bug: an undersized budget for this archive is an
                // expected, silent-fallback case, not an error condition. The
                // caller (`archive::open_decoded_with`) already frames this
                // as a fallback, so this message states only the reason.
                return Err(io::Error::other(format!(
                    "the {ceiling}-byte inflate budget would be exceeded"
                )));
            }
            Err(e) => {
                return Err(io::Error::other(format!(
                    "whole-buffer inflate failed for {}: {e:?}",
                    path.display()
                )));
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;

    fn gz(body: &[u8]) -> Vec<u8> {
        let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
        e.write_all(body).unwrap();
        e.finish().unwrap()
    }

    /// Deliberately poorly-compressible bytes, so a gzip of `n` of these
    /// stays close to `n` bytes on disk instead of collapsing via run-length
    /// matches the way `vec![b'a'; n]` would.
    fn incompressible(n: usize) -> Vec<u8> {
        let mut s: u64 = 0x2545F4914F6CDD1D;
        (0..n)
            .map(|_| {
                s = s.wrapping_mul(6364136223846793005).wrapping_add(1);
                (s >> 33) as u8
            })
            .collect()
    }

    fn tmp(bytes: &[u8]) -> (tempfile::TempDir, std::path::PathBuf) {
        let d = tempfile::tempdir().unwrap();
        let p = d.path().join("x.gz");
        std::fs::write(&p, bytes).unwrap();
        (d, p)
    }

    #[test]
    fn reads_the_isize_trailer() {
        let body = vec![b'a'; 5000];
        let (_d, p) = tmp(&gz(&body));
        assert_eq!(isize_hint(&p), Some(5000));
    }

    #[test]
    fn isize_of_a_too_short_file_is_none() {
        let (_d, p) = tmp(b"tiny");
        assert_eq!(isize_hint(&p), None);
    }

    /// Measured 2026-10-23 on x86-64, arm64 and the neh-app container: with
    /// zlib-ng streaming, the buffered path never used less total CPU, and on
    /// archives inflating to 42 MB it used 8-59% more. So `auto` streams even
    /// when the archive would fit the budget.
    #[test]
    fn auto_streams_even_when_the_hint_fits_the_budget() {
        let (_d, p) = tmp(&gz(&vec![b'a'; 5000]));
        assert!(matches!(
            choose(crate::cli::InflateMode::Auto, &p, 64 * 1024 * 1024),
            Strategy::Stream
        ));
    }

    #[test]
    fn explicit_modes_override_the_heuristic() {
        let (_d, p) = tmp(&gz(&vec![b'a'; 5000]));
        assert!(matches!(choose(crate::cli::InflateMode::Stream, &p, u64::MAX), Strategy::Stream));
        assert!(matches!(choose(crate::cli::InflateMode::Buffer, &p, 0), Strategy::Buffer { .. }));
    }

    /// The trailer is a lower bound on the output, so a trailer over the
    /// ceiling means the attempt cannot fit: refuse before decompressing
    /// anything. The body here is corrupt, so only a refusal that never
    /// touched it can name the budget rather than the bad data.
    #[test]
    fn a_trailer_over_the_ceiling_is_refused_without_decompressing() {
        let mut bytes = gz(&incompressible(200_000));
        bytes[12] ^= 0xff;
        bytes[13] ^= 0xff;
        let err = inflate_bytes(&bytes, Path::new("x.gz"), 1 << 16, 100_000).unwrap_err();
        assert!(err.to_string().contains("budget"), "decompressed despite the trailer: {err}");
    }

    #[test]
    fn buffered_inflate_grows_when_the_hint_was_too_small() {
        // ISIZE lies (concatenated stream); inflate must still succeed. The
        // body must exceed the 65_536-byte capacity floor in
        // `inflate_buffered`, or the grow/retry loop never actually runs and
        // this test would pass without exercising it.
        let (_d, p) = tmp(&gz(&vec![b'z'; 200_000]));
        let out = inflate_buffered(&p, 16, 1 << 20).unwrap();
        assert_eq!(out.len(), 200_000);
    }

    // --- Fix round 1: --inflate-budget is a hard ceiling in every mode ---

    #[test]
    fn buffer_mode_capacity_and_ceiling_never_exceed_the_budget() {
        let (_d, p) = tmp(&gz(&vec![b'a'; 500_000]));
        let budget = 10_000u64;
        match choose(InflateMode::Buffer, &p, budget) {
            Strategy::Buffer { capacity, ceiling } => {
                let bound = budget.max(1 << 16);
                assert!(capacity as u64 <= bound, "capacity {capacity} exceeds budget-derived bound {bound}");
                assert!(ceiling as u64 <= bound, "ceiling {ceiling} exceeds budget-derived bound {bound}");
            }
            Strategy::Stream => {} // also an acceptable outcome per the compressed-size guard
        }
    }

    #[test]
    fn a_compressed_file_larger_than_the_budget_selects_stream() {
        // Explicit `buffer` mode never runs the `auto` heuristic's
        // hint-vs-compressed "implausible" check, so this is the case that
        // isolates the compressed-file-size guard itself: a large,
        // incompressible body keeps the on-disk gzip well above the budget's
        // effective ceiling even though the mode is forced.
        let (_d, p) = tmp(&gz(&incompressible(200_000)));
        assert!(matches!(choose(InflateMode::Buffer, &p, 1024), Strategy::Stream));
    }

    #[test]
    fn inflate_budget_zero_still_means_never_buffer_in_auto() {
        let (_d, p) = tmp(&gz(&vec![b'a'; 5000]));
        assert!(matches!(choose(InflateMode::Auto, &p, 0), Strategy::Stream));
    }
}