mzstd 0.1.0

//! Frame headers and per-frame decoding state (RFC 8878 section 3.1.1).
//!
//! [`FrameHeader::parse`] reads a frame header after its magic number.
//! [`FrameDecoder`] validates it against the caller's limits, then
//! regenerates the frame's blocks one at a time, tracking the produced size
//! against `Frame_Content_Size` and hashing the output for the optional
//! content checksum.

use alloc::vec::Vec;

use super::block::{BlockDecoder, BlockHeader};
use super::sequences::BlockContext;
use crate::decompressor::{DecodeError, DecodeLimits, InvalidDataKind};
use crate::shared::format::{
    BLOCK_SIZE_MAX, CHECKSUM_SIZE, FRAME_HEADER_SIZE_MAX, WINDOW_LOG_ABSOLUTE_MIN,
};
use crate::shared::xxhash::Xxh64;

/// Size of a frame magic number.
pub(crate) const MAGIC_SIZE: usize = 4;

// The largest header (window byte, four-byte dictionary ID, eight-byte
// content size) and its magic number make the format's maximum.
const _: () = assert!(MAGIC_SIZE + header_size(0xc3) == FRAME_HEADER_SIZE_MAX);

/// A parsed frame header.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct FrameHeader {
    /// Largest back-reference distance, and the history a decoder keeps.
    pub(crate) window_size: u64,
    /// Declared decompressed size, when present.
    pub(crate) content_size: Option<u64>,
    /// Declared dictionary; zero means none.
    pub(crate) dictionary_id: u32,
    /// Whether a content checksum follows the last block.
    pub(crate) checksum: bool,
}

/// Size of a frame header, from its descriptor byte to its last field
/// (the magic number excluded).
pub(crate) const fn header_size(descriptor: u8) -> usize {
    let single_segment = descriptor & 0x20 != 0;
    let window = if single_segment { 0 } else { 1 };
    let dictionary = match descriptor & 3 {
        0 => 0,
        1 => 1,
        2 => 2,
        _ => 4,
    };
    let content_size = match descriptor >> 6 {
        0 => {
            if single_segment {
                1
            } else {
                0
            }
        }
        1 => 2,
        2 => 4,
        _ => 8,
    };
    1 + window + dictionary + content_size
}

/// Reads a little-endian field of up to eight bytes.
fn read_le(bytes: &[u8]) -> u64 {
    bytes
        .iter()
        .rev()
        .fold(0u64, |value, &byte| (value << 8) | u64::from(byte))
}

impl FrameHeader {
    /// Parses a header of exactly [`header_size`] bytes.
    pub(crate) fn parse(bytes: &[u8]) -> Result<Self, InvalidDataKind> {
        let invalid = InvalidDataKind::FrameHeader;
        let (&descriptor, mut fields) = bytes.split_first().ok_or(invalid)?;
        if bytes.len() != header_size(descriptor) || descriptor & 0x08 != 0 {
            return Err(invalid);
        }
        let single_segment = descriptor & 0x20 != 0;
        let mut window_size = 0;
        if !single_segment {
            let (&window, rest) = fields.split_first().ok_or(invalid)?;
            let log = WINDOW_LOG_ABSOLUTE_MIN + u32::from(window >> 3);
            let base = 1u64 << log;
            window_size = base + (base / 8) * u64::from(window & 7);
            fields = rest;
        }
        let dictionary_bytes = match descriptor & 3 {
            0 => 0,
            1 => 1,
            2 => 2,
            _ => 4,
        };
        let (dictionary, fields) = fields.split_at_checked(dictionary_bytes).ok_or(invalid)?;
        let content_size = match fields.len() {
            0 => None,
            2 => Some(read_le(fields) + 256),
            _ => Some(read_le(fields)),
        };
        if single_segment {
            window_size = content_size.unwrap_or(0);
        }
        Ok(Self {
            window_size,
            content_size,
            dictionary_id: read_le(dictionary) as u32,
            checksum: descriptor & 0x04 != 0,
        })
    }
}

/// Decoding state of the frame in progress.
#[derive(Debug)]
pub(crate) struct FrameDecoder {
    header: Option<FrameHeader>,
    block_max: usize,
    produced: u64,
    hasher: Xxh64,
    blocks: BlockDecoder,
}

impl Default for FrameDecoder {
    fn default() -> Self {
        Self {
            header: None,
            block_max: 0,
            produced: 0,
            hasher: Xxh64::new(0),
            blocks: BlockDecoder::default(),
        }
    }
}

