libfw_core/

compress.rs

//! Streaming compression abstractions backed by [`zrip`] (zstd).
//!
//! # Design
//!
//! Both [`Compressor`] and [`Decompressor`] are *stream-oriented*: they
//! accept bounded input chunks and produce bounded output chunks, so the
//! heap footprint of processing a file does **not** grow with the file
//! size. The [`STREAM_BUF_SIZE`](crate::STREAM_BUF_SIZE) sliding window
//! keeps memory constant (~64 KiB typical, capped well below 4 MiB even
//! for adversarial frames).
//!
//! # Wire format
//!
//! The compressed body is a concatenation of independent zstd *frames*:
//! each [`Compressor::compress`] call emits one complete frame when the
//! next call arrives (or on [`Compressor::finish`]).
//!
//! The decompressor buffers compressed bytes until a complete frame is
//! present (frame boundaries are detected by a small built-in zstd frame
//! parser), then decodes exactly that frame with
//! [`zrip::decompress_with_limit`]. A frame may therefore span several
//! `decompress()` inputs, and one input may contain several frames.
//!
//! # Constant memory
//!
//! - **Compressor**: buffers at most one input chunk internally plus the
//!   zrip encoder workspace (~150 KiB).
//! - **Decompressor**: holds at most one *compressed* frame
//!   ([`MAX_PENDING_FRAME`]) and one *decompressed* frame
//!   ([`MAX_FRAME_OUTPUT`]) at a time.
//!
//! The server should feed [`Compressor::compress`] in
//! [`STREAM_BUF_SIZE`](crate::STREAM_BUF_SIZE) windows (64 KiB), keeping
//! frames ~64 KiB and the client's transient buffers small.

use std::io::Write;

use crate::error::{CompressError, DecompressError};
use crate::{CHUNK_SIZE, STREAM_BUF_SIZE};

/// Compression formats understood by libfw.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompressionFormat {
    /// No compression; the body passes through verbatim.
    None,
    /// zstd via the [`zrip`] codec (levels -8..=4).
    Zrip,
}

impl CompressionFormat {
    /// Wire representation used in
    /// [`HEADER_COMPRESS`](crate::HEADER_COMPRESS) / `Content-Encoding`.
    pub fn as_str(self) -> &'static str {
        match self {
            CompressionFormat::None => "identity",
            CompressionFormat::Zrip => "zrip",
        }
    }

    /// Parses a header value (`identity`, `zrip`, …).
    pub fn parse_header(s: &str) -> Option<Self> {
        match s.trim().to_ascii_lowercase().as_str() {
            "" | "identity" | "none" => Some(CompressionFormat::None),
            "zrip" | "zstd" => Some(CompressionFormat::Zrip),
            _ => None,
        }
    }
}

/// Safety ceiling for the decompressed size of a single frame.
///
/// The protocol keeps frames ≤ [`CHUNK_SIZE`](crate::CHUNK_SIZE); this
/// limit is a guard against buggy or hostile peers.
pub const MAX_FRAME_OUTPUT: usize = CHUNK_SIZE as usize;

/// Maximum compressed bytes buffered while waiting for a frame boundary.
///
/// The worst-case incompressible [`CHUNK_SIZE`](crate::CHUNK_SIZE) chunk
/// compresses to ≈ chunk size + framing overhead, hence the slack.
pub const MAX_PENDING_FRAME: usize = CHUNK_SIZE as usize + STREAM_BUF_SIZE;

/// Default zrip compression level (Fast strategy, good for network transfer).
pub const ZRIP_DEFAULT_LEVEL: i32 = 1;

/// Streaming compressor. Feed it bounded input chunks; it appends the
/// compressed bytes produced so far to the provided output buffer.
pub trait Compressor: Send {
    /// The format this compressor produces.
    fn format(&self) -> CompressionFormat;

    /// Compress `input`, appending compressed bytes to `out`.
    ///
    /// `input` should be kept ≤ [`CHUNK_SIZE`](crate::CHUNK_SIZE) to keep
    /// the output frames small and memory constant. When the next chunk
    /// arrives, the previous chunk's frame is finalized and emitted.
    fn compress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), CompressError>;

    /// Finalize the stream, appending the last frame (if any) to `out`.
    ///
    /// Must be called exactly once; afterwards the compressor is spent.
    fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), CompressError>;
}

