primase-foundations 0.3.0

//! The compression codecs games store their data in, behind two traits: [`Decompressor`] and [`Compressor`].
//!
//! | codec | type | decode | encode |
//! | --- | --- | --- | --- |
//! | LZ4 raw block | [`Lz4`] | yes | yes |
//! | raw deflate (RFC 1951) | [`Deflate`] | yes | yes |
//! | zlib (RFC 1950) | [`Zlib`] | yes | yes |
//! | zlib with a missing end marker or checksum | [`ZlibUnterminated`] | yes | no |
//! | Zstandard frame (RFC 8878) | [`Zstd`] | yes | no |
//! | LZO2A | [`Lzo2a`] | yes | no |
//! | Oodle (Kraken, Mermaid, Selkie, Leviathan) | [`oodle::Oodle`] | yes | Kraken, with the game's library only |
//!
//! The adapters are strict. A decode takes the exact output size the container states and refuses a stream that
//! produces more or fewer bytes, has input left over, ends early, or (where the format carries one) fails its
//! checksum; the error names the codec and the input offset. There is no format autodetection: the caller says which
//! codec a stream is. A size claimed by a container is untrusted, so what a decode allocates for it is bounded by the
//! input: LZ4, deflate, zlib and LZO2A refuse, before allocating anything, a claim above their format's ratio
//! ([`LZ4_MAX_RATIO`], [`DEFLATE_MAX_RATIO`], [`LZO2A_MAX_RATIO`]) times the input's length, and zstd grows its output
//! block by block as the stream produces it. Oodle's ratio is not bounded here: its decode allocates the claimed size
//! up front (fallibly, as [`Error::Limit`]), so bound the size before asking. The tighter bound a container knows is
//! the caller's to apply.
//!
//! # Oodle
//!
//! Oodle is a commercial codec and its library is not part of this crate. [`oodle::Oodle::load`] loads the game's own
//! copy of the library from a path the caller gives (an `unsafe` call, because loading a native library runs its
//! initialization). Off Windows, [`oodle::Oodle::open`] with no path decodes with the pure-Rust oozextract crate
//! instead, which cannot compress. On Windows there is no built-in decoder, so a path is required.
//!
//! The library and oozextract decode a stream without saying whether bytes were left after its last block, so
//! [`oodle::Oodle`]'s decode does not check that. [`oodle_framing::walk`] does, by reading the stream's block and
//! quantum headers without decoding it: a container that knows its payload's exact extent calls it first.
#![deny(clippy::disallowed_methods, clippy::disallowed_types)]

use std::io::Write;

mod lzo2a;
pub use lzo2a::Lzo2a;
#[allow(unsafe_code)]
pub mod oodle;
pub mod oodle_framing;
#[cfg(any(feature = "reference-decoder", not(windows)))]
pub mod reference;

