Consolidate the pending bug fixes: 8 open PRs plus the fz-branch findings

.gitignore

@@ -11,4 +11,4 @@ install
 .vscode
 sz3_install
 sz3_build
-test
+/test

include/SZ3/api/impl/SZAlgoBioMD.hpp

@@ -1,14 +1,15 @@
 #ifndef SZ3_SZ_BIOMD_HPP
 #define SZ3_SZ_BIOMD_HPP
 
+#include "SZ3/compressor/SZGenericCompressor.hpp"
 #include "SZ3/decomposition/SZBioMDDecomposition.hpp"
 #include "SZ3/decomposition/SZBioMDXtcDecomposition.hpp"
 #include "SZ3/def.hpp"
+#include "SZ3/encoder/HuffmanEncoder.hpp"
+#include "SZ3/encoder/HuffmanEncoderV2.hpp"
 #include "SZ3/encoder/XtcBasedEncoder.hpp"
 #include "SZ3/lossless/Lossless_bypass.hpp"
 #include "SZ3/lossless/Lossless_zstd.hpp"
-#include "SZ3/encoder/HuffmanEncoderV2.hpp"
-#include "SZ3/encoder/HuffmanEncoder.hpp"
 #include "SZ3/quantizer/LinearQuantizer.hpp"
 #include "SZ3/utils/Config.hpp"
 #include "SZ3/utils/Statistic.hpp"

include/SZ3/api/impl/SZAlgoInterp.hpp

@@ -1,6 +1,8 @@
 #ifndef SZ3_SZALGO_INTERP_HPP
 #define SZ3_SZALGO_INTERP_HPP
 
+#include <memory>
+
 #include "SZ3/api/impl/SZAlgoLorenzoReg.hpp"
 #include "SZ3/decomposition/BlockwiseDecomposition.hpp"
 #include "SZ3/decomposition/InterpolationDecomposition.hpp"

include/SZ3/api/impl/SZDispatcher.hpp

@@ -1,10 +1,12 @@
 #ifndef SZ3_IMPL_SZDISPATCHER_HPP
 #define SZ3_IMPL_SZDISPATCHER_HPP
 
+#include <memory>
+
+#include "SZ3/api/impl/SZAlgoBioMD.hpp"
 #include "SZ3/api/impl/SZAlgoInterp.hpp"
 #include "SZ3/api/impl/SZAlgoLorenzoReg.hpp"
 #include "SZ3/api/impl/SZAlgoNopred.hpp"
-#include "SZ3/api/impl/SZAlgoBioMD.hpp"
 #include "SZ3/utils/Config.hpp"
 #include "SZ3/utils/Statistic.hpp"
 
@@ -63,14 +65,15 @@ size_t SZ_compress_dispatcher(Config &conf, const T *data, uchar *cmpData, size_
         auto zstd = Lossless_zstd();
         auto zstdCmpCap = ZSTD_compressBound(conf.num * sizeof(T)) + sizeof(size_t);
         auto zstdCmpData = static_cast<uchar *>(malloc(zstdCmpCap));
+        // RAII: zstd.compress can throw, which would leak this buffer with a bare free() at the end.
+        std::unique_ptr<uchar, void (*)(void *)> zstd_cmp_data_owner(zstdCmpData, &free);
         size_t zstdCmpSize =
             zstd.compress(reinterpret_cast<const uchar *>(data), conf.num * sizeof(T), zstdCmpData, zstdCmpCap);
         if (zstdCmpSize < cmpSize && zstdCmpSize <= cmpCap) {
             conf.cmprAlgo = ALGO_LOSSLESS;
             memcpy(cmpData, zstdCmpData, zstdCmpSize);
             cmpSize = zstdCmpSize;
         }
-        free(zstdCmpData);
     }
     return cmpSize;
 }

include/SZ3/api/impl/SZImplOMP.hpp

@@ -2,7 +2,9 @@
 #define SZ3_IMPL_SZDISPATCHER_OMP_HPP
 
 #include <cmath>
+#include <cstdlib>
 #include <memory>
