Skip to content

Commit

Permalink
Fix zlib (de)compressor memory leaks
Browse files Browse the repository at this point in the history
  • Loading branch information
savilli authored Jun 16, 2022
1 parent 622d857 commit ae55546
Showing 1 changed file with 16 additions and 5 deletions.
21 changes: 16 additions & 5 deletions src/serialization.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ with this program; if not, write to the Free Software Foundation, Inc.,
#include <zstd.h>

/* report a zlib or i/o error */
void zerr(int ret)
static void zerr(int ret)
{
dstream<<"zerr: ";
switch (ret) {
Expand All @@ -52,6 +52,17 @@ void zerr(int ret)
}
}

// Make sure that z is deleted in case of exception
template <int (*F)(z_stream*)>
class ZlibAutoDeleter {
public:
ZlibAutoDeleter(z_stream *ptr) : ptr_(ptr) {}
~ZlibAutoDeleter() { F(ptr_); }

private:
z_stream *ptr_;
};

void compressZlib(const u8 *data, size_t data_size, std::ostream &os, int level)
{
z_stream z;
Expand All @@ -68,6 +79,8 @@ void compressZlib(const u8 *data, size_t data_size, std::ostream &os, int level)
if(ret != Z_OK)
throw SerializationError("compressZlib: deflateInit failed");

ZlibAutoDeleter<deflateEnd> deleter(&z);

// Point zlib to our input buffer
z.next_in = (Bytef*)&data[0];
z.avail_in = data_size;
Expand All @@ -91,8 +104,6 @@ void compressZlib(const u8 *data, size_t data_size, std::ostream &os, int level)
if(status == Z_STREAM_END)
break;
}

deflateEnd(&z);
}

void compressZlib(const std::string &data, std::ostream &os, int level)
Expand All @@ -119,6 +130,8 @@ void decompressZlib(std::istream &is, std::ostream &os, size_t limit)
if(ret != Z_OK)
throw SerializationError("dcompressZlib: inflateInit failed");

ZlibAutoDeleter<inflateEnd> deleter(&z);

z.avail_in = 0;

for(;;)
Expand Down Expand Up @@ -180,8 +193,6 @@ void decompressZlib(std::istream &is, std::ostream &os, size_t limit)
break;
}
}

inflateEnd(&z);
}

struct ZSTD_Deleter {
Expand Down

0 comments on commit ae55546

Please sign in to comment.