zstd-rs/rust/src/zstd_compress_frame.rs

#![allow(non_camel_case_types)]
#![allow(non_snake_case)]
#![allow(clippy::missing_safety_doc)]

//! Frame-header and frame-trailer writing leaves.
//!
//! The high-level compression context remains in C.  These leaves accept the
//! five scalar frame parameters they need, so no `ZSTD_CCtx_params` layout
//! crosses the language boundary.

use crate::errors::{ZstdErrorCode, ERROR};
use std::ffi::c_void;
use std::os::raw::{c_int, c_uint};
use std::ptr;

const ZSTD_MAGICNUMBER: u32 = 0xFD2F_B528;
const ZSTD_MAGIC_SKIPPABLE_START: u32 = 0x184D_2A50;
const ZSTD_FRAMEHEADERSIZE_MAX: usize = 18;
const ZSTD_SKIPPABLEHEADERSIZE: usize = 8;
const ZSTD_BLOCKHEADERSIZE: usize = 3;
const ZSTD_WINDOWLOG_ABSOLUTEMIN: u32 = 10;
const ZSTD_CONTENTSIZE_UNKNOWN: u64 = u64::MAX;
const ZSTD_F_ZSTD1: c_int = 0;

#[inline]
unsafe fn write_le16(dst: *mut u8, value: u16) {
    let bytes = value.to_le_bytes();
    unsafe { ptr::copy_nonoverlapping(bytes.as_ptr(), dst, bytes.len()) };
}

#[inline]
unsafe fn write_le24(dst: *mut u8, value: u32) {
    let bytes = value.to_le_bytes();
    unsafe { ptr::copy_nonoverlapping(bytes.as_ptr(), dst, 3) };
}

#[inline]
unsafe fn write_le32(dst: *mut u8, value: u32) {
    let bytes = value.to_le_bytes();
    unsafe { ptr::copy_nonoverlapping(bytes.as_ptr(), dst, bytes.len()) };
}

#[inline]
unsafe fn write_le64(dst: *mut u8, value: u64) {
    let bytes = value.to_le_bytes();
    unsafe { ptr::copy_nonoverlapping(bytes.as_ptr(), dst, bytes.len()) };
}

/// Rust implementation of the private `ZSTD_writeFrameHeader()` leaf.
///
/// `no_dict_id_flag`, `checksum_flag`, `content_size_flag`, `format`, and
/// `window_log` are extracted by a C wrapper from `ZSTD_CCtx_params`.  The
/// wrapper preserves the existing static C function signature, avoiding any
/// dependency on the full context-parameter layout here.
///
/// The caller must provide a writable buffer of at least
/// `ZSTD_FRAMEHEADERSIZE_MAX` bytes, or this returns `dstSize_tooSmall`.  As
/// in the C routine, `pledged_src_size` must not be
/// `ZSTD_CONTENTSIZE_UNKNOWN` when `content_size_flag` is nonzero.
#[no_mangle]
pub unsafe extern "C" fn ZSTD_rust_writeFrameHeader(
    dst: *mut c_void,
    dst_capacity: usize,
    no_dict_id_flag: c_int,
    checksum_flag: c_int,
    content_size_flag: c_int,
    format: c_int,
    window_log: c_uint,
    pledged_src_size: u64,
    dict_id: u32,
) -> usize {
    if dst_capacity < ZSTD_FRAMEHEADERSIZE_MAX {
        return ERROR(ZstdErrorCode::DstSizeTooSmall);
    }

    let dst = dst.cast::<u8>();
    let dict_id_size_code_length =
        u32::from(dict_id > 0) + u32::from(dict_id >= 256) + u32::from(dict_id >= 65_536);
    let dict_id_size_code = if no_dict_id_flag != 0 {
        0
    } else {
        dict_id_size_code_length
    };
    let checksum_flag = u32::from(checksum_flag > 0);
    let content_size_flag = content_size_flag != 0;

    debug_assert!(!(content_size_flag && pledged_src_size == ZSTD_CONTENTSIZE_UNKNOWN));

    /* Parameters are validated before this leaf is reached.  Avoid a Rust
     * shift panic if a malformed caller bypasses that validation. */
    let window_size = 1u64.checked_shl(window_log).unwrap_or(0);
    let single_segment = u32::from(content_size_flag && window_size >= pledged_src_size);
    let window_log_byte = window_log
        .wrapping_sub(ZSTD_WINDOWLOG_ABSOLUTEMIN)
        .wrapping_shl(3) as u8;
    let fcs_code = if content_size_flag {
        u32::from(pledged_src_size >= 256)
            + u32::from(pledged_src_size >= 65_536 + 256)
            + u32::from(pledged_src_size >= 0xFFFF_FFFF)
    } else {
        0
    };
    let frame_header_description_byte =
        (dict_id_size_code + (checksum_flag << 2) + (single_segment << 5) + (fcs_code << 6)) as u8;

    let mut pos = 0usize;
    if format == ZSTD_F_ZSTD1 {
        unsafe { write_le32(dst, ZSTD_MAGICNUMBER) };
        pos = 4;
    }
    unsafe { dst.add(pos).write(frame_header_description_byte) };
    pos += 1;
    if single_segment == 0 {
        unsafe { dst.add(pos).write(window_log_byte) };
        pos += 1;
    }

    match dict_id_size_code {
        0 => {}
        1 => {
            unsafe { dst.add(pos).write(dict_id as u8) };
            pos += 1;
        }
        2 => {
            unsafe { write_le16(dst.add(pos), dict_id as u16) };
            pos += 2;
        }
        3 => {
            unsafe { write_le32(dst.add(pos), dict_id) };
            pos += 4;
        }
        _ => unreachable!("dictionary ID size code is bounded to 0..=3"),
    }

    match fcs_code {
        0 => {
            if single_segment != 0 {
                unsafe { dst.add(pos).write(pledged_src_size as u8) };
                pos += 1;
            }
        }
        1 => {
            unsafe { write_le16(dst.add(pos), (pledged_src_size - 256) as u16) };
            pos += 2;
        }
        2 => {
            unsafe { write_le32(dst.add(pos), pledged_src_size as u32) };
            pos += 4;
        }
        3 => {
            unsafe { write_le64(dst.add(pos), pledged_src_size) };
            pos += 8;
        }
        _ => unreachable!("frame content size code is bounded to 0..=3"),
    }
    pos
}

