Fix 64-bit safety issue in BZ2Compressor and BZ2Decompressor.

1 files changed, 35 insertions, 1 deletions

diff --git a/Lib/test/test_bz2.py b/Lib/test/test_bz2.py
index cee38e0..3567b36 100644
--- a/Lib/test/test_bz2.py
+++ b/Lib/test/test_bz2.py
@@ -1,10 +1,11 @@
 #!/usr/bin/env python3
 from test import support
-from test.support import TESTFN
+from test.support import TESTFN, precisionbigmemtest, _4G
 
 import unittest
 from io import BytesIO
 import os
+import random
 import subprocess
 import sys
 
@@ -415,6 +416,23 @@ class BZ2CompressorTest(BaseTest):
         data += bz2c.flush()
         self.assertEqual(self.decompress(data), self.TEXT)
 
+    @precisionbigmemtest(size=_4G + 100, memuse=2)
+    def testCompress4G(self, size):
+        # "Test BZ2Compressor.compress()/flush() with >4GiB input"
+        bz2c = BZ2Compressor()
+        data = b"x" * size
+        try:
+            compressed = bz2c.compress(data)
+            compressed += bz2c.flush()
+        finally:
+            data = None  # Release memory
+        data = bz2.decompress(compressed)
+        try:
+            self.assertEqual(len(data), size)
+            self.assertEqual(len(data.strip(b"x")), 0)
+        finally:
+            data = None
+
 class BZ2DecompressorTest(BaseTest):
     def test_Constructor(self):
         self.assertRaises(TypeError, BZ2Decompressor, 42)
@@ -453,6 +471,22 @@ class BZ2DecompressorTest(BaseTest):
         text = bz2d.decompress(self.DATA)
         self.assertRaises(EOFError, bz2d.decompress, b"anything")
 
+    @precisionbigmemtest(size=_4G + 100, memuse=3)
+    def testDecompress4G(self, size):
+        # "Test BZ2Decompressor.decompress() with >4GiB input"
+        blocksize = 10 * 1024 * 1024
+        block = random.getrandbits(blocksize * 8).to_bytes(blocksize, 'little')
+        try:
+            data = block * (size // blocksize + 1)
+            compressed = bz2.compress(data)
+            bz2d = BZ2Decompressor()
+            decompressed = bz2d.decompress(compressed)
+            self.assertTrue(decompressed == data)
+        finally:
+            data = None
+            compressed = None
+            decompressed = None
+
 
 class FuncTest(BaseTest):
     "Test module functions"