From 459782ad55e0161cfde8bb9066db85d9ddbade7d Mon Sep 17 00:00:00 2001 From: ddidderr Date: Sat, 18 Jul 2026 05:49:37 +0200 Subject: [PATCH] feat(compress): move frame epilogue serialization to Rust Port the scalar frame-finalization policy and byte serialization into the Rust compression leaf while keeping checksum finalization, tracing, and the compression context in C. Preserve the init-to-ongoing transition on partial failure and reset the stage only after a successful epilogue. Test Plan: - cargo test --manifest-path rust/Cargo.toml --no-default-features --features compression - cargo clippy --manifest-path rust/Cargo.toml - cargo clippy --manifest-path rust/Cargo.toml --benches - cargo clippy --manifest-path rust/Cargo.toml --tests - make -B -C lib -j2 lib - make -B -C tests -j2 test-zstream (84 deterministic, 6737 and 9012 fuzzer cases) --- lib/compress/zstd_compress.c | 49 +++------ rust/src/zstd_compress_frame.rs | 174 +++++++++++++++++++++++++++++++- 2 files changed, 189 insertions(+), 34 deletions(-) diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index bf12e3614..35f3093a8 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -37,6 +37,10 @@ size_t ZSTD_rust_writeFrameHeader(void* dst, size_t dstCapacity, int contentSizeFlag, int format, U32 windowLog, U64 pledgedSrcSize, U32 dictID); +size_t ZSTD_rust_writeEpilogue(void* dst, size_t dstCapacity, int* stage, + int noDictIDFlag, int checksumFlag, + int contentSizeFlag, int format, U32 windowLog, + U32 checksum); size_t ZSTD_rust_resetCCtxForSimpleCompression(void* cctx); size_t ZSTD_rust_prepareCCtxForSimpleCompression(void* cctx, size_t srcSize, @@ -3938,41 +3942,20 @@ size_t ZSTD_compressBegin(ZSTD_CCtx* cctx, int compressionLevel) * @return : nb of bytes written into dst (or an error code) */ static size_t ZSTD_writeEpilogue(ZSTD_CCtx* cctx, void* dst, size_t dstCapacity) { - BYTE* const ostart = (BYTE*)dst; - BYTE* op = ostart; - + U32 const checksum = cctx->appliedParams.fParams.checksumFlag + ? (U32)XXH64_digest(&cctx->xxhState) + : 0; DEBUGLOG(4, "ZSTD_writeEpilogue"); - RETURN_ERROR_IF(cctx->stage == ZSTDcs_created, stage_wrong, "init missing"); - - /* special case : empty frame */ - if (cctx->stage == ZSTDcs_init) { - size_t fhSize = ZSTD_writeFrameHeader(dst, dstCapacity, &cctx->appliedParams, 0, 0); - FORWARD_IF_ERROR(fhSize, "ZSTD_writeFrameHeader failed"); - dstCapacity -= fhSize; - op += fhSize; - cctx->stage = ZSTDcs_ongoing; - } - - if (cctx->stage != ZSTDcs_ending) { - /* write one last empty block, make it the "last" block */ - U32 const cBlockHeader24 = 1 /* last block */ + (((U32)bt_raw)<<1) + 0; - ZSTD_STATIC_ASSERT(ZSTD_BLOCKHEADERSIZE == 3); - RETURN_ERROR_IF(dstCapacity<3, dstSize_tooSmall, "no room for epilogue"); - MEM_writeLE24(op, cBlockHeader24); - op += ZSTD_blockHeaderSize; - dstCapacity -= ZSTD_blockHeaderSize; - } - - if (cctx->appliedParams.fParams.checksumFlag) { - U32 const checksum = (U32) XXH64_digest(&cctx->xxhState); - RETURN_ERROR_IF(dstCapacity<4, dstSize_tooSmall, "no room for checksum"); + if (cctx->appliedParams.fParams.checksumFlag) DEBUGLOG(4, "ZSTD_writeEpilogue: write checksum : %08X", (unsigned)checksum); - MEM_writeLE32(op, checksum); - op += 4; - } - - cctx->stage = ZSTDcs_created; /* return to "created but no init" status */ - return (size_t)(op-ostart); + return ZSTD_rust_writeEpilogue( + dst, dstCapacity, (int*)&cctx->stage, + cctx->appliedParams.fParams.noDictIDFlag, + cctx->appliedParams.fParams.checksumFlag, + cctx->appliedParams.fParams.contentSizeFlag, + (int)cctx->appliedParams.format, + cctx->appliedParams.cParams.windowLog, + checksum); } void ZSTD_CCtx_trace(ZSTD_CCtx* cctx, size_t extraCSize) diff --git a/rust/src/zstd_compress_frame.rs b/rust/src/zstd_compress_frame.rs index b3e6d9c0c..6819705c9 100644 --- a/rust/src/zstd_compress_frame.rs +++ b/rust/src/zstd_compress_frame.rs @@ -8,7 +8,7 @@ //! five scalar frame parameters they need, so no `ZSTD_CCtx_params` layout //! crosses the language boundary. -use crate::errors::{ZstdErrorCode, ERROR}; +use crate::errors::{ERR_isError, ZstdErrorCode, ERROR}; use crate::zstd_presplit::ZSTD_splitBlock; use std::ffi::c_void; use std::os::raw::{c_int, c_longlong, c_uint}; @@ -26,6 +26,10 @@ const ZSTD_BLOCK_SIZE: usize = 128 << 10; const ZSTD_FAST: c_int = 1; const ZSTD_BTULTRA2: c_int = 9; const SPLIT_LEVELS: [c_int; 10] = [0, 0, 1, 2, 2, 3, 3, 4, 4, 4]; +const ZSTDCS_CREATED: c_int = 0; +const ZSTDCS_INIT: c_int = 1; +const ZSTDCS_ONGOING: c_int = 2; +const ZSTDCS_ENDING: c_int = 3; #[inline] unsafe fn write_le16(dst: *mut u8, value: u16) { @@ -250,6 +254,82 @@ pub unsafe extern "C" fn ZSTD_rust_optimalBlockSize( unsafe { ZSTD_splitBlock(src, block_size_max, split_level, workspace, workspace_size) } } +/// Rust implementation of the private `ZSTD_writeEpilogue()` serializer. +/// +/// The C caller owns the compression context and passes its stage by scalar +/// pointer so the transition to `ongoing` after an empty-frame header remains +/// visible even when a later write fails. The checksum is already finalized +/// by C; Rust only serializes its low 32 bits. +#[no_mangle] +pub unsafe extern "C" fn ZSTD_rust_writeEpilogue( + dst: *mut c_void, + dst_capacity: usize, + stage: *mut c_int, + no_dict_id_flag: c_int, + checksum_flag: c_int, + content_size_flag: c_int, + format: c_int, + window_log: c_uint, + checksum: u32, +) -> usize { + debug_assert!(!stage.is_null()); + if stage.is_null() { + return ERROR(ZstdErrorCode::StageWrong); + } + + let mut stage_value = unsafe { stage.read() }; + if stage_value == ZSTDCS_CREATED { + return ERROR(ZstdErrorCode::StageWrong); + } + + let start = dst.cast::(); + let mut output = start; + let mut capacity = dst_capacity; + + if stage_value == ZSTDCS_INIT { + let header_size = unsafe { + ZSTD_rust_writeFrameHeader( + output.cast(), + capacity, + no_dict_id_flag, + checksum_flag, + content_size_flag, + format, + window_log, + 0, + 0, + ) + }; + if ERR_isError(header_size) { + return header_size; + } + output = unsafe { output.add(header_size) }; + capacity -= header_size; + stage_value = ZSTDCS_ONGOING; + unsafe { stage.write(stage_value) }; + } + + if stage_value != ZSTDCS_ENDING { + let block_size = unsafe { ZSTD_writeLastEmptyBlock(output.cast(), capacity) }; + if ERR_isError(block_size) { + return block_size; + } + output = unsafe { output.add(block_size) }; + capacity -= block_size; + } + + if checksum_flag != 0 { + if capacity < 4 { + return ERROR(ZstdErrorCode::DstSizeTooSmall); + } + unsafe { write_le32(output, checksum) }; + output = unsafe { output.add(4) }; + } + + unsafe { stage.write(ZSTDCS_CREATED) }; + output as usize - start as usize +} + #[cfg(test)] mod tests { use super::*; @@ -530,4 +610,96 @@ mod tests { assert!(result > 0); assert!(result <= ZSTD_BLOCK_SIZE); } + + #[test] + fn epilogue_rejects_created_stage() { + let mut output = [0u8; ZSTD_FRAMEHEADERSIZE_MAX]; + let mut stage = ZSTDCS_CREATED; + let result = unsafe { + ZSTD_rust_writeEpilogue( + output.as_mut_ptr().cast(), + output.len(), + &mut stage, + 0, + 0, + 1, + ZSTD_F_ZSTD1, + 20, + 0, + ) + }; + assert_eq!(ERR_getErrorCode(result), ZstdErrorCode::StageWrong as i32); + assert_eq!(stage, ZSTDCS_CREATED); + } + + #[test] + fn epilogue_writes_empty_frame_and_resets_stage() { + let mut output = [0u8; 18]; + let mut stage = ZSTDCS_INIT; + let result = unsafe { + ZSTD_rust_writeEpilogue( + output.as_mut_ptr().cast(), + output.len(), + &mut stage, + 0, + 0, + 1, + ZSTD_F_ZSTD1, + 20, + 0, + ) + }; + assert_eq!(result, 9); + assert_eq!( + &output[..result], + &[0x28, 0xB5, 0x2F, 0xFD, 0x20, 0x00, 1, 0, 0] + ); + assert_eq!(stage, ZSTDCS_CREATED); + } + + #[test] + fn epilogue_ending_stage_writes_only_checksum() { + let mut output = [0u8; 4]; + let mut stage = ZSTDCS_ENDING; + let result = unsafe { + ZSTD_rust_writeEpilogue( + output.as_mut_ptr().cast(), + output.len(), + &mut stage, + 0, + 1, + 0, + ZSTD_F_ZSTD1, + 10, + 0x1234_5678, + ) + }; + assert_eq!(result, 4); + assert_eq!(output, [0x78, 0x56, 0x34, 0x12]); + assert_eq!(stage, ZSTDCS_CREATED); + } + + #[test] + fn epilogue_preserves_ongoing_stage_when_checksum_does_not_fit() { + let mut output = [0u8; 6]; + let mut stage = ZSTDCS_ONGOING; + let result = unsafe { + ZSTD_rust_writeEpilogue( + output.as_mut_ptr().cast(), + output.len(), + &mut stage, + 0, + 1, + 0, + ZSTD_F_ZSTD1, + 10, + 0x1234_5678, + ) + }; + assert_eq!( + ERR_getErrorCode(result), + ZstdErrorCode::DstSizeTooSmall as i32 + ); + assert_eq!(stage, ZSTDCS_ONGOING); + } }