crabka-compression 0.3.9

//! Xerial-snappy framing over `snap` raw blocks.
//!
//! Kafka does not use Google's official Snappy stream format. It uses
//! "xerial-snappy", a Java-library convention with this layout:
//!
//! ```text
//! [\x82 'S' 'N' 'A' 'P' 'P' 'Y' \x00]                 # 8-byte magic
//! [\x00 \x00 \x00 \x01]                               # version       (BE u32)
//! [\x00 \x00 \x00 \x01]                               # minCompatibleVersion (BE u32)
//! ( [BE u32 chunk length] [raw snappy block ...] )*   # zero or more chunks
//! ```
//!
//! There is no end-of-stream marker; chunks run until EOF.
//!
//! JVM xerial-snappy byte equality is verified by the differential test
//! suite once the oracle gains a `compress` op.

use bytes::{BufMut, Bytes, BytesMut};

use crate::CompressionError;

/// Xerial framing header. 16 bytes total.
const XERIAL_HEADER: [u8; 16] = [
    0x82, b'S', b'N', b'A', b'P', b'P', b'Y', 0x00, 0x00, 0x00, 0x00, 0x01, // version = 1
    0x00, 0x00, 0x00, 0x01, // minCompatibleVersion = 1
];

/// Largest single chunk Kafka writes. Kafka's `SnappyOutputStream` writes
/// chunks up to 32 KiB by default; using the same size keeps our output
/// byte-identical with the JVM for differential-equal cases.
const XERIAL_CHUNK_SIZE: usize = 32 * 1024;

pub fn compress(data: &[u8]) -> Result<Bytes, CompressionError> {
    let mut out = BytesMut::with_capacity(XERIAL_HEADER.len() + data.len());
    out.put_slice(&XERIAL_HEADER);

    let mut encoder = snap::raw::Encoder::new();
    for chunk in data.chunks(XERIAL_CHUNK_SIZE) {
        let max = snap::raw::max_compress_len(chunk.len());
        let mut buf = vec![0u8; max];
        let n = encoder
            .compress(chunk, &mut buf)
            .map_err(|e| CompressionError::InvalidData(format!("snappy encode: {e}")))?;
        out.put_u32(u32::try_from(n).expect("chunk size fits u32"));
        out.put_slice(&buf[..n]);
    }
    Ok(out.freeze())
}

pub fn decompress(data: &[u8], max_output: usize) -> Result<Bytes, CompressionError> {
    if data.len() < XERIAL_HEADER.len() {
        return Err(CompressionError::InvalidData(
            "snappy payload too short for xerial header".into(),
        ));
    }
    if data[..8] != XERIAL_HEADER[..8] {
        return Err(CompressionError::InvalidData(
            "snappy missing xerial magic".into(),
        ));
    }
    // Ignore version fields (bytes 8..16); Kafka never bumped them.
    let mut rest = &data[XERIAL_HEADER.len()..];

    let mut out = BytesMut::with_capacity(data.len().saturating_mul(2).min(max_output));
    let mut decoder = snap::raw::Decoder::new();
    while !rest.is_empty() {
        if rest.len() < 4 {
            return Err(CompressionError::InvalidData(
                "snappy chunk header truncated".into(),
            ));
        }
        let len = u32::from_be_bytes([rest[0], rest[1], rest[2], rest[3]]) as usize;
        rest = &rest[4..];
        if rest.len() < len {
            return Err(CompressionError::InvalidData(
                "snappy chunk body truncated".into(),
            ));
        }
        let (block, tail) = rest.split_at(len);
        rest = tail;

        let max_out = snap::raw::decompress_len(block)
            .map_err(|e| CompressionError::InvalidData(format!("snappy decode_len: {e}")))?;
        // Reject before allocating this chunk if it would push us past the cap.
        // `decompress_len` reads the block's stored uncompressed size, so this
        // bounds allocation without materializing the oversized output.
        if out.len().saturating_add(max_out) > max_output {
            return Err(CompressionError::TooLarge { limit: max_output });
        }
        let mut buf = vec![0u8; max_out];
        let n = decoder
            .decompress(block, &mut buf)
            .map_err(|e| CompressionError::InvalidData(format!("snappy decode: {e}")))?;
        out.put_slice(&buf[..n]);
    }
    Ok(out.freeze())
}

