From d6cda74b84d813b3f51264baedc51ad6973e3ffe Mon Sep 17 00:00:00 2001
From: ddidderr <ddidderr@paul.network>
Date: Mon, 20 Jul 2026 18:25:52 +0200
Subject: [PATCH] refactor(compress): move end directive validation to Rust

Project the scalar ZSTD_compressStream2 end directive into Rust and keep C responsible for the private streaming context, existing parameter error, and diagnostic text. The Rust range policy preserves the original unsigned enum validation for continue, flush, and end.

Test Plan: ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo clippy --manifest-path rust/Cargo.toml --all-targets -- -D warnings; cargo +nightly fmt --manifest-path rust/Cargo.toml --all -- --check; make -j1; ./tests/rustLibSmoke; make -j1 -C tests test (all shell tests, large streaming tests, native tester, fuzzer phases, and zstream tester passed).
---
 lib/compress/zstd_compress.c | 14 ++++++++-
 rust/src/zstd_compress.rs    | 57 ++++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c
index f42148342..0e485ee66 100644
--- a/lib/compress/zstd_compress.c
+++ b/lib/compress/zstd_compress.c
@@ -769,6 +769,15 @@ size_t ZSTD_compressStream2_c(ZSTD_CCtx* cctx,
                               ZSTD_outBuffer* output,
                               ZSTD_inBuffer* input,
                               ZSTD_EndDirective endOp);
+typedef struct {
+    int endOp;
+} ZSTD_rust_compressStream2PolicyState;
+int ZSTD_rust_compressStream2Policy(
+        const ZSTD_rust_compressStream2PolicyState* state);
+typedef char ZSTD_rust_compress_stream2_policy_state_layout[
+        (offsetof(ZSTD_rust_compressStream2PolicyState, endOp) == 0
+         && sizeof(ZSTD_rust_compressStream2PolicyState) == sizeof(int))
+        ? 1 : -1];
 int ZSTD_rust_simpleCompress2Level(const void* cctx, size_t srcSize);
 typedef int (*ZSTD_rust_simpleCompress2Level_f)(const void* cctx, size_t srcSize);
 typedef struct {
@@ -8266,7 +8275,10 @@ size_t ZSTD_compressStream2_c( ZSTD_CCtx* cctx,
     /* check conditions */
     RETURN_ERROR_IF(output->pos > output->size, dstSize_tooSmall, "invalid output buffer");
     RETURN_ERROR_IF(input->pos  > input->size, srcSize_wrong, "invalid input buffer");
-    RETURN_ERROR_IF((U32)endOp > (U32)ZSTD_e_end, parameter_outOfBound, "invalid endDirective");
+    {   ZSTD_rust_compressStream2PolicyState const state = {(int)endOp};
+        RETURN_ERROR_IF(!ZSTD_rust_compressStream2Policy(&state),
+                        parameter_outOfBound, "invalid endDirective");
+    }
     assert(cctx != NULL);
 
     /* transparent initialization stage */
diff --git a/rust/src/zstd_compress.rs b/rust/src/zstd_compress.rs
index 73dafab93..0c3acab9d 100644
--- a/rust/src/zstd_compress.rs
+++ b/rust/src/zstd_compress.rs
@@ -163,6 +163,44 @@ const ZSTD_CHUNKSIZE_MAX: usize = u32::MAX as usize - ZSTD_CURRENT_MAX;
 const ZSTD_E_END: c_int = 2;
 const ZSTD_E_CONTINUE: c_int = 0;
 const ZSTD_E_FLUSH: c_int = 1;
+
+/// Scalar projection for the public `ZSTD_compressStream2` directive policy.
+///
+/// The C fallback retains the private context and streaming implementation;
+/// Rust owns only the public range decision; C retains error encoding and
+/// diagnostics.
+#[repr(C)]
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub struct ZSTD_rust_compressStream2PolicyState {
+    pub end_op: c_int,
+}
+
+const _: () = {
+    assert!(offset_of!(ZSTD_rust_compressStream2PolicyState, end_op) == 0);
+    assert!(size_of::<ZSTD_rust_compressStream2PolicyState>() == size_of::<c_int>());
+};
+
+#[inline]
+fn compress_stream2_policy(end_op: c_int) -> c_int {
+    if (ZSTD_E_CONTINUE..=ZSTD_E_END).contains(&end_op) {
+        1
+    } else {
+        0
+    }
+}
+
+/// Decide whether the public stream end directive is in range before entering
+/// C's private fallback state machine.
+#[no_mangle]
+pub unsafe extern "C" fn ZSTD_rust_compressStream2Policy(
+    state: *const ZSTD_rust_compressStream2PolicyState,
+) -> c_int {
+    let Some(state) = (unsafe { state.as_ref() }) else {
+        return 0;
+    };
+    compress_stream2_policy(state.end_op)
+}
+
 const ZSTD_C_WINDOW_LOG: c_int = 101;
 const ZSTD_C_HASH_LOG: c_int = 102;
 const ZSTD_C_CHAIN_LOG: c_int = 103;
@@ -14033,6 +14071,25 @@ mod tests {
         assert_eq!(context.calls, 0);
     }
 
+    #[test]
+    fn compress_stream2_policy_accepts_all_public_directives() {
+        for end_op in [ZSTD_E_CONTINUE, ZSTD_E_FLUSH, ZSTD_E_END] {
+            let state = ZSTD_rust_compressStream2PolicyState { end_op };
+            let result = unsafe { ZSTD_rust_compressStream2Policy(&state) };
+            assert_eq!(result, 1);
+        }
+    }
+
+    #[test]
+    fn compress_stream2_policy_rejects_out_of_range_directives() {
+        for end_op in [c_int::MIN, -1, ZSTD_E_END + 1, c_int::MAX] {
+            let state = ZSTD_rust_compressStream2PolicyState { end_op };
+            let result = unsafe { ZSTD_rust_compressStream2Policy(&state) };
+            assert_eq!(result, 0);
+        }
+        assert_eq!(unsafe { ZSTD_rust_compressStream2Policy(ptr::null()) }, 0);
+    }
+
     #[test]
     fn reduce_table_applies_threshold_and_wrapping_subtraction() {
         let mut table = [0, 1, 2, 3, 4, 5, 6, u32::MAX, 0, 0, 0, 0, 0, 0, 0, 0];