feat(compress): move CCtx stage guard into Rust

Project the private ZSTD_CCtx compression-stage value into the Rust copy
bridge and reject non-init sources before any destination mutation. Keep the
C private context layout and reset/table operations behind the existing
callbacks while preserving the stage error contract.

Test Plan:
- ulimit -v 41943040 && CARGO_BUILD_JOBS=1 cargo clippy --manifest-path rust/Cargo.toml --all-targets -- -D warnings
- ulimit -v 41943040 && make -j1
- ulimit -v 41943040 && CARGO_BUILD_JOBS=1 cargo test --manifest-path rust/Cargo.toml
- ulimit -v 41943040 && make -j1 -C tests test-zstream ZSTREAM_TESTTIME=-T2s
- ulimit -v 41943040 && make -j1 -C tests test-fuzzer FUZZERTEST=-T3s FUZZER_FLAGS=--no-big-tests

lib/compress/zstd_compress.c

@@ -457,8 +457,6 @@ typedef char ZSTD_rust_reset_cctx_state_layout[
                 == 6 * sizeof(void*)
          && sizeof(ZSTD_rust_resetCCtxState) == 7 * sizeof(void*))
         ? 1 : -1];
-typedef size_t (*ZSTD_rust_copyCCtxCheckStage_f)(
-        void* context, const void* srcCCtx);
 typedef size_t (*ZSTD_rust_copyCCtxReset_f)(
         void* context, const void* srcCCtx,
         const ZSTD_frameParameters* fParams,
@@ -471,7 +469,7 @@ typedef struct {
     const void* srcCCtx;
     const ZSTD_frameParameters* fParams;
     U64 pledgedSrcSize;
-    ZSTD_rust_copyCCtxCheckStage_f checkStage;
+    const int* sourceStage;
     void* destinationCustomMem;
     const void* sourceCustomMem;
     ZSTD_rust_copyCCtxReset_f reset;
@@ -531,7 +529,7 @@ typedef char ZSTD_rust_copy_cctx_internal_state_layout[
                 == 2 * sizeof(void*)
          && offsetof(ZSTD_rust_copyCCtxInternalState, pledgedSrcSize)
                 == 3 * sizeof(void*)
-         && offsetof(ZSTD_rust_copyCCtxInternalState, checkStage)
+         && offsetof(ZSTD_rust_copyCCtxInternalState, sourceStage)
                 == 3 * sizeof(void*) + sizeof(U64)
          && offsetof(ZSTD_rust_copyCCtxInternalState, zbuff)
                 == 3 * sizeof(void*) + sizeof(U64)
@@ -4744,17 +4742,6 @@ static size_t ZSTD_resetCCtx_usingCDict(ZSTD_CCtx* cctx,
     return ZSTD_rust_resetCCtxUsingCDict(&state);
 }
 
-static size_t ZSTD_rust_copyCCtx_check_stage(
-        void* context, const void* srcCCtx)
-{
-    const ZSTD_CCtx* const src = (const ZSTD_CCtx*)srcCCtx;
-    (void)context;
-    RETURN_ERROR_IF(src->stage != ZSTDcs_init, stage_wrong,
-                    "Can't copy a ctx that's not in init stage.");
-    DEBUGLOG(5, "ZSTD_copyCCtx_internal");
-    return 0;
-}
-
 static size_t ZSTD_rust_copyCCtx_reset(
         void* context, const void* srcCCtx,
         const ZSTD_frameParameters* fParams,
@@ -4832,11 +4819,12 @@ static size_t ZSTD_rust_copyCCtx_internal_callback(
         U64 pledgedSrcSize, int zbuff)
 {
     ZSTD_rust_copyCCtxInternalState state;
+    int const sourceStage = (int)((const ZSTD_CCtx*)srcCCtx)->stage;
     state.callbackContext = context;
     state.srcCCtx = srcCCtx;
     state.fParams = fParams;
     state.pledgedSrcSize = pledgedSrcSize;
-    state.checkStage = ZSTD_rust_copyCCtx_check_stage;
+    state.sourceStage = &sourceStage;
     state.destinationCustomMem =
             &((ZSTD_CCtx*)context)->customMem;
     state.sourceCustomMem =

rust/src/zstd_compress.rs

@@ -1672,7 +1672,6 @@ pub unsafe extern "C" fn ZSTD_rust_copyCCtx(state: *const ZSTD_rust_copyCCtxStat
     }
 }
 
-type CopyCCtxCheckStageFn = unsafe extern "C" fn(*mut c_void, *const c_void) -> usize;
 type CopyCCtxResetFn = unsafe extern "C" fn(
     *mut c_void,
     *const c_void,
@@ -1729,7 +1728,7 @@ pub struct ZSTD_rust_copyCCtxInternalState {
     src_cctx: *const c_void,
     f_params: *const ZSTD_frameParameters,
     pledged_src_size: u64,
-    check_stage: Option<CopyCCtxCheckStageFn>,
+    source_stage: *const c_int,
     destination_custom_mem: *mut c_void,
     source_custom_mem: *const c_void,
     reset: Option<CopyCCtxResetFn>,
@@ -1752,7 +1751,6 @@ pub struct ZSTD_rust_copyCCtxInternalState {
 }
 
 const _: () = {
-    assert!(size_of::<CopyCCtxCheckStageFn>() == size_of::<usize>());
     assert!(size_of::<CopyCCtxResetFn>() == size_of::<usize>());
     assert!(size_of::<CopyCCtxMarkTablesFn>() == size_of::<usize>());
     assert!(size_of::<CopyCCtxCopyTablesFn>() == size_of::<usize>());
@@ -1763,7 +1761,7 @@ const _: () = {
         offset_of!(ZSTD_rust_copyCCtxInternalState, pledged_src_size) == 3 * size_of::<usize>()
     );
     assert!(
-        offset_of!(ZSTD_rust_copyCCtxInternalState, check_stage)
+        offset_of!(ZSTD_rust_copyCCtxInternalState, source_stage)
             == 3 * size_of::<usize>() + size_of::<u64>()
     );
     assert!(
@@ -1788,25 +1786,18 @@ pub unsafe extern "C" fn ZSTD_rust_copyCCtxInternal(
         return ERROR(ZstdErrorCode::Generic);
     }
     let state = unsafe { &*state };
-    let (
+    let (Some(reset), Some(mark_tables_dirty), Some(copy_tables), Some(mark_tables_clean)) = (
-        Some(check_stage),
-        Some(reset),
-        Some(mark_tables_dirty),
-        Some(copy_tables),
-        Some(mark_tables_clean),
-    ) = (
-        state.check_stage,
         state.reset,
         state.mark_tables_dirty,
         state.copy_tables,
         state.mark_tables_clean,
-    )
+    ) else {
-    else {
         return ERROR(ZstdErrorCode::Generic);
     };
     if state.callback_context.is_null()
         || state.src_cctx.is_null()
         || state.f_params.is_null()
+        || state.source_stage.is_null()
         || state.destination_custom_mem.is_null()
         || state.source_custom_mem.is_null()
         || state.destination_window.is_null()
@@ -1825,9 +1816,8 @@ pub unsafe extern "C" fn ZSTD_rust_copyCCtxInternal(
         return ERROR(ZstdErrorCode::Generic);
     }
 
-    let stage_error = unsafe { check_stage(state.callback_context, state.src_cctx) };
+    if unsafe { *state.source_stage } != ZSTD_COMPRESSION_STAGE_INIT {
-    if ERR_isError(stage_error) {
+        return ERROR(ZstdErrorCode::StageWrong);
-        return stage_error;
     }
 
     unsafe {
@@ -13900,7 +13890,6 @@ mod tests {
     #[derive(Default)]
     struct CopyCCtxInternalTestContext {
         events: Vec<&'static str>,
-        stage_result: usize,
         reset_result: usize,
         frame_params: ZSTD_frameParameters,
         pledged_src_size: u64,
@@ -13913,15 +13902,6 @@ mod tests {
         unsafe { &mut *context.cast::<CopyCCtxInternalTestContext>() }
     }
 
-    unsafe extern "C" fn copy_cctx_internal_test_check_stage(
-        context: *mut c_void,
-        _src_cctx: *const c_void,
-    ) -> usize {
-        let context = unsafe { copy_cctx_internal_test_context(context) };
-        context.events.push("check");
-        context.stage_result
-    }
-
     unsafe extern "C" fn copy_cctx_internal_test_reset(
         context: *mut c_void,
         _src_cctx: *const c_void,
@@ -13961,6 +13941,7 @@ mod tests {
         context: &mut CopyCCtxInternalTestContext,
         src_cctx: *const c_void,
         f_params: &ZSTD_frameParameters,
+        source_stage: *const c_int,
         destination_custom_mem: *mut c_void,
         source_custom_mem: *const c_void,
         destination_window: *mut c_void,
@@ -13981,7 +13962,7 @@ mod tests {
             src_cctx,
             f_params,
             pledged_src_size: 123,
-            check_stage: Some(copy_cctx_internal_test_check_stage),
+            source_stage,
             destination_custom_mem,
             source_custom_mem,
             reset: Some(copy_cctx_internal_test_reset),
@@ -14012,6 +13993,7 @@ mod tests {
             checksumFlag: 1,
             noDictIDFlag: 1,
         };
+        let source_stage = ZSTD_COMPRESSION_STAGE_INIT;
         let source_custom_mem = [0x11usize, 0x22, 0x33];
         let mut destination_custom_mem = [0usize; 3];
         let mut source_block_state =
@@ -14051,6 +14033,7 @@ mod tests {
             &mut context,
             0x5000usize as *const c_void,
             &f_params,
+            &source_stage,
             (&mut destination_custom_mem as *mut [usize; 3]).cast(),
             (&source_custom_mem as *const [usize; 3]).cast(),
             (&mut destination_window as *mut ZSTD_rust_copyWindowState).cast(),
@@ -14070,10 +14053,7 @@ mod tests {
         let result = unsafe { ZSTD_rust_copyCCtxInternal(&state) };
 
         assert_eq!(result, 0);
-        assert_eq!(
+        assert_eq!(context.events, ["reset", "dirty", "tables", "clean"]);
-            context.events,
-            ["check", "reset", "dirty", "tables", "clean"]
-        );
         assert_eq!(context.frame_params, f_params);
         assert_eq!(context.pledged_src_size, 123);
         assert_eq!(context.zbuff, 7);
@@ -14096,11 +14076,9 @@ mod tests {
 
     #[test]
     fn copy_cctx_internal_rejects_stage_before_mutation() {
-        let mut context = CopyCCtxInternalTestContext {
+        let mut context = CopyCCtxInternalTestContext::default();
-            stage_result: ERROR(ZstdErrorCode::StageWrong),
-            ..CopyCCtxInternalTestContext::default()
-        };
         let f_params = ZSTD_frameParameters::default();
+        let source_stage = ZSTD_COMPRESSION_STAGE_ONGOING;
         let source_custom_mem = [0usize; 3];
         let mut destination_custom_mem = [0usize; 3];
         let mut destination_block_state =
@@ -14117,6 +14095,7 @@ mod tests {
             &mut context,
             0x5000usize as *const c_void,
             &f_params,
+            &source_stage,
             (&mut destination_custom_mem as *mut [usize; 3]).cast(),
             (&source_custom_mem as *const [usize; 3]).cast(),
             ptr::dangling_mut(),
@@ -14136,7 +14115,7 @@ mod tests {
         let result = unsafe { ZSTD_rust_copyCCtxInternal(&state) };
 
         assert_eq!(result, ERROR(ZstdErrorCode::StageWrong));
-        assert_eq!(context.events, ["check"]);
+        assert!(context.events.is_empty());
     }
 
     #[test]
@@ -14146,6 +14125,7 @@ mod tests {
             ..CopyCCtxInternalTestContext::default()
         };
         let f_params = ZSTD_frameParameters::default();
+        let source_stage = ZSTD_COMPRESSION_STAGE_INIT;
         let source_custom_mem = [0usize; 3];
         let mut destination_custom_mem = [0usize; 3];
         let mut destination_block_state =
@@ -14162,6 +14142,7 @@ mod tests {
             &mut context,
             0x5000usize as *const c_void,
             &f_params,
+            &source_stage,
             (&mut destination_custom_mem as *mut [usize; 3]).cast(),
             (&source_custom_mem as *const [usize; 3]).cast(),
             ptr::dangling_mut(),
@@ -14181,7 +14162,7 @@ mod tests {
         let result = unsafe { ZSTD_rust_copyCCtxInternal(&state) };
 
         assert_eq!(result, ERROR(ZstdErrorCode::MemoryAllocation));
-        assert_eq!(context.events, ["check", "reset"]);
+        assert_eq!(context.events, ["reset"]);
     }
 
     #[derive(Default)]