+#include <new>
 
 #include "SZ3/api/impl/SZDispatcher.hpp"
 
@@ -70,8 +72,15 @@ size_t SZ_compress_OMP(Config& conf, const T* data, uchar* cmpData, size_t cmpCa
 
         conf_t[tid] = conf;
         conf_t[tid].setDims(dims_t.begin(), dims_t.end());
-        size_t cmp_size_cap = ZSTD_compressBound(conf_t[tid].num * sizeof(T));
-        compressed_t[tid] = static_cast<uchar*>(malloc(cmp_size_cap));
+        // Reserve room for the size header that Lossless_zstd::compress writes in front of the zstd stream,
+        // otherwise the direct lossless path in SZ_compress_dispatcher throws for poorly compressible chunks.
+        size_t cmp_size_cap = sizeof(size_t) + ZSTD_compressBound(conf_t[tid].num * sizeof(T));
+        // The buffer is owned so that it is released even if the compression below throws.
+        std::unique_ptr<uchar, void (*)(void*)> compressed_owner(static_cast<uchar*>(malloc(cmp_size_cap)), &free);
+        if (!compressed_owner) {
+            throw std::bad_alloc();
+        }
+        compressed_t[tid] = compressed_owner.get();
         // we have to use conf_t[tid].N instead of N since each chunk may be a slice of the original data
         if (conf_t[tid].N == 1) {
             cmp_size_t[tid] = SZ_compress_dispatcher<T, 1>(conf_t[tid], data_t, compressed_t[tid], cmp_size_cap);
@@ -105,7 +114,6 @@ size_t SZ_compress_OMP(Config& conf, const T* data, uchar* cmpData, size_t cmpCa
         }
 
         memcpy(buffer_pos + cmp_start_t[tid], compressed_t[tid], cmp_size_t[tid]);
-        free(compressed_t[tid]);
     }
 
     return buffer_pos - cmpData + cmp_start_t[nThreads];
@@ -121,14 +129,22 @@ void SZ_decompress_OMP(Config& conf, const uchar* cmpData, size_t cmpSize, T* de
 #ifdef _OPENMP
 
     auto cmpr_data_pos = cmpData;
+    const uchar* const cmp_end = cmpData + cmpSize;
     int nThreads = 1;
+    // Everything below is read from untrusted data; bound every read against the end of the buffer.
+    if (static_cast<size_t>(cmp_end - cmpr_data_pos) < sizeof(nThreads))
+        throw std::out_of_range("SZ3 OMP: truncated thread count");
     read(nThreads, cmpr_data_pos);
+    // Each per-thread config and size entry occupies at least one byte, so the thread count can not exceed
+    // the size of the compressed buffer.
+    if (nThreads <= 0 || static_cast<size_t>(nThreads) > cmpSize)
+        throw std::out_of_range("SZ3 OMP: invalid thread count");
     omp_set_num_threads(nThreads);
     printf("OpenMP enabled for decompression, threads = %d\n", nThreads);
 
     std::vector<Config> conf_t(nThreads);
     for (int i = 0; i < nThreads; i++) {
-        conf_t[i].load(cmpr_data_pos);
+        conf_t[i].load(cmpr_data_pos, static_cast<size_t>(cmp_end - cmpr_data_pos));
     }
 
     if (conf_t[0].sz3MagicNumber != SZ3_MAGIC_NUMBER) {
@@ -145,12 +161,19 @@ void SZ_decompress_OMP(Config& conf, const uchar* cmpData, size_t cmpSize, T* de
 
     std::vector<size_t> cmp_start_t, cmp_size_t;
     cmp_size_t.resize(nThreads);
+    if (static_cast<size_t>(cmp_end - cmpr_data_pos) < static_cast<size_t>(nThreads) * sizeof(size_t))
+        throw std::out_of_range("SZ3 OMP: truncated per-thread sizes");
     read(cmp_size_t.data(), nThreads, cmpr_data_pos);
     auto cmpr_data_p = cmpr_data_pos;
 
     cmp_start_t.resize(nThreads + 1);
     cmp_start_t[0] = 0;
+    // The per-thread payloads follow back-to-back and must all fit in the remaining buffer. Build the running
+    // offsets with an overflow-safe bound so a crafted size can not point a thread's slice out of bounds.
+    const size_t payload_avail = static_cast<size_t>(cmp_end - cmpr_data_p);
     for (int i = 1; i <= nThreads; i++) {
+        if (cmp_size_t[i - 1] > payload_avail - cmp_start_t[i - 1])
+            throw std::out_of_range("SZ3 OMP: per-thread compressed sizes exceed the buffer");
         cmp_start_t[i] = cmp_start_t[i - 1] + cmp_size_t[i - 1];
     }
 
@@ -199,8 +222,9 @@ size_t SZ_compress_size_bound_omp(const Config& conf) {
     }
     size_t chunk_size = conf.dims[0] / static_cast<size_t>(nThreads) * (conf.num / conf.dims[0]);
     size_t last_chunk_size = (conf.dims[0] - conf.dims[0] / nThreads * (nThreads - 1)) * (conf.num / conf.dims[0]);
-    //for each thread, we save conf, compressed size, and compressed data
-    return sizeof(int) + nThreads * conf.size_est() + nThreads * sizeof(size_t) +
+    // for each thread, we save conf, compressed size, and compressed data
+    // the per-chunk compressed data may carry the size header written by Lossless_zstd::compress
+    return sizeof(int) + nThreads * conf.size_est() + 2 * nThreads * sizeof(size_t) +
            (nThreads - 1) * ZSTD_compressBound(chunk_size * sizeof(T)) +
            ZSTD_compressBound(last_chunk_size * sizeof(T));
 #else

include/SZ3/api/sz.hpp

@@ -21,10 +21,11 @@
 #ifndef SZ3_SZ_HPP
 #define SZ3_SZ_HPP
 
+#include <memory>
+
 #include "SZ3/api/impl/SZImpl.hpp"
 #include "SZ3/version.hpp"
 
-
 /**
  * Compresses the input data using the provided configuration and stores the result in a pre-allocated buffer.
  * @tparam T The data type of the source data.
@@ -95,10 +96,10 @@ char* SZ_compress(const SZ3::Config& config, const T* data, size_t& cmpSize) {
     using namespace SZ3;
 
     size_t bufferLen = SZ_compress_size_bound<T>(config);
-    auto buffer = new char[bufferLen];
-    cmpSize = SZ_compress(config, data, buffer, bufferLen);
+    std::unique_ptr<char[]> buffer(new char[bufferLen]);
+    cmpSize = SZ_compress(config, data, buffer.get(), bufferLen);
 
-    return buffer;
+    return buffer.release();
 }
 
 /**
@@ -119,6 +120,11 @@ void SZ_decompress(SZ3::Config& config, const char* cmpData, size_t cmpSize, T*&
 
     auto cmpDataPos = reinterpret_cast<const uchar*>(cmpData);
 
+    // Header layout: magic number (4) + data version (4) + compressed payload size (8) = 16 bytes.
+    if (cmpSize < 16) {
+        throw std::out_of_range("SZ3: compressed data is smaller than the header");
+    }
+
     read(config.sz3MagicNumber, cmpDataPos);
     if (config.sz3MagicNumber != SZ3_MAGIC_NUMBER) {
         throw std::invalid_argument("magic number mismatch, the input data is not compressed by SZ3");
@@ -137,8 +143,12 @@ void SZ_decompress(SZ3::Config& config, const char* cmpData, size_t cmpSize, T*&
     uint64_t cmpDataSize = 0;
     read(cmpDataSize, cmpDataPos);
 
+    // The compressed payload is followed by the serialized config; both must fit in the remaining bytes.
+    if (cmpDataSize > cmpSize - 16) {
+        throw std::out_of_range("SZ3: compressed payload size exceeds the buffer");
+    }
     auto cmpConfPos = cmpDataPos + cmpDataSize;
-    config.load(cmpConfPos);
+    config.load(cmpConfPos, cmpSize - 16 - cmpDataSize);
 
     if (decData == nullptr) {
         decData = new T[config.num];

include/SZ3/compressor/SZGenericCompressor.hpp

@@ -1,7 +1,10 @@
 #ifndef SZ3_COMPRESSOR_TYPE_ONE_HPP
 #define SZ3_COMPRESSOR_TYPE_ONE_HPP
 
+#include <cstdlib>
 #include <cstring>
+#include <memory>
+#include <type_traits>
 
 #include "SZ3/compressor/Compressor.hpp"
 #include "SZ3/decomposition/Decomposition.hpp"
@@ -13,6 +16,16 @@
 #include "SZ3/utils/Timer.hpp"
 
 namespace SZ3 {
+
+/// Detects the optional (non-virtual) `set_decode_bound()` an encoder may expose so the compressor can
+/// hand it the exact number of bytes its encoded stream may read. Encoders without it keep whatever
+/// bound their own `load()` recorded.
+template <class E, class = void>
+struct encoder_has_decode_bound : std::false_type {};
+template <class E>
+struct encoder_has_decode_bound<E, std::void_t<decltype(std::declval<E &>().set_decode_bound(size_t{}))>>
+    : std::true_type {};
+
 /**
  * SZGenericCompressor glues together decomposition, encoder, and lossless modules to form the compressor.
  * It only takes Decomposition, not Predictor.
@@ -46,38 +59,62 @@ class SZGenericCompressor : public concepts::CompressorInterface<T> {
             1000, 2 * (decomposition.size_est() + encoder.size_est() + sizeof(T) * quant_inds.size()));
 
         auto buffer = static_cast<uchar *>(malloc(bufferSize));
+        // Own the scratch buffer with RAII so it is released on every path: the encoder and the lossless
+        // layer below can throw (e.g. Lossless_zstd::compress throws std::length_error when the destination
+        // capacity is too small for poorly-compressible data), and the caller catches and continues, so a
+        // bare free() at the end leaks the buffer on each failed compression.
+        std::unique_ptr<uchar, void (*)(void *)> buffer_owner(buffer, &free);
         uchar *buffer_pos = buffer;
 
         decomposition.save(buffer_pos);
         encoder.save(buffer_pos);
 
-        //store the size of quant_inds is necessary as it is not always equal to conf.num
+        // store the size of quant_inds is necessary as it is not always equal to conf.num
         write<size_t>(quant_inds.size(), buffer_pos);
         encoder.encode(quant_inds, buffer_pos);
         encoder.postprocess_encode();
 
         auto cmpSize = lossless.compress(buffer, buffer_pos - buffer, cmpData, cmpCap);
-        free(buffer);
 
         return cmpSize;
     }
 
     T *decompress(const Config &conf, uchar const *cmpData, size_t cmpSize, T *decData) override {
         uchar *buffer = nullptr;
+        // No bound is passed to the lossless layer here. The internal buffer compress() produced is
+        // sized max(1000, 2 * (decomposition.size_est() + encoder.size_est() + sizeof(Q) * bins)), which
+        // for a wide bin type exceeds any bound derivable from conf alone -- bounding it by
+        // SZ_compress_size_bound rejects valid streams. Every module in this tree emits int bins so the
+        // bound happens to hold here, but an out-of-tree module with 64-bit bins would be rejected.
+        // A corrupted declared size is still caught by the zstd frame check and the size comparison in
+        // Lossless_zstd::decompress, after the allocation.
         size_t bufferSize = 0;
         lossless.decompress(cmpData, cmpSize, buffer, bufferSize);
 
+        // The lossless layer allocated `buffer` with malloc. Own it with RAII so it is released on every path
+        // below - including the parsing steps that operate on untrusted data and can throw before we are done
+        // with it - instead of being leaked. decompress() is reached repeatedly for corrupted blocks (fuzzing).
+        std::unique_ptr<uchar, void (*)(void *)> buffer_owner(buffer, &free);
+
         uchar const *bufferPos = buffer;
 
         decomposition.load(bufferPos, bufferSize);
         encoder.load(bufferPos, bufferSize);
 
         size_t quant_inds_size = 0;
-        read(quant_inds_size, bufferPos);
+        // Read the count with the bounded overload so a truncated buffer can not be read past its end.
+        read(quant_inds_size, bufferPos, bufferSize);
+        // The count field sits between the encoder's tree and its encoded stream, so the bound load()
+        // recorded for decode() is that many bytes too large. Hand the encoder the exact remaining length
+        // now that the field has been consumed. Optional: encoders without the hook keep load()'s bound.
+        if constexpr (encoder_has_decode_bound<Encoder>::value) {
+            encoder.set_decode_bound(bufferSize);
+        }
         auto quant_inds = encoder.decode(bufferPos, quant_inds_size);
         encoder.postprocess_decode();
 
-        free(buffer);
+        // The remaining work uses `quant_inds` and `decData` only, so release the internal buffer now.
+        buffer_owner.reset();
 
         decomposition.decompress(conf, quant_inds, decData);
         return decData;

include/SZ3/compressor/specialized/SZExaaltCompressor.hpp

@@ -1,6 +1,11 @@
 #ifndef SZ3_EXAALT_COMPRESSSOR_HPP
 #define SZ3_EXAALT_COMPRESSSOR_HPP
 
+#include <iostream>
+#include <limits>
+#include <memory>
+
+#include "SZ3/compressor/Compressor.hpp"
 #include "SZ3/def.hpp"
 #include "SZ3/encoder/Encoder.hpp"
 #include "SZ3/lossless/Lossless.hpp"

include/SZ3/compressor/specialized/SZTruncateCompressor.hpp

@@ -2,6 +2,7 @@
 #define SZ3_Truncate_COMPRESSOR_HPP
 
 #include <cstring>
+#include <memory>
 
 #include "SZ3/compressor/Compressor.hpp"
 #include "SZ3/decomposition/Decomposition.hpp"
@@ -28,14 +29,16 @@ class SZTruncateCompressor : public concepts::CompressorInterface<T> {
 
     size_t compress(const Config &conf, T *data, uchar *cmpData, size_t cmpCap) override {
         auto buffer = static_cast<uchar *>(malloc(conf.num * sizeof(T)));
+        // RAII: the lossless layer below can throw (std::length_error when the destination capacity is too
+        // small), which would leak this scratch buffer with a bare free() at the end.
+        std::unique_ptr<uchar, void (*)(void *)> buffer_owner(buffer, &free);
         auto buffer_pos = buffer;
 
         //            Timer timer(true);
         truncateArray(data, conf.num, byteLen, buffer_pos);
         //            timer.stop("Prediction & Quantization");
 
         auto cmpSize = lossless.compress(buffer, buffer_pos - buffer, cmpData, cmpCap);
-        free(buffer);
         return cmpSize;
         //            lossless.postcompress_data(buffer);
         //            return lossless_data;

include/SZ3/decomposition/BlockwiseDecomposition.hpp

@@ -2,15 +2,17 @@
 #define SZ3_BLOCKWISE_DECOMPOSITION_HPP
 
 #include <cstring>
+#include <limits>
+#include <memory>
 
 #include "Decomposition.hpp"
 #include "SZ3/def.hpp"
 #include "SZ3/predictor/LorenzoPredictor.hpp"
 #include "SZ3/predictor/Predictor.hpp"
 #include "SZ3/quantizer/LinearQuantizer.hpp"
+#include "SZ3/utils/BlockwiseIterator.hpp"
 #include "SZ3/utils/Config.hpp"
 #include "SZ3/utils/FileUtil.hpp"
-#include "SZ3/utils/BlockwiseIterator.hpp"
 #include "SZ3/utils/Timer.hpp"
 
 namespace SZ3 {

include/SZ3/decomposition/Decomposition.hpp

@@ -1,9 +1,11 @@
 #ifndef SZ3_DECOMPOSITION_INTERFACE
 #define SZ3_DECOMPOSITION_INTERFACE
 
+#include <utility>
 #include <vector>
 
 #include "SZ3/def.hpp"
+#include "SZ3/utils/Config.hpp"
 
 namespace SZ3::concepts {