diff --git a/include/SZ3/api/impl/SZAlgoInterp.hpp b/include/SZ3/api/impl/SZAlgoInterp.hpp index 914756ea..4dd6b32a 100644 --- a/include/SZ3/api/impl/SZAlgoInterp.hpp +++ b/include/SZ3/api/impl/SZAlgoInterp.hpp @@ -12,6 +12,8 @@ #include "SZ3/utils/Sample.hpp" #include "SZ3/utils/Statistic.hpp" +#include + namespace SZ3 { template size_t SZ_compress_Interp(Config &conf, T *data, uchar *cmpData, size_t cmpCap) { @@ -61,6 +63,9 @@ double interp_compress_test( std::max(1000, 1.2 * (sz.size_est() + encoder.size_est() + sizeof(T) * total_quant_bins.size())); auto buffer = static_cast(malloc(bufferSize)); + // RAII: the encoder and the lossless layer below can throw (std::length_error when the destination + // capacity is too small), which would leak this scratch buffer with a bare free() at the end. + std::unique_ptr buffer_owner(buffer, &free); uchar *buffer_pos = buffer; sz.save(buffer_pos); encoder.save(buffer_pos); @@ -70,7 +75,6 @@ double interp_compress_test( encoder.encode(total_quant_bins, buffer_pos); encoder.postprocess_encode(); auto cmpSize = lossless.compress(buffer, buffer_pos - buffer, cmpData, cmpCap); - free(buffer); auto compression_ratio = conf.num * sampled_blocks.size() * sizeof(T) * 1.0 / cmpSize; return compression_ratio; } @@ -100,6 +104,9 @@ double lorenzo_compress_test( size_t bufferSize = std::max(1000, 1.2 * (encoder.size_est() + sizeof(T) * total_quant_bins.size())); auto buffer = static_cast(malloc(bufferSize)); + // RAII: the encoder and the lossless layer below can throw (std::length_error when the destination + // capacity is too small), which would leak this scratch buffer with a bare free() at the end. + std::unique_ptr buffer_owner(buffer, &free); uchar *buffer_pos = buffer; sz.save(buffer_pos); encoder.save(buffer_pos); @@ -109,7 +116,6 @@ double lorenzo_compress_test( encoder.encode(total_quant_bins, buffer_pos); encoder.postprocess_encode(); auto cmpSize = lossless.compress(buffer, buffer_pos - buffer, cmpData, cmpCap); - free(buffer); auto compression_ratio = conf.num * sampled_blocks.size() * sizeof(T) * 1.0 / cmpSize; return compression_ratio; // } @@ -179,6 +185,9 @@ size_t SZ_compress_Interp_lorenzo(Config &conf, T *data, uchar *cmpData, size_t double best_lorenzo_ratio = 0, best_interp_ratio = 0, ratio; size_t bufferCap = conf.num * sizeof(T); auto buffer = static_cast(malloc(bufferCap)); + // RAII: the tuning calls below run full compression rounds that can throw, which would leak this + // scratch buffer with a bare free() at the end. + std::unique_ptr buffer_owner(buffer, &free); Config lorenzo_config = conf; { @@ -281,7 +290,6 @@ size_t SZ_compress_Interp_lorenzo(Config &conf, T *data, uchar *cmpData, size_t cmpSize = SZ_compress_LorenzoReg(conf, data, cmpData, cmpCap); } - free(buffer); return cmpSize; } } // namespace SZ3 diff --git a/include/SZ3/api/impl/SZDispatcher.hpp b/include/SZ3/api/impl/SZDispatcher.hpp index 9de7aa13..da43c658 100644 --- a/include/SZ3/api/impl/SZDispatcher.hpp +++ b/include/SZ3/api/impl/SZDispatcher.hpp @@ -8,6 +8,8 @@ #include "SZ3/utils/Config.hpp" #include "SZ3/utils/Statistic.hpp" +#include + namespace SZ3 { template size_t SZ_compress_dispatcher(Config &conf, const T *data, uchar *cmpData, size_t cmpCap) { @@ -63,6 +65,8 @@ size_t SZ_compress_dispatcher(Config &conf, const T *data, uchar *cmpData, size_ auto zstd = Lossless_zstd(); auto zstdCmpCap = ZSTD_compressBound(conf.num * sizeof(T)) + sizeof(size_t); auto zstdCmpData = static_cast(malloc(zstdCmpCap)); + // RAII: zstd.compress can throw, which would leak this buffer with a bare free() at the end. + std::unique_ptr zstd_cmp_data_owner(zstdCmpData, &free); size_t zstdCmpSize = zstd.compress(reinterpret_cast(data), conf.num * sizeof(T), zstdCmpData, zstdCmpCap); if (zstdCmpSize < cmpSize && zstdCmpSize <= cmpCap) { @@ -70,7 +74,6 @@ size_t SZ_compress_dispatcher(Config &conf, const T *data, uchar *cmpData, size_ memcpy(cmpData, zstdCmpData, zstdCmpSize); cmpSize = zstdCmpSize; } - free(zstdCmpData); } return cmpSize; } diff --git a/include/SZ3/api/impl/SZImplOMP.hpp b/include/SZ3/api/impl/SZImplOMP.hpp index 664a14dc..666656fd 100644 --- a/include/SZ3/api/impl/SZImplOMP.hpp +++ b/include/SZ3/api/impl/SZImplOMP.hpp @@ -2,7 +2,9 @@ #define SZ3_IMPL_SZDISPATCHER_OMP_HPP #include +#include #include +#include #include "SZ3/api/impl/SZDispatcher.hpp" @@ -70,8 +72,15 @@ size_t SZ_compress_OMP(Config& conf, const T* data, uchar* cmpData, size_t cmpCa conf_t[tid] = conf; conf_t[tid].setDims(dims_t.begin(), dims_t.end()); - size_t cmp_size_cap = ZSTD_compressBound(conf_t[tid].num * sizeof(T)); - compressed_t[tid] = static_cast(malloc(cmp_size_cap)); + // Reserve room for the size header that Lossless_zstd::compress writes in front of the zstd stream, + // otherwise the direct lossless path in SZ_compress_dispatcher throws for poorly compressible chunks. + size_t cmp_size_cap = sizeof(size_t) + ZSTD_compressBound(conf_t[tid].num * sizeof(T)); + // The buffer is owned so that it is released even if the compression below throws. + std::unique_ptr compressed_owner(static_cast(malloc(cmp_size_cap)), &free); + if (!compressed_owner) { + throw std::bad_alloc(); + } + compressed_t[tid] = compressed_owner.get(); // we have to use conf_t[tid].N instead of N since each chunk may be a slice of the original data if (conf_t[tid].N == 1) { cmp_size_t[tid] = SZ_compress_dispatcher(conf_t[tid], data_t, compressed_t[tid], cmp_size_cap); @@ -105,7 +114,6 @@ size_t SZ_compress_OMP(Config& conf, const T* data, uchar* cmpData, size_t cmpCa } memcpy(buffer_pos + cmp_start_t[tid], compressed_t[tid], cmp_size_t[tid]); - free(compressed_t[tid]); } return buffer_pos - cmpData + cmp_start_t[nThreads]; @@ -200,7 +208,8 @@ 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) + + //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..e2d0f4a1 100644 --- a/include/SZ3/api/sz.hpp +++ b/include/SZ3/api/sz.hpp @@ -24,6 +24,8 @@ #include "SZ3/api/impl/SZImpl.hpp" #include "SZ3/version.hpp" +#include + /** * Compresses the input data using the provided configuration and stores the result in a pre-allocated buffer. @@ -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(); } /** diff --git a/include/SZ3/compressor/SZGenericCompressor.hpp b/include/SZ3/compressor/SZGenericCompressor.hpp index 4a3f471a..9f4bf44c 100644 --- a/include/SZ3/compressor/SZGenericCompressor.hpp +++ b/include/SZ3/compressor/SZGenericCompressor.hpp @@ -46,6 +46,11 @@ class SZGenericCompressor : public concepts::CompressorInterface { 1000, 2 * (decomposition.size_est() + encoder.size_est() + sizeof(T) * quant_inds.size())); auto buffer = static_cast(malloc(bufferSize)); + // Own the scratch buffer with RAII so it is released on every path: the encoder and the lossless + // layer below can throw (e.g. Lossless_zstd::compress throws std::length_error when the destination + // capacity is too small for poorly-compressible data), and the caller catches and continues, so a + // bare free() at the end leaks the buffer on each failed compression. + std::unique_ptr buffer_owner(buffer, &free); uchar *buffer_pos = buffer; decomposition.save(buffer_pos); @@ -57,7 +62,6 @@ class SZGenericCompressor : public concepts::CompressorInterface { encoder.postprocess_encode(); auto cmpSize = lossless.compress(buffer, buffer_pos - buffer, cmpData, cmpCap); - free(buffer); return cmpSize; } diff --git a/include/SZ3/compressor/specialized/SZExaaltCompressor.hpp b/include/SZ3/compressor/specialized/SZExaaltCompressor.hpp index a62e33f6..43753e17 100644 --- a/include/SZ3/compressor/specialized/SZExaaltCompressor.hpp +++ b/include/SZ3/compressor/specialized/SZExaaltCompressor.hpp @@ -1,6 +1,8 @@ #ifndef SZ3_EXAALT_COMPRESSSOR_HPP #define SZ3_EXAALT_COMPRESSSOR_HPP +#include + #include "SZ3/def.hpp" #include "SZ3/encoder/Encoder.hpp" #include "SZ3/lossless/Lossless.hpp" @@ -90,6 +92,9 @@ class SZExaaltCompressor : public SZ3::concepts::CompressorInterface { quantizer.postcompress_data(); auto buffer = static_cast(malloc(4 * conf.num * sizeof(T))); + // RAII: the encoder and the lossless layer below can throw (std::length_error when the destination + // capacity is too small), which would leak this scratch buffer with a bare free() at the end. + std::unique_ptr buffer_owner(buffer, &free); uchar *buffer_pos = buffer; quantizer.save(buffer_pos); // quantizer.print(); @@ -112,7 +117,6 @@ class SZExaaltCompressor : public SZ3::concepts::CompressorInterface { encoder.postprocess_encode(); auto cmpSize = lossless.compress(buffer, buffer_pos - buffer, cmpData, cmpCap); - free(buffer); return cmpSize; } @@ -121,6 +125,8 @@ class SZExaaltCompressor : public SZ3::concepts::CompressorInterface { uchar *buffer = nullptr; size_t bufferSize = 0; lossless.decompress(cmpData, cmpSize, buffer, bufferSize); + // RAII: the parsing steps below operate on untrusted data and can throw before the buffer is freed. + std::unique_ptr buffer_owner(buffer, &free); size_t remaining_length = cmpSize; uchar const *buffer_pos = buffer; @@ -134,8 +140,6 @@ class SZExaaltCompressor : public SZ3::concepts::CompressorInterface { auto pred_inds = encoder.decode(buffer_pos, pred_inds_num); encoder.postprocess_decode(); - free(buffer); - quantizer.predecompress_data(); auto l = pred_inds[0] - level_num; diff --git a/include/SZ3/compressor/specialized/SZTruncateCompressor.hpp b/include/SZ3/compressor/specialized/SZTruncateCompressor.hpp index f4d3b766..d6a6f3f2 100644 --- a/include/SZ3/compressor/specialized/SZTruncateCompressor.hpp +++ b/include/SZ3/compressor/specialized/SZTruncateCompressor.hpp @@ -2,6 +2,7 @@ #define SZ3_Truncate_COMPRESSOR_HPP #include +#include #include "SZ3/compressor/Compressor.hpp" #include "SZ3/decomposition/Decomposition.hpp" @@ -28,6 +29,9 @@ class SZTruncateCompressor : public concepts::CompressorInterface { size_t compress(const Config &conf, T *data, uchar *cmpData, size_t cmpCap) override { auto buffer = static_cast(malloc(conf.num * sizeof(T))); + // RAII: the lossless layer below can throw (std::length_error when the destination capacity is too + // small), which would leak this scratch buffer with a bare free() at the end. + std::unique_ptr buffer_owner(buffer, &free); auto buffer_pos = buffer; // Timer timer(true); @@ -35,7 +39,6 @@ class SZTruncateCompressor : public concepts::CompressorInterface { // timer.stop("Prediction & Quantization"); auto cmpSize = lossless.compress(buffer, buffer_pos - buffer, cmpData, cmpCap); - free(buffer); return cmpSize; // lossless.postcompress_data(buffer); // return lossless_data;