Add test for ZSTD_compressSuperBlock and fix UB

lib/compress/hist.rs

@@ -89,6 +89,11 @@ unsafe fn HIST_count_parallel_wksp(
     check: CheckInput,
     workSpace: &mut [u32; 1024],
 ) -> size_t {
+    // TEMP: this can probably be removed if we make this function fully safe
+    // Some callers reuse `workSpace`'s memory as `count`: to prevent aliasing issues
+    // with the `&mut` reference, skip writes to `count`.
+    let aliasesWorkSpace = count as *mut u8 == workSpace.as_mut_ptr().cast::<u8>();
+
     let mut ip = source as *const u8;
     let iend = ip.add(sourceSize);
     let countSize = (usize::from(*maxSymbolValuePtr) + 1) * size_of::<core::ffi::c_uint>();
@@ -101,7 +106,9 @@ unsafe fn HIST_count_parallel_wksp(
 
     // safety checks
     if sourceSize == 0 {
-        ptr::write_bytes(count as *mut u8, 0, countSize);
+        if !aliasesWorkSpace {
+            ptr::write_bytes(count as *mut u8, 0, countSize);
+        }
         *maxSymbolValuePtr = 0;
         return 0;
     }
@@ -166,7 +173,9 @@ unsafe fn HIST_count_parallel_wksp(
         return Error::maxSymbolValue_tooSmall.to_error_code();
     }
     *maxSymbolValuePtr = maxSymbolValue;
-    core::ptr::copy(workSpace.as_ptr().cast::<u8>(), count as *mut u8, countSize);
+    if !aliasesWorkSpace {
+        core::ptr::copy(workSpace.as_ptr().cast::<u8>(), count as *mut u8, countSize);
+    }
 
     max as size_t
 }

lib/compress/zstd_fast.rs

@@ -374,8 +374,8 @@ unsafe fn ZSTD_compressBlock_fast_noDict_generic<const MLS: u32, const USE_CMOV:
                     // advance to next positions
                     ip0 = ip1;
                     ip1 = ip2;
-                    ip2 = ip0.add(step);
-                    ip3 = ip1.add(step);
+                    ip2 = ip0.wrapping_add(step);
+                    ip3 = ip1.wrapping_add(step);
 
                     // calculate step
                     if ip2 >= nextStep {

lib/compress/zstd_lazy.rs

@@ -1845,7 +1845,7 @@ unsafe fn ZSTD_compressBlock_lazy_generic<
                     if matchLength < 4 {
                         let step =
                             (ip.offset_from_unsigned(anchor) >> kSearchStrength).wrapping_add(1); // jump faster over incompressible sections
-                        ip = ip.add(step);
+                        ip = ip.wrapping_add(step);
                         // Enter the lazy skipping mode once we are skipping more than 8 bytes at a time.
                         // In this mode we stop inserting every position into our tables, and only insert
                         // positions that we search, which is one in step positions.

test-libzstd-rs-sys/src/compress.rs

@@ -140,6 +140,75 @@ mod compress2_strats {
     }
 }
 
+mod target_cblock_size {
+    use crate::assert_eq_rs_c;
+    use std::ffi::c_void;
+
+    #[cfg(miri)]
+    const INPUT: &[u8] = include_bytes!("../test-data/compress-input-tiny.dat");
+    #[cfg(not(miri))]
+    const INPUT: &[u8] = include_bytes!("../test-data/compress-input-small.dat");
+
+    const INCOMPRESSIBLE_INPUT: &[u8] = include_bytes!("../test-data/random-input.dat");
+
+    macro_rules! compress_target {
+        ($strategy:expr, $target_size:expr, $input:expr) => {{
+            let cctx = ZSTD_createCCtx();
+            assert!(!cctx.is_null());
+
+            let err = ZSTD_CCtx_setParameter(cctx, ZSTD_cParameter::ZSTD_c_strategy, $strategy);
+            assert_eq!(ZSTD_isError(err), 0);
+
+            // set targetCBlockSize to use ZSTD_compressSuperBlock
+            let err = ZSTD_CCtx_setParameter(
+                cctx,
+                ZSTD_cParameter::ZSTD_c_targetCBlockSize,
+                $target_size,
+            );
+            assert_eq!(ZSTD_isError(err), 0);
+
+            let bound = ZSTD_compressBound($input.len());
+            let mut dst = vec![0u8; bound];
+
+            let written = ZSTD_compress2(
+                cctx,
+                dst.as_mut_ptr() as *mut c_void,
+                dst.len(),
+                $input.as_ptr() as *const c_void,
+                $input.len(),
+            );
+            assert_eq!(ZSTD_isError(written), 0);
+            dst.truncate(written);
+
+            ZSTD_freeCCtx(cctx);
+
+            dst
+        }};
+    }
+
+    #[cfg(not(miri))]
+    const STRATEGIES: [i32; 4] = [1, 3, 6, 9];
+    #[cfg(miri)]
+    const STRATEGIES: [i32; 1] = [1];
+
+    #[test]
+    fn compressible_input() {
+        for strategy in STRATEGIES {
+            // test both ZSTD_TARGETCBLOCKSIZE_MIN and ZSTD_TARGETCBLOCKSIZE_MAX
+            for target_size in [1340, 131072] {
+                assert_eq_rs_c!({ compress_target!(strategy, target_size, INPUT) });
+            }
+        }
+    }
+
+    #[test]
+    fn incompressible_input() {
+        for strategy in STRATEGIES {
+            assert_eq_rs_c!({ compress_target!(strategy, 1340, INCOMPRESSIBLE_INPUT) });
+        }
+    }
+}
+
 #[test]
 #[cfg_attr(miri, ignore = "slow")]
 fn test_compress_stream_2() {