Bshuf_decompress_zstd_block consistent with bshuf_decompress_lz4_block

.github/workflows/main.yml

@@ -55,4 +55,4 @@ jobs:
         python setup.py install --h5plugin --h5plugin-dir ~/hdf5/lib --zstd
 
     - name: Run tests
-      run: pytest .
+      run: pytest -v .

src/bitshuffle.c

@@ -188,6 +188,8 @@ int64_t bshuf_decompress_zstd_block(ioc_chain *C_ptr,
         free(tmp_buf);
         return -91;
     }
+
+    nbytes = nbytes_from_header;
     count = bshuf_untrans_bit_elem(tmp_buf, out, size, elem_size);
     CHECK_ERR_FREE(count, tmp_buf);
     nbytes += 4;

src/bitshuffle.h

@@ -98,11 +98,6 @@ int64_t bshuf_compress_lz4(const void* in, void* out, const size_t size, const s
  * To properly unshuffle bitshuffled data, *size*, *elem_size* and *block_size*
  * must patch the parameters used to compress the data.
  *
- * NOT TO BE USED WITH UNTRUSTED DATA: This routine uses the function
- * LZ4_decompress_fast from LZ4, which does not protect against maliciously
- * formed datasets. By modifying the compressed data, this function could be
- * coerced into leaving the boundaries of the input buffer.
- *
  * Parameters
  * ----------
  *  in : input buffer
@@ -184,11 +179,6 @@ int64_t bshuf_compress_zstd(const void* in, void* out, const size_t size, const
  * To properly unshuffle bitshuffled data, *size*, *elem_size* and *block_size*
  * must patch the parameters used to compress the data.
  *
- * NOT TO BE USED WITH UNTRUSTED DATA: This routine uses the function
- * ZSTD_decompress_fast from ZSTD, which does not protect against maliciously
- * formed datasets. By modifying the compressed data, this function could be
- * coerced into leaving the boundaries of the input buffer.
- *
  * Parameters
  * ----------
  *  in : input buffer

tests/test_ext.py

@@ -6,7 +6,7 @@
 import numpy as np
 from numpy import random
 
-from bitshuffle import ext
+from bitshuffle import ext, __zstd__
 
 
 # If we are doing timeings by what factor to increase workload.
@@ -379,6 +379,22 @@ def test_10d_decompress_64(self):
         )
         self.check_data = pre_trans
 
+    @unittest.skipUnless(__zstd__, "ZSTD support not included")
+    def test_10c_compress_z64(self):
+        self.case = "compress zstd  64"
+        self.data = self.data.view(np.float64)
+        self.fun = lambda x: ext.compress_zstd(x, BLOCK)
+
+    @unittest.skipUnless(__zstd__, "ZSTD support not included")
+    def test_10d_decompress_z64(self):
+        self.case = "decompress zstd 64"
+        pre_trans = self.data.view(np.float64)
+        self.data = ext.compress_zstd(pre_trans, BLOCK)
+        self.fun = lambda x: ext.decompress_zstd(
+            x, pre_trans.shape, pre_trans.dtype, BLOCK
+        )
+        self.check_data = pre_trans
+
 
 """
 Commented out to prevent nose from finding them.
@@ -539,6 +555,23 @@ def test_circle_with_compression(self):
                 self.assertTrue(out.dtype is data.dtype)
                 self.assertTrue(np.all(data.view(np.uint8) == out.view(np.uint8)))
 
+    @unittest.skipUnless(__zstd__, "ZSTD support not included")
+    def test_circle_with_zstd_compression(self):
+        nmax = 100000
+        reps = 20
+        for dtype in TEST_DTYPES:
+            itemsize = np.dtype(dtype).itemsize
+            nbyte_max = nmax * itemsize
+            dbuf = random.randint(0, 255, nbyte_max).astype(np.uint8)
+            dbuf = dbuf.view(dtype)
+            for ii in range(reps):
+                n = random.randint(0, nmax, 1)[0]
+                data = dbuf[:n]
+                shuff = ext.compress_zstd(data)
+                out = ext.decompress_zstd(shuff, data.shape, data.dtype)
+                self.assertTrue(out.dtype is data.dtype)
+                self.assertTrue(np.all(data.view(np.uint8) == out.view(np.uint8)))
+
 
 # Python implementations for checking results.