Issue #11564: Avoid crashes when trying to pickle huge objects or containers

(more than 2**31 items).  Instead, in most cases, an OverflowError is raised.

3 files changed, 111 insertions, 9 deletions

diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py
index e4ab0dd..2b1fdd2 100644
--- a/Lib/test/pickletester.py
+++ b/Lib/test/pickletester.py
@@ -2,10 +2,14 @@ import io
 import unittest
 import pickle
 import pickletools
+import sys
 import copyreg
 from http.cookies import SimpleCookie
 
-from test.support import TestFailed, TESTFN, run_with_locale
+from test.support import (
+    TestFailed, TESTFN, run_with_locale,
+    _2G, _4G, precisionbigmemtest,
+    )
 
 from pickle import bytes_types
 
@@ -14,6 +18,8 @@ from pickle import bytes_types
 # kind of outer loop.
 protocols = range(pickle.HIGHEST_PROTOCOL + 1)
 
+character_size = 4 if sys.maxunicode > 0xFFFF else 2
+
 
 # Return True if opcode code appears in the pickle, else False.
 def opcode_in_pickle(code, pickle):
@@ -1098,6 +1104,100 @@ class AbstractPickleTests(unittest.TestCase):
         empty = self.loads(b'\x80\x03U\x00q\x00.', encoding='koi8-r')
         self.assertEqual(empty, '')
 
+    def check_negative_32b_binXXX(self, dumped):
+        if sys.maxsize > 2**32:
+            self.skipTest("test is only meaningful on 32-bit builds")
+        # XXX Pure Python pickle reads lengths as signed and passes
+        # them directly to read() (hence the EOFError)
+        with self.assertRaises((pickle.UnpicklingError, EOFError,
+                                ValueError, OverflowError)):
+            self.loads(dumped)
+
+    def test_negative_32b_binbytes(self):
+        # On 32-bit builds, a BINBYTES of 2**31 or more is refused
+        self.check_negative_32b_binXXX(b'\x80\x03B\xff\xff\xff\xffxyzq\x00.')
+
+    def test_negative_32b_binunicode(self):
+        # On 32-bit builds, a BINUNICODE of 2**31 or more is refused
+        self.check_negative_32b_binXXX(b'\x80\x03X\xff\xff\xff\xffxyzq\x00.')
+
+
+class BigmemPickleTests(unittest.TestCase):
+
+    # Binary protocols can serialize longs of up to 2GB-1
+
+    @precisionbigmemtest(size=_2G, memuse=1 + 1, dry_run=False)
+    def test_huge_long_32b(self, size):
+        data = 1 << (8 * size)
+        try:
+            for proto in protocols:
+                if proto < 2:
+                    continue
+                with self.assertRaises((ValueError, OverflowError)):
+                    self.dumps(data, protocol=proto)
+        finally:
+            data = None
+
+    # Protocol 3 can serialize up to 4GB-1 as a bytes object
+    # (older protocols don't have a dedicated opcode for bytes and are
+    # too inefficient)
+
+    @precisionbigmemtest(size=_2G, memuse=1 + 1, dry_run=False)
+    def test_huge_bytes_32b(self, size):
+        data = b"abcd" * (size // 4)
+        try:
+            for proto in protocols:
+                if proto < 3:
+                    continue
+                try:
+                    pickled = self.dumps(data, protocol=proto)
+                    self.assertTrue(b"abcd" in pickled[:15])
+                    self.assertTrue(b"abcd" in pickled[-15:])
+                finally:
+                    pickled = None
+        finally:
+            data = None
+
+    @precisionbigmemtest(size=_4G, memuse=1 + 1, dry_run=False)
+    def test_huge_bytes_64b(self, size):
+        data = b"a" * size
+        try:
+            for proto in protocols:
+                if proto < 3:
+                    continue
+                with self.assertRaises((ValueError, OverflowError)):
+                    self.dumps(data, protocol=proto)
+        finally:
+            data = None
+
+    # All protocols use 1-byte per printable ASCII character; we add another
+    # byte because the encoded form has to be copied into the internal buffer.
+
+    @precisionbigmemtest(size=_2G, memuse=2 + character_size, dry_run=False)
+    def test_huge_str_32b(self, size):
+        data = "abcd" * (size // 4)
+        try:
+            for proto in protocols:
+                try:
+                    pickled = self.dumps(data, protocol=proto)
+                    self.assertTrue(b"abcd" in pickled[:15])
+                    self.assertTrue(b"abcd" in pickled[-15:])
+                finally:
+                    pickled = None
+        finally:
+            data = None
+
+    @precisionbigmemtest(size=_4G, memuse=1 + character_size, dry_run=False)
+    def test_huge_str_64b(self, size):
+        data = "a" * size
+        try:
+            for proto in protocols:
+                with self.assertRaises((ValueError, OverflowError)):
+                    self.dumps(data, protocol=proto)
+        finally:
+            data = None
+
+
 # Test classes for reduce_ex
 
 class REX_one(object):
diff --git a/Lib/test/support.py b/Lib/test/support.py
index e82a4f9..8e6ca2a 100644
--- a/Lib/test/support.py
+++ b/Lib/test/support.py
@@ -1089,7 +1089,7 @@ def bigmemtest(minsize, memuse):
         return wrapper
     return decorator
 
-def precisionbigmemtest(size, memuse):
+def precisionbigmemtest(size, memuse, dry_run=True):
     def decorator(f):
         def wrapper(self):
             size = wrapper.size
@@ -1099,10 +1099,11 @@ def precisionbigmemtest(size, memuse):
             else:
                 maxsize = size
 
-                if real_max_memuse and real_max_memuse < maxsize * memuse:
-                    raise unittest.SkipTest(
-                        "not enough memory: %.1fG minimum needed"
-                        % (size * memuse / (1024 ** 3)))
+            if ((real_max_memuse or not dry_run)
+                and real_max_memuse < maxsize * memuse):
+                raise unittest.SkipTest(
+                    "not enough memory: %.1fG minimum needed"
+                    % (size * memuse / (1024 ** 3)))
 
             return f(self, maxsize)
         wrapper.size = size
diff --git a/Lib/test/test_pickle.py b/Lib/test/test_pickle.py
index a3878fd..9da2cae 100644
--- a/Lib/test/test_pickle.py
+++ b/Lib/test/test_pickle.py
@@ -7,6 +7,7 @@ from test.pickletester import AbstractPickleTests
 from test.pickletester import AbstractPickleModuleTests
 from test.pickletester import AbstractPersistentPicklerTests
 from test.pickletester import AbstractPicklerUnpicklerObjectTests
+from test.pickletester import BigmemPickleTests
 
 try:
     import _pickle
@@ -37,13 +38,13 @@ class PyPicklerTests(AbstractPickleTests):
         return u.load()
 
 
-class InMemoryPickleTests(AbstractPickleTests):
+class InMemoryPickleTests(AbstractPickleTests, BigmemPickleTests):
 
     pickler = pickle._Pickler
     unpickler = pickle._Unpickler
 
-    def dumps(self, arg, proto=None):
-        return pickle.dumps(arg, proto)
+    def dumps(self, arg, protocol=None):
+        return pickle.dumps(arg, protocol)
 
     def loads(self, buf, **kwds):
         return pickle.loads(buf, **kwds)