From 0c3d96ae1dccc081fbc2a8c50ba6e41fbdc384de Mon Sep 17 00:00:00 2001
From: Nadeem Vawda <nadeem.vawda@gmail.com>
Date: Sun, 15 May 2011 00:19:50 +0200
Subject: [PATCH] Issue #8650: Make zlib.[de]compressobj().[de]compress()
 64-bit clean.

Raise an OverflowError if the input data is too large, instead of silently
truncating the input and returning an incorrect result.
---
 Lib/test/test_zlib.py | 11 +++++++++++
 Misc/NEWS             |  5 +++++
 Modules/zlibmodule.c  | 37 +++++++++++++++++++++++--------------
 3 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/Lib/test/test_zlib.py b/Lib/test/test_zlib.py
index 6bc386ed857..0909aa17b21 100644
--- a/Lib/test/test_zlib.py
+++ b/Lib/test/test_zlib.py
@@ -523,6 +523,17 @@ class CompressObjectTestCase(BaseCompressTestCase, unittest.TestCase):
         decompress = lambda s: d.decompress(s) + d.flush()
         self.check_big_decompress_buffer(size, decompress)
 
+    @precisionbigmemtest(size=_4G + 100, memuse=1)
+    def test_length_overflow(self, size):
+        if size < _4G + 100:
+            self.skipTest("not enough free memory, need at least 4 GB")
+        data = b'x' * size
+        try:
+            self.assertRaises(OverflowError, zlib.compress, data, 1)
+            self.assertRaises(OverflowError, zlib.decompress, data)
+        finally:
+            data = None
+
 
 def genblock(seed, length, step=1024, generator=random):
     """length-byte stream of random data from a seed (in step-byte blocks)."""
diff --git a/Misc/NEWS b/Misc/NEWS
index 06c1d623549..8618363e2dd 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -69,6 +69,11 @@ Core and Builtins
 Library
 -------
 
+- Issue #8650: Make zlib module 64-bit clean. compress(), decompress() and
+  their incremental counterparts now raise OverflowError if given an input
+  larger than 4GB, instead of silently truncating the input and returning
+  an incorrect result.
+
 - Issue #12050: zlib.decompressobj().decompress() now clears the unconsumed_tail
   attribute when called without a max_length argument.
 
diff --git a/Modules/zlibmodule.c b/Modules/zlibmodule.c
index fa07739a16f..ba0e59ce064 100644
--- a/Modules/zlibmodule.c
+++ b/Modules/zlibmodule.c
@@ -420,22 +420,26 @@ PyDoc_STRVAR(comp_compress__doc__,
 static PyObject *
 PyZlib_objcompress(compobject *self, PyObject *args)
 {
-    int err, inplen;
+    int err;
+    unsigned int inplen;
     Py_ssize_t length = DEFAULTALLOC;
-    PyObject *RetVal;
+    PyObject *RetVal = NULL;
     Py_buffer pinput;
     Byte *input;
     unsigned long start_total_out;
 
     if (!PyArg_ParseTuple(args, "y*:compress", &pinput))
         return NULL;
+    if (pinput.len > UINT_MAX) {
+        PyErr_SetString(PyExc_OverflowError,
+                        "Size does not fit in an unsigned int");
+        goto error_outer;
+    }
     input = pinput.buf;
     inplen = pinput.len;
 
-    if (!(RetVal = PyBytes_FromStringAndSize(NULL, length))) {
-        PyBuffer_Release(&pinput);
-        return NULL;
-    }
+    if (!(RetVal = PyBytes_FromStringAndSize(NULL, length)))
+        goto error_outer;
 
     ENTER_ZLIB(self);
 
@@ -484,6 +488,7 @@ PyZlib_objcompress(compobject *self, PyObject *args)
 
  error:
     LEAVE_ZLIB(self);
+ error_outer:
     PyBuffer_Release(&pinput);
     return RetVal;
 }
@@ -502,9 +507,10 @@ PyDoc_STRVAR(decomp_decompress__doc__,
 static PyObject *
 PyZlib_objdecompress(compobject *self, PyObject *args)
 {
-    int err, inplen, max_length = 0;
+    int err, max_length = 0;
+    unsigned int inplen;
     Py_ssize_t old_length, length = DEFAULTALLOC;
-    PyObject *RetVal;
+    PyObject *RetVal = NULL;
     Py_buffer pinput;
     Byte *input;
     unsigned long start_total_out;
@@ -512,22 +518,24 @@ PyZlib_objdecompress(compobject *self, PyObject *args)
     if (!PyArg_ParseTuple(args, "y*|i:decompress", &pinput,
                           &max_length))
         return NULL;
+    if (pinput.len > UINT_MAX) {
+        PyErr_SetString(PyExc_OverflowError,
+                        "Size does not fit in an unsigned int");
+        goto error_outer;
+    }
     input = pinput.buf;
     inplen = pinput.len;
     if (max_length < 0) {
-        PyBuffer_Release(&pinput);
         PyErr_SetString(PyExc_ValueError,
                         "max_length must be greater than zero");
-        return NULL;
+        goto error_outer;
     }
 
     /* limit amount of data allocated to max_length */
     if (max_length && length > max_length)
         length = max_length;
-    if (!(RetVal = PyBytes_FromStringAndSize(NULL, length))) {
-        PyBuffer_Release(&pinput);
-        return NULL;
-    }
+    if (!(RetVal = PyBytes_FromStringAndSize(NULL, length)))
+        goto error_outer;
 
     ENTER_ZLIB(self);
 
@@ -621,6 +629,7 @@ PyZlib_objdecompress(compobject *self, PyObject *args)
 
  error:
     LEAVE_ZLIB(self);
+ error_outer:
     PyBuffer_Release(&pinput);
     return RetVal;
 }