Browse Source

Fix zlib (de)compressor memory leaks

savilli 1 year ago
parent
commit
ae555465ba
1 changed files with 16 additions and 5 deletions
  1. 16 5
      src/serialization.cpp

+ 16 - 5
src/serialization.cpp

@@ -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) {
@@ -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;
@@ -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;
@@ -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)
@@ -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(;;)
@@ -180,8 +193,6 @@ void decompressZlib(std::istream &is, std::ostream &os, size_t limit)
 			break;
 		}
 	}
-
-	inflateEnd(&z);
 }
 
 struct ZSTD_Deleter {