1 files changed, 54 insertions, 44 deletions

diff --git a/Modules/zlibmodule.c b/Modules/zlibmodule.c
index 3e9fdf4..9e1c2ae 100644
--- a/Modules/zlibmodule.c
+++ b/Modules/zlibmodule.c
@@ -561,6 +561,49 @@ PyZlib_objcompress(compobject *self, PyObject *args)
     return RetVal;
 }
 
+/* Helper for objdecompress() and unflush(). Saves any unconsumed input data in
+   self->unused_data or self->unconsumed_tail, as appropriate. */
+static int
+save_unconsumed_input(compobject *self, int err)
+{
+    if (err == Z_STREAM_END) {
+        /* The end of the compressed data has been reached. Store the leftover
+           input data in self->unused_data. */
+        if (self->zst.avail_in > 0) {
+            Py_ssize_t old_size = PyBytes_GET_SIZE(self->unused_data);
+            Py_ssize_t new_size;
+            PyObject *new_data;
+            if (self->zst.avail_in > PY_SSIZE_T_MAX - old_size) {
+                PyErr_NoMemory();
+                return -1;
+            }
+            new_size = old_size + self->zst.avail_in;
+            new_data = PyBytes_FromStringAndSize(NULL, new_size);
+            if (new_data == NULL)
+                return -1;
+            Py_MEMCPY(PyBytes_AS_STRING(new_data),
+                      PyBytes_AS_STRING(self->unused_data), old_size);
+            Py_MEMCPY(PyBytes_AS_STRING(new_data) + old_size,
+                      self->zst.next_in, self->zst.avail_in);
+            Py_DECREF(self->unused_data);
+            self->unused_data = new_data;
+            self->zst.avail_in = 0;
+        }
+    }
+    if (self->zst.avail_in > 0 || PyBytes_GET_SIZE(self->unconsumed_tail)) {
+        /* This code handles two distinct cases:
+           1. Output limit was reached. Save leftover input in unconsumed_tail.
+           2. All input data was consumed. Clear unconsumed_tail. */
+        PyObject *new_data = PyBytes_FromStringAndSize(
+                (char *)self->zst.next_in, self->zst.avail_in);
+        if (new_data == NULL)
+            return -1;
+        Py_DECREF(self->unconsumed_tail);
+        self->unconsumed_tail = new_data;
+    }
+    return 0;
+}
+
 PyDoc_STRVAR(decomp_decompress__doc__,
 "decompress(data, max_length) -- Return a string containing the decompressed\n"
 "version of the data.\n"
@@ -668,61 +711,21 @@ PyZlib_objdecompress(compobject *self, PyObject *args)
         Py_END_ALLOW_THREADS
     }
 
-    if(max_length) {
-        /* Not all of the compressed data could be accommodated in a buffer of
-           the specified size. Return the unconsumed tail in an attribute. */
-        Py_DECREF(self->unconsumed_tail);
-        self->unconsumed_tail = PyBytes_FromStringAndSize((char *)self->zst.next_in,
-                                                           self->zst.avail_in);
-    }
-    else if (PyBytes_GET_SIZE(self->unconsumed_tail) > 0) {
-        /* All of the compressed data was consumed. Clear unconsumed_tail. */
-        Py_DECREF(self->unconsumed_tail);
-        self->unconsumed_tail = PyBytes_FromStringAndSize("", 0);
-    }
-    if (self->unconsumed_tail == NULL) {
+    if (save_unconsumed_input(self, err) < 0) {
         Py_DECREF(RetVal);
         RetVal = NULL;
         goto error;
     }
 
-    /* The end of the compressed data has been reached, so set the
-       unused_data attribute to a string containing the remainder of the
-       data in the string.  Note that this is also a logical place to call
-       inflateEnd, but the old behaviour of only calling it on flush() is
-       preserved.
-    */
     if (err == Z_STREAM_END) {
-        if (self->zst.avail_in > 0) {
-            /* Append the leftover data to the existing value of unused_data. */
-            Py_ssize_t old_size = PyBytes_GET_SIZE(self->unused_data);
-            Py_ssize_t new_size = old_size + self->zst.avail_in;
-            PyObject *new_data;
-            if (new_size <= old_size) {  /* Check for overflow. */
-                PyErr_NoMemory();
-                Py_DECREF(RetVal);
-                RetVal = NULL;
-                goto error;
-            }
-            new_data = PyBytes_FromStringAndSize(NULL, new_size);
-            if (new_data == NULL) {
-                Py_DECREF(RetVal);
-                RetVal = NULL;
-                goto error;
-            }
-            Py_MEMCPY(PyBytes_AS_STRING(new_data),
-                      PyBytes_AS_STRING(self->unused_data), old_size);
-            Py_MEMCPY(PyBytes_AS_STRING(new_data) + old_size,
-                      self->zst.next_in, self->zst.avail_in);
-            Py_DECREF(self->unused_data);
-            self->unused_data = new_data;
-        }
+        /* This is the logical place to call inflateEnd, but the old behaviour
+           of only calling it on flush() is preserved. */
         self->eof = 1;
+    } else if (err != Z_OK && err != Z_BUF_ERROR) {
         /* We will only get Z_BUF_ERROR if the output buffer was full
            but there wasn't more output when we tried again, so it is
            not an error condition.
         */
-    } else if (err != Z_OK && err != Z_BUF_ERROR) {
         zlib_error(self->zst, err, "while decompressing data");
         Py_DECREF(RetVal);
         RetVal = NULL;
@@ -996,6 +999,12 @@ PyZlib_unflush(compobject *self, PyObject *args)
         Py_END_ALLOW_THREADS
     }
 
+    if (save_unconsumed_input(self, err) < 0) {
+        Py_DECREF(retval);
+        retval = NULL;
+        goto error;
+    }
+
     /* If at end of stream, clean up any memory allocated by zlib. */
     if (err == Z_STREAM_END) {
         self->eof = 1;
@@ -1008,6 +1017,7 @@ PyZlib_unflush(compobject *self, PyObject *args)
             goto error;
         }
     }
+
     if (_PyBytes_Resize(&retval, self->zst.total_out - start_total_out) < 0) {
         Py_DECREF(retval);
         retval = NULL;