#[cfg(test)]
mod tests {
    use assert2::assert;

    use super::*;

    const HELLO: &[u8] = b"hello kafka, this is a moderately repetitive payload to compress";
    const BIG_CAP: usize = 256 * 1024 * 1024;

    #[test]
    fn roundtrip() {
        let z = compress(HELLO).unwrap();
        let back = decompress(&z, BIG_CAP).unwrap();
        assert!(back.as_ref() == HELLO);
    }

    #[test]
    fn decompress_truncated_header() {
        assert!(matches!(
            decompress(&XERIAL_HEADER[..4], BIG_CAP),
            Err(CompressionError::InvalidData(_))
        ));
    }

    #[test]
    fn decompress_missing_magic() {
        let bytes = [0u8; 20];
        assert!(matches!(
            decompress(&bytes, BIG_CAP),
            Err(CompressionError::InvalidData(_))
        ));
    }

    #[test]
    fn decompress_truncated_chunk() {
        let mut bytes = XERIAL_HEADER.to_vec();
        bytes.extend_from_slice(&[0, 0, 0, 100]); // claim 100-byte chunk
        bytes.push(0); // only 1 byte present
        assert!(matches!(
            decompress(&bytes, BIG_CAP),
            Err(CompressionError::InvalidData(_))
        ));
    }

    #[test]
    fn decompression_bomb_rejected() {
        let bomb = vec![0u8; 64 * 1024 * 1024];
        let z = compress(&bomb).unwrap();
        assert!(matches!(
            decompress(&z, 1024),
            Err(CompressionError::TooLarge { limit: 1024 })
        ));
        let back = decompress(&z, BIG_CAP).unwrap();
        assert!(back.len() == bomb.len());
    }

    #[test]
    fn decompress_at_exact_cap_succeeds() {
        let z = compress(HELLO).unwrap();
        // The cumulative per-chunk cap check is `out + chunk > max_output`
        // (not `>=`), so a cap equal to the exact output size must pass.
        let back = decompress(&z, HELLO.len()).unwrap();
        assert!(back.as_ref() == HELLO);
        // One byte under the exact size is rejected.
        assert!(matches!(
            decompress(&z, HELLO.len() - 1),
            Err(CompressionError::TooLarge { limit }) if limit == HELLO.len() - 1
        ));
    }

    #[test]
    fn compress_uses_32kib_chunks() {
        // A payload between 1 KiB and 32 KiB must serialize as a SINGLE xerial
        // chunk (chunk size is 32 * 1024). A smaller chunk size would split it
        // into several length-prefixed chunks.
        let payload = vec![0x5Au8; 4096];
        let z = compress(&payload).unwrap();
        let mut rest = &z[XERIAL_HEADER.len()..];
        let mut chunks = 0;
        while rest.len() >= 4 {
            let len = u32::from_be_bytes([rest[0], rest[1], rest[2], rest[3]]) as usize;
            rest = &rest[4 + len..];
            chunks += 1;
        }
        assert!(chunks == 1, "expected one 32 KiB chunk, got {chunks}");
    }

    #[test]
    fn decompress_bare_header_is_empty() {
        // compress("") emits exactly the 16-byte xerial header and no chunks;
        // decompressing it must succeed and yield empty. Boundary on the
        // `data.len() < header` guard: a length equal to the header is valid.
        let z = compress(b"").unwrap();
        assert!(z.len() == XERIAL_HEADER.len());
        let back = decompress(&z, 1024).unwrap();
        assert!(back.is_empty());
    }

    #[test]
    fn decompress_four_byte_chunk_header_reads_length() {
        // With exactly a 4-byte chunk header present and no body, we must get
        // PAST the "header truncated" guard (`rest.len() < 4` is false at 4)
        // and fail on the missing *body* instead. Pins the `< 4` boundary.
        let mut bytes = XERIAL_HEADER.to_vec();
        bytes.extend_from_slice(&[0, 0, 0, 5]); // claim a 5-byte chunk, supply none
        let Err(CompressionError::InvalidData(msg)) = decompress(&bytes, 1024) else {
            panic!("expected InvalidData");
        };
        assert!(msg.contains("body"), "msg={msg}");
    }
}