diff --git a/lib/decompress/zstd_decompress.c b/lib/decompress/zstd_decompress.c index 41c735471..1966f2605 100644 --- a/lib/decompress/zstd_decompress.c +++ b/lib/decompress/zstd_decompress.c @@ -5,7 +5,8 @@ * optional members vary with the build configuration and the block decoder * shares that object. This translation unit is therefore a deliberately * narrow ABI adapter. It projects field addresses to Rust while retaining - * the private layout and configuration-bound legacy and trace leaves. + * the private layout and configuration-bound legacy leaves. Weak trace-symbol + * discovery also stays here; Rust receives only a narrow trace projection. */ #define ZSTD_STATIC_LINKING_ONLY @@ -105,7 +106,23 @@ typedef struct { size_t dctx_size; } ZSTD_rustDctxView; +typedef U64 ZSTD_rustTraceCtx; +typedef ZSTD_rustTraceCtx (*ZSTD_rustTraceBeginFn)(const void* dctx); +typedef void (*ZSTD_rustTraceEndFn)(ZSTD_rustTraceCtx traceCtx, + const void* trace); + +typedef struct { + ZSTD_rustTraceCtx* trace_ctx; + ZSTD_rustTraceBeginFn begin; + ZSTD_rustTraceEndFn end; + const void* dctx; + const void* ddict; + int dictionary_is_cold; +} ZSTD_rustDctxTraceView; + void ZSTD_rust_dctx_view(ZSTD_DCtx* dctx, ZSTD_rustDctxView* out); +void ZSTD_rust_dctx_trace_view(ZSTD_DCtx* dctx, + ZSTD_rustDctxTraceView* out); size_t ZSTD_rust_dctx_sizeof(void); void ZSTD_rust_dctx_init_platform(ZSTD_DCtx* dctx); size_t ZSTD_rust_dctx_default_max_window_size(void); @@ -117,11 +134,25 @@ ZSTD_DDict* ZSTD_rust_create_ddict(const void* dict, size_t dictSize, ZSTD_dictLoadMethod_e dictLoadMethod, ZSTD_dictContentType_e dictContentType, ZSTD_customMem customMem); -void ZSTD_rust_dctx_trace_begin(ZSTD_DCtx* dctx); -void ZSTD_rust_dctx_trace_end(ZSTD_DCtx* dctx, U64 uncompressedSize, - U64 compressedSize, int streaming); unsigned ZSTD_rust_legacy_support(void); +#if ZSTD_TRACE +static ZSTD_rustTraceCtx ZSTD_rust_trace_begin_callback(const void* dctx) +{ + return (ZSTD_trace_decompress_begin != NULL) + ? ZSTD_trace_decompress_begin((const ZSTD_DCtx*)dctx) : 0; +} + +static void ZSTD_rust_trace_end_callback(ZSTD_rustTraceCtx traceCtx, + const void* trace) +{ + if (ZSTD_trace_decompress_end != NULL) { + ZSTD_trace_decompress_end((ZSTD_TraceCtx)traceCtx, + (const ZSTD_Trace*)trace); + } +} +#endif + void ZSTD_rust_dctx_view(ZSTD_DCtx* dctx, ZSTD_rustDctxView* out) { ZSTD_memset(out, 0, sizeof(*out)); @@ -260,40 +291,19 @@ ZSTD_DDict* ZSTD_rust_create_ddict(const void* dict, size_t dictSize, dictContentType, customMem); } -void ZSTD_rust_dctx_trace_begin(ZSTD_DCtx* dctx) +void ZSTD_rust_dctx_trace_view(ZSTD_DCtx* dctx, + ZSTD_rustDctxTraceView* out) { + ZSTD_memset(out, 0, sizeof(*out)); + out->dctx = dctx; #if ZSTD_TRACE - dctx->traceCtx = (ZSTD_trace_decompress_begin != NULL) - ? ZSTD_trace_decompress_begin(dctx) : 0; -#else - (void)dctx; -#endif -} - -void ZSTD_rust_dctx_trace_end(ZSTD_DCtx* dctx, U64 uncompressedSize, - U64 compressedSize, int streaming) -{ -#if ZSTD_TRACE - if (dctx->traceCtx && ZSTD_trace_decompress_end != NULL) { - ZSTD_Trace trace; - ZSTD_memset(&trace, 0, sizeof(trace)); - trace.version = ZSTD_VERSION_NUMBER; - trace.streaming = streaming; - if (dctx->ddict) { - trace.dictionaryID = ZSTD_getDictID_fromDDict(dctx->ddict); - trace.dictionarySize = ZSTD_DDict_dictSize(dctx->ddict); - trace.dictionaryIsCold = dctx->ddictIsCold; - } - trace.uncompressedSize = (size_t)uncompressedSize; - trace.compressedSize = (size_t)compressedSize; - trace.dctx = dctx; - ZSTD_trace_decompress_end(dctx->traceCtx, &trace); - } -#else - (void)dctx; - (void)uncompressedSize; - (void)compressedSize; - (void)streaming; + out->trace_ctx = &dctx->traceCtx; + out->begin = (ZSTD_trace_decompress_begin != NULL) + ? ZSTD_rust_trace_begin_callback : NULL; + out->end = (ZSTD_trace_decompress_end != NULL) + ? ZSTD_rust_trace_end_callback : NULL; + out->ddict = dctx->ddict; + out->dictionary_is_cold = dctx->ddictIsCold; #endif } diff --git a/rust/src/zstd_decompress.rs b/rust/src/zstd_decompress.rs index 767fe531e..f4415fef7 100644 --- a/rust/src/zstd_decompress.rs +++ b/rust/src/zstd_decompress.rs @@ -299,8 +299,40 @@ struct ZSTD_rustDctxView { dctx_size: usize, } +type ZSTD_rustTraceBeginFn = unsafe extern "C" fn(*const c_void) -> u64; +type ZSTD_rustTraceEndFn = unsafe extern "C" fn(u64, *const c_void); + +#[repr(C)] +#[derive(Clone, Copy)] +struct ZSTD_rustDctxTraceView { + trace_ctx: *mut u64, + begin: Option, + end: Option, + dctx: *const c_void, + ddict: *const c_void, + dictionary_is_cold: c_int, +} + +#[repr(C)] +struct ZSTD_rustTrace { + version: c_uint, + streaming: c_int, + dictionary_id: c_uint, + dictionary_is_cold: c_int, + dictionary_size: usize, + uncompressed_size: usize, + compressed_size: usize, + params: *const c_void, + cctx: *const c_void, + dctx: *const c_void, +} + +/* Keep this in sync with ZSTD_VERSION_NUMBER in lib/zstd.h. */ +const ZSTD_TRACE_VERSION: c_uint = 10_507; + unsafe extern "C" { fn ZSTD_rust_dctx_view(dctx: *mut ZSTD_DCtx, out: *mut ZSTD_rustDctxView); + fn ZSTD_rust_dctx_trace_view(dctx: *mut ZSTD_DCtx, out: *mut ZSTD_rustDctxTraceView); fn ZSTD_rust_dctx_sizeof() -> usize; fn ZSTD_rust_dctx_init_platform(dctx: *mut ZSTD_DCtx); fn ZSTD_rust_dctx_default_max_window_size() -> usize; @@ -319,13 +351,6 @@ unsafe extern "C" { dict_content_type: c_int, custom_mem: ZSTD_customMem, ) -> *mut ZSTD_DDict; - fn ZSTD_rust_dctx_trace_end( - dctx: *mut ZSTD_DCtx, - uncompressed_size: u64, - compressed_size: u64, - streaming: c_int, - ); - fn ZSTD_rust_dctx_trace_begin(dctx: *mut ZSTD_DCtx); #[cfg(all( not(test), any( @@ -366,6 +391,76 @@ unsafe fn dctx_view(dctx: *mut ZSTD_DCtx) -> ZSTD_rustDctxView { unsafe { view.assume_init() } } +unsafe fn dctx_trace_view(dctx: *mut ZSTD_DCtx) -> ZSTD_rustDctxTraceView { + let mut view = MaybeUninit::::zeroed(); + unsafe { ZSTD_rust_dctx_trace_view(dctx, view.as_mut_ptr()) }; + unsafe { view.assume_init() } +} + +unsafe fn trace_begin_with_view(view: &ZSTD_rustDctxTraceView) { + if view.trace_ctx.is_null() { + return; + } + let trace_ctx = view + .begin + .map(|begin| unsafe { begin(view.dctx) }) + .unwrap_or(0); + unsafe { view.trace_ctx.write(trace_ctx) }; +} + +unsafe fn trace_end_with_view( + view: &ZSTD_rustDctxTraceView, + uncompressed_size: u64, + compressed_size: u64, + streaming: c_int, +) { + if view.trace_ctx.is_null() { + return; + } + let trace_ctx = unsafe { view.trace_ctx.read() }; + if trace_ctx == 0 { + return; + } + let Some(end) = view.end else { + return; + }; + + let mut trace = ZSTD_rustTrace { + version: ZSTD_TRACE_VERSION, + streaming, + dictionary_id: 0, + dictionary_is_cold: 0, + dictionary_size: 0, + uncompressed_size: uncompressed_size as usize, + compressed_size: compressed_size as usize, + params: ptr::null(), + cctx: ptr::null(), + dctx: view.dctx, + }; + if !view.ddict.is_null() { + let ddict = view.ddict.cast::(); + trace.dictionary_id = unsafe { ZSTD_getDictID_fromDDict(ddict) }; + trace.dictionary_size = unsafe { ZSTD_DDict_dictSize(ddict) }; + trace.dictionary_is_cold = view.dictionary_is_cold; + } + unsafe { end(trace_ctx, (&mut trace as *mut ZSTD_rustTrace).cast()) }; +} + +unsafe fn trace_begin(dctx: *mut ZSTD_DCtx) { + let view = unsafe { dctx_trace_view(dctx) }; + unsafe { trace_begin_with_view(&view) }; +} + +unsafe fn trace_end( + dctx: *mut ZSTD_DCtx, + uncompressed_size: u64, + compressed_size: u64, + streaming: c_int, +) { + let view = unsafe { dctx_trace_view(dctx) }; + unsafe { trace_end_with_view(&view, uncompressed_size, compressed_size, streaming) }; +} + #[inline] unsafe fn copy_dctx_prefix(dst: *mut u8, src: *const u8, in_buff: *const c_void) { let to_copy = (in_buff as usize).wrapping_sub(dst as usize); @@ -2607,6 +2702,171 @@ mod dctx_copy_tests { } } +#[cfg(test)] +mod dctx_trace_tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + struct TraceProbe { + begin_calls: AtomicUsize, + end_calls: AtomicUsize, + last_context: AtomicUsize, + last_version: AtomicUsize, + last_streaming: AtomicUsize, + last_dictionary_id: AtomicUsize, + last_dictionary_is_cold: AtomicUsize, + last_dictionary_size: AtomicUsize, + last_uncompressed_size: AtomicUsize, + last_compressed_size: AtomicUsize, + } + + impl TraceProbe { + fn new() -> Self { + Self { + begin_calls: AtomicUsize::new(0), + end_calls: AtomicUsize::new(0), + last_context: AtomicUsize::new(0), + last_version: AtomicUsize::new(0), + last_streaming: AtomicUsize::new(0), + last_dictionary_id: AtomicUsize::new(0), + last_dictionary_is_cold: AtomicUsize::new(0), + last_dictionary_size: AtomicUsize::new(0), + last_uncompressed_size: AtomicUsize::new(0), + last_compressed_size: AtomicUsize::new(0), + } + } + } + + unsafe extern "C" fn trace_begin_probe(dctx: *const c_void) -> u64 { + let probe = dctx.cast::(); + if probe.is_null() { + return 0; + } + unsafe { &*probe } + .begin_calls + .fetch_add(1, Ordering::SeqCst); + 0xD1CE + } + + unsafe extern "C" fn trace_end_probe(trace_ctx: u64, trace_ptr: *const c_void) { + if trace_ptr.is_null() { + return; + } + let trace = unsafe { &*trace_ptr.cast::() }; + let probe = trace.dctx.cast::(); + if probe.is_null() { + return; + } + let probe = unsafe { &*probe }; + probe.end_calls.fetch_add(1, Ordering::SeqCst); + probe + .last_context + .store(trace_ctx as usize, Ordering::SeqCst); + probe + .last_version + .store(trace.version as usize, Ordering::SeqCst); + probe + .last_streaming + .store(trace.streaming as usize, Ordering::SeqCst); + probe + .last_dictionary_id + .store(trace.dictionary_id as usize, Ordering::SeqCst); + probe + .last_dictionary_is_cold + .store(trace.dictionary_is_cold as usize, Ordering::SeqCst); + probe + .last_dictionary_size + .store(trace.dictionary_size, Ordering::SeqCst); + probe + .last_uncompressed_size + .store(trace.uncompressed_size, Ordering::SeqCst); + probe + .last_compressed_size + .store(trace.compressed_size, Ordering::SeqCst); + } + + #[test] + fn trace_enabled_preserves_record_policy_and_context() { + let probe = TraceProbe::new(); + let mut trace_ctx = 0u64; + let view = ZSTD_rustDctxTraceView { + trace_ctx: std::ptr::addr_of_mut!(trace_ctx), + begin: Some(trace_begin_probe), + end: Some(trace_end_probe), + dctx: (&probe as *const TraceProbe).cast(), + ddict: std::ptr::null(), + dictionary_is_cold: 1, + }; + + unsafe { trace_begin_with_view(&view) }; + assert_eq!(trace_ctx, 0xD1CE); + unsafe { + trace_end_with_view(&view, 0x1_0000_0001, 0x2_0000_0002, 1); + } + + assert_eq!(probe.begin_calls.load(Ordering::SeqCst), 1); + assert_eq!(probe.end_calls.load(Ordering::SeqCst), 1); + assert_eq!(probe.last_context.load(Ordering::SeqCst), 0xD1CE); + assert_eq!(probe.last_version.load(Ordering::SeqCst), 10_507); + assert_eq!(probe.last_streaming.load(Ordering::SeqCst), 1); + assert_eq!(probe.last_dictionary_id.load(Ordering::SeqCst), 0); + assert_eq!(probe.last_dictionary_is_cold.load(Ordering::SeqCst), 0); + assert_eq!(probe.last_dictionary_size.load(Ordering::SeqCst), 0); + assert_eq!( + probe.last_uncompressed_size.load(Ordering::SeqCst), + 0x1_0000_0001u64 as usize + ); + assert_eq!( + probe.last_compressed_size.load(Ordering::SeqCst), + 0x2_0000_0002u64 as usize + ); + assert_eq!(trace_ctx, 0xD1CE); + } + + #[test] + fn trace_disabled_projection_is_a_noop() { + let probe = TraceProbe::new(); + let trace_ctx = 0xBEEFu64; + let view = ZSTD_rustDctxTraceView { + trace_ctx: std::ptr::null_mut(), + begin: Some(trace_begin_probe), + end: Some(trace_end_probe), + dctx: (&probe as *const TraceProbe).cast(), + ddict: std::ptr::null(), + dictionary_is_cold: 1, + }; + + unsafe { + trace_begin_with_view(&view); + trace_end_with_view(&view, 1, 2, 0); + } + + assert_eq!(trace_ctx, 0xBEEF); + assert_eq!(probe.begin_calls.load(Ordering::SeqCst), 0); + assert_eq!(probe.end_calls.load(Ordering::SeqCst), 0); + } + + #[test] + fn trace_null_callbacks_follow_begin_and_end_policy() { + let mut trace_ctx = 0xBEEFu64; + let view = ZSTD_rustDctxTraceView { + trace_ctx: std::ptr::addr_of_mut!(trace_ctx), + begin: None, + end: None, + dctx: std::ptr::null(), + ddict: std::ptr::null(), + dictionary_is_cold: 0, + }; + + unsafe { trace_begin_with_view(&view) }; + assert_eq!(trace_ctx, 0); + + trace_ctx = 0xBEEF; + unsafe { trace_end_with_view(&view, 1, 2, 0) }; + assert_eq!(trace_ctx, 0xBEEF); + } +} + 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() { @@ -2667,7 +2927,7 @@ unsafe fn decode_frame_header( unsafe fn decompress_begin(view: &ZSTD_rustDctxView) -> usize { unsafe { - ZSTD_rust_dctx_trace_begin(view.dctx.cast()); + trace_begin(view.dctx.cast()); let format = field::(view.format); set_field(view.expected, frame_header_prefix(format)); set_field(view.stage, ZSTDDS_GET_FRAME_HEADER_SIZE); @@ -2986,7 +3246,7 @@ unsafe fn decompress_frame( remaining -= 4; } unsafe { - ZSTD_rust_dctx_trace_end( + trace_end( dctx, decoded as u64, (ip as usize).wrapping_sub(istart as usize) as u64, @@ -3475,12 +3735,7 @@ pub unsafe extern "C" fn ZSTD_decompressContinue( set_field(view.expected, 4usize); set_field(view.stage, ZSTDDS_CHECK_CHECKSUM); } else { - ZSTD_rust_dctx_trace_end( - dctx, - decoded_total, - field::(view.processed_c_size), - 1, - ); + trace_end(dctx, decoded_total, field::(view.processed_c_size), 1); set_field(view.expected, 0usize); set_field(view.stage, ZSTDDS_GET_FRAME_HEADER_SIZE); } @@ -3506,7 +3761,7 @@ pub unsafe extern "C" fn ZSTD_decompressContinue( } } unsafe { - ZSTD_rust_dctx_trace_end( + trace_end( dctx, field::(view.decoded_size), field::(view.processed_c_size),