impl FrameDecoder {
    /// Validates `header` against `limits` and starts its frame.
    pub(crate) fn begin(
        &mut self,
        header: FrameHeader,
        limits: &DecodeLimits,
    ) -> Result<(), DecodeError> {
        if header.dictionary_id != 0 {
            return Err(DecodeError::DictionaryUnsupported {
                id: header.dictionary_id,
            });
        }
        let limit = limits.max_window_size();
        if header.window_size > limit {
            return Err(DecodeError::WindowLimitExceeded {
                window_size: header.window_size,
                limit,
            });
        }
        // At most 128 KiB, so the conversion is lossless.
        self.block_max = header.window_size.min(BLOCK_SIZE_MAX as u64) as usize;
        self.header = Some(header);
        self.produced = 0;
        self.hasher.reset(0);
        self.blocks.reset();
        Ok(())
    }

    /// Heap bytes held between blocks.
    pub(crate) fn retained_bytes(&self) -> usize {
        self.blocks.retained_bytes()
    }

    /// Drops the entropy tables and buffers, keeping the header.
    pub(crate) fn release(&mut self) {
        self.blocks = BlockDecoder::default();
    }

    /// The header of the frame in progress.
    pub(crate) const fn header(&self) -> Option<&FrameHeader> {
        self.header.as_ref()
    }

    /// Largest block of the frame in progress.
    pub(crate) const fn block_max(&self) -> usize {
        self.block_max
    }

    /// Whether the frame in progress ends with a content checksum.
    pub(crate) fn has_checksum(&self) -> bool {
        self.header.is_some_and(|header| header.checksum)
    }

    /// Regenerates one block onto `out`, whose bytes from `history_start`
    /// are this frame's retained history.
    #[cfg_attr(all(feature = "hotpath", not(feature = "no_std")), hotpath::measure)]
    pub(crate) fn decode_block(
        &mut self,
        block: &BlockHeader,
        content: &[u8],
        out: &mut Vec<u8>,
        history_start: usize,
    ) -> Result<(), DecodeError> {
        let header = self.header.ok_or(DecodeError::InvalidState)?;
        out.try_reserve(self.block_max)
            .map_err(|_| DecodeError::AllocationFailed)?;
        let start = out.len();
        let context = BlockContext {
            history_start,
            window: header.window_size,
            block_max: self.block_max,
        };
        self.blocks.decode(block, content, out, context)?;
        let produced = &out[start..];
        if header.checksum {
            self.hasher.update(produced);
        }
        self.produced += produced.len() as u64;
        if let Some(declared) = header.content_size
            && (self.produced > declared || (block.last && self.produced != declared))
        {
            return Err(DecodeError::ContentSizeMismatch {
                declared,
                decoded: self.produced,
            });
        }
        Ok(())
    }