/// Rust implementation of the public `ZSTD_writeSkippableFrame()` ABI.
///
/// Integration removes the C function body, allowing this direct export to
/// provide the existing public symbol without a wrapper.
#[no_mangle]
pub unsafe extern "C" fn ZSTD_writeSkippableFrame(
    dst: *mut c_void,
    dst_capacity: usize,
    src: *const c_void,
    src_size: usize,
    magic_variant: c_uint,
) -> usize {
    let total_size = match src_size.checked_add(ZSTD_SKIPPABLEHEADERSIZE) {
        Some(size) => size,
        None if src_size > u32::MAX as usize => return ERROR(ZstdErrorCode::SrcSizeWrong),
        /* A frame whose payload fills the `size_t` address space cannot have
         * room for its eight-byte header. */
        None => return ERROR(ZstdErrorCode::DstSizeTooSmall),
    };
    if dst_capacity < total_size {
        return ERROR(ZstdErrorCode::DstSizeTooSmall);
    }
    if src_size > u32::MAX as usize {
        return ERROR(ZstdErrorCode::SrcSizeWrong);
    }
    if magic_variant > 15 {
        return ERROR(ZstdErrorCode::ParameterOutOfBound);
    }

    let dst = dst.cast::<u8>();
    let src = src.cast::<u8>();
    unsafe { write_le32(dst, ZSTD_MAGIC_SKIPPABLE_START + magic_variant) };
    unsafe { write_le32(dst.add(4), src_size as u32) };
    if src_size != 0 {
        unsafe { ptr::copy_nonoverlapping(src, dst.add(ZSTD_SKIPPABLEHEADERSIZE), src_size) };
    }
    total_size
}

