From 154bdf92fce957f331a57a54e3cdcaeefeba25ff Mon Sep 17 00:00:00 2001
From: Nadeem Vawda <nadeem.vawda@gmail.com>
Date: Sat, 14 May 2011 23:07:36 +0200
Subject: Fix memory leaks in zlib.compress() and .decompress().

Also, make sure that test_zlib tests decompress() for overly-large inputs.
---
 Lib/test/test_zlib.py |  1 +
 Modules/zlibmodule.c  | 25 +++++++++++--------------
 2 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/Lib/test/test_zlib.py b/Lib/test/test_zlib.py
index 852857b..48d9f58 100644
--- a/Lib/test/test_zlib.py
+++ b/Lib/test/test_zlib.py
@@ -193,6 +193,7 @@ class CompressTestCase(BaseCompressTestCase, unittest.TestCase):
         data = b'x' * size
         try:
             self.assertRaises(OverflowError, zlib.compress, data, 1)
+            self.assertRaises(OverflowError, zlib.decompress, data)
         finally:
             data = None
 
diff --git a/Modules/zlibmodule.c b/Modules/zlibmodule.c
index 969db4a..fa07739 100644
--- a/Modules/zlibmodule.c
+++ b/Modules/zlibmodule.c
@@ -116,7 +116,7 @@ PyZlib_compress(PyObject *self, PyObject *args)
 {
     PyObject *ReturnVal = NULL;
     Py_buffer pinput;
-    Byte *input, *output;
+    Byte *input, *output = NULL;
     unsigned int length;
     int level=Z_DEFAULT_COMPRESSION, err;
     z_stream zst;
@@ -127,20 +127,19 @@ PyZlib_compress(PyObject *self, PyObject *args)
 
     if (pinput.len > UINT_MAX) {
         PyErr_SetString(PyExc_OverflowError,
-            "size does not fit in an unsigned int");
-        return NULL;
+                        "Size does not fit in an unsigned int");
+        goto error;
     }
-    length = pinput.len;
     input = pinput.buf;
+    length = pinput.len;
 
     zst.avail_out = length + length/1000 + 12 + 1;
 
     output = (Byte*)malloc(zst.avail_out);
     if (output == NULL) {
-        PyBuffer_Release(&pinput);
         PyErr_SetString(PyExc_MemoryError,
                         "Can't allocate memory to compress data");
-        return NULL;
+        goto error;
     }
 
     /* Past the point of no return.  From here on out, we need to make sure
@@ -203,7 +202,7 @@ PyDoc_STRVAR(decompress__doc__,
 static PyObject *
 PyZlib_decompress(PyObject *self, PyObject *args)
 {
-    PyObject *result_str;
+    PyObject *result_str = NULL;
     Py_buffer pinput;
     Byte *input;
     unsigned int length;
@@ -218,11 +217,11 @@ PyZlib_decompress(PyObject *self, PyObject *args)
 
     if (pinput.len > UINT_MAX) {
         PyErr_SetString(PyExc_OverflowError,
-            "size does not fit in an unsigned int");
-        return NULL;
+                        "Size does not fit in an unsigned int");
+        goto error;
     }
-    length = pinput.len;
     input = pinput.buf;
+    length = pinput.len;
 
     if (r_strlen <= 0)
         r_strlen = 1;
@@ -230,10 +229,8 @@ PyZlib_decompress(PyObject *self, PyObject *args)
     zst.avail_in = length;
     zst.avail_out = r_strlen;
 
-    if (!(result_str = PyBytes_FromStringAndSize(NULL, r_strlen))) {
-        PyBuffer_Release(&pinput);
-        return NULL;
-    }
+    if (!(result_str = PyBytes_FromStringAndSize(NULL, r_strlen)))
+        goto error;
 
     zst.zalloc = (alloc_func)NULL;
     zst.zfree = (free_func)Z_NULL;
-- 
cgit v0.12