From 8848c7a37f929b471267bd893f91c3b818fafce0 Mon Sep 17 00:00:00 2001
From: Victor Stinner <victor.stinner@haypocalc.com>
Date: Tue, 4 Jan 2011 02:07:36 +0000
Subject: [PATCH] Issue #8650: zlib.compress() and zlib.decompress() raise an
 OverflowError if the input buffer length doesn't fit into an unsigned int
 (length bigger than 2^32-1 bytes).

---
 Lib/test/test_zlib.py | 12 +++++++++++-
 Misc/NEWS             |  4 ++++
 Modules/zlibmodule.c  | 22 ++++++++++++++++++----
 3 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/Lib/test/test_zlib.py b/Lib/test/test_zlib.py
index 0f5a1cac88b..5615c2d2396 100644
--- a/Lib/test/test_zlib.py
+++ b/Lib/test/test_zlib.py
@@ -2,7 +2,7 @@ import unittest
 from test import support
 import binascii
 import random
-from test.support import precisionbigmemtest, _1G
+from test.support import precisionbigmemtest, _1G, _4G
 
 zlib = support.import_module('zlib')
 
@@ -158,6 +158,16 @@ class CompressTestCase(BaseCompressTestCase, unittest.TestCase):
     def test_big_decompress_buffer(self, size):
         self.check_big_decompress_buffer(size, zlib.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)
+        finally:
+            data = None
+
 
 class CompressObjectTestCase(BaseCompressTestCase, unittest.TestCase):
     # Test compression object
diff --git a/Misc/NEWS b/Misc/NEWS
index d7610b3db72..1b77950249b 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -30,6 +30,10 @@ Core and Builtins
 Library
 -------
 
+- Issue #8650: zlib.compress() and zlib.decompress() raise an OverflowError if
+  the input buffer length doesn't fit into an unsigned int (length bigger than
+  2^32-1 bytes).
+
 - Issue #6643: Reinitialize locks held within the threading module after fork
   to avoid a potential rare deadlock or crash on some platforms.
 
diff --git a/Modules/zlibmodule.c b/Modules/zlibmodule.c
index 54ab9a14995..a03045ed169 100644
--- a/Modules/zlibmodule.c
+++ b/Modules/zlibmodule.c
@@ -117,14 +117,21 @@ PyZlib_compress(PyObject *self, PyObject *args)
     PyObject *ReturnVal = NULL;
     Py_buffer pinput;
     Byte *input, *output;
-    int length, level=Z_DEFAULT_COMPRESSION, err;
+    unsigned int length;
+    int level=Z_DEFAULT_COMPRESSION, err;
     z_stream zst;
 
     /* require Python string object, optional 'level' arg */
     if (!PyArg_ParseTuple(args, "y*|i:compress", &pinput, &level))
         return NULL;
-    input = pinput.buf;
+
+    if (pinput.len > UINT_MAX) {
+        PyErr_SetString(PyExc_OverflowError,
+            "size does not fit in an unsigned int");
+        return NULL;
+    }
     length = pinput.len;
+    input = pinput.buf;
 
     zst.avail_out = length + length/1000 + 12 + 1;
 
@@ -199,7 +206,8 @@ PyZlib_decompress(PyObject *self, PyObject *args)
     PyObject *result_str;
     Py_buffer pinput;
     Byte *input;
-    int length, err;
+    unsigned int length;
+    int err;
     int wsize=DEF_WBITS;
     Py_ssize_t r_strlen=DEFAULTALLOC;
     z_stream zst;
@@ -207,8 +215,14 @@ PyZlib_decompress(PyObject *self, PyObject *args)
     if (!PyArg_ParseTuple(args, "y*|in:decompress",
                           &pinput, &wsize, &r_strlen))
         return NULL;
-    input = pinput.buf;
+
+    if (pinput.len > UINT_MAX) {
+        PyErr_SetString(PyExc_OverflowError,
+            "size does not fit in an unsigned int");
+        return NULL;
+    }
     length = pinput.len;
+    input = pinput.buf;
 
     if (r_strlen <= 0)
         r_strlen = 1;