/// Rust implementation of the private `ZSTD_writeLastEmptyBlock()` leaf.
///
/// Integration removes the C body and declares this symbol before the C call
/// sites, so the direct export needs no context or layout bridge.
#[no_mangle]
pub unsafe extern "C" fn ZSTD_writeLastEmptyBlock(dst: *mut c_void, dst_capacity: usize) -> usize {
    if dst_capacity < ZSTD_BLOCKHEADERSIZE {
        return ERROR(ZstdErrorCode::DstSizeTooSmall);
    }
    /* lastBlock = 1, block type = bt_raw, and block size = 0. */
    unsafe { write_le24(dst.cast::<u8>(), 1) };
    ZSTD_BLOCKHEADERSIZE
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::errors::{ERR_getErrorCode, ERR_isError};

    #[test]
    fn writes_single_segment_empty_frame_header() {
        let mut output = [0u8; ZSTD_FRAMEHEADERSIZE_MAX];
        let written = unsafe {
            ZSTD_rust_writeFrameHeader(
                output.as_mut_ptr().cast(),
                output.len(),
                0,
                0,
                1,
                ZSTD_F_ZSTD1,
                20,
                0,
                0,
            )
        };

        assert_eq!(written, 6);
        assert_eq!(&output[..written], &[0x28, 0xB5, 0x2F, 0xFD, 0x20, 0x00]);
    }

    #[test]
    fn writes_magicless_unknown_size_header() {
        let mut output = [0u8; ZSTD_FRAMEHEADERSIZE_MAX];
        let written = unsafe {
            ZSTD_rust_writeFrameHeader(
                output.as_mut_ptr().cast(),
                output.len(),
                0,
                0,
                0,
                1,
                10,
                ZSTD_CONTENTSIZE_UNKNOWN,
                0,
            )
        };

        assert_eq!(written, 2);
        assert_eq!(&output[..written], &[0x00, 0x00]);
    }

    #[test]
    fn writes_maximal_frame_header() {
        let mut output = [0u8; ZSTD_FRAMEHEADERSIZE_MAX];
        let content_size = 0x1_0000_0000u64;
        let written = unsafe {
            ZSTD_rust_writeFrameHeader(
                output.as_mut_ptr().cast(),
                output.len(),
                0,
                1,
                1,
                ZSTD_F_ZSTD1,
                20,
                content_size,
                0x1234_5678,
            )
        };

        assert_eq!(written, ZSTD_FRAMEHEADERSIZE_MAX);
        assert_eq!(
            &output[..written],
            &[
                0x28, 0xB5, 0x2F, 0xFD, 0xC7, 0x50, 0x78, 0x56, 0x34, 0x12, 0x00, 0x00, 0x00, 0x00,
                0x01, 0x00, 0x00, 0x00,
            ]
        );
    }

    #[test]
    fn encodes_short_dictionary_and_content_size_fields() {
        let mut output = [0u8; ZSTD_FRAMEHEADERSIZE_MAX];
        let written = unsafe {
            ZSTD_rust_writeFrameHeader(
                output.as_mut_ptr().cast(),
                output.len(),
                0,
                0,
                1,
                ZSTD_F_ZSTD1,
                10,
                256,
                255,
            )
        };

        assert_eq!(written, 8);
        assert_eq!(
            &output[..written],
            &[0x28, 0xB5, 0x2F, 0xFD, 0x61, 0xFF, 0x00, 0x00]
        );
    }

    #[test]
    fn ignores_dictionary_id_when_requested() {
        let mut output = [0u8; ZSTD_FRAMEHEADERSIZE_MAX];
        let written = unsafe {
            ZSTD_rust_writeFrameHeader(
                output.as_mut_ptr().cast(),
                output.len(),
                1,
                0,
                1,
                ZSTD_F_ZSTD1,
                10,
                256,
                0x1234_5678,
            )
        };

        assert_eq!(written, 7);
        assert_eq!(
            &output[..written],
            &[0x28, 0xB5, 0x2F, 0xFD, 0x60, 0x00, 0x00]
        );
    }

    #[test]
    fn frame_header_requires_worst_case_capacity() {
        let mut output = [0u8; ZSTD_FRAMEHEADERSIZE_MAX - 1];
        let result = unsafe {
            ZSTD_rust_writeFrameHeader(
                output.as_mut_ptr().cast(),
                output.len(),
                0,
                0,
                1,
                ZSTD_F_ZSTD1,
                20,
                0,
                0,
            )
        };

        assert!(ERR_isError(result));
        assert_eq!(
            ERR_getErrorCode(result),
            ZstdErrorCode::DstSizeTooSmall as i32
        );
    }

    #[test]
    fn writes_skippable_frame_and_validates_parameters() {
        let mut output = [0u8; 11];
        let written = unsafe {
            ZSTD_writeSkippableFrame(
                output.as_mut_ptr().cast(),
                output.len(),
                b"abc".as_ptr().cast(),
                3,
                2,
            )
        };
        assert_eq!(written, output.len());
        assert_eq!(
            output,
            [0x52, 0x2A, 0x4D, 0x18, 0x03, 0x00, 0x00, 0x00, b'a', b'b', b'c']
        );

        let too_small = unsafe {
            ZSTD_writeSkippableFrame(
                output.as_mut_ptr().cast(),
                output.len() - 1,
                b"abc".as_ptr().cast(),
                3,
                2,
            )
        };
        assert_eq!(
            ERR_getErrorCode(too_small),
            ZstdErrorCode::DstSizeTooSmall as i32
        );

        let invalid_variant = unsafe {
            ZSTD_writeSkippableFrame(
                output.as_mut_ptr().cast(),
                output.len(),
                b"abc".as_ptr().cast(),
                3,
                16,
            )
        };
        assert_eq!(
            ERR_getErrorCode(invalid_variant),
            ZstdErrorCode::ParameterOutOfBound as i32
        );
    }

    #[test]
    fn writes_last_empty_block() {
        let mut output = [0u8; ZSTD_BLOCKHEADERSIZE];
        assert_eq!(
            unsafe { ZSTD_writeLastEmptyBlock(output.as_mut_ptr().cast(), output.len()) },
            ZSTD_BLOCKHEADERSIZE
        );
        assert_eq!(output, [1, 0, 0]);

        let result = unsafe { ZSTD_writeLastEmptyBlock(output.as_mut_ptr().cast(), 2) };
        assert_eq!(
            ERR_getErrorCode(result),
            ZstdErrorCode::DstSizeTooSmall as i32
        );
    }
}