diff --git a/.github/workflows/integration_test.yml b/.github/workflows/integration_test.yml index 222911c0..4e97c455 100644 --- a/.github/workflows/integration_test.yml +++ b/.github/workflows/integration_test.yml @@ -10,7 +10,25 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - dataset: [exaalt-small, hurricane, miranda, cesm-atm, scale-letkf, hacc, exaalt-copper, exaalt-helium] + # A dataset entry may name the fields to run, so one too slow to be a single job is spread + # over several. hacc took 2h09 as one job, scale-letkf 1h00 and exaalt-helium 46m; the rest + # were already under half an hour. Split to keep every job near 35 minutes. + dataset: + - exaalt-small + - cesm-atm + - miranda + - hurricane + - exaalt-copper + - "hacc:vx.f32" + - "hacc:vz.f32" + - "hacc:vy.f32" + - "hacc:zz.f32" + - "hacc:xx.f32" + - "hacc:yy.f32" + - "exaalt-helium:dataset1-7852x1037.x.f32.dat,dataset1-7852x1037.y.f32.dat,dataset1-7852x1037.z.f32.dat" + - "exaalt-helium:dataset2-2338x106711.x.f32.dat,dataset2-2338x106711.y.f32.dat,dataset2-2338x106711.z.f32.dat" + - "scale-letkf:T-98x1200x1200.f32,PRES-98x1200x1200.f32,U-98x1200x1200.f32,V-98x1200x1200.f32,QC-98x1200x1200.f32,QS-98x1200x1200.f32" + - "scale-letkf:QG-98x1200x1200.f32,QV-98x1200x1200.f32,W-98x1200x1200.f32,QI-98x1200x1200.f32,RH-98x1200x1200.f32,QR-98x1200x1200.f32" fail-fast: false name: Test ${{ matrix.dataset }} diff --git a/.gitignore b/.gitignore index 0569725c..5a8e409e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,8 @@ test.dat compressed.out perf.data* -build +# Anchored to the repo root, and glob so build-release and friends are covered too +/build* debug cmake-build-debug cmake-build-release @@ -9,6 +10,14 @@ cmake-build-release install .DS_Store .vscode +.cache +.claude +tmp +docs/site/html sz3_install sz3_build -test \ No newline at end of file +/test + +# Python bytecode +__pycache__/ +*.pyc diff --git a/CMakeLists.txt b/CMakeLists.txt index 7a9cc03e..695a1e09 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,5 +1,5 @@ cmake_minimum_required(VERSION 3.18) -project(SZ3 VERSION 3.3.2) +project(SZ3 VERSION 3.3.3) #data version defines the version of the compressed data format #it is not always equal to the program version (e.g., SZ3 v3.1.0 and SZ3 v.3.1.1 may use the same data version of v.3.1.0) diff --git a/README.md b/README.md index 0b93bd35..53cb0b21 100644 --- a/README.md +++ b/README.md @@ -79,6 +79,7 @@ Version New features * SZ 3.3.0 Add key QoZ v1 and v2 features to improve compression speed and data quality. The full QoZ is available from **a separate branch** (https://github.com/szcompressor/SZ3/tree/QoZ). * SZ 3.3.1: SZ3 Windows support for both Visual Studio and MinGW toolchains. pySZ v1 released and available via `pip install pysz`. Bio algorithms added. * SZ 3.3.2: bugfix for compressed format. +* SZ 3.3.3: decompression is bounds-checked against corrupted input, and several bugs reachable on valid data are fixed. The compressed format is unchanged. ## 3rd party libraries/tools * [Zstandard](https://facebook.github.io/zstd/) v1.4.5 will be fetched if libzstd can not be found by pkg-config. 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..7c0121a3 100644 --- a/include/SZ3/api/impl/SZDispatcher.hpp +++ b/include/SZ3/api/impl/SZDispatcher.hpp @@ -1,10 +1,13 @@ #ifndef SZ3_IMPL_SZDISPATCHER_HPP #define SZ3_IMPL_SZDISPATCHER_HPP +#include +#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" @@ -62,15 +65,14 @@ size_t SZ_compress_dispatcher(Config &conf, const T *data, uchar *cmpData, size_ if (conf.num * sizeof(T) / 1.0 / cmpSize < 3) { auto zstd = Lossless_zstd(); auto zstdCmpCap = ZSTD_compressBound(conf.num * sizeof(T)) + sizeof(size_t); - auto zstdCmpData = static_cast(malloc(zstdCmpCap)); - size_t zstdCmpSize = - zstd.compress(reinterpret_cast(data), conf.num * sizeof(T), zstdCmpData, zstdCmpCap); + std::unique_ptr zstdCmpData(new uchar[zstdCmpCap]); + size_t zstdCmpSize = zstd.compress(reinterpret_cast(data), conf.num * sizeof(T), + zstdCmpData.get(), zstdCmpCap); if (zstdCmpSize < cmpSize && zstdCmpSize <= cmpCap) { conf.cmprAlgo = ALGO_LOSSLESS; - memcpy(cmpData, zstdCmpData, zstdCmpSize); + memcpy(cmpData, zstdCmpData.get(), zstdCmpSize); cmpSize = zstdCmpSize; } - free(zstdCmpData); } return cmpSize; } @@ -79,9 +81,8 @@ template void SZ_decompress_dispatcher(Config &conf, const uchar *cmpData, size_t cmpSize, T *decData) { if (conf.cmprAlgo == ALGO_LOSSLESS) { auto zstd = Lossless_zstd(); - size_t decDataSize = 0; auto decDataPos = reinterpret_cast(decData); - zstd.decompress(cmpData, cmpSize, decDataPos, decDataSize); + size_t decDataSize = zstd.decompress(cmpData, cmpSize, decDataPos, conf.num * sizeof(T)); if (decDataSize != conf.num * sizeof(T)) { throw std::runtime_error("Decompressed data size does not match the original data size"); } diff --git a/include/SZ3/api/impl/SZImplOMP.hpp b/include/SZ3/api/impl/SZImplOMP.hpp index 664a14dc..2f1d563e 100644 --- a/include/SZ3/api/impl/SZImplOMP.hpp +++ b/include/SZ3/api/impl/SZImplOMP.hpp @@ -2,13 +2,16 @@ #define SZ3_IMPL_SZDISPATCHER_OMP_HPP #include +#include #include +#include #include "SZ3/api/impl/SZDispatcher.hpp" #ifdef _OPENMP #include +#include #endif namespace SZ3 { @@ -70,8 +73,10 @@ 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)); + // Room for the size header Lossless_zstd::compress writes ahead of the zstd stream. + size_t cmp_size_cap = sizeof(size_t) + ZSTD_compressBound(conf_t[tid].num * sizeof(T)); + std::unique_ptr compressed_owner(new uchar[cmp_size_cap]); + 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 +110,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 +125,21 @@ 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; + 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 thread contributes at least a config and a size, so the count cannot exceed the buffer size. + 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); + size_t confRemaining = static_cast(cmp_end - cmpr_data_pos); + conf_t[i].load(cmpr_data_pos, confRemaining); } if (conf_t[0].sz3MagicNumber != SZ3_MAGIC_NUMBER) { @@ -145,12 +156,18 @@ 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 payloads follow back-to-back; build the offsets without overflowing so no slice points out. + 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 +216,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..b4a927ba 100644 --- a/include/SZ3/api/sz.hpp +++ b/include/SZ3/api/sz.hpp @@ -21,10 +21,12 @@ #ifndef SZ3_SZ_HPP #define SZ3_SZ_HPP +#include +#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 +97,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 +121,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 +144,13 @@ 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); + size_t confRemaining = cmpSize - 16 - cmpDataSize; + config.load(cmpConfPos, confRemaining); if (decData == nullptr) { decData = new T[config.num]; diff --git a/include/SZ3/compressor/SZGenericCompressor.hpp b/include/SZ3/compressor/SZGenericCompressor.hpp index 4a3f471a..650ebc1b 100644 --- a/include/SZ3/compressor/SZGenericCompressor.hpp +++ b/include/SZ3/compressor/SZGenericCompressor.hpp @@ -1,7 +1,11 @@ #ifndef SZ3_COMPRESSOR_TYPE_ONE_HPP #define SZ3_COMPRESSOR_TYPE_ONE_HPP +#include #include +#include +#include +#include #include "SZ3/compressor/Compressor.hpp" #include "SZ3/decomposition/Decomposition.hpp" @@ -13,6 +17,7 @@ #include "SZ3/utils/Timer.hpp" namespace SZ3 { + /** * SZGenericCompressor glues together decomposition, encoder, and lossless modules to form the compressor. * It only takes Decomposition, not Predictor. @@ -45,27 +50,30 @@ class SZGenericCompressor : public concepts::CompressorInterface { size_t bufferSize = std::max( 1000, 2 * (decomposition.size_est() + encoder.size_est() + sizeof(T) * quant_inds.size())); - auto buffer = static_cast(malloc(bufferSize)); + // Owned, because the encoder and the lossless layer below can throw. + std::unique_ptr buffer_owner(new uchar[bufferSize]); + uchar *const buffer = buffer_owner.get(); 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; - lossless.decompress(cmpData, cmpSize, buffer, bufferSize); + size_t bufferSize = lossless.decompress(cmpData, cmpSize, buffer, 0); + + // malloc'd by the lossless layer, hence free(). Owned, because the parsing below can throw. + std::unique_ptr buffer_owner(buffer, &free); uchar const *bufferPos = buffer; @@ -73,11 +81,12 @@ class SZGenericCompressor : public concepts::CompressorInterface { encoder.load(bufferPos, bufferSize); size_t quant_inds_size = 0; - read(quant_inds_size, bufferPos); - auto quant_inds = encoder.decode(bufferPos, quant_inds_size); + read(quant_inds_size, bufferPos, bufferSize); + auto quant_inds = encoder.decode(bufferPos, quant_inds_size, bufferSize); 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..529a12d8 100644 --- a/include/SZ3/compressor/specialized/SZExaaltCompressor.hpp +++ b/include/SZ3/compressor/specialized/SZExaaltCompressor.hpp @@ -1,6 +1,12 @@ #ifndef SZ3_EXAALT_COMPRESSSOR_HPP #define SZ3_EXAALT_COMPRESSSOR_HPP +#include +#include +#include +#include + +#include "SZ3/compressor/Compressor.hpp" #include "SZ3/def.hpp" #include "SZ3/encoder/Encoder.hpp" #include "SZ3/lossless/Lossless.hpp" @@ -119,19 +125,19 @@ class SZExaaltCompressor : public SZ3::concepts::CompressorInterface { // T *decompress(uchar const *lossless_compressed_data, const size_t length) { T *decompress(const Config &conf, uchar const *cmpData, size_t cmpSize, T *dec_data) override { uchar *buffer = nullptr; - size_t bufferSize = 0; - lossless.decompress(cmpData, cmpSize, buffer, bufferSize); - size_t remaining_length = cmpSize; + size_t bufferSize = lossless.decompress(cmpData, cmpSize, buffer, 0); + // The parsing below walks the decompressed buffer, so bufferSize is its bound, not cmpSize. + size_t remaining_length = bufferSize; uchar const *buffer_pos = buffer; quantizer.load(buffer_pos, remaining_length); encoder.load(buffer_pos, remaining_length); - auto quant_inds = encoder.decode(buffer_pos, conf.num); + auto quant_inds = encoder.decode(buffer_pos, conf.num, remaining_length); encoder.postprocess_decode(); encoder.load(buffer_pos, remaining_length); auto pred_inds_num = (timestep_op == 1) ? conf.dims[1] : conf.num; - auto pred_inds = encoder.decode(buffer_pos, pred_inds_num); + auto pred_inds = encoder.decode(buffer_pos, pred_inds_num, remaining_length); encoder.postprocess_decode(); free(buffer); diff --git a/include/SZ3/compressor/specialized/SZTruncateCompressor.hpp b/include/SZ3/compressor/specialized/SZTruncateCompressor.hpp index f4d3b766..1e935653 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" @@ -27,7 +28,8 @@ 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))); + std::unique_ptr buffer_owner(new uchar[conf.num * sizeof(T)]); + uchar *const buffer = buffer_owner.get(); auto buffer_pos = buffer; // Timer timer(true); @@ -35,7 +37,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; @@ -43,8 +44,7 @@ class SZTruncateCompressor : public concepts::CompressorInterface { T *decompress(const Config &conf, uchar const *cmpData, size_t cmpSize, T *decData) override { uchar *buffer = nullptr; - size_t bufferSize = 0; - lossless.decompress(cmpData, cmpSize, buffer, bufferSize); + size_t bufferSize = lossless.decompress(cmpData, cmpSize, buffer, 0); // size_t remaining_length = bufferCap; uchar const *buffer_pos = buffer; diff --git a/include/SZ3/decomposition/BlockwiseDecomposition.hpp b/include/SZ3/decomposition/BlockwiseDecomposition.hpp index 88a90a0e..9daa6b66 100644 --- a/include/SZ3/decomposition/BlockwiseDecomposition.hpp +++ b/include/SZ3/decomposition/BlockwiseDecomposition.hpp @@ -2,15 +2,18 @@ #define SZ3_BLOCKWISE_DECOMPOSITION_HPP #include +#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 { @@ -46,6 +49,9 @@ class BlockwiseDecomposition : public concepts::DecompositionInterface &quant_inds, T *dec_data) override { + if (quant_inds.size() < conf.num) { + throw std::out_of_range("SZ3 blockwise: fewer bins than the grid consumes"); + } int *quant_inds_pos = &quant_inds[0]; auto data_with_padding = 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..5fd45efa 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,8 +26,22 @@ class InterpolationDecomposition : public concepts::DecompositionInterface &quant_inds, T *dec_data) override { + // load() read original_dimensions from the payload, and the grid walk below is sized by it while + // dec_data and quant_inds are sized by conf. A tampered value runs off both. + 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(); + if (quant_inds.size() < num_elements) { + throw std::out_of_range("SZ3 interpolation: fewer bins than the grid consumes"); + } this->quant_inds = quant_inds.data(); double eb = quantizer.get_eb(); @@ -146,6 +162,11 @@ class InterpolationDecomposition : public concepts::DecompositionInterface(interp_id) >= interpolators.size()) { + throw std::out_of_range("SZ3 interpolation: interpolator id is out of range"); + } num_elements = 1; interp_level = -1; bool use_anchor = false; @@ -190,6 +217,11 @@ class InterpolationDecomposition : public concepts::DecompositionInterface 0 && (anchor_stride & (anchor_stride - 1)) != 0) { + throw std::out_of_range("SZ3 interpolation: anchor stride must be a power of two"); + } if (anchor_stride > 0) { int max_interpolation_level = static_cast(log2(anchor_stride)) + 1; if (max_interpolation_level <= interp_level) { @@ -210,6 +242,9 @@ class InterpolationDecomposition : public concepts::DecompositionInterface(direction_sequence_id) >= dim_sequences.size()) { + throw std::out_of_range("SZ3 interpolation: direction sequence id is out of range"); + } } void build_anchor_grid(T *data) { // store anchor points. steplength: anchor_stride on each dimension diff --git a/include/SZ3/decomposition/NoPredictionDecomposition.hpp b/include/SZ3/decomposition/NoPredictionDecomposition.hpp index 34348dd3..a06b09c0 100644 --- a/include/SZ3/decomposition/NoPredictionDecomposition.hpp +++ b/include/SZ3/decomposition/NoPredictionDecomposition.hpp @@ -1,6 +1,8 @@ #ifndef SZ3_NO_PREDICTION_DECOMPOSITION_HPP #define SZ3_NO_PREDICTION_DECOMPOSITION_HPP +#include + #include "Decomposition.hpp" #include "SZ3/def.hpp" #include "SZ3/quantizer/Quantizer.hpp" @@ -16,6 +18,9 @@ class NoPredictionDecomposition : public concepts::DecompositionInterface &quant_inds, T *dec_data) override { + if (quant_inds.size() < conf.num) { + throw std::out_of_range("SZ3 no-prediction: fewer bins than the data has elements"); + } for (size_t i = 0; i < conf.num; i++) { dec_data[i] = quantizer.recover(0, quant_inds[i]); } @@ -32,6 +37,8 @@ class NoPredictionDecomposition : public concepts::DecompositionInterface #include +#include #include "Decomposition.hpp" #include "SZ3/utils/Config.hpp" @@ -180,8 +182,8 @@ class SZBioMDXtcDecomposition : public concepts::DecompositionInterface +#include +#include + #include "Decomposition.hpp" #include "SZ3/def.hpp" #include "SZ3/predictor/LorenzoPredictor.hpp" @@ -34,6 +38,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 +51,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 +65,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); @@ -73,6 +83,9 @@ class TimeSeriesDecomposition : public concepts::DecompositionInterface& quant_inds, T* dec_data) override { + if (quant_inds.size() < num_elements) { + throw std::out_of_range("SZ3 time series: fewer bins than the grid consumes"); + } int const* quant_inds_pos = quant_inds.data(); // std::array intra_block_dims; // auto dec_data = new T[num_elements]; diff --git a/include/SZ3/encoder/ArithmeticEncoder.hpp b/include/SZ3/encoder/ArithmeticEncoder.hpp index 5167faea..ea4353e8 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" @@ -505,6 +506,12 @@ class ArithmeticEncoder : public concepts::EncoderInterface { put_codes_to_output(buf, pending_bits + 1, &bytes, &lackBits, &outSize); } bytes += 1; + // decode() opens with an eight-byte read and then keeps its cursor a few bytes ahead of what it + // has consumed, so it always ran past the end of what this wrote. Pad by that much. + for (size_t i = 0; i < 8; i++) { + *bytes++ = 0; + outSize++; + } return outSize; } @@ -517,7 +524,10 @@ class ArithmeticEncoder : public concepts::EncoderInterface { * @param int *out (output) : the result (type array decompressed from the stream 's') * * */ - std::vector decode(const uchar *&bytes, size_t targetLength) override { + std::vector decode(const uchar *&bytes, size_t targetLength, size_t &remaining_length) override { + // The reads below are not individually bounded, so check what they consumed before charging it: + // subtracting more than is left would wrap remaining_length and unbound everything parsed after. + const uchar *decode_start = bytes; std::vector out(targetLength); // void ari_decode(AriCoder *ariCoder, unsigned char *s, size_t s_len, size_t targetLength, int *out) { @@ -527,7 +537,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++) { @@ -581,6 +591,11 @@ class ArithmeticEncoder : public concepts::EncoderInterface { } } bytes += s_counter; + const size_t consumed = static_cast(bytes - decode_start); + if (consumed > remaining_length) { + throw std::out_of_range("SZ3 arithmetic encoder: decode read past the end of the compressed buffer"); + } + remaining_length -= consumed; return out; } diff --git a/include/SZ3/encoder/BypassEncoder.hpp b/include/SZ3/encoder/BypassEncoder.hpp index 9c3ef8e6..6bbe95d5 100644 --- a/include/SZ3/encoder/BypassEncoder.hpp +++ b/include/SZ3/encoder/BypassEncoder.hpp @@ -2,6 +2,7 @@ #define SZ3_BYPASS_ENCODER_HPP #include +#include #include #include "Encoder.hpp" @@ -25,10 +26,14 @@ class BypassEncoder : public concepts::EncoderInterface { void preprocess_decode() override {} - std::vector decode(const uchar *&bytes, size_t targetLength) override { + std::vector decode(const uchar *&bytes, size_t targetLength, size_t &remaining_length) override { + if (targetLength > remaining_length / sizeof(T)) { + throw std::out_of_range("SZ3 bypass encoder: more bins requested than the buffer holds"); + } std::vector bins(targetLength); memcpy(bins.data(), bytes, sizeof(T) * targetLength); bytes += sizeof(T) * targetLength; + remaining_length -= sizeof(T) * targetLength; return bins; } diff --git a/include/SZ3/encoder/Encoder.hpp b/include/SZ3/encoder/Encoder.hpp index 822502dd..077ef3e1 100644 --- a/include/SZ3/encoder/Encoder.hpp +++ b/include/SZ3/encoder/Encoder.hpp @@ -1,12 +1,14 @@ #ifndef SZ3_ENCODER_HPP #define SZ3_ENCODER_HPP +#include #include #include "SZ3/def.hpp" namespace SZ3 { namespace concepts { + /** * Encoder changes the input to a more compact representative * Usually this step is lossless instead of lossy @@ -37,11 +39,16 @@ class EncoderInterface { /** * reverse of encode() + * + * `targetLength` says when to stop producing, `remaining_length` says how far the reads may go; + * an entropy-coded stream has no terminator, so neither number substitutes for the other. + * * @param bytes input in byte stream * @param targetLength size of the output vector + * @param remaining_length bytes readable from `bytes`; decremented by what is consumed * @return output in vector */ - virtual std::vector decode(const uchar *&bytes, size_t targetLength) = 0; + virtual std::vector decode(const uchar *&bytes, size_t targetLength, size_t &remaining_length) = 0; /** * serialize the encoder and store it to a buffer diff --git a/include/SZ3/encoder/HuffmanEncoder.hpp b/include/SZ3/encoder/HuffmanEncoder.hpp index 69645a0b..82e3e7a1 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 @@ -222,21 +224,26 @@ class HuffmanEncoder : public concepts::EncoderInterface { void preprocess_decode() override {} // perform decoding - std::vector decode(const uchar *&bytes, size_t targetLength) override { + std::vector decode(const uchar *&bytes, size_t targetLength, size_t &remaining_length) override { node t = treeRoot; std::vector out(targetLength); size_t i = 0, byteIndex = 0, count = 0; int r; node n = treeRoot; size_t encodedLength = 0; - read(encodedLength, bytes); + read(encodedLength, bytes, remaining_length); 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; } - for (i = 0; count < targetLength; i++) { + if (encodedLength > remaining_length) + throw std::out_of_range("SZ3 Huffman: encoded length exceeds compressed buffer"); + + // Walk at most the bits the stream holds, and stop once targetLength symbols are out. + const size_t maxBits = targetLength > 0 ? encodedLength * 8 : 0; + for (i = 0; i < maxBits; i++) { byteIndex = i >> 3; // i/8 r = i % 8; if (((bytes[byteIndex] >> (7 - r)) & 0x01) == 0) @@ -247,10 +254,12 @@ class HuffmanEncoder : public concepts::EncoderInterface { if (n->t) { out[count] = n->c + offset; n = t; - count++; + if (++count == targetLength) break; } } + if (count < targetLength) throw std::out_of_range("SZ3 Huffman: corrupted encoded stream"); bytes += encodedLength; + remaining_length -= encodedLength; return out; } @@ -260,8 +269,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; + // The stored state count is skipped: it sizes the encode-side code tables, which decoding never + // touches. nodeCount is bounded before it sizes anything, or the encodeStartIndex arithmetic + // below overflows and the tree is 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,9 +286,14 @@ class HuffmanEncoder : public concepts::EncoderInterface { encodeStartIndex = 1 + 2 * nodeCount * sizeof(unsigned int) + nodeCount * sizeof(unsigned char) + nodeCount * sizeof(T); - huffmanTree = createHuffmanTree(stateNum); + 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 pool is 4x what is asked for, and unpad_tree builds each of the nodeCount nodes once. + huffmanTree = createHuffmanTree(nodeCount); 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; } @@ -290,14 +309,10 @@ class HuffmanEncoder : public concepts::EncoderInterface { node reconstruct_HuffTree_from_bytes_anyStates(const unsigned char *bytes, uint nodeCount) { if (nodeCount <= 256) { - unsigned char *L = static_cast(malloc(nodeCount * sizeof(unsigned char))); - memset(L, 0, nodeCount * sizeof(unsigned char)); - unsigned char *R = static_cast(malloc(nodeCount * sizeof(unsigned char))); - memset(R, 0, nodeCount * sizeof(unsigned char)); - T *C = static_cast(malloc(nodeCount * sizeof(T))); - memset(C, 0, nodeCount * sizeof(T)); - unsigned char *t = static_cast(malloc(nodeCount * sizeof(unsigned char))); - memset(t, 0, nodeCount * sizeof(unsigned char)); + std::vector L(nodeCount); + std::vector R(nodeCount); + std::vector C(nodeCount); + std::vector t(nodeCount); // TODO: Endian type // unsigned char cmpSysEndianType = bytes[0]; // if(cmpSysEndianType!=(unsigned char)sysEndianType) @@ -314,27 +329,21 @@ class HuffmanEncoder : public concepts::EncoderInterface { // break; // } // } - memcpy(L, bytes + 1, nodeCount * sizeof(unsigned char)); - memcpy(R, bytes + 1 + nodeCount * sizeof(unsigned char), nodeCount * sizeof(unsigned char)); - memcpy(C, bytes + 1 + 2 * nodeCount * sizeof(unsigned char), nodeCount * sizeof(T)); - memcpy(t, bytes + 1 + 2 * nodeCount * sizeof(unsigned char) + nodeCount * sizeof(T), + memcpy(L.data(), bytes + 1, nodeCount * sizeof(unsigned char)); + memcpy(R.data(), bytes + 1 + nodeCount * sizeof(unsigned char), nodeCount * sizeof(unsigned char)); + memcpy(C.data(), bytes + 1 + 2 * nodeCount * sizeof(unsigned char), nodeCount * sizeof(T)); + memcpy(t.data(), bytes + 1 + 2 * nodeCount * sizeof(unsigned char) + nodeCount * sizeof(T), nodeCount * sizeof(unsigned char)); + std::vector seen(nodeCount, false); + seen[0] = true; node root = this->new_node2(C[0], t[0]); - this->unpad_tree(L, R, C, t, 0, root); - free(L); - free(R); - free(C); - free(t); + this->unpad_tree(L.data(), R.data(), C.data(), t.data(), 0, root, nodeCount, seen); return root; } else if (nodeCount <= 65536) { - unsigned short *L = static_cast(malloc(nodeCount * sizeof(unsigned short))); - memset(L, 0, nodeCount * sizeof(unsigned short)); - unsigned short *R = static_cast(malloc(nodeCount * sizeof(unsigned short))); - memset(R, 0, nodeCount * sizeof(unsigned short)); - T *C = static_cast(malloc(nodeCount * sizeof(T))); - memset(C, 0, nodeCount * sizeof(T)); - unsigned char *t = static_cast(malloc(nodeCount * sizeof(unsigned char))); - memset(t, 0, nodeCount * sizeof(unsigned char)); + std::vector L(nodeCount); + std::vector R(nodeCount); + std::vector C(nodeCount); + std::vector t(nodeCount); // TODO: Endian type // unsigned char cmpSysEndianType = bytes[0]; @@ -353,30 +362,24 @@ class HuffmanEncoder : public concepts::EncoderInterface { // } // } - memcpy(L, bytes + 1, nodeCount * sizeof(unsigned short)); - memcpy(R, bytes + 1 + nodeCount * sizeof(unsigned short), nodeCount * sizeof(unsigned short)); - memcpy(C, bytes + 1 + 2 * nodeCount * sizeof(unsigned short), nodeCount * sizeof(T)); + memcpy(L.data(), bytes + 1, nodeCount * sizeof(unsigned short)); + memcpy(R.data(), bytes + 1 + nodeCount * sizeof(unsigned short), nodeCount * sizeof(unsigned short)); + memcpy(C.data(), bytes + 1 + 2 * nodeCount * sizeof(unsigned short), nodeCount * sizeof(T)); - memcpy(t, bytes + 1 + 2 * nodeCount * sizeof(unsigned short) + nodeCount * sizeof(T), + memcpy(t.data(), bytes + 1 + 2 * nodeCount * sizeof(unsigned short) + nodeCount * sizeof(T), nodeCount * sizeof(unsigned char)); + std::vector seen(nodeCount, false); + seen[0] = true; node root = this->new_node2(0, 0); - this->unpad_tree(L, R, C, t, 0, root); - free(L); - free(R); - free(C); - free(t); + this->unpad_tree(L.data(), R.data(), C.data(), t.data(), 0, root, nodeCount, seen); return root; } else // nodeCount>65536 { - unsigned int *L = static_cast(malloc(nodeCount * sizeof(unsigned int))); - memset(L, 0, nodeCount * sizeof(unsigned int)); - unsigned int *R = static_cast(malloc(nodeCount * sizeof(unsigned int))); - memset(R, 0, nodeCount * sizeof(unsigned int)); - T *C = static_cast(malloc(nodeCount * sizeof(T))); - memset(C, 0, nodeCount * sizeof(T)); - unsigned char *t = static_cast(malloc(nodeCount * sizeof(unsigned char))); - memset(t, 0, nodeCount * sizeof(unsigned char)); + std::vector L(nodeCount); + std::vector R(nodeCount); + std::vector C(nodeCount); + std::vector t(nodeCount); // TODO: Endian type // unsigned char cmpSysEndianType = bytes[0]; // if(cmpSysEndianType!=(unsigned char)sysEndianType) @@ -394,19 +397,17 @@ class HuffmanEncoder : public concepts::EncoderInterface { // } // } - memcpy(L, bytes + 1, nodeCount * sizeof(unsigned int)); - memcpy(R, bytes + 1 + nodeCount * sizeof(unsigned int), nodeCount * sizeof(unsigned int)); - memcpy(C, bytes + 1 + 2 * nodeCount * sizeof(unsigned int), nodeCount * sizeof(T)); + memcpy(L.data(), bytes + 1, nodeCount * sizeof(unsigned int)); + memcpy(R.data(), bytes + 1 + nodeCount * sizeof(unsigned int), nodeCount * sizeof(unsigned int)); + memcpy(C.data(), bytes + 1 + 2 * nodeCount * sizeof(unsigned int), nodeCount * sizeof(T)); - memcpy(t, bytes + 1 + 2 * nodeCount * sizeof(unsigned int) + nodeCount * sizeof(T), + memcpy(t.data(), bytes + 1 + 2 * nodeCount * sizeof(unsigned int) + nodeCount * sizeof(T), nodeCount * sizeof(unsigned char)); + std::vector seen(nodeCount, false); + seen[0] = true; node root = this->new_node2(0, 0); - this->unpad_tree(L, R, C, t, 0, root); - free(L); - free(R); - free(C); - free(t); + this->unpad_tree(L.data(), R.data(), C.data(), t.data(), 0, root, nodeCount, seen); return root; } } @@ -479,11 +480,13 @@ 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 single-symbol tree gives the root a zero-length code, and shifting by 64 is undefined. + (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); + // len >= 128 would shift by >= 64, and such a code does 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 +536,11 @@ class HuffmanEncoder : public concepts::EncoderInterface { } } + // The state table is sized by the bin range rather than 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 +587,35 @@ 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, + std::vector &seen) { // root->c = C[i]; if (root->t == 0) { T1 l, r; l = L[i]; if (l != 0) { + // pad_tree gives a child a higher index than its parent, so a valid index satisfies i < l < nodeCount. + // Enforcing it keeps L/R/C/t reads inside the pool and rules out a cycle. + if (l <= i || l >= nodeCount) throw std::out_of_range("SZ3 Huffman: invalid left child index in tree"); + // Increasing indices rule out a cycle but not two parents naming one child, which would + // expand the tree exponentially instead of building nodeCount nodes. + if (seen[l]) throw std::out_of_range("SZ3 Huffman: tree node reached twice"); + seen[l] = true; 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, seen); } r = R[i]; if (r != 0) { + if (r <= i || r >= nodeCount) throw std::out_of_range("SZ3 Huffman: invalid right child index in tree"); + if (seen[r]) throw std::out_of_range("SZ3 Huffman: tree node reached twice"); + seen[r] = true; 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, seen); + } + if (root->left == nullptr || root->right == nullptr) { + throw std::out_of_range("SZ3 Huffman: internal tree node is missing a child"); } } } diff --git a/include/SZ3/encoder/HuffmanEncoderV2.hpp b/include/SZ3/encoder/HuffmanEncoderV2.hpp index fc041c3d..af0042da 100644 --- a/include/SZ3/encoder/HuffmanEncoderV2.hpp +++ b/include/SZ3/encoder/HuffmanEncoderV2.hpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include "SZ3/def.hpp" @@ -434,12 +435,20 @@ class HuffmanEncoderV2 : public concepts::EncoderInterface { void preprocess_decode() override { } - std::vector decode(const uchar*& bytes, size_t targetLength) override { + std::vector decode(const uchar*& bytes, size_t targetLength, size_t& remaining_length) override { + // The reads below are not individually bounded, so check what they consumed before charging it: + // subtracting more than is left would wrap remaining_length and unbound everything parsed after. + const uchar* decode_start = bytes; if (tree.maxval == 1) { size_t len = bytesToInt64_bigEndian(bytes) ^ 0x1234abcd; bytes += 8; // assert(len==targetLength); + const size_t consumed = static_cast(bytes - decode_start); + if (consumed > remaining_length) { + throw std::out_of_range("SZ3 HuffmanEncoderV2: decode read past the end of the compressed buffer"); + } + remaining_length -= consumed; return std::vector(len, tree.offset); } @@ -516,6 +525,11 @@ class HuffmanEncoderV2 : public concepts::EncoderInterface { bytes += (len + 7) >> 3; + const size_t consumed = static_cast(bytes - decode_start); + if (consumed > remaining_length) { + throw std::out_of_range("SZ3 HuffmanEncoderV2: decode read past the end of the compressed buffer"); + } + remaining_length -= consumed; return out; } @@ -669,6 +683,11 @@ class HuffmanEncoderV2 : public concepts::EncoderInterface { // timer.stop("decode"); + const size_t consumed = static_cast(bytes - decode_start); + if (consumed > remaining_length) { + throw std::out_of_range("SZ3 HuffmanEncoderV2: decode read past the end of the compressed buffer"); + } + remaining_length -= consumed; return out; } @@ -1037,6 +1056,12 @@ class HuffmanEncoderV2 : public concepts::EncoderInterface { void loadAsDFSOrder(const uchar*& bytes, size_t& remaining_length) { tree.init(); + // The tree is a fixed-size header plus a DFS bitstream, all of it untrusted; every read below is + // bounded and every allocation is sized from a checked field. + 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 +1073,32 @@ 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); + + // Each node costs at least one bit of the DFS stream, so the count can not exceed the bits available. + // tree.n is an int holding a 64-bit read, so check the sign before comparing. + 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 sizes the dense tables below. preprocess_encode leaves usemp == 0 only under 1 << 28, so a + // larger value here is inconsistent and would 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); } } + remaining_length -= header_size; + if (tree.n == 0) { tree.setConstructed(); return; @@ -1090,6 +1129,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 +1141,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 +1164,8 @@ class HuffmanEncoderV2 : public concepts::EncoderInterface { bytes += (i + 7) >> 3; + 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..fa890488 100644 --- a/include/SZ3/encoder/RunlengthEncoder.hpp +++ b/include/SZ3/encoder/RunlengthEncoder.hpp @@ -1,6 +1,7 @@ #ifndef SZ3_RUNLENGTH_ENCODER_HPP #define SZ3_RUNLENGTH_ENCODER_HPP +#include #include #include "Encoder.hpp" @@ -12,7 +13,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; @@ -37,13 +40,16 @@ class RunlengthEncoder : public concepts::EncoderInterface { void preprocess_decode() override {} - std::vector decode(const uchar *&bytes, size_t targetLength) override { + std::vector decode(const uchar *&bytes, size_t targetLength, size_t &remaining_length) override { std::vector bins(targetLength, 0); T value; int cnt; for (size_t i = 0; i < bins.size();) { - read(value, bytes); - read(cnt, bytes); + read(value, bytes, remaining_length); + read(cnt, bytes, remaining_length); + if (cnt < 0) { + throw std::out_of_range("SZ3 runlength encoder: negative run length"); + } if (i + cnt > bins.size()) { throw std::runtime_error("Decoded length exceeds targetLength"); } @@ -60,6 +66,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..8fecde6c 100644 --- a/include/SZ3/encoder/XtcBasedEncoder.hpp +++ b/include/SZ3/encoder/XtcBasedEncoder.hpp @@ -9,10 +9,14 @@ #define _SZ_XTC3_ENCODER_HPP #include +#include +#include +#include #include #include "SZ3/def.hpp" #include "SZ3/encoder/Encoder.hpp" +#include "SZ3/utils/Config.hpp" // #define DEBUG_OUTPUT @@ -313,12 +317,17 @@ class XtcBasedEncoder : public concepts::EncoderInterface { size_t bufferSize = size3 * 1.2; struct DataBuffer buffer; - int *intBufferPoiner = reinterpret_cast(malloc(size3 * sizeof(*intBufferPoiner))); - buffer.data = reinterpret_cast(malloc(bufferSize * sizeof(int))); - if (buffer.data == nullptr) { - fprintf(stderr, "malloc failed\n"); - exit(1); + std::unique_ptr index_owner( + static_cast(malloc(size3 * sizeof(int))), &free); + // Zeroed: the bit packing below leaves untouched the bits it does not set, and buffer.index bytes + // of this go into the compressed output. + std::unique_ptr buffer_owner( + static_cast(calloc(bufferSize, sizeof(int))), &free); + if (index_owner == nullptr || buffer_owner == nullptr) { + throw std::runtime_error("SZ3 Xtc: can not allocate the compression buffer"); } + int *intBufferPoiner = index_owner.get(); + buffer.data = buffer_owner.get(); buffer.index = 0; buffer.lastbits = 0; buffer.lastbyte = 0; @@ -431,14 +440,17 @@ class XtcBasedEncoder : public concepts::EncoderInterface { } *intOutputPtr++ = smallIdx; - int maxIdx = std::min(LASTIDX, smallIdx + CHAR_BIT); + // LASTIDX is one past the last entry, and the loop above stops there when no entry reaches minDiff, + // which is every input with fewer than two triplets. The decoder clamps the same way. + const int smallLookup = std::min(smallIdx, LASTIDX - 1); + int maxIdx = std::min(LASTIDX - 1, smallIdx + CHAR_BIT); int minIdx = maxIdx - CHAR_BIT; /* often this equal smallIdx */ - int smaller = magicInts[std::max(FIRSTIDX, smallIdx - 1)] / 2; - int smallNum = magicInts[smallIdx] / 2; + int smaller = magicInts[std::max(FIRSTIDX, smallLookup - 1)] / 2; + int smallNum = magicInts[smallLookup] / 2; unsigned int sizeSmall[3]; - sizeSmall[0] = magicInts[smallIdx]; - sizeSmall[1] = magicInts[smallIdx]; - sizeSmall[2] = magicInts[smallIdx]; + sizeSmall[0] = magicInts[smallLookup]; + sizeSmall[1] = magicInts[smallLookup]; + sizeSmall[2] = magicInts[smallLookup]; int larger = magicInts[maxIdx] / 2; size_t i = 0; unsigned int *localUnsignedIntBufferPointer = reinterpret_cast(intBufferPoiner); @@ -563,8 +575,6 @@ class XtcBasedEncoder : public concepts::EncoderInterface { remain -= batchSize; } while (remain > 0); - free(buffer.data); - free(intBufferPoiner); size_t outputSize = charOutputPtr - bytes; @@ -586,7 +596,10 @@ class XtcBasedEncoder : public concepts::EncoderInterface { * this routine decompresses a large number of compressed 3d coordinates. * */ - std::vector decode(const unsigned char *&bytes, size_t targetLength) override { + std::vector decode(const unsigned char *&bytes, size_t targetLength, size_t &remaining_length) override { + // The reads below are not individually bounded, so check what they consumed before charging it: + // subtracting more than is left would wrap remaining_length and unbound everything parsed after. + const unsigned char *decode_start = bytes; #ifdef DEBUG_OUTPUT printf("\nDecoding, targetLength: %ld\n", targetLength); #endif @@ -597,10 +610,8 @@ 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"); - } + // Allocated below, once size3 is known. + buffer.data = nullptr; buffer.index = 0; buffer.lastbits = 0; buffer.lastbyte = 0; @@ -637,20 +648,32 @@ class XtcBasedEncoder : public concepts::EncoderInterface { } int smallIdx = *inputIntPtr++; + // The encoder writes LASTIDX when no table entry reaches minDiff, and clamps its own lookups. + if (smallIdx < 0 || smallIdx > LASTIDX) throw std::out_of_range("SZ3 Xtc: small index out of range"); + const int smallLookup = std::min(smallIdx, LASTIDX - 1); - int smaller = magicInts[std::max(FIRSTIDX, smallIdx - 1)] / 2; - int smallNum = magicInts[smallIdx] / 2; + int smaller = magicInts[std::max(FIRSTIDX, smallLookup - 1)] / 2; + int smallNum = magicInts[smallLookup] / 2; unsigned int sizeSmall[3]; - sizeSmall[0] = magicInts[smallIdx]; - sizeSmall[1] = magicInts[smallIdx]; - sizeSmall[2] = magicInts[smallIdx]; + sizeSmall[0] = magicInts[smallLookup]; + sizeSmall[1] = magicInts[smallLookup]; + sizeSmall[2] = magicInts[smallLookup]; size_t size3 = targetLength; bufferSize = size3 * 1.2; - buffer.data = reinterpret_cast(malloc(bufferSize * sizeof(int))); + std::unique_ptr buffer_owner( + static_cast(malloc(bufferSize * sizeof(int))), &free); + buffer.data = buffer_owner.get(); + 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 the byte count the memcpy loop below copies into buffer.data. + 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); @@ -669,7 +692,12 @@ class XtcBasedEncoder : public concepts::EncoderInterface { int run = 0; size_t i = 0; - int *intBufferPoiner = reinterpret_cast(malloc(size3 * sizeof(*intBufferPoiner))); + std::unique_ptr index_owner( + static_cast(malloc(size3 * sizeof(int))), &free); + if (index_owner == nullptr) { + throw std::runtime_error("SZ3 Xtc: can not allocate the index buffer"); + } + int *intBufferPoiner = index_owner.get(); int *localIntBufferPointer = intBufferPoiner; unsigned char *charOutputPtr = reinterpret_cast(quantData.data()); int *intOutputPtr = reinterpret_cast(charOutputPtr); @@ -756,8 +784,6 @@ class XtcBasedEncoder : public concepts::EncoderInterface { } sizeSmall[0] = sizeSmall[1] = sizeSmall[2] = magicInts[smallIdx]; } - free(buffer.data); - free(intBufferPoiner); #ifdef DEBUG_OUTPUT printf("Decoded %llu triplets.\n", numTriplets); @@ -773,6 +799,11 @@ class XtcBasedEncoder : public concepts::EncoderInterface { quantData[quantData.size() - 1] = reminder1; quantData[quantData.size() - 2] = reminder2; } + const size_t consumed = static_cast(bytes - decode_start); + if (consumed > remaining_length) { + throw std::out_of_range("SZ3 Xtc: decode read past the end of the compressed buffer"); + } + remaining_length -= consumed; return quantData; } diff --git a/include/SZ3/lossless/Lossless.hpp b/include/SZ3/lossless/Lossless.hpp index acf29630..5bb77d3e 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 { /** @@ -28,12 +32,13 @@ class LosslessInterface { /** * reverse of compress(), decompress the data with lossless compressors * @param src data to be decompressed - * @param srcLen length (in bytes) of the data to be decompressed (as input) or the data decompressed (as output). - * @param dst decompressed data - * @param dstLen length (in bytes) of the decompressed data + * @param srcLen length (in bytes) of that data + * @param dst buffer to decompress into; when null on entry the callee allocates it with malloc() + * and the caller frees it + * @param dstCap the capacity of dst, ignored when dst is null * @return length (in bytes) of the data decompressed */ - virtual size_t decompress(const uchar *src, const size_t srcLen, uchar *&dst, size_t &dstLen) = 0; + virtual size_t decompress(const uchar *src, size_t srcLen, uchar *&dst, size_t dstCap) = 0; }; } // namespace SZ3::concepts diff --git a/include/SZ3/lossless/Lossless_bypass.hpp b/include/SZ3/lossless/Lossless_bypass.hpp index c281dc52..9121834e 100644 --- a/include/SZ3/lossless/Lossless_bypass.hpp +++ b/include/SZ3/lossless/Lossless_bypass.hpp @@ -6,25 +6,51 @@ #define SZ3_LOSSLESS_BYPASS_HPP #include +#include + #include "SZ3/def.hpp" #include "SZ3/lossless/Lossless.hpp" namespace SZ3 { class Lossless_bypass : public concepts::LosslessInterface { public: + /** + * compress data with lossless compressors + * @param src data to be compressed + * @param srcLen length (in bytes) of the data to be compressed + * @param dst compressed data + * @param dstCap capacity (in bytes) for storing the compressed data + * @return length (in bytes) of the data compressed + */ size_t compress(const uchar *src, size_t srcLen, uchar *dst, size_t dstCap) override { + 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 { - dstLen = srcLen; + /** + * reverse of compress(), decompress the data with lossless compressors + * @param src data to be decompressed + * @param srcLen length (in bytes) of that data + * @param dst buffer to decompress into; when null on entry the callee allocates it with malloc() + * and the caller frees it + * @param dstCap the capacity of dst, ignored when dst is null + * @return length (in bytes) of the data decompressed + */ + size_t decompress(const uchar *src, size_t srcLen, uchar *&dst, size_t dstCap) override { + // malloc, because the caller frees what it gets back with free(). if (dst == nullptr) { - dst = static_cast(malloc(dstLen)); + dst = static_cast(malloc(srcLen)); + if (dst == nullptr) { + throw std::runtime_error("SZ3 bypass lossless: can not allocate the decompression buffer"); + } + } else if (srcLen > dstCap) { + throw std::out_of_range("SZ3 bypass lossless: payload exceeds the allowed capacity"); } - std::memcpy(dst, src, dstLen); - return dstLen; + std::memcpy(dst, src, srcLen); + return srcLen; } }; } // namespace SZ3 diff --git a/include/SZ3/lossless/Lossless_zstd.hpp b/include/SZ3/lossless/Lossless_zstd.hpp index 3cb9953a..c915b8f7 100644 --- a/include/SZ3/lossless/Lossless_zstd.hpp +++ b/include/SZ3/lossless/Lossless_zstd.hpp @@ -5,6 +5,7 @@ #ifndef SZ3_LOSSLESS_ZSTD_HPP #define SZ3_LOSSLESS_ZSTD_HPP +#include #include #include "SZ3/def.hpp" @@ -26,6 +27,14 @@ class Lossless_zstd : public concepts::LosslessInterface { * This behavior is not desirable in SZ, as we need the whole compressed data for decompression. * Therefore, we need to check if the dst buffer (dstCap) is large enough for zstd */ + /** + * compress data with lossless compressors + * @param src data to be compressed + * @param srcLen length (in bytes) of the data to be compressed + * @param dst compressed data + * @param dstCap capacity (in bytes) for storing the compressed data + * @return length (in bytes) of the data compressed + */ size_t compress(const uchar *src, size_t srcLen, uchar *dst, size_t dstCap) override { write(srcLen, dst); dstCap -= sizeof(size_t); // reserve space for srcLen @@ -36,12 +45,44 @@ class Lossless_zstd : public concepts::LosslessInterface { return dstLen + sizeof(size_t); } - size_t decompress(const uchar *src, const size_t srcLen, uchar *&dst, size_t &dstLen) override { + /** + * reverse of compress(), decompress the data with lossless compressors + * @param src data to be decompressed + * @param srcLen length (in bytes) of that data + * @param dst buffer to decompress into; when null on entry the callee allocates it with malloc() + * and the caller frees it + * @param dstCap the capacity of dst, ignored when dst is null + * @return length (in bytes) of the data decompressed + */ + size_t decompress(const uchar *src, size_t srcLen, uchar *&dst, size_t dstCap) override { + // The stream is a decompressed-size field followed by the zstd frame, all untrusted. + if (srcLen < sizeof(size_t)) { + throw std::out_of_range("SZ3 lossless: compressed data is smaller than the size header"); + } + size_t dstLen = 0; read(dstLen, src); + + // malloc, because the caller frees what it gets back with free(). + std::unique_ptr owner(nullptr, &free); if (dst == nullptr) { - dst = static_cast(malloc(dstLen)); + owner.reset(static_cast(malloc(dstLen))); + if (owner == nullptr) { + throw std::runtime_error("SZ3 lossless: can not allocate the decompression buffer"); + } + } else if (dstLen > dstCap) { + throw std::out_of_range("SZ3 lossless: declared decompressed size exceeds the allowed capacity"); + } + uchar *out = (dst != nullptr) ? dst : owner.get(); + + // A short frame would leave the tail of the output uninitialized for the caller to read. + size_t res = ZSTD_decompress(out, dstLen, src, srcLen - sizeof(size_t)); + if (ZSTD_isError(res) || res != dstLen) { + throw std::runtime_error("SZ3 lossless: stream does not decompress to the size it declares"); } - return ZSTD_decompress(dst, dstLen, src, srcLen - sizeof(dstLen)); + + dst = out; + owner.release(); + return res; } private: diff --git a/include/SZ3/predictor/ComposedPredictor.hpp b/include/SZ3/predictor/ComposedPredictor.hpp index 13f4f6f4..d30de641 100644 --- a/include/SZ3/predictor/ComposedPredictor.hpp +++ b/include/SZ3/predictor/ComposedPredictor.hpp @@ -4,6 +4,7 @@ #include #include #include +#include #include "SZ3/encoder/HuffmanEncoder.hpp" #include "SZ3/predictor/Predictor.hpp" @@ -45,7 +46,12 @@ class ComposedPredictor : public concepts::PredictorInterface { } bool predecompress(const block_iter &block) override { + // selection and its entries are untrusted, and 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,7 +78,7 @@ class ComposedPredictor : public concepts::PredictorInterface { if (selection_size > 0) { HuffmanEncoder selection_encoder; selection_encoder.load(c, remaining_length); - this->selection = selection_encoder.decode(c, selection_size); + this->selection = selection_encoder.decode(c, selection_size, remaining_length); selection_encoder.postprocess_decode(); } } diff --git a/include/SZ3/predictor/LorenzoPredictor.hpp b/include/SZ3/predictor/LorenzoPredictor.hpp index e61942fa..80c509a0 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,8 @@ 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` keeps + // the step negative instead of wrapping it into an out-of-bounds pointer. 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..64294934 100644 --- a/include/SZ3/predictor/RegressionPredictor.hpp +++ b/include/SZ3/predictor/RegressionPredictor.hpp @@ -2,6 +2,7 @@ #define SZ3_REGRESSION_PREDICTOR_HPP #include +#include #include "SZ3/encoder/HuffmanEncoder.hpp" #include "SZ3/predictor/Predictor.hpp" @@ -114,9 +115,8 @@ class RegressionPredictor : public concepts::PredictorInterface { quantizer_liner.load(c, remaining_length); HuffmanEncoder encoder = HuffmanEncoder(); encoder.load(c, remaining_length); - regression_coeff_quant_inds = encoder.decode(c, coeff_size); + regression_coeff_quant_inds = encoder.decode(c, coeff_size, remaining_length); encoder.postprocess_decode(); - remaining_length -= coeff_size * sizeof(int); std::fill(current_coeffs.begin(), current_coeffs.end(), 0); regression_coeff_index = 0; } @@ -155,6 +155,8 @@ class RegressionPredictor : public concepts::PredictorInterface { } void pred_and_recover_coefficients() { + 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..f1fc1a00 100644 --- a/include/SZ3/quantizer/LinearQuantizer.hpp +++ b/include/SZ3/quantizer/LinearQuantizer.hpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include "SZ3/def.hpp" @@ -28,7 +29,6 @@ class LinearQuantizer : public concepts::QuantizerInterface { assert(eb != 0); } - double get_eb() const { return error_bound; } void set_eb(double eb) { @@ -38,12 +38,13 @@ class LinearQuantizer : public concepts::QuantizerInterface { std::pair get_out_range() const override { return std::make_pair(0, radius * 2); } - // quantize the data with a prediction value, and returns the quantization index and the decompressed data - // 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) { + // NaN data makes this product NaN and infinities push it past the int64_t range; casting either is + // undefined, so the range test is in floating point. What is not representable falls through to unpred. + double scaled = fabs(diff) * this->error_bound_reciprocal; + if (scaled < this->radius * 2 - 1) { + int64_t quant_index = static_cast(scaled) + 1; quant_index >>= 1; int half_index = quant_index; quant_index <<= 1; @@ -55,19 +56,14 @@ class LinearQuantizer : public concepts::QuantizerInterface { 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; - } 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 @@ -80,10 +76,15 @@ class LinearQuantizer : public concepts::QuantizerInterface { } ALWAYS_INLINE T recover_pred(T pred, int quant_index) { - return pred + 2 * (quant_index - this->radius) * this->error_bound; + // quant_index comes from the stream; in int, 2 * (quant_index - radius) overflows past INT_MAX/2. + // Exact for every index a valid stream carries. + return pred + 2 * (static_cast(quant_index) - this->radius) * this->error_bound; } - ALWAYS_INLINE T recover_unpred() { return unpred[index++]; } + ALWAYS_INLINE T recover_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 +116,9 @@ class LinearQuantizer : public concepts::QuantizerInterface { size_t unpred_size = 0; read(unpred_size, c, remaining_length); if (unpred_size > 0) { + // resize() below is sized from the stream, so check the count against the bytes that exist first. + 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..8adf8651 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 { diff --git a/include/SZ3/utils/BlockwiseIterator.hpp b/include/SZ3/utils/BlockwiseIterator.hpp index 3221d4f1..9387728d 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,19 @@ class block_data : public std::enable_shared_from_this> { } } + /** + * @brief The block's `num` values in unpadded layout. With padding this is an internal copy, + * materialized on the first call and valid until this object is destroyed. + */ + 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 +289,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/ByteUtil.hpp b/include/SZ3/utils/ByteUtil.hpp index 99f85441..ff7333ee 100644 --- a/include/SZ3/utils/ByteUtil.hpp +++ b/include/SZ3/utils/ByteUtil.hpp @@ -253,7 +253,7 @@ std::vector bytes2vector(const unsigned char *&c, uint8_t bit_width, size_t n size_t byte_index = bit_index / 8; size_t bit_offset = bit_index % 8; - value |= ((c[byte_index] >> bit_offset) & 1) << j; + value |= static_cast((c[byte_index] >> bit_offset) & 1) << j; } data[i] = value; } diff --git a/include/SZ3/utils/Config.hpp b/include/SZ3/utils/Config.hpp index 81febd8c..bfbacac4 100644 --- a/include/SZ3/utils/Config.hpp +++ b/include/SZ3/utils/Config.hpp @@ -12,11 +12,14 @@ #ifndef SZ3_Config_HPP #define SZ3_Config_HPP +#include #include #include #include +#include #include #include +#include #include #include "SZ3/def.hpp" @@ -337,6 +340,8 @@ class Config { } else if (errorBoundMode == EB_ABS_AND_REL) { write(absErrorBound, c); write(relErrorBound, c); + } else { + throw std::invalid_argument("SZ3 Config::save: unknown error bound mode"); } uint8_t boolvals = (lorenzo & 1) << 7 | (lorenzo2 & 1) << 6 | (regression & 1) << 5 | (regression2 & 1) << 4 | @@ -357,59 +362,74 @@ class Config { * @brief Deserialize the configuration from a byte array. * * @param c Pointer to the byte array. + * @param remaining_length bytes readable from `c`; decremented by what is consumed. */ - void load(const unsigned char*& c) { + void load(const unsigned char*& c, size_t& remaining_length) { + const unsigned char* const c0 = c; + const unsigned char* const cend = c + remaining_length; uchar confSize = 0; - read(confSize, c); - auto c1 = c + confSize; + read(confSize, c, remaining_length); + // Where the optional fields at the end stop. A writer that declares more than it sends stops + // at the buffer instead. + const unsigned char* const c1 = std::min(c0 + confSize, cend); - read(N, c); + read(N, c, remaining_length); + if (N > 4) throw std::out_of_range("SZ3 Config::load: invalid number of dimensions"); uint8_t bitWidth; - read(bitWidth, c); + read(bitWidth, c, remaining_length); + if (bitWidth > 64) throw std::out_of_range("SZ3 Config::load: invalid dimension bit width"); + const size_t dim_bytes = (static_cast(N) * bitWidth + 7) / 8; + if (dim_bytes > remaining_length) + throw std::out_of_range("SZ3 Config::load: dimensions exceed the buffer"); dims = bytes2vector(c, bitWidth, N); - // dims.resize(N); - // read(dims.data(), N, c); - read(num, c); - read(cmprAlgo, c); + remaining_length -= dim_bytes; + read(num, c, remaining_length); + // num must equal the product of the dimensions, or the predictor walks more grid positions than + // were allocated. No dimensions means no elements: the HDF5 filter loads a config from cd_values + // before it knows the dataset shape. + size_t dims_product = dims.empty() ? 0 : 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"); + read(cmprAlgo, c, remaining_length); - read(errorBoundMode, c); + read(errorBoundMode, c, remaining_length); if (errorBoundMode == EB_ABS) { - read(absErrorBound, c); + read(absErrorBound, c, remaining_length); } else if (errorBoundMode == EB_REL) { - read(relErrorBound, c); + read(relErrorBound, c, remaining_length); } else if (errorBoundMode == EB_PSNR) { - read(psnrErrorBound, c); + read(psnrErrorBound, c, remaining_length); } else if (errorBoundMode == EB_L2NORM) { - read(l2normErrorBound, c); - } else if (errorBoundMode == EB_ABS_OR_REL) { - read(absErrorBound, c); - read(relErrorBound, c); - } else if (errorBoundMode == EB_ABS_AND_REL) { - read(absErrorBound, c); - read(relErrorBound, c); + read(l2normErrorBound, c, remaining_length); + } else if (errorBoundMode == EB_ABS_OR_REL || errorBoundMode == EB_ABS_AND_REL) { + read(absErrorBound, c, remaining_length); + read(relErrorBound, c, remaining_length); + } else { + // save() always writes a bound here, so an unhandled mode would shift every field below it. + throw std::invalid_argument("SZ3 Config::load: unknown error bound mode"); } if (c < c1) { uint8_t boolvals; - read(boolvals, c); + read(boolvals, c, remaining_length); lorenzo = (boolvals >> 7) & 1; lorenzo2 = (boolvals >> 6) & 1; regression = (boolvals >> 5) & 1; regression2 = (boolvals >> 4) & 1; openmp = (boolvals >> 3) & 1; } - if (c < c1) { - read(dataType, c); - } - if (c < c1) { - read(quantbinCnt, c); - } - if (c < c1) { - read(blockSize, c); - } - if (c < c1) { - read(predDim, c); - } + if (c < c1) read(dataType, c, remaining_length); + if (c < c1) read(quantbinCnt, c, remaining_length); + if (c < c1) read(blockSize, c, remaining_length); + if (c < c1) read(predDim, c, remaining_length); } /** 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..6a198302 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,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) { - assert(num_elements * sizeof(T1) <= remaining_length); + if (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++) { @@ -109,7 +110,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..60ee81d7 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 @@ -65,7 +66,8 @@ herr_t get_SZ3_conf_from_H5(const hid_t propertyList, SZ3::Config& conf) { // if not empty, load cd_values into config if (cd_nelmts > 0) { auto buffer = reinterpret_cast(cd_values.data()); - conf.load(buffer); + size_t cd_bytes = cd_nelmts * sizeof(unsigned int); + conf.load(buffer, cd_bytes); } } return 1; @@ -159,7 +161,9 @@ 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; + // The bound assumes the payload fits in the raw size, so leave headroom on top of it for + // algorithms whose output can reach or exceed that. + 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); @@ -186,7 +190,8 @@ static size_t H5Z_filter_sz3(unsigned int flags, size_t cd_nelmts, const unsigne SZ3::Config conf; auto buffer = reinterpret_cast(cd_values); - conf.load(buffer); + size_t cd_bytes = cd_nelmts * sizeof(unsigned int); + conf.load(buffer, cd_bytes); // conf.print(); if (conf.num < 20) return nbytes; diff --git a/tools/H5Z-SZ3/test/cdvalueHelper.cpp b/tools/H5Z-SZ3/test/cdvalueHelper.cpp index e5442b29..796b466c 100644 --- a/tools/H5Z-SZ3/test/cdvalueHelper.cpp +++ b/tools/H5Z-SZ3/test/cdvalueHelper.cpp @@ -90,7 +90,8 @@ int main(int argc, char* argv[]) { SZ3::Config conf; auto buffer = reinterpret_cast(cd_values.data()); - conf.load(buffer); + size_t cd_bytes = cd_values.size() * sizeof(unsigned int); + conf.load(buffer, cd_bytes); std::ofstream file(outPath); if (!file.is_open()) { diff --git a/tools/H5Z-SZ3/test/cdvalueHelper.py b/tools/H5Z-SZ3/test/cdvalueHelper.py index 021313ba..35a9ee62 100644 --- a/tools/H5Z-SZ3/test/cdvalueHelper.py +++ b/tools/H5Z-SZ3/test/cdvalueHelper.py @@ -29,7 +29,7 @@ def __init__(self, algo="ALGO_INTERP_LORENZO", absolute=None, relative=None, psn # placeholder values will be overwritten when the HDF5 filter calls the set_local function, serialized = bytearray() - serialized.extend(struct.pack(' #include #include +#include #include "SZ3/compressor/specialized/SZExaaltCompressor.hpp" #include "SZ3/decomposition/BlockwiseDecomposition.hpp" @@ -136,7 +137,9 @@ float *VQ(Config conf, size_t ts, T *data, size_t &compressed_size, bool decom, Timer timer(true); compressed_size = 2 * conf.num * sizeof(T); auto compressed = static_cast(malloc(compressed_size)); - sz->compress(conf, data, compressed, compressed_size); + // compress() returns what it actually wrote; compressed_size went in as the capacity, and the + // decompression below reads it as the stream length. + compressed_size = sz->compress(conf, data, compressed, compressed_size); total_compress_time += timer.stop("Compression"); if (!decom) { free(compressed); @@ -163,7 +166,9 @@ float *MT(Config conf, size_t ts, T *data, size_t &compressed_size, bool decom, Timer timer(true); compressed_size = 2 * conf.num * sizeof(T); auto compressed = static_cast(malloc(compressed_size)); - sz->compress(conf, data, compressed, compressed_size); + // compress() returns what it actually wrote; compressed_size went in as the capacity, and the + // decompression below reads it as the stream length. + compressed_size = sz->compress(conf, data, compressed, compressed_size); total_compress_time += timer.stop("Compression"); if (!decom) { free(compressed); @@ -191,7 +196,9 @@ float *SZ2(Config conf, size_t ts, T *data, size_t &compressed_size, bool decom) Timer timer(true); compressed_size = 2 * conf.num * sizeof(T); auto compressed = static_cast(malloc(compressed_size)); - sz->compress(conf, data, compressed, compressed_size); + // compress() returns what it actually wrote; compressed_size went in as the capacity, and the + // decompression below reads it as the stream length. + compressed_size = sz->compress(conf, data, compressed, compressed_size); total_compress_time += timer.stop("Compression"); if (!decom) { @@ -300,7 +307,8 @@ uchar *LAMMPS_compress(Config conf, T *data, int method, size_t &compressed_size sz = make_sz(conf); // sz->compress(conf, data, compressed_size); } - sz->compress(conf, data, compressed_data, compressed_size); + // compressed_size is the caller's out-parameter for the stream length, not the capacity it went in as. + compressed_size = sz->compress(conf, data, compressed_data, compressed_size); // auto ratio = conf.num * sizeof(T) * 1.0 / compressed_size; // std::cout << "Compression Ratio = " << ratio << std::endl; // std::cout << "Compressed size = " << compressed_size << std::endl; diff --git a/tools/sz3/sz3.cpp b/tools/sz3/sz3.cpp index be0a844c..2055f0c0 100644 --- a/tools/sz3/sz3.cpp +++ b/tools/sz3/sz3.cpp @@ -130,7 +130,8 @@ template void compress(char *inPath, char *cmpPath, SZ3::Config &conf) { T *data = new T[conf.num]; SZ3::readfile(inPath, conf.num, data); - size_t bytesCap = 2 * conf.num * sizeof(T); + // SZ_compress refuses a capacity below its own bound, which a small input falls under. + size_t bytesCap = SZ3::SZ_compress_size_bound(conf); auto bytes = new char[bytesCap]; SZ3::Timer timer(true); diff --git a/tools/test/deprecated/SZBlockInterpolationCompressor.hpp b/tools/test/deprecated/SZBlockInterpolationCompressor.hpp index cd7ec88e..6573b4df 100644 --- a/tools/test/deprecated/SZBlockInterpolationCompressor.hpp +++ b/tools/test/deprecated/SZBlockInterpolationCompressor.hpp @@ -35,8 +35,7 @@ class SZBlockInterpolationCompressor { T *decompress(const Config &conf, uchar const *cmpData, size_t cmpSize, T *decData) { uchar *buffer = nullptr; - size_t bufferSize = 0; - lossless.decompress(cmpData, cmpSize, buffer, bufferSize); + size_t bufferSize = lossless.decompress(cmpData, cmpSize, buffer, 0); size_t remaining_length = bufferSize; uchar const *buffer_pos = buffer; diff --git a/tools/test/integration/integration_test_driver.py b/tools/test/integration/integration_test_driver.py index 101d63e1..207ed712 100644 --- a/tools/test/integration/integration_test_driver.py +++ b/tools/test/integration/integration_test_driver.py @@ -123,11 +123,22 @@ def main(): sys.exit(1) if len(sys.argv) > 2: - selected_dataset = sys.argv[2] + # "name" runs the whole dataset, "name:a.f32,b.f32" runs those fields of it, so a dataset + # too slow to be one job can be spread over several. + selected_dataset, _, selected_fields = sys.argv[2].partition(":") if selected_dataset not in datasets: print(f"Dataset {selected_dataset} not found in {datasets_json}") sys.exit(1) - datasets = {selected_dataset: datasets[selected_dataset]} + dataset_info = datasets[selected_dataset] + if selected_fields: + wanted = [f for f in selected_fields.split(",") if f] + missing = [f for f in wanted if f not in dataset_info["fields"]] + if missing: + print(f"Fields {missing} not found in dataset {selected_dataset}") + sys.exit(1) + dataset_info = dict(dataset_info) + dataset_info["fields"] = {f: dataset_info["fields"][f] for f in wanted} + datasets = {selected_dataset: dataset_info} script_dir = os.path.dirname(os.path.abspath(__file__)) project_source_dir = os.path.abspath(os.path.join(script_dir, "..", "..", "..")) 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 diff --git a/tools/test/modules/test_encoder.cpp b/tools/test/modules/test_encoder.cpp index c7319acb..9770cc0b 100644 --- a/tools/test/modules/test_encoder.cpp +++ b/tools/test/modules/test_encoder.cpp @@ -33,7 +33,7 @@ void runFunctionalTest() { const SZ3::uchar *buffer_conf_pos = buffer_conf.data(); Encoder coder; coder.load(buffer_conf_pos, conf_len); - auto dataDecoded = coder.decode(buffer_data_pos, N); + auto dataDecoded = coder.decode(buffer_data_pos, N, data_len); for (int i = 0; i < N; i++) { EXPECT_EQ(data[i], dataDecoded[i]); } diff --git a/tools/test/modules/test_lossless.cpp b/tools/test/modules/test_lossless.cpp index 00e23cef..eed0c359 100644 --- a/tools/test/modules/test_lossless.cpp +++ b/tools/test/modules/test_lossless.cpp @@ -22,8 +22,7 @@ void runFunctionalTest() { std::vector decompressed(N); SZ3::uchar* decompressed_pos = decompressed.data(); - size_t decompressedSize; - lossless.decompress(dst.data(), compressedSize, decompressed_pos, decompressedSize); + size_t decompressedSize = lossless.decompress(dst.data(), compressedSize, decompressed_pos, decompressed.size()); EXPECT_EQ(decompressedSize, src.size()); EXPECT_EQ(std::vector(decompressed.data(), decompressed.data() + decompressedSize), src);