    /// Checks the stored content checksum against the frame's output.
    pub(crate) fn verify_checksum(&self, stored: [u8; CHECKSUM_SIZE]) -> Result<(), DecodeError> {
        let expected = u32::from_le_bytes(stored);
        let computed = self.hasher.digest() as u32;
        if expected == computed {
            Ok(())
        } else {
            Err(DecodeError::ChecksumMismatch { expected, computed })
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::decompressor::core::block::BlockKind;

    #[test]
    fn header_sizes_follow_the_descriptor() {
        assert_eq!(header_size(0x00), 2);
        assert_eq!(header_size(0x20), 2);
        assert_eq!(header_size(0x03), 6);
        assert_eq!(header_size(0x40), 4);
        assert_eq!(header_size(0x80), 6);
        assert_eq!(header_size(0xe2), 11);
    }

    #[test]
    fn windows_content_sizes_and_dictionaries_parse() {
        // Window exponent 1, mantissa 3: 2 KiB + 3 * 256.
        let header = FrameHeader::parse(&[0x04, 0x0b]).expect("valid");
        assert_eq!(header.window_size, 2048 + 3 * 256);
        assert_eq!(header.content_size, None);
        assert!(header.checksum);
        // Single segment with a one-byte size: the window is the size.
        let header = FrameHeader::parse(&[0x20, 200]).expect("valid");
        assert_eq!(header.window_size, 200);
        assert_eq!(header.content_size, Some(200));
        // Two-byte sizes are offset by 256; a two-byte dictionary ID.
        let header = FrameHeader::parse(&[0x62, 0x34, 0x12, 0x01, 0x00]).expect("valid");
        assert_eq!(header.dictionary_id, 0x1234);
        assert_eq!(header.content_size, Some(257));
        let header = FrameHeader::parse(&[0xa0, 1, 0, 0, 0]).expect("valid");
        assert_eq!(header.content_size, Some(1));
        let header = FrameHeader::parse(&[0xe1, 7, 0, 0, 0, 0, 0, 0, 0, 1]).expect("valid");
        assert_eq!(header.dictionary_id, 7);
        assert_eq!(header.content_size, Some(1 << 56));
    }

    #[test]
    fn reserved_bits_and_wrong_lengths_are_rejected() {
        assert_eq!(
            FrameHeader::parse(&[0x08, 0]),
            Err(InvalidDataKind::FrameHeader)
        );
        assert_eq!(
            FrameHeader::parse(&[0x00]),
            Err(InvalidDataKind::FrameHeader)
        );
        assert_eq!(FrameHeader::parse(&[]), Err(InvalidDataKind::FrameHeader));
        // The unused bit is ignored.
        assert!(FrameHeader::parse(&[0x10, 0]).is_ok());
    }

    #[test]
    fn begin_enforces_dictionary_and_window_limits() {
        let mut frame = FrameDecoder::default();
        let limits = DecodeLimits::default();
        let with_dictionary = FrameHeader::parse(&[0x01, 0x00, 9]).expect("valid");
        assert!(matches!(
            frame.begin(with_dictionary, &limits),
            Err(DecodeError::DictionaryUnsupported { id: 9 })
        ));
        // Exponent 18: a 2^28-byte window exceeds the default 2^27.
        let wide = FrameHeader::parse(&[0x00, 18 << 3]).expect("valid");
        assert!(matches!(
            frame.begin(wide, &limits),
            Err(DecodeError::WindowLimitExceeded {
                window_size: 0x1000_0000,
                limit: 0x0800_0000
            })
        ));
        let empty = FrameHeader::parse(&[0x20, 0]).expect("valid");
        frame.begin(empty, &limits).expect("valid");
        assert_eq!(frame.block_max(), 0);
        assert!(!frame.has_checksum());
        assert_eq!(frame.header(), Some(&empty));
    }

    #[test]
    fn blocks_are_checked_against_the_declared_size_and_checksum() {
        let limits = DecodeLimits::default();
        let mut frame = FrameDecoder::default();
        let mut out = Vec::new();
        let raw = |last, size| BlockHeader {
            last,
            kind: BlockKind::Raw,
            size,
        };
        assert!(matches!(
            frame.decode_block(&raw(true, 1), b"a", &mut out, 0),
            Err(DecodeError::InvalidState)
        ));
        frame
            .begin(FrameHeader::parse(&[0x24, 3]).expect("valid"), &limits)
            .expect("valid");
        frame
            .decode_block(&raw(false, 2), b"ab", &mut out, 0)
            .expect("fits");
        assert!(matches!(
            frame.decode_block(&raw(true, 0), b"", &mut out, 0),
            Err(DecodeError::ContentSizeMismatch {
                declared: 3,
                decoded: 2
            })
        ));
        frame
            .begin(FrameHeader::parse(&[0x24, 3]).expect("valid"), &limits)
            .expect("valid");
        out.clear();
        frame
            .decode_block(&raw(true, 3), b"abc", &mut out, 0)
            .expect("fits");
        // XXH64("abc") = 0x44BC2CF5AD770999.
        frame
            .verify_checksum(0xAD77_0999u32.to_le_bytes())
            .expect("matches");
        assert!(matches!(
            frame.verify_checksum([0; 4]),
            Err(DecodeError::ChecksumMismatch {
                expected: 0,
                computed: 0xAD77_0999
            })
        ));
        // Producing more than declared fails before the last block.
        frame
            .begin(FrameHeader::parse(&[0x20, 1]).expect("valid"), &limits)
            .expect("valid");
        assert!(matches!(
            frame.decode_block(&raw(false, 1), b"a", &mut out, 0),
            Ok(())
        ));
        assert!(matches!(
            frame.decode_block(&raw(false, 1), b"b", &mut out, 0),
            Err(DecodeError::ContentSizeMismatch {
                declared: 1,
                decoded: 2
            })
        ));
    }
}