/// Streaming decompressor. Feed it arbitrary compressed bytes; decoded
/// output is appended to the provided buffer.
pub trait Decompressor: Send {
    /// The format this decompressor consumes.
    fn format(&self) -> CompressionFormat;

    /// Decompress `input`, appending decoded bytes to `out`.
    ///
    /// Input may split frames arbitrarily. `out` should be drained by the
    /// caller after each call to keep memory bounded.
    fn decompress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), DecompressError>;

    /// Signal the end of the stream.
    ///
    /// Flushes any complete trailing frames and verifies the stream is
    /// well-formed (no truncated frames). Call exactly once.
    fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), DecompressError>;
}

/// Construct a compressor for `format`.
pub fn compressor(format: CompressionFormat) -> Result<Box<dyn Compressor>, CompressError> {
    match format {
        CompressionFormat::None => Ok(Box::new(PassthroughCompressor)),
        CompressionFormat::Zrip => Ok(Box::new(ZripCompressor::new(ZRIP_DEFAULT_LEVEL)?)),
    }
}

/// Construct a decompressor for `format`.
pub fn decompressor(format: CompressionFormat) -> Box<dyn Decompressor> {
    match format {
        CompressionFormat::None => Box::new(PassthroughDecompressor),
        CompressionFormat::Zrip => Box::new(ZripDecompressor::new()),
    }
}

// ---------------------------------------------------------------------------
// Passthrough (identity)
// ---------------------------------------------------------------------------

struct PassthroughCompressor;

impl Compressor for PassthroughCompressor {
    fn format(&self) -> CompressionFormat {
        CompressionFormat::None
    }

    fn compress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), CompressError> {
        out.extend_from_slice(input);
        Ok(())
    }

    fn finish(&mut self, _out: &mut Vec<u8>) -> Result<(), CompressError> {
        Ok(())
    }
}

struct PassthroughDecompressor;

impl Decompressor for PassthroughDecompressor {
    fn format(&self) -> CompressionFormat {
        CompressionFormat::None
    }

    fn decompress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), DecompressError> {
        out.extend_from_slice(input);
        Ok(())
    }

    fn finish(&mut self, _out: &mut Vec<u8>) -> Result<(), DecompressError> {
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// zrip (zstd) implementation
// ---------------------------------------------------------------------------

/// zstd compressor emitting one independent frame per input chunk.
pub struct ZripCompressor {
    encoder: Option<zrip::FrameEncoder<Vec<u8>>>,
    /// Whether data has been written since the last frame boundary.
    dirty: bool,
}

impl ZripCompressor {
    /// Create a compressor at `level` (-8..=4; 0 = library default).
    pub fn new(level: i32) -> Result<Self, CompressError> {
        Ok(ZripCompressor {
            encoder: Some(zrip::FrameEncoder::new(Vec::new(), level)?),
            dirty: false,
        })
    }
}

impl Compressor for ZripCompressor {
    fn format(&self) -> CompressionFormat {
        CompressionFormat::Zrip
    }

    fn compress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), CompressError> {
        if input.is_empty() {
            return Ok(());
        }
        let encoder = self
            .encoder
            .as_mut()
            .ok_or_else(|| std::io::Error::other("compressor already finished"))?;
        // Finalize the previous chunk's frame and hand its compressed bytes back.
        if self.dirty {
            let finished = encoder.reset(Vec::new())?;
            out.extend_from_slice(&finished);
        }
        encoder.write_all(input)?;
        self.dirty = true;
        Ok(())
    }

    fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), CompressError> {
        let encoder = self
            .encoder
            .take()
            .ok_or_else(|| std::io::Error::other("compressor already finished"))?;
        let tail = encoder.finish()?;
        out.extend_from_slice(&tail);
        self.dirty = false;
        Ok(())
    }
}

/// zstd decompressor reassembling frames across arbitrary chunk boundaries.
pub struct ZripDecompressor {
    /// Compressed bytes not yet turned into complete frames.
    pending: Vec<u8>,
    finished: bool,
}

impl ZripDecompressor {
    /// Create a decompressor.
    pub fn new() -> Self {
        ZripDecompressor {
            pending: Vec::with_capacity(STREAM_BUF_SIZE),
            finished: false,
        }
    }

