Fix memory leaks in zlib.compress() and .decompress().

Also, make sure that test_zlib tests decompress() for overly-large inputs.
This commit is contained in:
Nadeem Vawda 2011-05-14 23:07:36 +02:00
parent d8eab60c00
commit 154bdf92fc
2 changed files with 12 additions and 14 deletions

View File

@ -193,6 +193,7 @@ class CompressTestCase(BaseCompressTestCase, unittest.TestCase):
data = b'x' * size data = b'x' * size
try: try:
self.assertRaises(OverflowError, zlib.compress, data, 1) self.assertRaises(OverflowError, zlib.compress, data, 1)
self.assertRaises(OverflowError, zlib.decompress, data)
finally: finally:
data = None data = None

View File

@ -116,7 +116,7 @@ PyZlib_compress(PyObject *self, PyObject *args)
{ {
PyObject *ReturnVal = NULL; PyObject *ReturnVal = NULL;
Py_buffer pinput; Py_buffer pinput;
Byte *input, *output; Byte *input, *output = NULL;
unsigned int length; unsigned int length;
int level=Z_DEFAULT_COMPRESSION, err; int level=Z_DEFAULT_COMPRESSION, err;
z_stream zst; z_stream zst;
@ -127,20 +127,19 @@ PyZlib_compress(PyObject *self, PyObject *args)
if (pinput.len > UINT_MAX) { if (pinput.len > UINT_MAX) {
PyErr_SetString(PyExc_OverflowError, PyErr_SetString(PyExc_OverflowError,
"size does not fit in an unsigned int"); "Size does not fit in an unsigned int");
return NULL; goto error;
} }
length = pinput.len;
input = pinput.buf; input = pinput.buf;
length = pinput.len;
zst.avail_out = length + length/1000 + 12 + 1; zst.avail_out = length + length/1000 + 12 + 1;
output = (Byte*)malloc(zst.avail_out); output = (Byte*)malloc(zst.avail_out);
if (output == NULL) { if (output == NULL) {
PyBuffer_Release(&pinput);
PyErr_SetString(PyExc_MemoryError, PyErr_SetString(PyExc_MemoryError,
"Can't allocate memory to compress data"); "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 /* 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 * static PyObject *
PyZlib_decompress(PyObject *self, PyObject *args) PyZlib_decompress(PyObject *self, PyObject *args)
{ {
PyObject *result_str; PyObject *result_str = NULL;
Py_buffer pinput; Py_buffer pinput;
Byte *input; Byte *input;
unsigned int length; unsigned int length;
@ -218,11 +217,11 @@ PyZlib_decompress(PyObject *self, PyObject *args)
if (pinput.len > UINT_MAX) { if (pinput.len > UINT_MAX) {
PyErr_SetString(PyExc_OverflowError, PyErr_SetString(PyExc_OverflowError,
"size does not fit in an unsigned int"); "Size does not fit in an unsigned int");
return NULL; goto error;
} }
length = pinput.len;
input = pinput.buf; input = pinput.buf;
length = pinput.len;
if (r_strlen <= 0) if (r_strlen <= 0)
r_strlen = 1; r_strlen = 1;
@ -230,10 +229,8 @@ PyZlib_decompress(PyObject *self, PyObject *args)
zst.avail_in = length; zst.avail_in = length;
zst.avail_out = r_strlen; zst.avail_out = r_strlen;
if (!(result_str = PyBytes_FromStringAndSize(NULL, r_strlen))) { if (!(result_str = PyBytes_FromStringAndSize(NULL, r_strlen)))
PyBuffer_Release(&pinput); goto error;
return NULL;
}
zst.zalloc = (alloc_func)NULL; zst.zalloc = (alloc_func)NULL;
zst.zfree = (free_func)Z_NULL; zst.zfree = (free_func)Z_NULL;