diff --git a/lib/decompress/zstd_decompress.c b/lib/decompress/zstd_decompress.c index 2d8246fd9..5d4ec0474 100644 --- a/lib/decompress/zstd_decompress.c +++ b/lib/decompress/zstd_decompress.c @@ -137,13 +137,56 @@ void ZSTD_dctx_init_platform(ZSTD_DCtx* dctx); void ZSTD_rust_dctx_init_platform(void* bmi2, int dynamic_bmi2); size_t ZSTD_rust_dctx_default_max_window_size(void); int ZSTD_rust_no_forward_progress_max(void); -int ZSTD_rust_heapmode(void); -int ZSTD_rust_decompress_stack_uses_heap(int heapmode); size_t ZSTD_rust_decompress_stack(void* dst, size_t dstCapacity, const void* src, size_t srcSize); size_t ZSTD_rust_decompress_stack_context(ZSTD_DCtx* dctx, void* dst, size_t dstCapacity, const void* src, size_t srcSize); + +typedef struct { + int heapmode; + void* workspace; + size_t workspace_size; + void* dst; + size_t dst_capacity; + const void* src; + size_t src_size; +} ZSTD_rustDecompressStackProjection; +typedef char ZSTD_rust_decompress_stack_projection_layout[ + (offsetof(ZSTD_rustDecompressStackProjection, heapmode) == 0 + && offsetof(ZSTD_rustDecompressStackProjection, workspace) == sizeof(void*) + && offsetof(ZSTD_rustDecompressStackProjection, workspace_size) + == 2 * sizeof(void*) + && offsetof(ZSTD_rustDecompressStackProjection, dst) + == 3 * sizeof(void*) + && offsetof(ZSTD_rustDecompressStackProjection, dst_capacity) + == 4 * sizeof(void*) + && offsetof(ZSTD_rustDecompressStackProjection, src) + == 5 * sizeof(void*) + && offsetof(ZSTD_rustDecompressStackProjection, src_size) + == 6 * sizeof(void*) + && sizeof(ZSTD_rustDecompressStackProjection) == 7 * sizeof(void*)) + ? 1 : -1]; + +typedef ZSTD_DCtx* (*ZSTD_rustInitStaticDCtxFn)(void* workspace, + size_t workspace_size); +typedef size_t (*ZSTD_rustDecompressStackContextFn)( + ZSTD_DCtx* dctx, void* dst, size_t dstCapacity, + const void* src, size_t srcSize); +typedef struct { + ZSTD_rustInitStaticDCtxFn init_static_dctx; + ZSTD_rustDecompressStackContextFn decompress_context; +} ZSTD_rustDecompressStackCallbacks; +typedef char ZSTD_rust_decompress_stack_callbacks_layout[ + (offsetof(ZSTD_rustDecompressStackCallbacks, init_static_dctx) == 0 + && offsetof(ZSTD_rustDecompressStackCallbacks, decompress_context) + == sizeof(void*) + && sizeof(ZSTD_rustDecompressStackCallbacks) == 2 * sizeof(void*)) + ? 1 : -1]; + +size_t ZSTD_rust_decompress_stack_action( + const ZSTD_rustDecompressStackProjection* projection, + const ZSTD_rustDecompressStackCallbacks* callbacks); unsigned ZSTD_rust_legacy_support(void); #if ZSTD_TRACE @@ -275,14 +318,21 @@ int ZSTD_rust_heapmode(void) size_t ZSTD_rust_decompress_stack(void* dst, size_t dstCapacity, const void* src, size_t srcSize) { - if (ZSTD_rust_decompress_stack_uses_heap(ZSTD_HEAPMODE)) { - return ERROR(GENERIC); - } ZSTD_DCtx dctx; - ZSTD_DCtx* const initialized = ZSTD_initStaticDCtx(&dctx, sizeof(dctx)); - if (initialized == NULL) return ERROR(memory_allocation); - return ZSTD_rust_decompress_stack_context(initialized, dst, dstCapacity, - src, srcSize); + ZSTD_rustDecompressStackProjection const projection = { + ZSTD_HEAPMODE, + &dctx, + sizeof(dctx), + dst, + dstCapacity, + src, + srcSize + }; + ZSTD_rustDecompressStackCallbacks const callbacks = { + ZSTD_initStaticDCtx, + ZSTD_rust_decompress_stack_context + }; + return ZSTD_rust_decompress_stack_action(&projection, &callbacks); } void ZSTD_rust_dctx_trace_view(ZSTD_DCtx* dctx, diff --git a/rust/src/zstd_decompress.rs b/rust/src/zstd_decompress.rs index bc273607b..849228541 100644 --- a/rust/src/zstd_decompress.rs +++ b/rust/src/zstd_decompress.rs @@ -411,6 +411,59 @@ mod dctx_platform_tests { type ZstdDecompressDCtxFn = unsafe extern "C" fn(*mut ZSTD_DCtx, *mut c_void, usize, *const c_void, usize) -> usize; +pub type ZSTD_rustInitStaticDCtxFn = unsafe extern "C" fn(*mut c_void, usize) -> *mut ZSTD_DCtx; +pub type ZSTD_rustDecompressStackContextFn = ZstdDecompressDCtxFn; + +/// Scalar inputs for the stack-decompression entry path. The C translation +/// unit keeps the local `ZSTD_DCtx` object and its private layout; Rust owns +/// only the heap-mode branch and the order in which the C leaves run. +#[repr(C)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct ZSTD_rustDecompressStackProjection { + pub heapmode: c_int, + pub workspace: *mut c_void, + pub workspace_size: usize, + pub dst: *mut c_void, + pub dst_capacity: usize, + pub src: *const c_void, + pub src_size: usize, +} + +const _: () = { + assert!(offset_of!(ZSTD_rustDecompressStackProjection, heapmode) == 0); + assert!( + offset_of!(ZSTD_rustDecompressStackProjection, workspace) + == if size_of::() == 8 { 8 } else { 4 } + ); + assert!( + offset_of!(ZSTD_rustDecompressStackProjection, workspace_size) == 2 * size_of::() + ); + assert!(offset_of!(ZSTD_rustDecompressStackProjection, dst) == 3 * size_of::()); + assert!(offset_of!(ZSTD_rustDecompressStackProjection, dst_capacity) == 4 * size_of::()); + assert!(offset_of!(ZSTD_rustDecompressStackProjection, src) == 5 * size_of::()); + assert!(offset_of!(ZSTD_rustDecompressStackProjection, src_size) == 6 * size_of::()); + assert!(size_of::() == 7 * size_of::()); +}; + +/// C-owned leaves used by [`ZSTD_rust_decompress_stack_action`]. +#[repr(C)] +#[derive(Clone, Copy, Debug, Default)] +pub struct ZSTD_rustDecompressStackCallbacks { + pub init_static_dctx: Option, + pub decompress_context: Option, +} + +const _: () = { + assert!(size_of::() == size_of::()); + assert!(size_of::>() == size_of::()); + assert!(size_of::() == size_of::()); + assert!(size_of::>() == size_of::()); + assert!(offset_of!(ZSTD_rustDecompressStackCallbacks, init_static_dctx) == 0); + assert!( + offset_of!(ZSTD_rustDecompressStackCallbacks, decompress_context) == size_of::() + ); + assert!(size_of::() == 2 * size_of::()); +}; #[inline] unsafe fn field(slot: *mut c_void) -> T { @@ -3001,6 +3054,169 @@ mod stack_context_tests { } } +#[cfg(test)] +mod stack_action_tests { + use super::*; + + #[repr(C)] + struct StackActionProbe { + events: [u8; 2], + event_count: usize, + init_calls: usize, + dispatch_calls: usize, + init_workspace_size: usize, + dispatch_dctx: usize, + dispatch_dst: usize, + dispatch_dst_capacity: usize, + dispatch_src: usize, + dispatch_src_size: usize, + init_returns_null: c_int, + result: usize, + } + + impl StackActionProbe { + fn new(init_returns_null: c_int, result: usize) -> Self { + Self { + events: [0; 2], + event_count: 0, + init_calls: 0, + dispatch_calls: 0, + init_workspace_size: 0, + dispatch_dctx: 0, + dispatch_dst: 0, + dispatch_dst_capacity: 0, + dispatch_src: 0, + dispatch_src_size: 0, + init_returns_null, + result, + } + } + + fn projection( + &mut self, + heapmode: c_int, + dst: *mut c_void, + dst_capacity: usize, + src: *const c_void, + src_size: usize, + ) -> ZSTD_rustDecompressStackProjection { + ZSTD_rustDecompressStackProjection { + heapmode, + workspace: (self as *mut Self).cast(), + workspace_size: size_of::(), + dst, + dst_capacity, + src, + src_size, + } + } + } + + unsafe fn record(probe: &mut StackActionProbe, event: u8) { + if probe.event_count < probe.events.len() { + probe.events[probe.event_count] = event; + } + probe.event_count += 1; + } + + unsafe extern "C" fn recording_init( + workspace: *mut c_void, + workspace_size: usize, + ) -> *mut ZSTD_DCtx { + let probe = unsafe { &mut *workspace.cast::() }; + unsafe { record(probe, b'I') }; + probe.init_calls += 1; + probe.init_workspace_size = workspace_size; + if probe.init_returns_null != 0 { + ptr::null_mut() + } else { + workspace.cast() + } + } + + unsafe extern "C" fn recording_decompress( + dctx: *mut ZSTD_DCtx, + dst: *mut c_void, + dst_capacity: usize, + src: *const c_void, + src_size: usize, + ) -> usize { + let probe = unsafe { &mut *dctx.cast::() }; + unsafe { record(probe, b'D') }; + probe.dispatch_calls += 1; + probe.dispatch_dctx = dctx as usize; + probe.dispatch_dst = dst as usize; + probe.dispatch_dst_capacity = dst_capacity; + probe.dispatch_src = src as usize; + probe.dispatch_src_size = src_size; + probe.result + } + + fn callbacks() -> ZSTD_rustDecompressStackCallbacks { + ZSTD_rustDecompressStackCallbacks { + init_static_dctx: Some(recording_init), + decompress_context: Some(recording_decompress), + } + } + + #[test] + fn stack_action_rejects_heap_mode_before_any_callback() { + let mut probe = StackActionProbe::new(0, 0x1234); + let projection = probe.projection(1, ptr::null_mut(), 0, ptr::null(), 0); + let callbacks = callbacks(); + + let result = unsafe { ZSTD_rust_decompress_stack_action(&projection, &callbacks) }; + + assert_eq!(result, ERROR(ZstdErrorCode::Generic)); + assert_eq!(probe.event_count, 0); + assert_eq!(probe.init_calls, 0); + assert_eq!(probe.dispatch_calls, 0); + } + + #[test] + fn stack_action_maps_initializer_failure_without_dispatch() { + let mut probe = StackActionProbe::new(1, 0x1234); + let projection = probe.projection(0, ptr::null_mut(), 0, ptr::null(), 0); + let callbacks = callbacks(); + + let result = unsafe { ZSTD_rust_decompress_stack_action(&projection, &callbacks) }; + + assert_eq!(result, ERROR(ZstdErrorCode::MemoryAllocation)); + assert_eq!(&probe.events, b"I\0"); + assert_eq!(probe.init_calls, 1); + assert_eq!(probe.init_workspace_size, size_of::()); + assert_eq!(probe.dispatch_calls, 0); + } + + #[test] + fn stack_action_initializes_then_dispatches_with_projected_arguments() { + let mut probe = StackActionProbe::new(0, 0x5a5a); + let mut dst = [0u8; 9]; + let src = [1u8, 2, 3, 4]; + let projection = probe.projection( + 0, + dst.as_mut_ptr().cast(), + dst.len(), + src.as_ptr().cast(), + src.len(), + ); + let callbacks = callbacks(); + + let result = unsafe { ZSTD_rust_decompress_stack_action(&projection, &callbacks) }; + + assert_eq!(result, probe.result); + assert_eq!(&probe.events, b"ID"); + assert_eq!(probe.init_calls, 1); + assert_eq!(probe.dispatch_calls, 1); + assert_eq!(probe.init_workspace_size, size_of::()); + assert_eq!(probe.dispatch_dctx, projection.workspace as usize); + assert_eq!(probe.dispatch_dst, dst.as_mut_ptr() as usize); + assert_eq!(probe.dispatch_dst_capacity, dst.len()); + assert_eq!(probe.dispatch_src, src.as_ptr() as usize); + assert_eq!(probe.dispatch_src_size, src.len()); + } +} + unsafe fn select_frame_ddict(view: &ZSTD_rustDctxView) { let set: *mut DDictHashSet = unsafe { field(view.ddict_set) }; if set.is_null() || unsafe { dctx_ddict(view) }.is_null() { @@ -3635,6 +3851,52 @@ unsafe fn decompress_stack_context( unsafe { dispatch(view.dctx.cast(), dst, dst_capacity, src, src_size) } } +unsafe fn decompress_stack_action( + projection: ZSTD_rustDecompressStackProjection, + callbacks: ZSTD_rustDecompressStackCallbacks, +) -> usize { + if decompress_stack_uses_heap(projection.heapmode) { + return ERROR(ZstdErrorCode::Generic); + } + + let Some(init_static_dctx) = callbacks.init_static_dctx else { + return ERROR(ZstdErrorCode::Generic); + }; + let Some(decompress_context) = callbacks.decompress_context else { + return ERROR(ZstdErrorCode::Generic); + }; + + let dctx = unsafe { init_static_dctx(projection.workspace, projection.workspace_size) }; + if dctx.is_null() { + return ERROR(ZstdErrorCode::MemoryAllocation); + } + unsafe { + decompress_context( + dctx, + projection.dst, + projection.dst_capacity, + projection.src, + projection.src_size, + ) + } +} + +/// Order the scalar stack-decompression policy while C retains the stack +/// object, private DCtx layout, static initializer, and context leaf. +#[no_mangle] +pub unsafe extern "C" fn ZSTD_rust_decompress_stack_action( + projection: *const ZSTD_rustDecompressStackProjection, + callbacks: *const ZSTD_rustDecompressStackCallbacks, +) -> usize { + let Some(projection) = (unsafe { projection.as_ref() }).copied() else { + return ERROR(ZstdErrorCode::Generic); + }; + let Some(callbacks) = (unsafe { callbacks.as_ref() }).copied() else { + return ERROR(ZstdErrorCode::Generic); + }; + unsafe { decompress_stack_action(projection, callbacks) } +} + #[inline] fn decompress_stack_uses_heap(heapmode: c_int) -> bool { heapmode >= 1