    /// Decode as many complete frames as `pending` holds.
    fn drain_frames(&mut self, out: &mut Vec<u8>) -> Result<(), DecompressError> {
        loop {
            let boundary = frame_boundary(&self.pending);
            match boundary {
                FrameBoundary::Complete { len, content_size } => {
                    if content_size.is_some_and(|cs| cs > MAX_FRAME_OUTPUT as u64) {
                        return Err(DecompressError::TooLarge {
                            limit: MAX_FRAME_OUTPUT,
                        });
                    }
                    let decoded = {
                        let frame = &self.pending[..len];
                        zrip::decompress_with_limit(frame, MAX_FRAME_OUTPUT).map_err(|e| {
                            DecompressError::Io(std::io::Error::new(
                                std::io::ErrorKind::InvalidData,
                                e,
                            ))
                        })?
                    };
                    out.extend_from_slice(&decoded);
                    self.pending.drain(..len);
                }
                FrameBoundary::Incomplete => return Ok(()),
                FrameBoundary::Invalid => {
                    return Err(DecompressError::Io(std::io::Error::new(
                        std::io::ErrorKind::InvalidData,
                        "invalid zstd frame data",
                    )))
                }
            }
        }
    }
}

impl Default for ZripDecompressor {
    fn default() -> Self {
        ZripDecompressor::new()
    }
}

impl Decompressor for ZripDecompressor {
    fn format(&self) -> CompressionFormat {
        CompressionFormat::Zrip
    }

    fn decompress(&mut self, input: &[u8], out: &mut Vec<u8>) -> Result<(), DecompressError> {
        if self.finished {
            return Err(DecompressError::Io(std::io::Error::other(
                "decompressor already finished",
            )));
        }
        if !input.is_empty() {
            self.pending.extend_from_slice(input);
        }
        self.drain_frames(out)?;
        if self.pending.len() > MAX_PENDING_FRAME {
            return Err(DecompressError::TooLarge {
                limit: MAX_PENDING_FRAME,
            });
        }
        Ok(())
    }

    fn finish(&mut self, out: &mut Vec<u8>) -> Result<(), DecompressError> {
        if self.finished {
            return Ok(());
        }
        self.finished = true;
        self.drain_frames(out)?;
        if !self.pending.is_empty() {
            return Err(DecompressError::Truncated(std::io::Error::new(
                std::io::ErrorKind::UnexpectedEof,
                format!("{} trailing compressed bytes", self.pending.len()),
            )));
        }
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Minimal zstd frame-boundary parser
// ---------------------------------------------------------------------------

const ZSTD_MAGIC: u32 = 0xFD2F_B528;
const SKIPPABLE_MASK: u32 = 0xFFFF_FFF0;
const SKIPPABLE_MAGIC: u32 = 0x184D_2A50;

enum FrameBoundary {
    /// A complete frame occupying `len` bytes, declaring `content_size`
    /// output bytes (`None` when the frame header does not state it).
    Complete { len: usize, content_size: Option<u64> },
    /// Need more bytes to know the boundary.
    Incomplete,
    /// Corrupt / unsupported.
    Invalid,
}

/// Determine the byte length of the next zstd frame in `buf` (if present).
fn frame_boundary(buf: &[u8]) -> FrameBoundary {
    if buf.len() < 4 {
        return FrameBoundary::Incomplete;
    }
    let magic = u32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]);
    if magic == ZSTD_MAGIC {
        zstd_frame_boundary(buf)
    } else if (magic & SKIPPABLE_MASK) == SKIPPABLE_MAGIC {
        skippable_frame_boundary(buf)
    } else {
        FrameBoundary::Invalid
    }
}

/// Skippable frames: `magic(4) | skip_size u32(4) | payload(skip_size)`.
fn skippable_frame_boundary(buf: &[u8]) -> FrameBoundary {
    if buf.len() < 8 {
        return FrameBoundary::Incomplete;
    }
    let skip = u32::from_le_bytes([buf[4], buf[5], buf[6], buf[7]]) as usize;
    let total = 8usize.saturating_add(skip);
    if buf.len() >= total {
        FrameBoundary::Complete {
            len: total,
            content_size: Some(0),
        }
    } else {
        FrameBoundary::Incomplete
    }
}