/// Why a stream was refused or a codec could not be used.
#[derive(Debug, thiserror::Error)]
pub enum Error {
    /// The stream is malformed or not what the adapter accepts.
    #[error("{codec} at input offset {offset}: {detail}")]
    Invalid {
        /// The codec that refused it.
        codec: &'static str,
        /// Where in the input the problem was found.
        offset: usize,
        /// What was wrong.
        detail: String,
    },
    /// The stream decoded to a different number of bytes than expected (`actual` is the amount reached when the
    /// decoder stopped, which for an overrun can be just past `expected`).
    #[error("output size {actual}, expected {expected}")]
    Size {
        /// The size the caller stated.
        expected: usize,
        /// The size the stream produced.
        actual: usize,
    },
    /// The output buffer could not be allocated, or a size does not fit the platform's integers.
    #[error("allocation/size limit exceeded")]
    Limit,
    /// An I/O error from an encoder, or from resolving a library path.
    #[error("I/O: {0}")]
    Io(#[from] std::io::Error),
    /// The Oodle library could not be loaded or lacks an export.
    #[error("Oodle library: {0}")]
    Library(#[from] libloading::Error),
}
/// A result with this module's [`Error`].
pub type Result<T> = std::result::Result<T, Error>;
pub(crate) fn invalid(codec: &'static str, offset: usize, detail: impl ToString) -> Error {
    Error::Invalid {
        codec,
        offset,
        detail: detail.to_string(),
    }
}
/// Refuses, before anything is allocated for it, a claimed output of `expected` bytes from `input_len` bytes of a format
/// that cannot expand more than `max_ratio` to 1.
pub(crate) fn check_claim(
    codec: &'static str,
    input_len: usize,
    expected: usize,
    max_ratio: usize,
) -> Result<()> {
    let most = input_len.saturating_mul(max_ratio);
    if expected > most {
        return Err(invalid(
            codec,
            0,
            format!(
                "the container claims {expected} output bytes, but {input_len} input bytes make at most {most} ({max_ratio}:1)"
            ),
        ));
    }
    Ok(())
}
/// How many input bytes the decompressor has consumed (a count past `usize` is beyond any input).
fn total_in(d: &zlib_rs::Inflate) -> usize {
    usize::try_from(d.total_in()).unwrap_or(usize::MAX)
}
/// What the inflater says of a stream it refused, in the words this module has always used for it.
fn inflate_error(e: zlib_rs::InflateError) -> &'static str {
    match e {
        zlib_rs::InflateError::NeedDict { .. } => "requires a dictionary",
        zlib_rs::InflateError::StreamError
        | zlib_rs::InflateError::DataError
        | zlib_rs::InflateError::MemError => "deflate decompression error",
    }
}
/// The window an inflater is made with: the largest deflate has, which every stream fits.
const WINDOW_BITS: u8 = 15;
/// What a raw DEFLATE stream's window holds before its first byte: zeros. Some streams copy from before their start
/// (GTA V Enhanced stores one: `x64a.rpf\textures\parachute_decals.ytd`, whose matches reach back past byte 0), which
/// a strict inflater refuses as "invalid distance too far back"; read against a zero-filled window that stream decodes
/// to exactly what flate2's miniz_oxide made of it, and to what zlib 1.3.1 makes of it with this window as its
/// dictionary. A stream that copies only from its own output decodes the same with or without the window.
static ZERO_WINDOW: [u8; 1 << WINDOW_BITS] = [0; 1 << WINDOW_BITS];
/// An inflater for a raw DEFLATE stream (`zlib` false: its window starts as [`ZERO_WINDOW`]) or a zlib one (`zlib`
/// true: a zlib stream names any dictionary it needs in its header, so its window starts empty, as RFC 1950 has it).
fn inflater(zlib: bool) -> Result<zlib_rs::Inflate> {
    let mut d = zlib_rs::Inflate::new(zlib, WINDOW_BITS);
    if !zlib {
        d.set_dictionary(&ZERO_WINDOW)
            .map_err(|e| invalid("deflate", 0, inflate_error(e)))?;
    }
    Ok(d)
}
/// The input offset a decoder reached when `left` bytes remained.
fn consumed(input: &[u8], left: usize) -> usize {
    input.len().saturating_sub(left)
}
pub(crate) fn buffer(size: usize) -> Result<Vec<u8>> {
    let mut out = Vec::new();
    out.try_reserve_exact(size).map_err(|_| Error::Limit)?;
    out.resize(size, 0);
    Ok(out)
}
pub(crate) fn check_size(actual: usize, expected: usize) -> Result<()> {
    if actual == expected {
        Ok(())
    } else {
        Err(Error::Size { expected, actual })
    }
}
/// Decodes one exactly framed payload to an output size the caller states.
/// Implementations reject output overrun; the caller must bound untrusted sizes.
pub trait Decompressor {
    /// The `expected` bytes `input` decodes to; any other outcome is an error.
    fn decompress(&self, input: &[u8], expected: usize) -> Result<Vec<u8>>;
}
/// Encodes one payload; framing belongs to the selected codec.
pub trait Compressor {
    /// `input` compressed.
    fn compress(&self, input: &[u8]) -> Result<Vec<u8>>;
}

/// The most bytes an LZ4 block can produce for one byte of input (a match length is extended by bytes of 255, each
/// adding 255 to it), rounded up: a claimed output size above this times the input is false.
pub const LZ4_MAX_RATIO: usize = 255;

/// Raw LZ4 block, no size prefix and no frame. All input bytes belong to the block; a final literal sequence terminates
/// it. The output size comes from the enclosing archive.
#[derive(Debug, Clone, Copy)]
pub struct Lz4;
impl Decompressor for Lz4 {
    fn decompress(&self, input: &[u8], expected: usize) -> Result<Vec<u8>> {
        if input.is_empty() {
            return Err(invalid("lz4", 0, "missing token: expected at least one byte"));
        }
        check_claim("lz4", input.len(), expected, LZ4_MAX_RATIO)?;
        let mut out = buffer(expected)?;
        let n = lz4_flex::block::decompress_into(input, &mut out).map_err(|e| invalid("lz4", 0, e))?;
        check_size(n, expected)?;
        Ok(out)
    }
}
impl Compressor for Lz4 {
    fn compress(&self, input: &[u8]) -> Result<Vec<u8>> {
        Ok(lz4_flex::block::compress(input))
    }
}

/// One zstd frame (RFC 8878), decoded exactly: no skippable frame, no dictionary, a stated content size (where the
/// frame states one) equal to the expected size, every input byte consumed, a stated content checksum verified, and
/// never more than the expected output (decoded a block at a time, refused as soon as it would exceed it).
#[derive(Debug, Clone, Copy)]
pub struct Zstd;
impl Decompressor for Zstd {
    fn decompress(&self, input: &[u8], expected: usize) -> Result<Vec<u8>> {
        use ruzstd::{BlockDecodingStrategy, FrameDecoder};
        use std::io::Read;
        let (frame, _) = ruzstd::frame::read_frame_header(input).map_err(|e| invalid("zstd", 0, e))?;
        if let Some(id) = frame.header.dictionary_id() {
            return Err(invalid("zstd", 0, format!("frame needs dictionary {id}")));
        }
        let descriptor = &frame.header.descriptor;
        let states_size = descriptor.single_segment_flag()
            || descriptor
                .frame_content_size_bytes()
                .map_err(|e| invalid("zstd", 0, e))?
                > 0;
        if states_size {
            let stated = usize::try_from(frame.header.frame_content_size()).map_err(|_| Error::Limit)?;
            check_size(stated, expected)?;
        }
        let mut d = FrameDecoder::new();
        let mut src = input;
        d.init(&mut src).map_err(|e| invalid("zstd", 0, e))?;
        // grown with what the stream produces, which the check below keeps within `expected`
        let mut out = Vec::new();
        loop {
            let before = src.len();
            d.decode_blocks(&mut src, BlockDecodingStrategy::UptoBlocks(1))
                .map_err(|e| invalid("zstd", consumed(input, before), e))?;
            let ready = d.can_collect();
            let total = out.len().checked_add(ready).ok_or(Error::Limit)?;
            if total > expected {
                return Err(Error::Size {
                    expected,
                    actual: total,
                });
            }
            let at = out.len();
            out.try_reserve(ready).map_err(|_| Error::Limit)?;
            out.resize(total, 0);
            let fresh = out.get_mut(at..).ok_or_else(|| {
                invalid(
                    "zstd",
                    consumed(input, src.len()),
                    "output shorter than collected",
                )
            })?;
            d.read_exact(fresh)?;
            if d.is_finished() && d.can_collect() == 0 {
                break;
            }
            if src.len() == before && ready == 0 {
                return Err(invalid(
                    "zstd",
                    consumed(input, src.len()),
                    "frame ends before its last block: expected another block",
                ));
            }
        }
        if !src.is_empty() {
            return Err(invalid(
                "zstd",
                consumed(input, src.len()),
                "trailing input after the frame: expected the frame to end here",
            ));
        }
        if let (Some(stored), Some(computed)) = (d.get_checksum_from_data(), d.get_calculated_checksum())
            && stored != computed
        {
            return Err(invalid(
                "zstd",
                input.len(),
                format!("content checksum {stored:08x}, computed {computed:08x}"),
            ));
        }
        check_size(out.len(), expected)?;
        Ok(out)
    }
}

/// RFC1951 raw DEFLATE. Requires the final block marker and exact input consumption.
#[derive(Debug, Clone, Copy)]
pub struct Deflate;
/// RFC1950 zlib: header, DEFLATE, Adler32 trailer. Requires StreamEnd and exact consumption.
#[derive(Debug, Clone, Copy)]
pub struct Zlib;
/// zlib as some archives' tables of contents store it: with a missing final block marker or Adler-32 checksum. Such a
/// stream is accepted only if all input is consumed, the decoder reports no error, and exactly the expected bytes are
/// produced. An incomplete stream cannot be checked against its Adler-32, so this is a size check, not an integrity
/// check. Complete streams are accepted too.
#[derive(Debug, Clone, Copy)]
pub struct ZlibUnterminated;
fn inflate(input: &[u8], expected: usize, zlib: bool, incomplete: bool) -> Result<Vec<u8>> {
    check_claim("deflate/zlib", input.len(), expected, DEFLATE_MAX_RATIO)?;
    let mut d = inflater(zlib)?;
    let mut out = buffer(expected.checked_add(1).ok_or(Error::Limit)?)?;
    let status = d
        .decompress(input, &mut out, zlib_rs::InflateFlush::NoFlush)
        .map_err(|e| invalid("deflate/zlib", total_in(&d), inflate_error(e)))?;
    check_size(usize::try_from(d.total_out()).unwrap_or(usize::MAX), expected)?;
    if total_in(&d) != input.len() {
        return Err(invalid(
            "deflate/zlib",
            total_in(&d),
            "trailing input: expected the stream to end here",
        ));
    }
    if status != zlib_rs::Status::StreamEnd && !incomplete {
        return Err(invalid("deflate/zlib", input.len(), "missing stream end"));
    }
    // Even incomplete empty zlib needs a valid two-byte header.
    if zlib && input.len() < 2 {
        return Err(invalid("zlib", input.len(), "short header"));
    }
    out.truncate(expected);
    Ok(out)
}
impl Decompressor for Deflate {
    fn decompress(&self, input: &[u8], expected: usize) -> Result<Vec<u8>> {
        inflate(input, expected, false, false)
    }
}
/// The most bytes a DEFLATE stream can produce for one byte of input (RFC 1951: a match of 258 bytes takes about two
/// bits), rounded up: a claimed output size above this times the input is false.
pub const DEFLATE_MAX_RATIO: usize = 1032;

/// The most bytes an LZO2A stream can produce for one byte of input, rounded up: a long match of length `9 + 255 z + n`
/// takes the two bytes of its distance, `z` zero bytes and the byte `n`, so its output stays under 255 times its input.
pub const LZO2A_MAX_RATIO: usize = 255;

impl Deflate {
    /// A complete raw DEFLATE stream whose output size is not stated, at most `max` bytes. The output grows with what
    /// the stream produces (at most [`DEFLATE_MAX_RATIO`] times its input), never with `max`, so a caller's bound
    /// taken from an untrusted header costs no memory by itself. Requires the final block marker and exact input
    /// consumption; more than `max` bytes of output is [`Error::Size`] (`actual` is then `max + 1`).
    pub fn decompress_up_to(&self, input: &[u8], max: usize) -> Result<Vec<u8>> {
        let mut d = inflater(false)?;
        let mut out: Vec<u8> = Vec::new();
        loop {
            if out.len() == out.capacity() {
                // double, but never past one byte more than `max`, which is what shows the stream overran it
                let want = out
                    .capacity()
                    .max(4096)
                    .saturating_mul(2)
                    .min(max.saturating_add(1));
                let more = want.saturating_sub(out.len());
                if more == 0 {
                    return Err(Error::Size {
                        expected: max,
                        actual: out.len(),
                    });
                }
                out.try_reserve_exact(more).map_err(|_| Error::Limit)?;
            }
            let taken = total_in(&d);
            let rest = input.get(taken..).unwrap_or_default();
            let before = (d.total_in(), d.total_out());
            // the inflater writes into initialised bytes: the room is zero-filled, and cut back to what it produced
            let at = out.len();
            out.resize(out.capacity(), 0);
            let result = match out.get_mut(at..) {
                Some(room) => d.decompress(rest, room, zlib_rs::InflateFlush::NoFlush),
                None => Ok(zlib_rs::Status::BufError),
            };
            let produced = usize::try_from(d.total_out().saturating_sub(before.1)).unwrap_or(usize::MAX);
            out.truncate(at.saturating_add(produced));
            let status = result.map_err(|e| invalid("deflate", taken, inflate_error(e)))?;
            if out.len() > max {
                return Err(Error::Size {
                    expected: max,
                    actual: out.len(),
                });
            }
            if status == zlib_rs::Status::StreamEnd {
                break;
            }
            if (d.total_in(), d.total_out()) == before && out.len() < out.capacity() {
                return Err(invalid("deflate", input.len(), "missing stream end"));
            }
        }
        if total_in(&d) != input.len() {
            return Err(invalid(
                "deflate",
                total_in(&d),
                "trailing input: expected the stream to end here",
            ));
        }
        Ok(out)
    }
}
impl Decompressor for Zlib {
    fn decompress(&self, input: &[u8], expected: usize) -> Result<Vec<u8>> {
        inflate(input, expected, true, false)
    }
}
impl Decompressor for ZlibUnterminated {
    fn decompress(&self, input: &[u8], expected: usize) -> Result<Vec<u8>> {
        inflate(input, expected, true, true)
    }
}
impl Compressor for Deflate {
    fn compress(&self, input: &[u8]) -> Result<Vec<u8>> {
        let mut e = flate2::write::DeflateEncoder::new(Vec::new(), flate2::Compression::default());
        e.write_all(input)?;
        Ok(e.finish()?)
    }
}
impl Compressor for Zlib {
    fn compress(&self, input: &[u8]) -> Result<Vec<u8>> {
        let mut e = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default());
        e.write_all(input)?;
        Ok(e.finish()?)
    }
}