From 3843cd8e86673ce8969c6e948cd15735168882d3 Mon Sep 17 00:00:00 2001
From: Antoine Pitrou <solipsis@pitrou.net>
Date: Fri, 7 May 2010 16:50:34 +0000
Subject: [PATCH] Issue #8571: Fix an internal error when compressing or
 decompressing a chunk larger than 1GB with the zlib module's compressor and
 decompressor objects.

---
 Lib/test/test_zlib.py | 62 ++++++++++++++++++++++++++++++++++++++++---
 Misc/NEWS             |  4 +++
 Modules/zlibmodule.c  |  7 ++---
 3 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/Lib/test/test_zlib.py b/Lib/test/test_zlib.py
index 6ffed99293d..cd021750d7e 100644
--- a/Lib/test/test_zlib.py
+++ b/Lib/test/test_zlib.py
@@ -2,6 +2,7 @@ import unittest
 from test import test_support
 import binascii
 import random
+from test.test_support import precisionbigmemtest, _1G
 
 zlib = test_support.import_module('zlib')
 
@@ -90,8 +91,39 @@ class ExceptionTestCase(unittest.TestCase):
         self.assertRaises(ValueError, zlib.decompressobj().flush, -1)
 
 
+class BaseCompressTestCase(object):
+    def check_big_compress_buffer(self, size, compress_func):
+        _1M = 1024 * 1024
+        fmt = "%%0%dx" % (2 * _1M)
+        # Generate 10MB worth of random, and expand it by repeating it.
+        # The assumption is that zlib's memory is not big enough to exploit
+        # such spread out redundancy.
+        data = ''.join([binascii.a2b_hex(fmt % random.getrandbits(8 * _1M))
+                        for i in range(10)])
+        data = data * (size // len(data) + 1)
+        try:
+            compress_func(data)
+        finally:
+            # Release memory
+            data = None
 
-class CompressTestCase(unittest.TestCase):
+    def check_big_decompress_buffer(self, size, decompress_func):
+        data = 'x' * size
+        try:
+            compressed = zlib.compress(data, 1)
+        finally:
+            # Release memory
+            data = None
+        data = decompress_func(compressed)
+        # Sanity check
+        try:
+            self.assertEqual(len(data), size)
+            self.assertEqual(len(data.strip('x')), 0)
+        finally:
+            data = None
+
+
+class CompressTestCase(BaseCompressTestCase, unittest.TestCase):
     # Test compression in one go (whole message compression)
     def test_speech(self):
         x = zlib.compress(HAMLET_SCENE)
@@ -103,10 +135,19 @@ class CompressTestCase(unittest.TestCase):
         x = zlib.compress(data)
         self.assertEqual(zlib.decompress(x), data)
 
+    # Memory use of the following functions takes into account overallocation
+
+    @precisionbigmemtest(size=_1G + 1024 * 1024, memuse=3)
+    def test_big_compress_buffer(self, size):
+        compress = lambda s: zlib.compress(s, 1)
+        self.check_big_compress_buffer(size, compress)
+
+    @precisionbigmemtest(size=_1G + 1024 * 1024, memuse=2)
+    def test_big_decompress_buffer(self, size):
+        self.check_big_decompress_buffer(size, zlib.decompress)
 
 
-
-class CompressObjectTestCase(unittest.TestCase):
+class CompressObjectTestCase(BaseCompressTestCase, unittest.TestCase):
     # Test compression object
     def test_pair(self):
         # straightforward compress/decompress objects
@@ -380,6 +421,21 @@ class CompressObjectTestCase(unittest.TestCase):
             d.flush()
             self.assertRaises(ValueError, d.copy)
 
+    # Memory use of the following functions takes into account overallocation
+
+    @precisionbigmemtest(size=_1G + 1024 * 1024, memuse=3)
+    def test_big_compress_buffer(self, size):
+        c = zlib.compressobj(1)
+        compress = lambda s: c.compress(s) + c.flush()
+        self.check_big_compress_buffer(size, compress)
+
+    @precisionbigmemtest(size=_1G + 1024 * 1024, memuse=2)
+    def test_big_decompress_buffer(self, size):
+        d = zlib.decompressobj()
+        decompress = lambda s: d.decompress(s) + d.flush()
+        self.check_big_decompress_buffer(size, decompress)
+
+
 def genblock(seed, length, step=1024, generator=random):
     """length-byte stream of random data from a seed (in step-byte blocks)."""
     if seed is not None:
diff --git a/Misc/NEWS b/Misc/NEWS
index 6d1423bfcd8..a29b76700d8 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -41,6 +41,10 @@ Core and Builtins
 Library
 -------
 
+- Issue #8571: Fix an internal error when compressing or decompressing a
+  chunk larger than 1GB with the zlib module's compressor and decompressor
+  objects.
+
 - Issue #8573: asyncore _strerror() function might throw ValueError.
 
 - Issue #8483: asyncore.dispatcher's __getattr__ method produced confusing 
diff --git a/Modules/zlibmodule.c b/Modules/zlibmodule.c
index 4f78dbcb15d..0cb1e64e03c 100644
--- a/Modules/zlibmodule.c
+++ b/Modules/zlibmodule.c
@@ -392,7 +392,8 @@ PyDoc_STRVAR(comp_compress__doc__,
 static PyObject *
 PyZlib_objcompress(compobject *self, PyObject *args)
 {
-    int err, inplen, length = DEFAULTALLOC;
+    int err, inplen;
+    Py_ssize_t length = DEFAULTALLOC;
     PyObject *RetVal;
     Byte *input;
     unsigned long start_total_out;
@@ -461,8 +462,8 @@ PyDoc_STRVAR(decomp_decompress__doc__,
 static PyObject *
 PyZlib_objdecompress(compobject *self, PyObject *args)
 {
-    int err, inplen, old_length, length = DEFAULTALLOC;
-    int max_length = 0;
+    int err, inplen, max_length = 0;
+    Py_ssize_t old_length, length = DEFAULTALLOC;
     PyObject *RetVal;
     Byte *input;
     unsigned long start_total_out;