Fix memory leaks in compression scratch buffers on exception

include/SZ3/api/impl/SZAlgoInterp.hpp

@@ -12,6 +12,8 @@
 #include "SZ3/utils/Sample.hpp"
 #include "SZ3/utils/Statistic.hpp"
 
+#include <memory>
+
 namespace SZ3 {
 template <class T, uint N>
 size_t SZ_compress_Interp(Config &conf, T *data, uchar *cmpData, size_t cmpCap) {
@@ -61,6 +63,9 @@ double interp_compress_test(
         std::max<size_t>(1000, 1.2 * (sz.size_est() + encoder.size_est() + sizeof(T) * total_quant_bins.size()));
 
     auto buffer = static_cast<uchar *>(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<uchar, void (*)(void *)> 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<size_t>(1000, 1.2 * (encoder.size_est() + sizeof(T) * total_quant_bins.size()));
 
     auto buffer = static_cast<uchar *>(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<uchar, void (*)(void *)> 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<uchar *>(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<uchar, void (*)(void *)> 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<T, N>(conf, data, cmpData, cmpCap);
     }
 
-    free(buffer);
     return cmpSize;
 }
 }  // namespace SZ3

include/SZ3/api/impl/SZDispatcher.hpp

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

include/SZ3/api/impl/SZImplOMP.hpp

@@ -2,7 +2,9 @@
 #define SZ3_IMPL_SZDISPATCHER_OMP_HPP
 
 #include <cmath>
+#include <cstdlib>
 #include <memory>
+#include <new>
 
 #include "SZ3/api/impl/SZDispatcher.hpp"
 
@@ -70,8 +72,15 @@ size_t SZ_compress_OMP(Config& conf, const T* data, uchar* cmpData, size_t cmpCa
 
         conf_t[tid] = conf;
         conf_t[tid].setDims(dims_t.begin(), dims_t.end());
-        size_t cmp_size_cap = ZSTD_compressBound(conf_t[tid].num * sizeof(T));
-        compressed_t[tid] = static_cast<uchar*>(malloc(cmp_size_cap));
+        // Reserve room for the size header that Lossless_zstd::compress writes in front of the zstd stream,
+        // otherwise the direct lossless path in SZ_compress_dispatcher throws for poorly compressible chunks.
+        size_t cmp_size_cap = sizeof(size_t) + ZSTD_compressBound(conf_t[tid].num * sizeof(T));
+        // The buffer is owned so that it is released even if the compression below throws.
+        std::unique_ptr<uchar, void (*)(void*)> compressed_owner(static_cast<uchar*>(malloc(cmp_size_cap)), &free);
+        if (!compressed_owner) {
+            throw std::bad_alloc();
+        }
+        compressed_t[tid] = compressed_owner.get();
         // we have to use conf_t[tid].N instead of N since each chunk may be a slice of the original data
         if (conf_t[tid].N == 1) {
             cmp_size_t[tid] = SZ_compress_dispatcher<T, 1>(conf_t[tid], data_t, compressed_t[tid], cmp_size_cap);
@@ -105,7 +114,6 @@ size_t SZ_compress_OMP(Config& conf, const T* data, uchar* cmpData, size_t cmpCa
         }
 
         memcpy(buffer_pos + cmp_start_t[tid], compressed_t[tid], cmp_size_t[tid]);
-        free(compressed_t[tid]);
     }
 
     return buffer_pos - cmpData + cmp_start_t[nThreads];
@@ -200,7 +208,8 @@ size_t SZ_compress_size_bound_omp(const Config& conf) {
     size_t chunk_size = conf.dims[0] / static_cast<size_t>(nThreads) * (conf.num / conf.dims[0]);
     size_t last_chunk_size = (conf.dims[0] - conf.dims[0] / nThreads * (nThreads - 1)) * (conf.num / conf.dims[0]);
     //for each thread, we save conf, compressed size, and compressed data
-    return sizeof(int) + nThreads * conf.size_est() + nThreads * sizeof(size_t) +
+    //the per-chunk compressed data may carry the size header written by Lossless_zstd::compress
+    return sizeof(int) + nThreads * conf.size_est() + 2 * nThreads * sizeof(size_t) +
            (nThreads - 1) * ZSTD_compressBound(chunk_size * sizeof(T)) +
            ZSTD_compressBound(last_chunk_size * sizeof(T));
 #else

include/SZ3/api/sz.hpp

@@ -24,6 +24,8 @@
 #include "SZ3/api/impl/SZImpl.hpp"
 #include "SZ3/version.hpp"
 
+#include <memory>
+
 
 /**
  * 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<T>(config);
-    auto buffer = new char[bufferLen];
-    cmpSize = SZ_compress(config, data, buffer, bufferLen);
+    std::unique_ptr<char[]> buffer(new char[bufferLen]);
+    cmpSize = SZ_compress(config, data, buffer.get(), bufferLen);
 
-    return buffer;
+    return buffer.release();
 }
 
 /**

include/SZ3/compressor/SZGenericCompressor.hpp

@@ -46,6 +46,11 @@ class SZGenericCompressor : public concepts::CompressorInterface<T> {
             1000, 2 * (decomposition.size_est() + encoder.size_est() + sizeof(T) * quant_inds.size()));
 
         auto buffer = static_cast<uchar *>(malloc(bufferSize));
+        // Own the scratch buffer with RAII so it is released on every path: the encoder and the lossless
+        // layer below can throw (e.g. Lossless_zstd::compress throws std::length_error when the destination
+        // capacity is too small for poorly-compressible data), and the caller catches and continues, so a
+        // bare free() at the end leaks the buffer on each failed compression.
+        std::unique_ptr<uchar, void (*)(void *)> buffer_owner(buffer, &free);
         uchar *buffer_pos = buffer;
 
         decomposition.save(buffer_pos);
@@ -57,7 +62,6 @@ class SZGenericCompressor : public concepts::CompressorInterface<T> {
         encoder.postprocess_encode();
 
         auto cmpSize = lossless.compress(buffer, buffer_pos - buffer, cmpData, cmpCap);
-        free(buffer);
 
         return cmpSize;
     }

include/SZ3/compressor/specialized/SZExaaltCompressor.hpp

@@ -1,6 +1,8 @@
 #ifndef SZ3_EXAALT_COMPRESSSOR_HPP
 #define SZ3_EXAALT_COMPRESSSOR_HPP
 
+#include <memory>
+
 #include "SZ3/def.hpp"
 #include "SZ3/encoder/Encoder.hpp"
 #include "SZ3/lossless/Lossless.hpp"
@@ -90,6 +92,9 @@ class SZExaaltCompressor : public SZ3::concepts::CompressorInterface<T> {
         quantizer.postcompress_data();
 
         auto buffer = static_cast<uchar *>(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<uchar, void (*)(void *)> buffer_owner(buffer, &free);
         uchar *buffer_pos = buffer;
         quantizer.save(buffer_pos);
         //            quantizer.print();
@@ -112,7 +117,6 @@ class SZExaaltCompressor : public SZ3::concepts::CompressorInterface<T> {
         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<T> {
         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<uchar, void (*)(void *)> buffer_owner(buffer, &free);
         size_t remaining_length = cmpSize;
         uchar const *buffer_pos = buffer;
 
@@ -134,8 +140,6 @@ class SZExaaltCompressor : public SZ3::concepts::CompressorInterface<T> {
         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;

include/SZ3/compressor/specialized/SZTruncateCompressor.hpp

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