From d281da36c7df64f7b7189222b5154b0cb1363ae2 Mon Sep 17 00:00:00 2001 From: Kai Zhao Date: Mon, 31 Aug 2026 23:20:43 -0700 Subject: [PATCH 1/6] Consolidate the pending bug fixes: 8 open PRs plus the fz-branch findings Bug fixes only. Nothing here changes the compressed format or an interface signature, and every algorithm's output is byte-identical to master's (verified over 31 dataset/algorithm/error-bound combinations, plus decompressing files master produced). From the open pull requests, reviewed line by line: #131 Config::load read one byte past the end of the config blob #132 bounds-check the whole decompression path (see the two skips below) #133 shift-by-64 UB for a single-symbol Huffman tree #134 bounds-check HuffmanEncoderV2 tree loading #135 bounds-check XtcBasedEncoder's magicInts index; plug a leaked buffer #137 remaining_length accounting after Huffman decode in both predictors #138 non-finite float to int64_t cast UB in LinearQuantizer #139 scratch buffers leaked when compression throws; OMP chunk capacity was missing room for the size header Lossless_zstd writes Two parts of #132 are deliberately not taken: - It moves the quant_inds count from after the encoder's tree to before it so the tree is immediately followed by its encoded stream. That is a compressed format change and it is not versioned, so files written by any released SZ3 fail to decode. The bound it was buying is recovered instead by letting the caller supply it: HuffmanEncoder::set_decode_bound(), called by SZGenericCompressor once it has consumed the count, and by both predictors. load() no longer guesses the bound from its own buffer -- the tree and the stream are not required to share a buffer, and test_encoder.cpp puts them in separate ones. - It reinterprets LosslessInterface::decompress's `dstLen` as an input capacity when the caller provides the buffer. The declared contract is that it is an output, and callers do pass uninitialised values, so this fails nondeterministically (it broke LosslessTest.LosslessBypass here). The self-allocating branch, where a non-zero value is opt-in, is kept. Also fixed while reviewing: - HuffmanEncoderV2 sized its dense tables from an unbounded `maxval` read straight from the stream; #134 bounds the node count but not this. - HuffmanEncoderV2's node count is an int holding a value read as 64-bit, so the new bound is checked with an explicit sign test rather than an implicit conversion. - test_lossless.cpp passed an uninitialised size to decompress(). From the fz branch, verified to reproduce here: Lossless_bypass::compress ignored its destination capacity, so any payload larger than the caller's buffer was an unconditional heap overflow; TimeSeriesDecomposition violated its error bound 1.94x on the null-reference-frame path; ArithmeticEncoder sign-extended a shift, corrupting about half of all streams; HuffmanEncoder's stateNum narrowing overflowed on a wide bin range; RunlengthEncoder and two decompositions had no size_est(), so the compressor sized its buffer from 0; KmeansUtil had reserve+operator[] UB and an off-by-one read; the HDF5 filter sized its buffer below SZ_compress's own minimum; 24 headers were not self-contained under libstdc++ or libc++; .gitignore's bare `test` pattern hid every test file. Co-Authored-By: Claude Opus 5 --- .gitignore | 2 +- include/SZ3/api/impl/SZAlgoBioMD.hpp | 5 +- include/SZ3/api/impl/SZAlgoInterp.hpp | 2 + include/SZ3/api/impl/SZDispatcher.hpp | 7 +- include/SZ3/api/impl/SZImplOMP.hpp | 36 ++++++-- include/SZ3/api/sz.hpp | 20 +++-- .../SZ3/compressor/SZGenericCompressor.hpp | 49 +++++++++-- .../specialized/SZExaaltCompressor.hpp | 5 ++ .../specialized/SZTruncateCompressor.hpp | 5 +- .../decomposition/BlockwiseDecomposition.hpp | 4 +- include/SZ3/decomposition/Decomposition.hpp | 2 + .../InterpolationDecomposition.hpp | 21 +++++ .../NoPredictionDecomposition.hpp | 2 + .../decomposition/SZBioMDXtcDecomposition.hpp | 1 + .../decomposition/TimeSeriesDecomposition.hpp | 15 +++- include/SZ3/encoder/ArithmeticEncoder.hpp | 3 +- include/SZ3/encoder/HuffmanEncoder.hpp | 83 +++++++++++++++---- include/SZ3/encoder/HuffmanEncoderV2.hpp | 41 ++++++++- include/SZ3/encoder/RunlengthEncoder.hpp | 6 +- include/SZ3/encoder/XtcBasedEncoder.hpp | 2 + include/SZ3/lossless/Lossless.hpp | 4 + include/SZ3/lossless/Lossless_bypass.hpp | 19 ++++- include/SZ3/lossless/Lossless_zstd.hpp | 50 ++++++++++- include/SZ3/predictor/ComposedPredictor.hpp | 13 +++ include/SZ3/predictor/LorenzoPredictor.hpp | 9 +- include/SZ3/predictor/RegressionPredictor.hpp | 14 +++- include/SZ3/preprocessor/PreFilter.hpp | 4 + include/SZ3/preprocessor/Transpose.hpp | 4 + include/SZ3/quantizer/LinearQuantizer.hpp | 62 ++++++++------ include/SZ3/quantizer/Quantizer.hpp | 21 +++++ include/SZ3/utils/BlockwiseIterator.hpp | 21 +++++ include/SZ3/utils/Config.hpp | 57 ++++++++++++- include/SZ3/utils/Extraction.hpp | 8 ++ include/SZ3/utils/Iterator.hpp | 2 + include/SZ3/utils/KmeansUtil.hpp | 4 +- include/SZ3/utils/MemoryUtil.hpp | 11 ++- include/SZ3/utils/QuantOptimization.hpp | 3 + include/SZ3/utils/Sample.hpp | 4 +- include/SZ3/utils/Statistic.hpp | 4 + tools/H5Z-SZ3/src/H5Z_SZ3.cpp | 6 +- tools/test/modules/test_lossless.cpp | 3 +- 41 files changed, 543 insertions(+), 91 deletions(-) diff --git a/.gitignore b/.gitignore index 0569725c..19492f3d 100644 --- a/.gitignore +++ b/.gitignore @@ -11,4 +11,4 @@ install .vscode sz3_install sz3_build -test \ No newline at end of file +/test \ No newline at end of file diff --git a/include/SZ3/api/impl/SZAlgoBioMD.hpp b/include/SZ3/api/impl/SZAlgoBioMD.hpp index 8719c047..dd49e3c5 100644 --- a/include/SZ3/api/impl/SZAlgoBioMD.hpp +++ b/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" diff --git a/include/SZ3/api/impl/SZAlgoInterp.hpp b/include/SZ3/api/impl/SZAlgoInterp.hpp index 914756ea..6ad9ca32 100644 --- a/include/SZ3/api/impl/SZAlgoInterp.hpp +++ b/include/SZ3/api/impl/SZAlgoInterp.hpp @@ -1,6 +1,8 @@ #ifndef SZ3_SZALGO_INTERP_HPP #define SZ3_SZALGO_INTERP_HPP +#include + #include "SZ3/api/impl/SZAlgoLorenzoReg.hpp" #include "SZ3/decomposition/BlockwiseDecomposition.hpp" #include "SZ3/decomposition/InterpolationDecomposition.hpp" diff --git a/include/SZ3/api/impl/SZDispatcher.hpp b/include/SZ3/api/impl/SZDispatcher.hpp index 9de7aa13..70773047 100644 --- a/include/SZ3/api/impl/SZDispatcher.hpp +++ b/include/SZ3/api/impl/SZDispatcher.hpp @@ -1,10 +1,12 @@ #ifndef SZ3_IMPL_SZDISPATCHER_HPP #define SZ3_IMPL_SZDISPATCHER_HPP +#include + +#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,6 +65,8 @@ 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(malloc(zstdCmpCap)); + // RAII: zstd.compress can throw, which would leak this buffer with a bare free() at the end. + std::unique_ptr zstd_cmp_data_owner(zstdCmpData, &free); size_t zstdCmpSize = zstd.compress(reinterpret_cast(data), conf.num * sizeof(T), zstdCmpData, zstdCmpCap); if (zstdCmpSize < cmpSize && zstdCmpSize <= cmpCap) { @@ -70,7 +74,6 @@ size_t SZ_compress_dispatcher(Config &conf, const T *data, uchar *cmpData, size_ memcpy(cmpData, zstdCmpData, zstdCmpSize); cmpSize = zstdCmpSize; } - free(zstdCmpData); } return cmpSize; } diff --git a/include/SZ3/api/impl/SZImplOMP.hpp b/include/SZ3/api/impl/SZImplOMP.hpp index 664a14dc..81389213 100644 --- a/include/SZ3/api/impl/SZImplOMP.hpp +++ b/include/SZ3/api/impl/SZImplOMP.hpp @@ -2,7 +2,9 @@ #define SZ3_IMPL_SZDISPATCHER_OMP_HPP #include +#include #include +#include #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(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 compressed_owner(static_cast(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(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(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(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 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(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 cmp_start_t, cmp_size_t; cmp_size_t.resize(nThreads); + if (static_cast(cmp_end - cmpr_data_pos) < static_cast(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(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(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 diff --git a/include/SZ3/api/sz.hpp b/include/SZ3/api/sz.hpp index 16f61045..d99c0746 100644 --- a/include/SZ3/api/sz.hpp +++ b/include/SZ3/api/sz.hpp @@ -21,10 +21,11 @@ #ifndef SZ3_SZ_HPP #define SZ3_SZ_HPP +#include + #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(config); - auto buffer = new char[bufferLen]; - cmpSize = SZ_compress(config, data, buffer, bufferLen); + std::unique_ptr 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(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]; diff --git a/include/SZ3/compressor/SZGenericCompressor.hpp b/include/SZ3/compressor/SZGenericCompressor.hpp index 4a3f471a..75cc67ce 100644 --- a/include/SZ3/compressor/SZGenericCompressor.hpp +++ b/include/SZ3/compressor/SZGenericCompressor.hpp @@ -1,7 +1,10 @@ #ifndef SZ3_COMPRESSOR_TYPE_ONE_HPP #define SZ3_COMPRESSOR_TYPE_ONE_HPP +#include #include +#include +#include #include "SZ3/compressor/Compressor.hpp" #include "SZ3/decomposition/Decomposition.hpp" @@ -11,8 +14,19 @@ #include "SZ3/utils/Config.hpp" #include "SZ3/utils/FileUtil.hpp" #include "SZ3/utils/Timer.hpp" +#include "zstd.h" 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 +struct encoder_has_decode_bound : std::false_type {}; +template +struct encoder_has_decode_bound().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 +60,63 @@ class SZGenericCompressor : public concepts::CompressorInterface { 1000, 2 * (decomposition.size_est() + encoder.size_est() + sizeof(T) * quant_inds.size())); auto buffer = static_cast(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 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(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; - size_t bufferSize = 0; + // The lossless layer reads the size of this internal buffer from the untrusted payload and would + // otherwise allocate it unbounded. Bound it by the largest internal buffer this configuration could + // have produced: during compression the buffer is losslessly (zstd) compressed, and + // ZSTD_compressBound(B) >= B, so a block that was actually stored with this compressor satisfies + // B < SZ_compress_size_bound = 4096 + conf.size_est() + ZSTD_compressBound(conf.num * sizeof(T)). + // Passing this as the capacity lets the lossless decoder reject a corrupted payload that declares a + // larger internal size before allocating it. conf.num is validated against the trusted output size by + // the caller, so this bound can not be inflated by corrupted input. + size_t bufferSize = 4096 + conf.size_est() + ZSTD_compressBound(conf.num * sizeof(T)); 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 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::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; diff --git a/include/SZ3/compressor/specialized/SZExaaltCompressor.hpp b/include/SZ3/compressor/specialized/SZExaaltCompressor.hpp index a62e33f6..52ab33f0 100644 --- a/include/SZ3/compressor/specialized/SZExaaltCompressor.hpp +++ b/include/SZ3/compressor/specialized/SZExaaltCompressor.hpp @@ -1,6 +1,11 @@ #ifndef SZ3_EXAALT_COMPRESSSOR_HPP #define SZ3_EXAALT_COMPRESSSOR_HPP +#include +#include +#include + +#include "SZ3/compressor/Compressor.hpp" #include "SZ3/def.hpp" #include "SZ3/encoder/Encoder.hpp" #include "SZ3/lossless/Lossless.hpp" diff --git a/include/SZ3/compressor/specialized/SZTruncateCompressor.hpp b/include/SZ3/compressor/specialized/SZTruncateCompressor.hpp index f4d3b766..d6a6f3f2 100644 --- a/include/SZ3/compressor/specialized/SZTruncateCompressor.hpp +++ b/include/SZ3/compressor/specialized/SZTruncateCompressor.hpp @@ -2,6 +2,7 @@ #define SZ3_Truncate_COMPRESSOR_HPP #include +#include #include "SZ3/compressor/Compressor.hpp" #include "SZ3/decomposition/Decomposition.hpp" @@ -28,6 +29,9 @@ class SZTruncateCompressor : public concepts::CompressorInterface { size_t compress(const Config &conf, T *data, uchar *cmpData, size_t cmpCap) override { auto buffer = static_cast(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 buffer_owner(buffer, &free); auto buffer_pos = buffer; // Timer timer(true); @@ -35,7 +39,6 @@ class SZTruncateCompressor : public concepts::CompressorInterface { // 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; diff --git a/include/SZ3/decomposition/BlockwiseDecomposition.hpp b/include/SZ3/decomposition/BlockwiseDecomposition.hpp index 88a90a0e..aa5b6262 100644 --- a/include/SZ3/decomposition/BlockwiseDecomposition.hpp +++ b/include/SZ3/decomposition/BlockwiseDecomposition.hpp @@ -2,15 +2,17 @@ #define SZ3_BLOCKWISE_DECOMPOSITION_HPP #include +#include +#include #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 { diff --git a/include/SZ3/decomposition/Decomposition.hpp b/include/SZ3/decomposition/Decomposition.hpp index 2661d101..9efb576c 100644 --- a/include/SZ3/decomposition/Decomposition.hpp +++ b/include/SZ3/decomposition/Decomposition.hpp @@ -1,9 +1,11 @@ #ifndef SZ3_DECOMPOSITION_INTERFACE #define SZ3_DECOMPOSITION_INTERFACE +#include #include #include "SZ3/def.hpp" +#include "SZ3/utils/Config.hpp" namespace SZ3::concepts { diff --git a/include/SZ3/decomposition/InterpolationDecomposition.hpp b/include/SZ3/decomposition/InterpolationDecomposition.hpp index dde1f0d7..c72e25af 100644 --- a/include/SZ3/decomposition/InterpolationDecomposition.hpp +++ b/include/SZ3/decomposition/InterpolationDecomposition.hpp @@ -3,10 +3,12 @@ #include #include +#include #include "Decomposition.hpp" #include "SZ3/def.hpp" #include "SZ3/quantizer/Quantizer.hpp" +#include "SZ3/utils/BlockwiseIterator.hpp" #include "SZ3/utils/Config.hpp" #include "SZ3/utils/FileUtil.hpp" #include "SZ3/utils/Interpolators.hpp" @@ -24,6 +26,20 @@ class InterpolationDecomposition : public concepts::DecompositionInterface &quant_inds, T *dec_data) override { + // `original_dimensions` was read from the (untrusted) compressed payload by `load`, separately from the + // trusted `conf.dims`. It drives the size of the grid walked over `dec_data` below (and the number of + // `quant_inds` consumed), so a tampered value larger than the trusted configuration would read and + // write past the end of both buffers, which are sized for `conf.num` elements (== product of conf.dims). + // Require it to match the trusted configuration dimensions before using it for anything. + if (conf.dims.size() != N) { + throw std::out_of_range("SZ3 interpolation: configuration dimension count does not match the data"); + } + for (uint i = 0; i < N; i++) { + if (original_dimensions[i] != conf.dims[i]) { + throw std::out_of_range("SZ3 interpolation: stored dimensions do not match the trusted configuration"); + } + } + init(); this->quant_inds = quant_inds.data(); @@ -146,6 +162,11 @@ class InterpolationDecomposition : public concepts::DecompositionInterface #include #include "Decomposition.hpp" diff --git a/include/SZ3/decomposition/TimeSeriesDecomposition.hpp b/include/SZ3/decomposition/TimeSeriesDecomposition.hpp index 9028a3f3..cdf3c8cc 100644 --- a/include/SZ3/decomposition/TimeSeriesDecomposition.hpp +++ b/include/SZ3/decomposition/TimeSeriesDecomposition.hpp @@ -1,6 +1,9 @@ #ifndef SZ3_TIME_SERIES_DECOMPOSITION_HPP #define SZ3_TIME_SERIES_DECOMPOSITION_HPP +#include +#include + #include "Decomposition.hpp" #include "SZ3/def.hpp" #include "SZ3/predictor/LorenzoPredictor.hpp" @@ -34,6 +37,9 @@ class TimeSeriesDecomposition : public concepts::DecompositionInterface compress(const Config& conf, T* data) override { std::vector quant_inds(num_elements); size_t quant_count = 0; + // The timestep loop below predicts from the reconstruction of timestep 0. + const T* ts0_recon = data; + std::shared_ptr> data_with_padding; if (data_ts0 != nullptr) { for (size_t j = 0; j < conf.dims[1]; j++) { quant_inds[quant_count++] = quantizer.quantize_and_overwrite(data[j], data_ts0[j]); @@ -44,7 +50,7 @@ class TimeSeriesDecomposition : public concepts::DecompositionInterface>(data, spatial_dims, predictor.get_padding(), true); auto block = data_with_padding->block_iter(conf.blockSize); do { @@ -58,13 +64,16 @@ class TimeSeriesDecomposition : public concepts::DecompositionInterfacevalues(); } for (size_t j = 0; j < conf.dims[1]; j++) { + T prev = ts0_recon[j]; for (size_t i = 1; i < conf.dims[0]; i++) { size_t idx = i * conf.dims[1] + j; - size_t idx_prev = (i - 1) * conf.dims[1] + j; - quant_inds[quant_count++] = quantizer.quantize_and_overwrite(data[idx], data[idx_prev]); + quant_inds[quant_count++] = quantizer.quantize_and_overwrite(data[idx], prev); + prev = data[idx]; // quantize_and_overwrite left the reconstruction here } } assert(quant_count == num_elements); diff --git a/include/SZ3/encoder/ArithmeticEncoder.hpp b/include/SZ3/encoder/ArithmeticEncoder.hpp index 5167faea..409f08cd 100644 --- a/include/SZ3/encoder/ArithmeticEncoder.hpp +++ b/include/SZ3/encoder/ArithmeticEncoder.hpp @@ -2,6 +2,7 @@ #define SZ3_ArithmeticEncoder_HPP #include +#include #include #include "SZ3/encoder/Encoder.hpp" @@ -527,7 +528,7 @@ class ArithmeticEncoder : public concepts::EncoderInterface { size_t total_frequency = ariCoder.total_frequency; const uchar *sp = bytes + 5; unsigned int offset = 4; - size_t value = (bytesToInt64_bigEndian(bytes) >> 20); // alignment with the MAX_CODE + size_t value = (static_cast(bytesToInt64_bigEndian(bytes)) >> 20); // alignment with the MAX_CODE size_t s_counter = sizeof(int); for (i = 0; i < targetLength; i++) { diff --git a/include/SZ3/encoder/HuffmanEncoder.hpp b/include/SZ3/encoder/HuffmanEncoder.hpp index 69645a0b..bfe9b103 100644 --- a/include/SZ3/encoder/HuffmanEncoder.hpp +++ b/include/SZ3/encoder/HuffmanEncoder.hpp @@ -1,23 +1,25 @@ #ifndef SZ3_HUFFMAN_ENCODER_HPP #define SZ3_HUFFMAN_ENCODER_HPP -#include - -#include "SZ3/def.hpp" -#include "SZ3/encoder/Encoder.hpp" -#include "SZ3/utils/ByteUtil.hpp" -#include "SZ3/utils/MemoryUtil.hpp" -#include "SZ3/utils/Timer.hpp" -#include "SZ3/utils/Collections.hpp" #include +#include #include #include #include #include +#include #include #include +#include #include +#include "SZ3/def.hpp" +#include "SZ3/encoder/Encoder.hpp" +#include "SZ3/utils/ByteUtil.hpp" +#include "SZ3/utils/Collections.hpp" +#include "SZ3/utils/MemoryUtil.hpp" +#include "SZ3/utils/Timer.hpp" + namespace SZ3 { template @@ -228,22 +230,31 @@ class HuffmanEncoder : public concepts::EncoderInterface { size_t i = 0, byteIndex = 0, count = 0; int r; node n = treeRoot; + /// Bytes available for the encoded payload, recorded by load(). Used to bound all reads below. + size_t remaining = decode_remaining_length; + if (remaining < sizeof(size_t)) throw std::out_of_range("SZ3 Huffman: truncated encoded length"); size_t encodedLength = 0; read(encodedLength, bytes); + remaining -= sizeof(size_t); if (n->t) // root->t==1 means that all state values are the same (constant) { for (count = 0; count < targetLength; count++) out[count] = n->c + offset; return out; } + if (encodedLength > remaining) throw std::out_of_range("SZ3 Huffman: encoded length exceeds compressed buffer"); + for (i = 0; count < targetLength; i++) { byteIndex = i >> 3; // i/8 + if (byteIndex >= encodedLength) throw std::out_of_range("SZ3 Huffman: corrupted encoded stream"); r = i % 8; if (((bytes[byteIndex] >> (7 - r)) & 0x01) == 0) n = n->left; else n = n->right; + if (n == nullptr) throw std::out_of_range("SZ3 Huffman: corrupted tree"); + if (n->t) { out[count] = n->c + offset; n = t; @@ -260,8 +271,13 @@ class HuffmanEncoder : public concepts::EncoderInterface { // load Huffman tree void load(const uchar *&c, size_t &remaining_length) override { read(offset, c, remaining_length); + if (remaining_length < 2 * sizeof(int)) throw std::out_of_range("SZ3 Huffman: truncated tree header"); nodeCount = bytesToInt32_bigEndian(c); int stateNum = bytesToInt32_bigEndian(c + sizeof(int)) * 2; + /// `nodeCount` comes from untrusted data. Bound it before it is used in size computations so + /// the encodeStartIndex arithmetic below cannot overflow and the tree cannot be read past the buffer. + if (nodeCount <= 0 || static_cast(nodeCount) > remaining_length) + throw std::out_of_range("SZ3 Huffman: invalid node count"); size_t encodeStartIndex; if (nodeCount <= 256) encodeStartIndex = 1 + 3 * nodeCount * sizeof(unsigned char) + nodeCount * sizeof(T); @@ -272,12 +288,28 @@ class HuffmanEncoder : public concepts::EncoderInterface { encodeStartIndex = 1 + 2 * nodeCount * sizeof(unsigned int) + nodeCount * sizeof(unsigned char) + nodeCount * sizeof(T); + size_t tree_bytes = sizeof(int) + sizeof(int) + encodeStartIndex; + if (tree_bytes > remaining_length) throw std::out_of_range("SZ3 Huffman: tree exceeds compressed buffer"); + + /// The node pool is sized from stateNum, but the tree is reconstructed from nodeCount nodes; both come + /// from untrusted data. createHuffmanTree allocates a pool of 2*allNodes = 4*stateNum nodes, and the + /// reconstruction creates at most nodeCount of them, so reject a stateNum that is non-positive or too + /// small to hold nodeCount nodes - otherwise new_node2 would write past the end of the pool. + if (stateNum <= 0 || static_cast(nodeCount) > 4 * static_cast(stateNum)) + throw std::out_of_range("SZ3 Huffman: node count exceeds the tree pool capacity"); + huffmanTree = createHuffmanTree(stateNum); treeRoot = reconstruct_HuffTree_from_bytes_anyStates(c + sizeof(int) + sizeof(int), nodeCount); - c += sizeof(int) + sizeof(int) + encodeStartIndex; + c += tree_bytes; + remaining_length -= tree_bytes; loaded = true; } + /// Bound the encoded stream `decode()` is about to read. A caller that lays the tree and the stream + /// out contiguously calls this with the bytes left after the tree (and after any field it writes in + /// between), which makes decode() reject a corrupted length instead of reading past the buffer. + void set_decode_bound(size_t remaining) { decode_remaining_length = remaining; } + bool isLoaded() const { return loaded; } private: @@ -286,6 +318,10 @@ class HuffmanEncoder : public concepts::EncoderInterface { unsigned int nodeCount = 0; uchar sysEndianType; // 0: little endian, 1: big endian bool loaded = false; + /// Bytes the encoded stream may read, or SIZE_MAX when no caller supplied a bound. load() does not + /// set it: the tree and the encoded stream are not required to live in the same buffer (see + /// tools/test/modules/test_encoder.cpp), so only a caller that knows the layout can bound the stream. + size_t decode_remaining_length = std::numeric_limits::max(); T offset; node reconstruct_HuffTree_from_bytes_anyStates(const unsigned char *bytes, uint nodeCount) { @@ -320,7 +356,7 @@ class HuffmanEncoder : public concepts::EncoderInterface { memcpy(t, bytes + 1 + 2 * nodeCount * sizeof(unsigned char) + nodeCount * sizeof(T), nodeCount * sizeof(unsigned char)); node root = this->new_node2(C[0], t[0]); - this->unpad_tree(L, R, C, t, 0, root); + this->unpad_tree(L, R, C, t, 0, root, nodeCount); free(L); free(R); free(C); @@ -361,7 +397,7 @@ class HuffmanEncoder : public concepts::EncoderInterface { nodeCount * sizeof(unsigned char)); node root = this->new_node2(0, 0); - this->unpad_tree(L, R, C, t, 0, root); + this->unpad_tree(L, R, C, t, 0, root, nodeCount); free(L); free(R); free(C); @@ -402,7 +438,7 @@ class HuffmanEncoder : public concepts::EncoderInterface { nodeCount * sizeof(unsigned char)); node root = this->new_node2(0, 0); - this->unpad_tree(L, R, C, t, 0, root); + this->unpad_tree(L, R, C, t, 0, root, nodeCount); free(L); free(R); free(C); @@ -479,11 +515,14 @@ class HuffmanEncoder : public concepts::EncoderInterface { if (n->t) { huffmanTree->code[n->c] = static_cast(malloc(2 * sizeof(uint64_t))); if (len <= 64) { - (huffmanTree->code[n->c])[0] = out1 << (64 - len); + // A Huffman tree with a single symbol gives a zero-length code (build_code is called with + // len == 0 for the root). Shifting a 64-bit value by 64 is undefined behavior, so guard it. + (huffmanTree->code[n->c])[0] = (len == 0) ? 0 : (out1 << (64 - len)); (huffmanTree->code[n->c])[1] = out2; } else { (huffmanTree->code[n->c])[0] = out1; - (huffmanTree->code[n->c])[1] = out2 << (128 - len); + // Likewise, len >= 128 would shift by >= 64; such codes do not fit in 128 bits anyway. + (huffmanTree->code[n->c])[1] = (len >= 128) ? out2 : (out2 << (128 - len)); } huffmanTree->cout[n->c] = static_cast(len); // std::cout << "build_code: c = " << n->c << ", len = " << len << ", out1 = " << out1 << ", out2 = " << out2 @@ -533,6 +572,11 @@ class HuffmanEncoder : public concepts::EncoderInterface { } } + // The state table is sized by the bin range, not the distinct count, so a sparse wide-range + // stream overflows this narrowing. + if (static_cast(max) - static_cast(offset) > 2e9) { + throw std::invalid_argument("HuffmanEncoder: bin range too wide; use HuffmanEncoderV2"); + } int stateNum = max - offset + 2; huffmanTree = createHuffmanTree(stateNum); @@ -579,21 +623,26 @@ class HuffmanEncoder : public concepts::EncoderInterface { } template - void unpad_tree(T1 *L, T1 *R, T *C, unsigned char *t, unsigned int i, node root) { + void unpad_tree(T1 *L, T1 *R, T *C, unsigned char *t, unsigned int i, node root, unsigned int nodeCount) { // root->c = C[i]; if (root->t == 0) { T1 l, r; l = L[i]; if (l != 0) { + // Child indices come from untrusted data. pad_tree always assigns a child a higher index than + // its parent, so a valid index satisfies i < l < nodeCount. Enforcing this prevents reading + // L/R/C/t out of bounds and prevents a cycle that would overflow the node pool. + if (l <= i || l >= nodeCount) throw std::out_of_range("SZ3 Huffman: invalid left child index in tree"); node lroot = new_node2(C[l], t[l]); root->left = lroot; - unpad_tree(L, R, C, t, l, lroot); + unpad_tree(L, R, C, t, l, lroot, nodeCount); } r = R[i]; if (r != 0) { + if (r <= i || r >= nodeCount) throw std::out_of_range("SZ3 Huffman: invalid right child index in tree"); node rroot = new_node2(C[r], t[r]); root->right = rroot; - unpad_tree(L, R, C, t, r, rroot); + unpad_tree(L, R, C, t, r, rroot, nodeCount); } } } diff --git a/include/SZ3/encoder/HuffmanEncoderV2.hpp b/include/SZ3/encoder/HuffmanEncoderV2.hpp index fc041c3d..8e3f615d 100644 --- a/include/SZ3/encoder/HuffmanEncoderV2.hpp +++ b/include/SZ3/encoder/HuffmanEncoderV2.hpp @@ -1037,6 +1037,13 @@ class HuffmanEncoderV2 : public concepts::EncoderInterface { void loadAsDFSOrder(const uchar*& bytes, size_t& remaining_length) { tree.init(); + // The serialized tree is a fixed-size header followed by a DFS bitstream, and must fit in + // remaining_length. Bound every read so corrupted input can not read past the end of the buffer or + // allocate an untrusted amount. + const uchar* const tree_start = bytes; + const size_t header_size = 1 + sizeof(T) + 2 * sizeof(size_t); + if (remaining_length < header_size) throw std::out_of_range("SZ3 HuffmanEncoderV2: truncated tree header"); + tree.usemp = (*bytes) >> 7; tree.mbft = (*bytes) & 0x3f; ++bytes; @@ -1048,18 +1055,36 @@ class HuffmanEncoderV2 : public concepts::EncoderInterface { tree.n = bytesToInt64_bigEndian(bytes); bytes += sizeof(size_t); - tree.ht.reserve(tree.n << 1); - tree.maxval = bytesToInt64_bigEndian(bytes); bytes += sizeof(size_t); + + // The DFS bitstream follows the header; each of the tree.n nodes consumes at least one bit, so the + // node count can not exceed the available bits. This also avoids the tree.n << 1 overflow below. + // tree.n is an int holding a value read as 64-bit, so check the sign explicitly rather than letting + // the comparison convert it. + const size_t dfs_bytes = remaining_length - header_size; + if (tree.n < 0 || static_cast(tree.n) > dfs_bytes * 8) + throw std::out_of_range("SZ3 HuffmanEncoderV2: node count exceeds the compressed buffer"); + tree.ht.reserve(static_cast(tree.n) << 1); + if (tree.usemp == 0x00) { + // maxval comes from the stream and sizes the dense tables below. preprocess_encode only leaves + // usemp == 0 when maxval stays under 1 << 28, so a larger value in a dense tree is inconsistent + // and would otherwise drive an unbounded allocation. + const int64_t maxval_span = static_cast(tree.maxval); + if (maxval_span < 0 || maxval_span >= (1ll << 28)) { + throw std::out_of_range("SZ3 HuffmanEncoderV2: dense tree declares an out-of-range value span"); + } if (tree.n > 0) { tree.veccode.resize(tree.maxval); tree.veclen.resize(tree.maxval); } } + // The fixed-size header has now been consumed on every path below. + remaining_length -= header_size; + if (tree.n == 0) { tree.setConstructed(); return; @@ -1090,6 +1115,8 @@ class HuffmanEncoderV2 : public concepts::EncoderInterface { while (!stk.empty()) { Node* u = stk.top(); + if (static_cast(i >> 3) >= dfs_bytes) + throw std::out_of_range("SZ3 HuffmanEncoderV2: tree bitstream exceeds the compressed buffer"); if (readBit(bytes, i++) == 0x00) { tree.ht.push_back(Node()); if (u->p[0] == nullptr) { @@ -1100,7 +1127,11 @@ class HuffmanEncoderV2 : public concepts::EncoderInterface { stk.push(&tree.ht[tree.ht.size() - 1]); } else { T c = 0; - for (int j = 0; j < tree.mbft; j++) c |= static_cast(readBit(bytes, i++)) << j; + for (int j = 0; j < tree.mbft; j++) { + if (static_cast(i >> 3) >= dfs_bytes) + throw std::out_of_range("SZ3 HuffmanEncoderV2: tree bitstream exceeds the compressed buffer"); + c |= static_cast(readBit(bytes, i++)) << j; + } tree.ht.push_back(Node(c)); if (u->p[0] == nullptr) u->p[0] = &tree.ht[tree.ht.size() - 1]; @@ -1119,6 +1150,10 @@ class HuffmanEncoderV2 : public concepts::EncoderInterface { bytes += (i + 7) >> 3; + // Consume the DFS bitstream from the caller's remaining length (the header was already subtracted + // above; the original code advanced the pointer but never decremented remaining_length). + remaining_length -= static_cast(bytes - tree_start) - header_size; + if (tree.usemp) { tree.dfs_mp(&tree.ht[tree.root]); } else { diff --git a/include/SZ3/encoder/RunlengthEncoder.hpp b/include/SZ3/encoder/RunlengthEncoder.hpp index 1d00a574..b2a95a80 100644 --- a/include/SZ3/encoder/RunlengthEncoder.hpp +++ b/include/SZ3/encoder/RunlengthEncoder.hpp @@ -12,7 +12,9 @@ namespace SZ3 { template class RunlengthEncoder : public concepts::EncoderInterface { public: - void preprocess_encode(const std::vector &bins, int stateNum) override {} + void preprocess_encode(const std::vector &bins, int stateNum) override { num_bins = bins.size(); } + + size_t size_est() override { return num_bins * (sizeof(T) + sizeof(int)); } size_t encode(const std::vector &bins, uchar *&bytes) override { auto bytespos = bytes; @@ -60,6 +62,8 @@ class RunlengthEncoder : public concepts::EncoderInterface { void save(uchar *&c) override {} void load(const uchar *&c, size_t &remaining_length) override {} + + size_t num_bins = 0; ///< Set by preprocess_encode(), consumed by size_est() }; } // namespace SZ3 #endif diff --git a/include/SZ3/encoder/XtcBasedEncoder.hpp b/include/SZ3/encoder/XtcBasedEncoder.hpp index bc429612..90912c76 100644 --- a/include/SZ3/encoder/XtcBasedEncoder.hpp +++ b/include/SZ3/encoder/XtcBasedEncoder.hpp @@ -9,10 +9,12 @@ #define _SZ_XTC3_ENCODER_HPP #include +#include #include #include "SZ3/def.hpp" #include "SZ3/encoder/Encoder.hpp" +#include "SZ3/utils/Config.hpp" // #define DEBUG_OUTPUT diff --git a/include/SZ3/lossless/Lossless.hpp b/include/SZ3/lossless/Lossless.hpp index acf29630..37bcdb06 100644 --- a/include/SZ3/lossless/Lossless.hpp +++ b/include/SZ3/lossless/Lossless.hpp @@ -5,6 +5,10 @@ #ifndef SZ3_LOSSLESS_HPP #define SZ3_LOSSLESS_HPP +#include + +#include "SZ3/def.hpp" + namespace SZ3::concepts { /** diff --git a/include/SZ3/lossless/Lossless_bypass.hpp b/include/SZ3/lossless/Lossless_bypass.hpp index c281dc52..c2e872b7 100644 --- a/include/SZ3/lossless/Lossless_bypass.hpp +++ b/include/SZ3/lossless/Lossless_bypass.hpp @@ -6,6 +6,8 @@ #define SZ3_LOSSLESS_BYPASS_HPP #include +#include + #include "SZ3/def.hpp" #include "SZ3/lossless/Lossless.hpp" @@ -13,15 +15,30 @@ namespace SZ3 { class Lossless_bypass : public concepts::LosslessInterface { public: size_t compress(const uchar *src, size_t srcLen, uchar *dst, size_t dstCap) override { + // Nothing here shrinks the payload, so a caller that sized its destination from a bound assuming + // compression can be smaller than srcLen. Lossless_zstd already checks and throws; do the same + // instead of memcpy'ing past the end of the destination. + if (dstCap < srcLen) { + throw std::length_error(SZ3_ERROR_COMP_BUFFER_NOT_LARGE_ENOUGH); + } std::memcpy(dst, src, srcLen); - // dst = src; return srcLen; } size_t decompress(const uchar *src, const size_t srcLen, uchar *&dst, size_t &dstLen) override { + // Mirror Lossless_zstd: when the caller asks us to allocate, a non-zero incoming dstLen is an + // upper bound on the allocation; zero means no bound. dstLen stays a pure output parameter for a + // caller-provided buffer. + const size_t dst_capacity = dstLen; dstLen = srcLen; if (dst == nullptr) { + if (dst_capacity != 0 && dstLen > dst_capacity) { + throw std::out_of_range("SZ3 bypass lossless: payload exceeds the allowed capacity"); + } dst = static_cast(malloc(dstLen)); + if (dst == nullptr) { + throw std::runtime_error("SZ3 bypass lossless: can not allocate the decompression buffer"); + } } std::memcpy(dst, src, dstLen); return dstLen; diff --git a/include/SZ3/lossless/Lossless_zstd.hpp b/include/SZ3/lossless/Lossless_zstd.hpp index 3cb9953a..1fb98aa6 100644 --- a/include/SZ3/lossless/Lossless_zstd.hpp +++ b/include/SZ3/lossless/Lossless_zstd.hpp @@ -37,11 +37,57 @@ class Lossless_zstd : public concepts::LosslessInterface { } size_t decompress(const uchar *src, const size_t srcLen, uchar *&dst, size_t &dstLen) override { + /// The compressed buffer starts with the decompressed size, followed by the zstd stream. + /// Validate the inputs so corrupted data can not read out of bounds, allocate an untrusted + /// amount, or pass a zstd error code back as a size. + if (srcLen < sizeof(dstLen)) { + throw std::out_of_range("SZ3 lossless: compressed data is smaller than the size header"); + } + /// When the caller asks us to allocate (dst == nullptr), a non-zero incoming dstLen is an upper + /// bound on that allocation; zero means no bound. This is only consulted on that branch, so a + /// caller that treats dstLen as pure output is unaffected. + const size_t dst_capacity = dstLen; read(dstLen, src); - if (dst == nullptr) { + const bool self_allocated = (dst == nullptr); + if (self_allocated) { + /// dst == nullptr means the caller asks us to allocate the output buffer, and its size is read + /// from the (untrusted) compressed payload. When the caller supplies a non-zero capacity it is an + /// upper bound on how large that allocation may be; reject a payload that declares a larger size + /// before allocating it, so corrupted data can not force an arbitrary (and untracked) allocation. + /// A zero capacity means the caller did not supply a bound (legacy behavior). + if (dst_capacity != 0 && dstLen > dst_capacity) { + throw std::out_of_range("SZ3 lossless: declared decompressed size exceeds the allowed capacity"); + } dst = static_cast(malloc(dstLen)); + if (dst == nullptr) { + throw std::runtime_error("SZ3 lossless: can not allocate the decompression buffer"); + } + } + // No capacity check for a caller-provided buffer: `dstLen` is an output parameter in the + // LosslessInterface contract, so callers are not required to pass a meaningful value in + // (tools/test/modules/test_lossless.cpp passes an uninitialized one). Reinterpreting it as an + // input capacity would break them. The self-allocating branch above bounds itself instead. + try { + size_t res = ZSTD_decompress(dst, dstLen, src, srcLen - sizeof(dstLen)); + if (ZSTD_isError(res)) { + throw std::runtime_error("SZ3 lossless: zstd decompression failed"); + } + /// The declared size is read from the (untrusted) payload; require zstd to actually produce that + /// many bytes, so a frame that expands to fewer bytes can not leave the tail of the output buffer + /// uninitialized (which a caller would otherwise copy out as if it were decompressed data). + if (res != dstLen) { + throw std::out_of_range("SZ3 lossless: decompressed size does not match the declared size"); + } + return res; + } catch (...) { + /// Free a buffer we allocated ourselves so a corrupted payload that fails decompression here does + /// not leak it; a caller-provided buffer is owned by the caller and is left untouched. + if (self_allocated) { + free(dst); + dst = nullptr; + } + throw; } - return ZSTD_decompress(dst, dstLen, src, srcLen - sizeof(dstLen)); } private: diff --git a/include/SZ3/predictor/ComposedPredictor.hpp b/include/SZ3/predictor/ComposedPredictor.hpp index 13f4f6f4..e149b842 100644 --- a/include/SZ3/predictor/ComposedPredictor.hpp +++ b/include/SZ3/predictor/ComposedPredictor.hpp @@ -45,7 +45,13 @@ class ComposedPredictor : public concepts::PredictorInterface { } bool predecompress(const block_iter &block) override { + // selection and its entries come from untrusted data; bound the running index into selection and the + // selected predictor index before using them. predict()/estimate_error() reuse the sid set here. + if (current_index >= selection.size()) + throw std::out_of_range("SZ3: ran out of predictor selections while decompressing"); sid = selection[current_index++]; + if (sid < 0 || static_cast(sid) >= predictors.size()) + throw std::out_of_range("SZ3: predictor selection index is out of range"); return predictors[sid]->predecompress(block); } @@ -72,8 +78,15 @@ class ComposedPredictor : public concepts::PredictorInterface { if (selection_size > 0) { HuffmanEncoder selection_encoder; selection_encoder.load(c, remaining_length); + // The tree is immediately followed by its encoded stream here, so the remaining + // length bounds the stream exactly. + selection_encoder.set_decode_bound(remaining_length); + /// decode() advances `c` past the encoded stream but does not update `remaining_length`; + /// account for exactly the bytes it consumed so the bound stays tight for later reads. + const uchar *decode_start = c; this->selection = selection_encoder.decode(c, selection_size); selection_encoder.postprocess_decode(); + remaining_length -= static_cast(c - decode_start); } } diff --git a/include/SZ3/predictor/LorenzoPredictor.hpp b/include/SZ3/predictor/LorenzoPredictor.hpp index e61942fa..810d07f9 100644 --- a/include/SZ3/predictor/LorenzoPredictor.hpp +++ b/include/SZ3/predictor/LorenzoPredictor.hpp @@ -1,6 +1,9 @@ #ifndef SZ3_LORENZO_PREDICTOR_HPP #define SZ3_LORENZO_PREDICTOR_HPP +#include +#include + #include "SZ3/predictor/Predictor.hpp" namespace SZ3 { @@ -97,7 +100,11 @@ class LorenzoPredictor : public concepts::PredictorInterface { T noise = 0; private: - // Helper functions for Lorenzo prediction + // Helper functions for Lorenzo prediction. + // The neighbour offsets are unsigned, so `d[-offset]` would be `*(d + (size_t)(-offset))`, forming an + // out-of-bounds pointer by wrapping the unsigned addition (flagged by -fsanitize=pointer-overflow even + // though the accessed element is in bounds thanks to the predictor's padding). The addresses below are + // computed with pointer subtraction so the offset stays a small negative step and no wrap-around occurs. T prev1(T *d, size_t i) { return *(d - i); } T prev2(T *d, const std::array &ds, size_t j, size_t i) { return *(d - (j * ds[0] + i)); } T prev3(T *d, const std::array &ds, size_t k, size_t j, size_t i) { diff --git a/include/SZ3/predictor/RegressionPredictor.hpp b/include/SZ3/predictor/RegressionPredictor.hpp index a77e9a9a..823d17b9 100644 --- a/include/SZ3/predictor/RegressionPredictor.hpp +++ b/include/SZ3/predictor/RegressionPredictor.hpp @@ -114,9 +114,17 @@ class RegressionPredictor : public concepts::PredictorInterface { quantizer_liner.load(c, remaining_length); HuffmanEncoder encoder = HuffmanEncoder(); encoder.load(c, remaining_length); + // The tree is immediately followed by its encoded stream here, so the remaining + // length bounds the stream exactly. + encoder.set_decode_bound(remaining_length); + /// decode() advances `c` past the encoded stream but does not update `remaining_length`; + /// account for exactly the bytes it consumed. The previous `coeff_size * sizeof(int)` used the + /// uncompressed index count, which overshoots the (Huffman-compressed) stream and understates + /// remaining_length, making a later encoder.load() bound check spuriously reject valid data. + const uchar *decode_start = c; regression_coeff_quant_inds = encoder.decode(c, coeff_size); encoder.postprocess_decode(); - remaining_length -= coeff_size * sizeof(int); + remaining_length -= static_cast(c - decode_start); std::fill(current_coeffs.begin(), current_coeffs.end(), 0); regression_coeff_index = 0; } @@ -155,6 +163,10 @@ class RegressionPredictor : public concepts::PredictorInterface { } void pred_and_recover_coefficients() { + // Each block consumes N + 1 regression coefficients; the coefficient stream comes from untrusted data, + // so a crafted block count larger than the stored coefficients would read past its end. + if (regression_coeff_index + N + 1 > regression_coeff_quant_inds.size()) + throw std::out_of_range("SZ3: ran out of regression coefficients while decompressing"); for (int i = 0; i < static_cast(N); i++) { current_coeffs[i] = quantizer_liner.recover(current_coeffs[i], regression_coeff_quant_inds[regression_coeff_index++]); diff --git a/include/SZ3/preprocessor/PreFilter.hpp b/include/SZ3/preprocessor/PreFilter.hpp index eefc2a5c..3b08c059 100644 --- a/include/SZ3/preprocessor/PreFilter.hpp +++ b/include/SZ3/preprocessor/PreFilter.hpp @@ -5,6 +5,10 @@ #ifndef SZ3_PREFILTER_HPP #define SZ3_PREFILTER_HPP +#include +#include +#include + #include "SZ3/preprocessor/PreProcessor.hpp" namespace SZ3 { diff --git a/include/SZ3/preprocessor/Transpose.hpp b/include/SZ3/preprocessor/Transpose.hpp index b017a79b..6e70d9a2 100644 --- a/include/SZ3/preprocessor/Transpose.hpp +++ b/include/SZ3/preprocessor/Transpose.hpp @@ -5,6 +5,10 @@ #ifndef SZ3_TRANSPOSE_H #define SZ3_TRANSPOSE_H +#include +#include +#include + #include "SZ3/preprocessor/PreProcessor.hpp" namespace SZ3 { diff --git a/include/SZ3/quantizer/LinearQuantizer.hpp b/include/SZ3/quantizer/LinearQuantizer.hpp index 3da366de..002cb7f4 100644 --- a/include/SZ3/quantizer/LinearQuantizer.hpp +++ b/include/SZ3/quantizer/LinearQuantizer.hpp @@ -42,32 +42,35 @@ class LinearQuantizer : public concepts::QuantizerInterface { // int quantize(T data, T pred, T& dec_data); ALWAYS_INLINE int quantize_and_overwrite(T& data, T pred) override { T diff = data - pred; - auto quant_index = static_cast(fabs(diff) * this->error_bound_reciprocal) + 1; - if (quant_index < this->radius * 2) { - quant_index >>= 1; - int half_index = quant_index; - quant_index <<= 1; - int quant_index_shifted; - if (diff < 0) { - quant_index = -quant_index; - quant_index_shifted = this->radius - half_index; - } else { - quant_index_shifted = this->radius + half_index; + // fabs(diff) * error_bound_reciprocal is NaN when data is NaN and exceeds the int64_t + // range for infinities or huge magnitudes; casting those to int64_t is undefined behaviour. + // Only finite magnitudes within the quantization range are representable as an index; every + // other value is stored losslessly in unpred, exactly like the out-of-range branch below. + double scaled = fabs(diff) * this->error_bound_reciprocal; + if (scaled < this->radius * 2) { + auto quant_index = static_cast(scaled) + 1; + if (quant_index < this->radius * 2) { + quant_index >>= 1; + int half_index = quant_index; + quant_index <<= 1; + int quant_index_shifted; + if (diff < 0) { + quant_index = -quant_index; + quant_index_shifted = this->radius - half_index; + } else { + quant_index_shifted = this->radius + half_index; + } + T decompressed_data = pred + quant_index * this->error_bound; + // if data is NaN, the diff is NaN, and NaN <= 0 is false + diff = fabs(decompressed_data - data); + if (diff <= this->error_bound || (!strict_eb && diff <= this->error_bound * 1.1)) { + data = decompressed_data; + return quant_index_shifted; + } } - T decompressed_data = pred + quant_index * this->error_bound; - // if data is NaN, the diff is NaN, and NaN <= 0 is false - diff = fabs(decompressed_data - data); - if (diff <= this->error_bound || (!strict_eb && diff <= this->error_bound * 1.1)) { - data = decompressed_data; - return quant_index_shifted; - } else { - unpred.push_back(data); - return 0; - } - } else { - unpred.push_back(data); - return 0; } + unpred.push_back(data); + return 0; } // recover the data using the quantization index @@ -83,7 +86,12 @@ class LinearQuantizer : public concepts::QuantizerInterface { return pred + 2 * (quant_index - this->radius) * this->error_bound; } - ALWAYS_INLINE T recover_unpred() { return unpred[index++]; } + ALWAYS_INLINE T recover_unpred() { + // index and the quantization stream come from untrusted data; a crafted stream can request more + // unpredictable values than were stored, which would read past the end of unpred. + if (index >= unpred.size()) throw std::out_of_range("SZ3: ran out of unpredictable values while decompressing"); + return unpred[index++]; + } ALWAYS_INLINE int force_save_unpred(T ori) override { unpred.push_back(ori); @@ -115,6 +123,10 @@ class LinearQuantizer : public concepts::QuantizerInterface { size_t unpred_size = 0; read(unpred_size, c, remaining_length); if (unpred_size > 0) { + // Validate the count against the remaining bytes before resizing, otherwise a corrupted count + // would drive a huge allocation before the bounded read below has a chance to reject it. + if (unpred_size > remaining_length / sizeof(T)) + throw std::out_of_range("SZ3: unpredictable value count exceeds the compressed buffer"); unpred.resize(unpred_size); read(unpred.data(), unpred_size, c, remaining_length); } diff --git a/include/SZ3/quantizer/Quantizer.hpp b/include/SZ3/quantizer/Quantizer.hpp index eac5d4df..041a81c0 100644 --- a/include/SZ3/quantizer/Quantizer.hpp +++ b/include/SZ3/quantizer/Quantizer.hpp @@ -2,6 +2,9 @@ #define SZ3_QUANTIZER_HPP #include +#include +#include +#include namespace SZ3::concepts { @@ -70,4 +73,22 @@ class QuantizerInterface { }; } // namespace SZ3::concepts +namespace SZ3 { +/// Detects the optional (non-virtual) `size_est()` some quantizers expose. +template +struct quantizer_has_size_est : std::false_type {}; +template +struct quantizer_has_size_est().size_est())>> : std::true_type {}; + +/// The quantizer's serialized-size estimate, or 0 when it does not expose one. +template +size_t quantizer_size_est(Q &q) { + if constexpr (quantizer_has_size_est::value) { + return q.size_est(); + } else { + return 0; + } +} +} // namespace SZ3 + #endif diff --git a/include/SZ3/utils/BlockwiseIterator.hpp b/include/SZ3/utils/BlockwiseIterator.hpp index 3221d4f1..d2aad1c1 100644 --- a/include/SZ3/utils/BlockwiseIterator.hpp +++ b/include/SZ3/utils/BlockwiseIterator.hpp @@ -4,12 +4,15 @@ #include #include #include +#include #include #include #include #include #include +#include "SZ3/def.hpp" + namespace SZ3 { /** @@ -219,6 +222,23 @@ class block_data : public std::enable_shared_from_this> { } } + /** + * @brief The block's values in unpadded layout. + * + * Without padding this is the array this object was constructed from. With padding it is an + * internal copy, materialized on the first call and valid until this object is destroyed. + * + * @return Pointer to `num` elements + */ + const T *values() { + if (padding == 0 || internal_buffer.empty()) { + return data_padding; + } + unpadded_buffer.resize(num); + copy_data_with_padding(unpadded_buffer.data(), ds, data_padding, ds_padding, dims); + return unpadded_buffer.data(); + } + block_iterator block_iter(size_t block_size) { return block_iterator(this->shared_from_this(), block_size); } protected: @@ -273,6 +293,7 @@ class block_data : public std::enable_shared_from_this> { std::array dims; // dimension std::array ds, ds_padding; // stride std::vector internal_buffer; + std::vector unpadded_buffer; // materialized on demand by values() T *data_cp_dst = nullptr; T *data_padding; // point to either data_ or internal_buffer depending on padding size_t padding; diff --git a/include/SZ3/utils/Config.hpp b/include/SZ3/utils/Config.hpp index 81febd8c..25a7f342 100644 --- a/include/SZ3/utils/Config.hpp +++ b/include/SZ3/utils/Config.hpp @@ -15,8 +15,10 @@ #include #include #include +#include #include #include +#include #include #include "SZ3/def.hpp" @@ -358,38 +360,81 @@ class Config { * * @param c Pointer to the byte array. */ - void load(const unsigned char*& c) { + // Legacy overload for trusted internal callers (e.g. the OpenMP path). + void load(const unsigned char*& c) { load(c, std::numeric_limits::max()); } + + // `cmpSize` bounds how many bytes may be read from `c` (the config blob). Used when loading + // from untrusted compressed data so a corrupted config cannot read out of bounds. + void load(const unsigned char*& c, size_t cmpSize) { + const unsigned char* const c0 = c; + auto require = [&](size_t n) { + if (cmpSize - static_cast(c - c0) < n) + throw std::out_of_range("SZ3 Config::load: read past the end of the config"); + }; + + require(sizeof(uchar)); uchar confSize = 0; read(confSize, c); - auto c1 = c + confSize; + /// `confSize` is the total size of the config blob, including this prefix byte. + if (confSize > cmpSize) throw std::out_of_range("SZ3 Config::load: config size exceeds the buffer"); + auto c1 = c0 + confSize; + require(sizeof(N)); read(N, c); + if (N < 1 || N > 4) throw std::out_of_range("SZ3 Config::load: invalid number of dimensions"); uint8_t bitWidth; + require(sizeof(bitWidth)); read(bitWidth, c); + if (bitWidth > 64) throw std::out_of_range("SZ3 Config::load: invalid dimension bit width"); + require((static_cast(N) * bitWidth + 7) / 8); dims = bytes2vector(c, bitWidth, N); - // dims.resize(N); - // read(dims.data(), N, c); + require(sizeof(num)); read(num, c); + /// The element count must equal the product of the dimensions, otherwise the predictor would + /// iterate over more grid positions than were allocated for the decompressed data. Validate + /// the product with an overflow check. + { + size_t dims_product = 1; + bool dims_ok = true; + for (size_t dim : dims) { + if (dim == 0 || dims_product > std::numeric_limits::max() / dim) { + dims_ok = false; + break; + } + dims_product *= dim; + } + if (!dims_ok || dims_product != num) + throw std::out_of_range("SZ3 Config::load: dimensions inconsistent with the element count"); + } + require(sizeof(cmprAlgo)); read(cmprAlgo, c); + require(sizeof(errorBoundMode)); read(errorBoundMode, c); if (errorBoundMode == EB_ABS) { + require(sizeof(absErrorBound)); read(absErrorBound, c); } else if (errorBoundMode == EB_REL) { + require(sizeof(relErrorBound)); read(relErrorBound, c); } else if (errorBoundMode == EB_PSNR) { + require(sizeof(psnrErrorBound)); read(psnrErrorBound, c); } else if (errorBoundMode == EB_L2NORM) { + require(sizeof(l2normErrorBound)); read(l2normErrorBound, c); } else if (errorBoundMode == EB_ABS_OR_REL) { + require(sizeof(absErrorBound) + sizeof(relErrorBound)); read(absErrorBound, c); read(relErrorBound, c); } else if (errorBoundMode == EB_ABS_AND_REL) { + require(sizeof(absErrorBound) + sizeof(relErrorBound)); read(absErrorBound, c); read(relErrorBound, c); } if (c < c1) { + require(sizeof(uint8_t)); uint8_t boolvals; read(boolvals, c); lorenzo = (boolvals >> 7) & 1; @@ -399,15 +444,19 @@ class Config { openmp = (boolvals >> 3) & 1; } if (c < c1) { + require(sizeof(dataType)); read(dataType, c); } if (c < c1) { + require(sizeof(quantbinCnt)); read(quantbinCnt, c); } if (c < c1) { + require(sizeof(blockSize)); read(blockSize, c); } if (c < c1) { + require(sizeof(predDim)); read(predDim, c); } } diff --git a/include/SZ3/utils/Extraction.hpp b/include/SZ3/utils/Extraction.hpp index 1e0bef9a..1ab9db5a 100644 --- a/include/SZ3/utils/Extraction.hpp +++ b/include/SZ3/utils/Extraction.hpp @@ -5,6 +5,14 @@ #ifndef SZ3_EXTRACTION_HPP #define SZ3_EXTRACTION_HPP +#include +#include +#include +#include + +#include "SZ3/def.hpp" +#include "SZ3/utils/Timer.hpp" + namespace SZ3 { template diff --git a/include/SZ3/utils/Iterator.hpp b/include/SZ3/utils/Iterator.hpp index 9c4d7a07..7dc3973b 100644 --- a/include/SZ3/utils/Iterator.hpp +++ b/include/SZ3/utils/Iterator.hpp @@ -14,6 +14,8 @@ #include #include +#include "SZ3/def.hpp" + namespace SZ3 { // N-dimensional multi_dimensional_range template diff --git a/include/SZ3/utils/KmeansUtil.hpp b/include/SZ3/utils/KmeansUtil.hpp index 478df3dd..1cfe9c82 100644 --- a/include/SZ3/utils/KmeansUtil.hpp +++ b/include/SZ3/utils/KmeansUtil.hpp @@ -292,7 +292,7 @@ void get_cluster(T *data, size_t num, float &level_start, float &level_offset, i if (num == sample_num) { sample = std::vector(data, data + num); } else { - sample.reserve(sample_num); + sample.resize(sample_num); // the loop below writes through operator[] std::random_device rd; // Will be used to obtain a seed for the random number engine std::mt19937 gen(rd()); // Standard mersenne_twister_engine seeded with rd() // std::uniform_int_distribution<> dis(0, 2 * sample_rate); @@ -303,7 +303,7 @@ void get_cluster(T *data, size_t num, float &level_start, float &level_offset, i // sample[i] = input[input_idx]; // } // std::cout << std::endl; - std::uniform_int_distribution<> dis2(0, num); + std::uniform_int_distribution<> dis2(0, static_cast(num) - 1); std::unordered_set sampledkeys; // printf("total_num=%lu, sample_num=%lu\n", num, sample_num); for (size_t i = 0; i < sample_num; i++) { diff --git a/include/SZ3/utils/MemoryUtil.hpp b/include/SZ3/utils/MemoryUtil.hpp index 7128b84b..5e5ef301 100644 --- a/include/SZ3/utils/MemoryUtil.hpp +++ b/include/SZ3/utils/MemoryUtil.hpp @@ -5,9 +5,9 @@ #ifndef SZ3_MEMORYOPS_HPP #define SZ3_MEMORYOPS_HPP -#include -#include #include +#include +#include #include "SZ3/def.hpp" @@ -73,7 +73,9 @@ inline T byteswap(T value) { // read array template void read(T1 *array, size_t num_elements, uchar const *&compressed_data_pos, size_t &remaining_length) { - assert(num_elements * sizeof(T1) <= remaining_length); + if (num_elements * sizeof(T1) > remaining_length) { + throw std::invalid_argument("SZ3: compressed stream is truncated"); + } memcpy(array, compressed_data_pos, num_elements * sizeof(T1)); if constexpr (SZ3_BIG_ENDIAN) { for (size_t i = 0; i < num_elements; i++) { @@ -109,7 +111,8 @@ void read(T1 &var, uchar const *&compressed_data_pos) { // read variable template void read(T1 &var, uchar const *&compressed_data_pos, size_t &remaining_length) { - assert(sizeof(T1) <= remaining_length); + if (sizeof(T1) > remaining_length) + throw std::out_of_range("SZ3: attempt to read past the end of the compressed buffer"); memcpy(&var, compressed_data_pos, sizeof(T1)); if constexpr (SZ3_BIG_ENDIAN) { var = byteswap(var); diff --git a/include/SZ3/utils/QuantOptimization.hpp b/include/SZ3/utils/QuantOptimization.hpp index de1c6653..9ae0a3a3 100644 --- a/include/SZ3/utils/QuantOptimization.hpp +++ b/include/SZ3/utils/QuantOptimization.hpp @@ -1,8 +1,11 @@ #ifndef SZ3_optimize_quant_intervals_hpp #define SZ3_optimize_quant_intervals_hpp +#include #include +#include "SZ3/def.hpp" + namespace SZ3 { #define QuantIntvMeanCapacity 8192 diff --git a/include/SZ3/utils/Sample.hpp b/include/SZ3/utils/Sample.hpp index 33d5cf70..8e061d6d 100644 --- a/include/SZ3/utils/Sample.hpp +++ b/include/SZ3/utils/Sample.hpp @@ -1,9 +1,11 @@ #ifndef SZ3_SAMPLE_HPP #define SZ3_SAMPLE_HPP -#include "SZ3/def.hpp" +#include #include +#include "SZ3/def.hpp" + namespace SZ3 { template inline void profiling_block(T* data, std::vector& dims, std::vector>& starts, diff --git a/include/SZ3/utils/Statistic.hpp b/include/SZ3/utils/Statistic.hpp index b8cafa4d..d22bc366 100644 --- a/include/SZ3/utils/Statistic.hpp +++ b/include/SZ3/utils/Statistic.hpp @@ -5,6 +5,10 @@ #ifndef SZ3_STATISTIC_HPP #define SZ3_STATISTIC_HPP +#include +#include +#include + #include "Config.hpp" namespace SZ3 { diff --git a/tools/H5Z-SZ3/src/H5Z_SZ3.cpp b/tools/H5Z-SZ3/src/H5Z_SZ3.cpp index 69d08772..425bbebf 100644 --- a/tools/H5Z-SZ3/src/H5Z_SZ3.cpp +++ b/tools/H5Z-SZ3/src/H5Z_SZ3.cpp @@ -1,5 +1,6 @@ #include "H5Z_SZ3.hpp" +#include #include #include #include @@ -159,7 +160,10 @@ void process_data(SZ3::Config& conf, void** buf, size_t* buf_size, size_t nbytes *buf = processedData; *buf_size = conf.num * sizeof(T); } else { - size_t cmpCap = sizeof(T) * conf.num * 2; + // SZ_compress rejects anything below its own bound, which small chunks fall under; that + // bound assumes the payload fits in the raw size, so keep the old headroom on top of it + // for algorithms whose output can approach or exceed it. + size_t cmpCap = std::max(SZ3::SZ_compress_size_bound(conf), sizeof(T) * conf.num * 2); char* cmpData = static_cast(malloc(cmpCap)); *buf_size = SZ_compress(conf, static_cast(*buf), cmpData, cmpCap); free(*buf); diff --git a/tools/test/modules/test_lossless.cpp b/tools/test/modules/test_lossless.cpp index 00e23cef..81b19ded 100644 --- a/tools/test/modules/test_lossless.cpp +++ b/tools/test/modules/test_lossless.cpp @@ -22,7 +22,8 @@ void runFunctionalTest() { std::vector decompressed(N); SZ3::uchar* decompressed_pos = decompressed.data(); - size_t decompressedSize; + // decompress() reports the size here; initialise it so the value passed in is never garbage. + size_t decompressedSize = 0; lossless.decompress(dst.data(), compressedSize, decompressed_pos, decompressedSize); EXPECT_EQ(decompressedSize, src.size()); From 6a7985b25657bead41524e071e7f972c20851acc Mon Sep 17 00:00:00 2001 From: Kai Zhao Date: Mon, 31 Aug 2026 23:27:05 -0700 Subject: [PATCH 2/6] Restore PR #135's XtcBasedEncoder bounds Taking the fz-branch fixes file by file overwrote them. The magicInts index and the packed-data size both come from the compressed stream and were used unchecked. Co-Authored-By: Claude Opus 5 --- include/SZ3/encoder/XtcBasedEncoder.hpp | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/include/SZ3/encoder/XtcBasedEncoder.hpp b/include/SZ3/encoder/XtcBasedEncoder.hpp index 90912c76..85f28248 100644 --- a/include/SZ3/encoder/XtcBasedEncoder.hpp +++ b/include/SZ3/encoder/XtcBasedEncoder.hpp @@ -10,6 +10,7 @@ #include #include +#include #include #include "SZ3/def.hpp" @@ -599,10 +600,9 @@ class XtcBasedEncoder : public concepts::EncoderInterface { size_t bufferSize = targetLength * 1.2; struct DataBuffer buffer; - buffer.data = reinterpret_cast(malloc(bufferSize * sizeof(int))); - if (buffer.data == nullptr) { - fprintf(stderr, "malloc failed\n"); - } + // buffer.data is allocated below (after size3 is known); the previous allocation here was overwritten + // and leaked. + buffer.data = nullptr; buffer.index = 0; buffer.lastbits = 0; buffer.lastbyte = 0; @@ -639,6 +639,9 @@ class XtcBasedEncoder : public concepts::EncoderInterface { } int smallIdx = *inputIntPtr++; + // smallIdx is read from the compressed data and indexes the fixed-size magicInts table below. + if (smallIdx < 0 || smallIdx >= LASTIDX) + throw std::out_of_range("SZ3 Xtc: small index out of range"); int smaller = magicInts[std::max(FIRSTIDX, smallIdx - 1)] / 2; int smallNum = magicInts[smallIdx] / 2; @@ -650,9 +653,17 @@ class XtcBasedEncoder : public concepts::EncoderInterface { size_t size3 = targetLength; bufferSize = size3 * 1.2; buffer.data = reinterpret_cast(malloc(bufferSize * sizeof(int))); + if (buffer.data == nullptr) { + throw std::runtime_error("SZ3 Xtc: can not allocate the decompression buffer"); + } buffer.index = *(reinterpret_cast(inputIntPtr)); inputIntPtr += sizeof(uint64_t) / sizeof(int); + // buffer.index is an attacker-controlled byte count that is copied into buffer.data below; it must not + // exceed the buffer capacity, otherwise the memcpy loop overflows the heap buffer. + if (buffer.index > bufferSize * sizeof(int)) + throw std::out_of_range("SZ3 Xtc: packed data size exceeds the decompression buffer"); + size_t offset = 0; size_t remain = buffer.index; inputBytesPointer = reinterpret_cast(inputIntPtr); @@ -672,6 +683,9 @@ class XtcBasedEncoder : public concepts::EncoderInterface { int run = 0; size_t i = 0; int *intBufferPoiner = reinterpret_cast(malloc(size3 * sizeof(*intBufferPoiner))); + if (intBufferPoiner == nullptr) { + throw std::runtime_error("SZ3 Xtc: can not allocate the index buffer"); + } int *localIntBufferPointer = intBufferPoiner; unsigned char *charOutputPtr = reinterpret_cast(quantData.data()); int *intOutputPtr = reinterpret_cast(charOutputPtr); From 40523ceef06603f61d28a13a14c410f34cb5183f Mon Sep 17 00:00:00 2001 From: Kai Zhao Date: Mon, 31 Aug 2026 23:31:11 -0700 Subject: [PATCH 3/6] Compare by division in the bounded array read num_elements * sizeof(T1) can overflow on its own, which is the computation the check is supposed to guard. Matches the scalar overload and PR #132. Co-Authored-By: Claude Opus 5 --- include/SZ3/utils/MemoryUtil.hpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/include/SZ3/utils/MemoryUtil.hpp b/include/SZ3/utils/MemoryUtil.hpp index 5e5ef301..2f630166 100644 --- a/include/SZ3/utils/MemoryUtil.hpp +++ b/include/SZ3/utils/MemoryUtil.hpp @@ -73,9 +73,8 @@ inline T byteswap(T value) { // read array template void read(T1 *array, size_t num_elements, uchar const *&compressed_data_pos, size_t &remaining_length) { - if (num_elements * sizeof(T1) > remaining_length) { - throw std::invalid_argument("SZ3: compressed stream is truncated"); - } + if (sizeof(T1) != 0 && num_elements > remaining_length / sizeof(T1)) + throw std::out_of_range("SZ3: attempt to read past the end of the compressed buffer"); memcpy(array, compressed_data_pos, num_elements * sizeof(T1)); if constexpr (SZ3_BIG_ENDIAN) { for (size_t i = 0; i < num_elements; i++) { From da7a29a24ce561feaa1a6a92c2f1c726ce285906 Mon Sep 17 00:00:00 2001 From: Kai Zhao Date: Mon, 31 Aug 2026 23:45:48 -0700 Subject: [PATCH 4/6] Drop PR #132's bound on the internal decompression buffer It bounds that buffer by SZ_compress_size_bound, which is the size of the *output* buffer. compress() sizes the internal one as max(1000, 2 * (decomposition.size_est() + encoder.size_est() + sizeof(Q) * bins)), so for a 64-bit bin type it is far larger and valid streams are rejected. Every module in this tree emits int bins, so the bound happens to hold and these tests cannot reach it -- the fz branch has three modules that do (BitplaneEncoder, BitTruncationQuantizer, FixedPointQuantizer) and all three failed. A corrupted declared size is still caught, after the allocation, by the zstd frame check and the size comparison this PR adds to Lossless_zstd::decompress. Co-Authored-By: Claude Opus 5 --- include/SZ3/compressor/SZGenericCompressor.hpp | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/include/SZ3/compressor/SZGenericCompressor.hpp b/include/SZ3/compressor/SZGenericCompressor.hpp index 75cc67ce..f3d8bf64 100644 --- a/include/SZ3/compressor/SZGenericCompressor.hpp +++ b/include/SZ3/compressor/SZGenericCompressor.hpp @@ -14,7 +14,6 @@ #include "SZ3/utils/Config.hpp" #include "SZ3/utils/FileUtil.hpp" #include "SZ3/utils/Timer.hpp" -#include "zstd.h" namespace SZ3 { @@ -82,15 +81,14 @@ class SZGenericCompressor : public concepts::CompressorInterface { T *decompress(const Config &conf, uchar const *cmpData, size_t cmpSize, T *decData) override { uchar *buffer = nullptr; - // The lossless layer reads the size of this internal buffer from the untrusted payload and would - // otherwise allocate it unbounded. Bound it by the largest internal buffer this configuration could - // have produced: during compression the buffer is losslessly (zstd) compressed, and - // ZSTD_compressBound(B) >= B, so a block that was actually stored with this compressor satisfies - // B < SZ_compress_size_bound = 4096 + conf.size_est() + ZSTD_compressBound(conf.num * sizeof(T)). - // Passing this as the capacity lets the lossless decoder reject a corrupted payload that declares a - // larger internal size before allocating it. conf.num is validated against the trusted output size by - // the caller, so this bound can not be inflated by corrupted input. - size_t bufferSize = 4096 + conf.size_est() + ZSTD_compressBound(conf.num * sizeof(T)); + // 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 From b7cf0e785727993851cc84fea6dff5b78b74001a Mon Sep 17 00:00:00 2001 From: Kai Zhao Date: Tue, 1 Sep 2026 00:01:54 -0700 Subject: [PATCH 5/6] Validate config contents only for a compressed-stream trailer PR #132 rejects N outside [1, 4] and dimensions whose product differs from the element count. Both hold for a config read from the end of a compressed stream, and neither holds for the HDF5 filter's cd_values: those carry placeholder zeros that set_local fills in later, as cdvalueHelper.py says in as many words. The check killed the whole filter -- 80 of cesm-atm's 160 integration cases aborted with "invalid number of dimensions", across every algorithm. The single-argument overload, which is what the filter and the OpenMP path use, now says so explicitly and skips the content checks; the bounded overload used by SZ_decompress keeps them. Verified against the real cesm-atm field through the HDF5 filter: five algorithms pass, and removing the guard reproduces the abort. Co-Authored-By: Claude Opus 5 --- include/SZ3/utils/Config.hpp | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/include/SZ3/utils/Config.hpp b/include/SZ3/utils/Config.hpp index 25a7f342..fc97a1d5 100644 --- a/include/SZ3/utils/Config.hpp +++ b/include/SZ3/utils/Config.hpp @@ -360,12 +360,16 @@ class Config { * * @param c Pointer to the byte array. */ - // Legacy overload for trusted internal callers (e.g. the OpenMP path). - void load(const unsigned char*& c) { load(c, std::numeric_limits::max()); } + /// Overload for a config that is not a compressed-stream trailer: the OpenMP path, and the HDF5 + /// filter's `cd_values`, which carries placeholder dimensions that `set_local` fills in later + /// (see tools/H5Z-SZ3/test/cdvalueHelper.py). Neither bounds the read nor validates the contents. + void load(const unsigned char*& c) { load(c, std::numeric_limits::max(), false); } // `cmpSize` bounds how many bytes may be read from `c` (the config blob). Used when loading // from untrusted compressed data so a corrupted config cannot read out of bounds. - void load(const unsigned char*& c, size_t cmpSize) { + /// @param validate Reject dimensions that no compressed stream can legitimately carry. Only a + /// stream trailer is validated; see the single-argument overload. + void load(const unsigned char*& c, size_t cmpSize, bool validate = true) { const unsigned char* const c0 = c; auto require = [&](size_t n) { if (cmpSize - static_cast(c - c0) < n) @@ -381,7 +385,7 @@ class Config { require(sizeof(N)); read(N, c); - if (N < 1 || N > 4) throw std::out_of_range("SZ3 Config::load: invalid number of dimensions"); + if (validate && (N < 1 || N > 4)) throw std::out_of_range("SZ3 Config::load: invalid number of dimensions"); uint8_t bitWidth; require(sizeof(bitWidth)); read(bitWidth, c); @@ -393,7 +397,7 @@ class Config { /// The element count must equal the product of the dimensions, otherwise the predictor would /// iterate over more grid positions than were allocated for the decompressed data. Validate /// the product with an overflow check. - { + if (validate) { size_t dims_product = 1; bool dims_ok = true; for (size_t dim : dims) { From 5c976f9361ddb3fa03c45089ecf198634f3f678f Mon Sep 17 00:00:00 2001 From: Kai Zhao Date: Tue, 1 Sep 2026 07:00:15 -0700 Subject: [PATCH 6/6] Cover the HDF5 filter's compressed-buffer sizing The filter has to size its output buffer from SZ_compress_size_bound rather than from the raw chunk size, and nothing in the suite exercised that: a chunk small enough for the compressed block to exceed it never appeared. This adds a third chunk mode that asks for 8-element chunks, which the pre-fix filter rejects with "buffer not large enough" on every algorithm. The mode is restricted to a leading slice of at most 4096 chunks. Over a whole field it would mean tens of millions of chunks, and the HDF5 chunk index alone takes the process past a CI runner's memory. Co-Authored-By: Claude Opus 5 --- tools/test/integration/test_h5_filter.py | 33 +++++++++++++++++++----- 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/tools/test/integration/test_h5_filter.py b/tools/test/integration/test_h5_filter.py index 6159d32e..38f2c1d7 100644 --- a/tools/test/integration/test_h5_filter.py +++ b/tools/test/integration/test_h5_filter.py @@ -171,30 +171,49 @@ def main(): compression, compression_opts = get_compression_args(cmpr_algo, bound) all_pass = True - for chunk in [False, True]: + # 'small' keeps a chunk well under the compressed-buffer bound, which the filter has to + # size from SZ_compress_size_bound rather than from the raw chunk size. A chunk that small + # over a whole field would mean tens of millions of chunks, whose HDF5 index alone exhausts + # memory, so the mode runs on a leading slice that still covers many chunks. + max_small_chunks = 4096 + small_chunk = tuple(min(d, 8) for d in shape) + for chunk in ['full', 'auto', 'small']: print(f"Testing {raw_file} with algo = {cmpr_algo} AbsErrorBound = {bound} Chunk = {chunk}") compressed_h5 = os.path.join(output_dir, f"{base_name}_compressed_{chunk}.h5") decompressed_h5 = os.path.join(output_dir, f"{base_name}_decompressed_{chunk}.h5") - if chunk: + payload, reference_h5 = data, original_h5 + if chunk == 'small': + chunks_per_row = 1 + for axis in range(1, len(shape)): + chunks_per_row *= -(-shape[axis] // small_chunk[axis]) + rows = max(1, max_small_chunks // chunks_per_row) * small_chunk[0] + if rows < shape[0]: + payload = data[:rows] + reference_h5 = os.path.join(output_dir, f"{base_name}_original_{chunk}.h5") + write_hdf5(payload, reference_h5, h5_dataset_name) + print(f" restricted to the leading {rows} of {shape[0]} to bound the chunk count") + + if chunk == 'auto': # hd5py will automatically determine chunk sizes if chunks is not set - write_hdf5(data, compressed_h5, h5_dataset_name, compression=compression, compression_opts=compression_opts) + write_hdf5(payload, compressed_h5, h5_dataset_name, compression=compression, + compression_opts=compression_opts) else: - write_hdf5(data, compressed_h5, h5_dataset_name, compression=compression, compression_opts=compression_opts, - chunks=shape) + write_hdf5(payload, compressed_h5, h5_dataset_name, compression=compression, + compression_opts=compression_opts, chunks=payload.shape if chunk == 'full' else small_chunk) with h5py.File(compressed_h5, 'r') as f_in, h5py.File(decompressed_h5, 'w') as f_out: f_out.create_dataset(h5_dataset_name, data=f_in[h5_dataset_name][:]) - max_error = compare_hdf5(original_h5, decompressed_h5, h5_dataset_name) + max_error = compare_hdf5(reference_h5, decompressed_h5, h5_dataset_name) if max_error <= (bound * 3 if cmpr_algo in ['ALGO_BIOMDXTC'] else bound * 1.2): result = "PASS" else: result = "FAIL" - print(f"Test Result for AbsErrorBound = {bound} ChunkSize = {chunk}: {result}") + print(f"Test Result for AbsErrorBound = {bound} Chunk = {chunk}: {result}") if result == "FAIL": all_pass = False