/// Walk the frame header + blocks of a standard zstd frame.
///
/// Layout: `magic(4) descriptor(1) [window(1)] [dict_id(0/1/2/4)]
/// [content_size(0/1/2/4/8)] block* checksum(0/4)`.
fn zstd_frame_boundary(buf: &[u8]) -> FrameBoundary {
    if buf.len() < 5 {
        return FrameBoundary::Incomplete;
    }
    let descriptor = buf[4];
    // Reserved bits (3..=4) must be zero.
    if descriptor & 0x18 != 0 {
        return FrameBoundary::Invalid;
    }
    let single_segment = descriptor & 0x20 != 0;
    let checksum = descriptor & 0x04 != 0;
    let dict_id_flag = descriptor & 0x03;
    let fcs_flag = (descriptor >> 6) & 0x03;

    let mut hdr_len = 5usize;
    if !single_segment {
        hdr_len += 1; // window descriptor
    }
    hdr_len += match dict_id_flag {
        0 => 0,
        1 => 1,
        2 => 2,
        3 => 4,
        _ => unreachable!(),
    };
    let fcs_size: usize = match fcs_flag {
        0 if single_segment => 1,
        0 => 0,
        1 => 2,
        2 => 4,
        3 => 8,
        _ => unreachable!(),
    };
    hdr_len += fcs_size;

    if buf.len() < hdr_len {
        return FrameBoundary::Incomplete;
    }
    let content_size = if fcs_size > 0 {
        let mut v = 0u64;
        for (i, &b) in buf[5..5 + fcs_size].iter().enumerate() {
            v |= (b as u64) << (8 * i);
        }
        Some(v)
    } else {
        None
    };

    // Walk blocks until the last-block flag.
    let mut off = hdr_len;
    loop {
        if buf.len() < off + 3 {
            return FrameBoundary::Incomplete;
        }
        // Standard zstd block header (as read by the reference decoder):
        // bit 0 = last_block, bits 1..=2 = block_type, bits 3..=23 = size.
        let block_header = u32::from_le_bytes([buf[off], buf[off + 1], buf[off + 2], 0]);
        let last = block_header & 0x01 != 0;
        let block_type = (block_header >> 1) & 0x03;
        let block_size = (block_header >> 3) as usize;
        if block_type == 3 {
            // Reserved block type.
            return FrameBoundary::Invalid;
        }
        off += 3 + block_size;
        if off > MAX_PENDING_FRAME {
            return FrameBoundary::Invalid;
        }
        if last {
            break;
        }
    }
    if checksum {
        off += 4;
    }
    if buf.len() < off {
        return FrameBoundary::Incomplete;
    }
    FrameBoundary::Complete {
        len: off,
        content_size,
    }
}

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

    fn roundtrip_chunks(data: &[u8], feed: &[usize]) {
        let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
        let mut compressed = Vec::new();
        let mut off = 0;
        for &n in feed {
            let end = (off + n).min(data.len());
            if end > off {
                c.compress(&data[off..end], &mut compressed).unwrap();
            }
            off = end;
        }
        c.finish(&mut compressed).unwrap();
        assert!(!compressed.is_empty());

        // Decode with arbitrarily split inputs, including mid-frame splits.
        let mut d = ZripDecompressor::new();
        let mut plain = Vec::new();
        let mut step = 1;
        let mut i = 0;
        while i < compressed.len() {
            let end = (i + step).min(compressed.len());
            d.decompress(&compressed[i..end], &mut plain).unwrap();
            i = end;
            step = step % 5 + 1; // 1..=5
        }
        d.finish(&mut plain).unwrap();
        assert_eq!(plain, data, "roundtrip mismatch with feed {feed:?}");
    }

    #[test]
    fn roundtrip_single_chunk() {
        let data: Vec<u8> = (0..100_000u32).map(|i| (i % 251) as u8).collect();
        roundtrip_chunks(&data, &[data.len()]);
    }

    #[test]
    fn roundtrip_multi_chunk_64k_windows() {
        // Server-style: feed in 64 KiB sliding windows.
        let data: Vec<u8> = (0..300_000u32).map(|i| (i / 7) as u8).collect();
        let feed: Vec<usize> = std::iter::repeat(STREAM_BUF_SIZE).take(5).collect();
        roundtrip_chunks(&data, &feed);
    }

    #[test]
    fn roundtrip_highly_compressible() {
        let data = b"libfw streaming compression test. ".repeat(10_000);
        // Feed sizes cycling 1KiB..16KiB until the whole input is consumed.
        let mut feed = Vec::new();
        let mut consumed = 0;
        for (_, &size) in [1024usize, 2048, 4096, 8192, 16384].iter().cycle().enumerate() {
            feed.push(size);
            consumed += size;
            if consumed >= data.len() {
                break;
            }
        }
        roundtrip_chunks(&data, &feed);
    }

    #[test]
    fn roundtrip_empty_stream() {
        let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
        let mut compressed = Vec::new();
        c.finish(&mut compressed).unwrap();

        let mut d = ZripDecompressor::new();
        let mut plain = Vec::new();
        d.decompress(&compressed, &mut plain).unwrap();
        d.finish(&mut plain).unwrap();
        assert!(plain.is_empty());
    }

    #[test]
    fn empty_input_chunks_are_noops() {
        let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
        let mut out = Vec::new();
        c.compress(&[], &mut out).unwrap();
        c.compress(b"hello", &mut out).unwrap();
        c.finish(&mut out).unwrap();

        let mut d = ZripDecompressor::new();
        let mut plain = Vec::new();
        d.decompress(&out, &mut plain).unwrap();
        d.finish(&mut plain).unwrap();
        assert_eq!(plain, b"hello");
    }

    #[test]
    fn format_header_roundtrip() {
        assert_eq!(CompressionFormat::parse_header("zrip"), Some(CompressionFormat::Zrip));
        assert_eq!(CompressionFormat::parse_header("ZSTD"), Some(CompressionFormat::Zrip));
        assert_eq!(CompressionFormat::parse_header("identity"), Some(CompressionFormat::None));
        assert_eq!(CompressionFormat::parse_header(""), Some(CompressionFormat::None));
        assert_eq!(CompressionFormat::parse_header("br"), None);
        assert_eq!(CompressionFormat::Zrip.as_str(), "zrip");
        assert_eq!(CompressionFormat::None.as_str(), "identity");
    }

    #[test]
    fn passthrough_roundtrip() {
        let mut c = compressor(CompressionFormat::None).unwrap();
        let mut d = decompressor(CompressionFormat::None);
        let mut compressed = Vec::new();
        let mut plain = Vec::new();
        c.compress(b"abc", &mut compressed).unwrap();
        c.finish(&mut compressed).unwrap();
        d.decompress(&compressed, &mut plain).unwrap();
        d.finish(&mut plain).unwrap();
        assert_eq!(plain, b"abc");
    }

    #[test]
    fn truncated_stream_is_detected() {
        let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
        let mut compressed = Vec::new();
        c.compress(&vec![7u8; 5000], &mut compressed).unwrap();
        c.finish(&mut compressed).unwrap();
        compressed.truncate(compressed.len() - 1); // chop one byte

        let mut d = ZripDecompressor::new();
        let mut plain = Vec::new();
        d.decompress(&compressed, &mut plain).unwrap();
        assert!(matches!(d.finish(&mut plain), Err(DecompressError::Truncated(_))));
    }

    #[test]
    fn corrupt_stream_is_detected() {
        let mut d = ZripDecompressor::new();
        let mut plain = Vec::new();
        let err = d.decompress(b"this is not a zstd frame at all", &mut plain);
        assert!(err.is_err());
    }

    #[test]
    fn zstd_compat_interop() {
        // Our frames are standard zstd: the reference `zstd` codec can
        // decode them, and we can decode its output.
        let data: Vec<u8> = (0..50_000u32).map(|i| (i % 31) as u8).collect();
        let mut c = ZripCompressor::new(ZRIP_DEFAULT_LEVEL).unwrap();
        let mut compressed = Vec::new();
        c.compress(&data, &mut compressed).unwrap();
        c.finish(&mut compressed).unwrap();
        let decoded = zstd::stream::decode_all(&compressed[..]).unwrap();
        assert_eq!(decoded, data);

        let zstd_enc = zstd::stream::encode_all(&data[..], 1).unwrap();
        let mut d = ZripDecompressor::new();
        let mut plain = Vec::new();
        d.decompress(&zstd_enc, &mut plain).unwrap();
        d.finish(&mut plain).unwrap();
        assert_eq!(plain, data);
    }
}