diff --git a/programs/fileio.c b/programs/fileio.c index 60798d43f..3d2a33c25 100644 --- a/programs/fileio.c +++ b/programs/fileio.c @@ -412,15 +412,15 @@ enum { typedef void (*FIO_rust_frame_progress_fn)(void* opaque, const char* srcFileName, U64 decodedSize); -int FIO_rust_decompressZstdFrame(void* fCtx, - void* dctx, - ReadPoolCtx_t* readCtx, - WritePoolCtx_t* writeCtx, - const char* srcFileName, - U64 alreadyDecoded, - U64* frameSize, - size_t* zstdError, - FIO_rust_frame_progress_fn progress); +int FIO_rust_decompressZstdFrames(void* fCtx, + void* dctx, + ReadPoolCtx_t* readCtx, + WritePoolCtx_t* writeCtx, + const char* srcFileName, + U64 alreadyDecoded, + U64* decodedSize, + size_t* zstdError, + FIO_rust_frame_progress_fn progress); void FIO_rust_displayCompressionParameters(const FIO_prefs_t* prefs); #ifdef ZSTD_LZ4COMPRESS int FIO_rust_LZ4_GetBlockSize_FromBlockId(int id); @@ -2235,20 +2235,20 @@ FIO_decompressZstdFrameProgress(void* const opaque, } static unsigned long long -FIO_decompressZstdFrame(FIO_ctx_t* const fCtx, dRess_t* ress, - const FIO_prefs_t* const prefs, - const char* srcFileName, - U64 alreadyDecoded) /* for multi-frames streams */ +FIO_decompressZstdFrames(FIO_ctx_t* const fCtx, dRess_t* ress, + const FIO_prefs_t* const prefs, + const char* srcFileName, + U64 alreadyDecoded) /* for multi-frames streams */ { - U64 frameSize = 0; + U64 decodedSize = 0; size_t zstdError = 0; - int const status = FIO_rust_decompressZstdFrame( + int const status = FIO_rust_decompressZstdFrames( fCtx, ress->dctx, ress->readCtx, ress->writeCtx, srcFileName, - alreadyDecoded, &frameSize, &zstdError, + alreadyDecoded, &decodedSize, &zstdError, FIO_decompressZstdFrameProgress); if (status == FIO_RUST_ZSTD_FRAME_OK) - return frameSize; + return decodedSize; if (status == FIO_RUST_ZSTD_FRAME_DECODING_ERROR) { DISPLAYLEVEL(1, "%s : Decoding error (36) : %s \n", srcFileName, ZSTD_getErrorName(zstdError)); @@ -2522,7 +2522,7 @@ static int FIO_decompressFrames(FIO_ctx_t* const fCtx, return 1; } if (ZSTD_isFrame(buf, ress.readCtx->srcBufferLoaded)) { - unsigned long long const frameSize = FIO_decompressZstdFrame(fCtx, &ress, prefs, srcFileName, filesize); + unsigned long long const frameSize = FIO_decompressZstdFrames(fCtx, &ress, prefs, srcFileName, filesize); if (frameSize == FIO_ERROR_FRAME_DECODING) return 1; filesize += frameSize; } else if (buf[0] == 31 && buf[1] == 139) { /* gz magic number */ diff --git a/rust/src/fileio_asyncio.rs b/rust/src/fileio_asyncio.rs index 920d9e682..d6b19ff1e 100644 --- a/rust/src/fileio_asyncio.rs +++ b/rust/src/fileio_asyncio.rs @@ -45,6 +45,7 @@ type FIO_zstd_decompress_fn = unsafe extern "C" fn( *mut crate::zstd_decompress::ZSTD_inBuffer, ) -> usize; type FIO_zstd_in_size_fn = extern "C" fn() -> usize; +type FIO_zstd_is_frame_fn = unsafe extern "C" fn(*const c_void, usize) -> c_uint; /// C's `FIO_prefs_t` from `programs/fileio_types.h`. /// @@ -1503,6 +1504,81 @@ unsafe fn decompress_zstd_frame_with( } } +/// Decompresses consecutive zstd frames until the next format boundary. +/// +/// The caller has already identified the first zstd frame. After each frame, +/// this helper fills only enough input to inspect the next four bytes and +/// leaves a non-zstd header or a short trailing buffer untouched for C's +/// mixed-format dispatcher. Completed-frame output remains accounted for if +/// a later frame reports an error. +unsafe fn decompress_zstd_frames_with( + f_ctx: *mut c_void, + dctx: *mut c_void, + read_ctx: *mut ReadPoolCtx_t, + write_ctx: *mut WritePoolCtx_t, + src_file_name: *const c_char, + already_decoded: u64, + decoded_size: *mut u64, + zstd_error: *mut usize, + progress: FIO_rust_frame_progress_fn, + reset: FIO_zstd_reset_fn, + decompress: FIO_zstd_decompress_fn, + dstream_in_size: FIO_zstd_in_size_fn, + is_frame: FIO_zstd_is_frame_fn, +) -> c_int { + assert!(!dctx.is_null()); + assert!(!read_ctx.is_null()); + assert!(!write_ctx.is_null()); + assert!(!src_file_name.is_null()); + assert!(!decoded_size.is_null()); + assert!(!zstd_error.is_null()); + + unsafe { + *decoded_size = 0; + *zstd_error = 0; + } + + loop { + unsafe { AIO_ReadPool_fillBuffer(read_ctx, 4) }; + let loaded = unsafe { read_buffer_loaded(read_ctx) }; + if loaded < 4 { + return FIO_RUST_ZSTD_FRAME_OK; + } + + let source = unsafe { read_buffer_ptr(read_ctx) }; + let is_zstd_frame = unsafe { is_frame(source.cast::(), loaded) != 0 }; + if !is_zstd_frame { + return FIO_RUST_ZSTD_FRAME_OK; + } + + let mut frame_size = 0; + let decoded_before_frame = unsafe { *decoded_size }; + let status = unsafe { + decompress_zstd_frame_with( + f_ctx, + dctx, + read_ctx, + write_ctx, + src_file_name, + already_decoded.wrapping_add(decoded_before_frame), + &mut frame_size, + zstd_error, + progress, + reset, + decompress, + dstream_in_size, + ) + }; + if status != FIO_RUST_ZSTD_FRAME_OK { + return status; + } + + unsafe { + *decoded_size = decoded_before_frame.wrapping_add(frame_size); + } + } +} + unsafe extern "C" fn fio_zstd_reset(dctx: *mut c_void, reset: c_int) -> usize { unsafe { crate::zstd_decompress::ZSTD_DCtx_reset( @@ -1556,6 +1632,37 @@ pub unsafe extern "C" fn FIO_rust_decompressZstdFrame( } } +#[no_mangle] +pub unsafe extern "C" fn FIO_rust_decompressZstdFrames( + f_ctx: *mut c_void, + dctx: *mut c_void, + read_ctx: *mut ReadPoolCtx_t, + write_ctx: *mut WritePoolCtx_t, + src_file_name: *const c_char, + already_decoded: u64, + decoded_size: *mut u64, + zstd_error: *mut usize, + progress: FIO_rust_frame_progress_fn, +) -> c_int { + unsafe { + decompress_zstd_frames_with( + f_ctx, + dctx, + read_ctx, + write_ctx, + src_file_name, + already_decoded, + decoded_size, + zstd_error, + progress, + fio_zstd_reset, + fio_zstd_decompress, + crate::zstd_decompress::ZSTD_DStreamInSize, + crate::zstd_decompress::ZSTD_isFrame, + ) + } +} + #[inline] unsafe fn read_buffer_ptr(ctx: *mut ReadPoolCtx_t) -> *const u8 { let context = ctx.cast::(); @@ -2038,6 +2145,7 @@ mod tests { struct ProgressState { calls: usize, last_decoded: u64, + decoded_values: Vec, } unsafe extern "C" fn record_progress( @@ -2050,11 +2158,12 @@ mod tests { let state = unsafe { &mut *opaque.cast::() }; state.calls += 1; state.last_decoded = decoded_size; + state.decoded_values.push(decoded_size); } - const MOCK_FRAME_HEADER_SIZE: usize = 5; + const MOCK_FRAME_HEADER_SIZE: usize = 8; const MOCK_DSTREAM_IN_SIZE: usize = 8; - const MOCK_FRAME_MAGIC: u8 = 0xA5; + const MOCK_FRAME_MAGIC: [u8; 4] = [0x28, 0xB5, 0x2F, 0xFD]; struct MockDecoder { header: [u8; MOCK_FRAME_HEADER_SIZE], @@ -2081,12 +2190,21 @@ mod tests { fn encode_frame(input: &[u8]) -> Vec { assert!(u32::try_from(input.len()).is_ok()); let mut frame = Vec::with_capacity(MOCK_FRAME_HEADER_SIZE + input.len()); - frame.push(MOCK_FRAME_MAGIC); + frame.extend_from_slice(&MOCK_FRAME_MAGIC); frame.extend_from_slice(&(input.len() as u32).to_le_bytes()); frame.extend_from_slice(input); frame } + unsafe extern "C" fn mock_is_frame(buffer: *const c_void, size: usize) -> c_uint { + if size < MOCK_FRAME_MAGIC.len() || buffer.is_null() { + return 0; + } + let magic = + unsafe { std::slice::from_raw_parts(buffer.cast::(), MOCK_FRAME_MAGIC.len()) }; + u32::from(magic == MOCK_FRAME_MAGIC) + } + unsafe extern "C" fn mock_reset(dctx: *mut c_void, _reset: c_int) -> usize { unsafe { (*dctx.cast::()).reset() }; 0 @@ -2111,16 +2229,18 @@ mod tests { if decoder.header_len < MOCK_FRAME_HEADER_SIZE { return (MOCK_FRAME_HEADER_SIZE - decoder.header_len).max(1); } - if decoder.header[0] != MOCK_FRAME_MAGIC { + if decoder.header[..MOCK_FRAME_MAGIC.len()] != MOCK_FRAME_MAGIC + || decoder.header[4..].iter().all(|&byte| byte == 0xFF) + { input.pos = call_start; return crate::errors::ERROR(crate::errors::ZstdErrorCode::Generic); } if decoder.remaining == 0 { let length = u32::from_le_bytes([ - decoder.header[1], - decoder.header[2], - decoder.header[3], decoder.header[4], + decoder.header[5], + decoder.header[6], + decoder.header[7], ]); decoder.remaining = length as usize; if decoder.remaining == 0 { @@ -2235,6 +2355,41 @@ mod tests { (status, frame_size, zstd_error) } + fn decompress_frames( + &mut self, + already_decoded: u64, + progress_state: Option<&mut ProgressState>, + ) -> (c_int, u64, usize) { + let mut decoded_size = 0; + let mut zstd_error = 0; + let progress_context = progress_state + .map(|state| state as *mut ProgressState as *mut c_void) + .unwrap_or(ptr::null_mut()); + let progress_callback = if progress_context.is_null() { + None + } else { + Some(record_progress as unsafe extern "C" fn(*mut c_void, *const c_char, u64)) + }; + let status = unsafe { + decompress_zstd_frames_with( + progress_context, + self.dctx, + self.read_ctx, + self.write_ctx, + c"frame-test.zst".as_ptr(), + already_decoded, + &mut decoded_size, + &mut zstd_error, + progress_callback, + mock_reset, + mock_decompress, + mock_dstream_in_size, + mock_is_frame, + ) + }; + (status, decoded_size, zstd_error) + } + fn unread_bytes(&self) -> Vec { let loaded = unsafe { read_buffer_loaded(self.read_ctx) }; let source = unsafe { read_buffer_ptr(self.read_ctx) }; @@ -2310,6 +2465,117 @@ mod tests { assert_eq!(progress.last_decoded, 37); } + #[test] + fn decompresses_one_frame_and_stops_at_eof() { + let input = b"one frame followed by eof".repeat(4_000); + let frame = encode_frame(&input); + + for async_io in [0, 1] { + let mut harness = FrameHarness::new(&frame, 113, async_io); + let (status, decoded_size, zstd_error) = harness.decompress_frames(0, None); + assert_eq!(status, FIO_RUST_ZSTD_FRAME_OK); + assert_eq!(decoded_size, input.len() as u64); + assert_eq!(zstd_error, 0); + assert!(harness.unread_bytes().is_empty()); + assert_eq!(harness.output_bytes(), input); + } + } + + #[test] + fn decompresses_concatenated_frames_and_accumulates_output() { + let first = b"first concatenated frame".repeat(1_000); + let second = b"second concatenated frame".repeat(1_000); + let mut stream = encode_frame(&first); + stream.extend_from_slice(&encode_frame(&second)); + + let mut harness = FrameHarness::new(&stream, 113, 1); + let (status, decoded_size, zstd_error) = harness.decompress_frames(0, None); + assert_eq!(status, FIO_RUST_ZSTD_FRAME_OK); + assert_eq!(decoded_size, (first.len() + second.len()) as u64); + assert_eq!(zstd_error, 0); + assert!(harness.unread_bytes().is_empty()); + + let mut expected = first; + expected.extend_from_slice(&second); + assert_eq!(harness.output_bytes(), expected); + } + + #[test] + fn stops_before_a_following_non_zstd_format_header() { + let input = b"zstd before gzip".repeat(100); + let frame = encode_frame(&input); + let following_format = [0x1F, 0x8B, 0x08, 0x00, 0xAA]; + let mut stream = frame; + stream.extend_from_slice(&following_format); + + let mut harness = FrameHarness::new(&stream, stream.len(), 0); + let (status, decoded_size, zstd_error) = harness.decompress_frames(0, None); + assert_eq!(status, FIO_RUST_ZSTD_FRAME_OK); + assert_eq!(decoded_size, input.len() as u64); + assert_eq!(zstd_error, 0); + assert_eq!(harness.unread_bytes(), following_format); + assert_eq!(harness.output_bytes(), input); + } + + #[test] + fn stops_before_one_to_three_trailing_bytes() { + let input = b"zstd before a short trailing header"; + let frame = encode_frame(input); + + for trailing_len in 1..=3 { + let trailing: Vec = (0..trailing_len).map(|byte| 0xC0 + byte).collect(); + let mut stream = frame.clone(); + stream.extend_from_slice(&trailing); + let mut harness = FrameHarness::new(&stream, stream.len(), 0); + + let (status, decoded_size, zstd_error) = harness.decompress_frames(0, None); + assert_eq!(status, FIO_RUST_ZSTD_FRAME_OK); + assert_eq!(decoded_size, input.len() as u64); + assert_eq!(zstd_error, 0); + assert_eq!(harness.unread_bytes(), trailing); + assert_eq!(harness.output_bytes(), input); + } + } + + #[test] + fn second_frame_error_preserves_its_input() { + let first = b"first frame is valid".repeat(100); + let mut invalid_second = MOCK_FRAME_MAGIC.to_vec(); + invalid_second.extend_from_slice(&u32::MAX.to_le_bytes()); + invalid_second.extend_from_slice(b"invalid second frame"); + let mut stream = encode_frame(&first); + stream.extend_from_slice(&invalid_second); + + let mut harness = FrameHarness::new(&stream, stream.len(), 0); + let (status, decoded_size, zstd_error) = harness.decompress_frames(0, None); + assert_eq!(status, FIO_RUST_ZSTD_FRAME_DECODING_ERROR); + assert_eq!(decoded_size, first.len() as u64); + assert_ne!(zstd_error, 0); + assert_eq!(harness.unread_bytes(), invalid_second); + assert_eq!(harness.output_bytes(), first); + } + + #[test] + fn consecutive_frame_progress_uses_the_nonzero_base() { + let first = b"first progress frame"; + let second = b"second progress frame"; + let mut stream = encode_frame(first); + stream.extend_from_slice(&encode_frame(second)); + + let mut harness = FrameHarness::new(&stream, stream.len(), 0); + let mut progress = ProgressState::default(); + let (status, decoded_size, zstd_error) = + harness.decompress_frames(41, Some(&mut progress)); + assert_eq!(status, FIO_RUST_ZSTD_FRAME_OK); + assert_eq!(decoded_size, (first.len() + second.len()) as u64); + assert_eq!(zstd_error, 0); + assert_eq!(progress.decoded_values, vec![41, 41 + first.len() as u64]); + assert_eq!( + harness.output_bytes(), + [first.as_slice(), second.as_slice()].concat() + ); + } + #[test] fn decompresses_concatenated_frames_without_crossing_boundary() { let first = b"first frame ".repeat(20_000);