Decompression bomb protection

PIL/Image.py

@@ -31,11 +31,18 @@
 import warnings
 
 
+class DecompressionBombWarning(RuntimeWarning):
+    pass
+
 class _imaging_not_installed:
     # module placeholder
     def __getattr__(self, id):
         raise ImportError("The _imaging C module is not installed")
 
+
+# Limit to around a quarter gigabyte for a 24 bit (3 bpp) image
+MAX_IMAGE_PIXELS = int(1024 * 1024 * 1024 / 4 / 3)
+
 try:
     # give Tk a chance to set up the environment, in case we're
     # using an _imaging module linked against libtcl/libtk (use
@@ -2172,6 +2179,20 @@ def fromarray(obj, mode=None):
 _fromarray_typemap[((1, 1), _ENDIAN + "f4")] = ("F", "F")
 
 
+def _decompression_bomb_check(size):
+    if MAX_IMAGE_PIXELS is None:
+        return
+
+    pixels = size[0] * size[1]
+
+    if pixels > MAX_IMAGE_PIXELS:
+        warnings.warn(
+            "Image size (%d pixels) exceeds limit of %d pixels, "
+            "could be decompression bomb DOS attack." %
+            (pixels, MAX_IMAGE_PIXELS),
+            DecompressionBombWarning)
+
+
 def open(fp, mode="r"):
     """
     Opens and identifies the given image file.
@@ -2209,7 +2230,9 @@ def open(fp, mode="r"):
             factory, accept = OPEN[i]
             if not accept or accept(prefix):
                 fp.seek(0)
-                return factory(fp, filename)
+                im = factory(fp, filename)
+                _decompression_bomb_check(im.size)
+                return im
         except (SyntaxError, IndexError, TypeError):
             # import traceback
             # traceback.print_exc()
@@ -2222,7 +2245,9 @@ def open(fp, mode="r"):
                 factory, accept = OPEN[i]
                 if not accept or accept(prefix):
                     fp.seek(0)
-                    return factory(fp, filename)
+                    im = factory(fp, filename)
+                    _decompression_bomb_check(im.size)
+                    return im
             except (SyntaxError, IndexError, TypeError):
                 # import traceback
                 # traceback.print_exc()

Tests/test_decompression_bomb.py

@@ -0,0 +1,45 @@
+from helper import unittest, PillowTestCase, tearDownModule
+
+from PIL import Image
+
+test_file = "Images/lena.ppm"
+
+ORIGINAL_LIMIT = Image.MAX_IMAGE_PIXELS
+
+
+class TestDecompressionBomb(PillowTestCase):
+
+    def tearDown(self):
+        Image.MAX_IMAGE_PIXELS = ORIGINAL_LIMIT
+
+    def test_no_warning_small_file(self):
+        # Implicit assert: no warning.
+        # A warning would cause a failure.
+        Image.open(test_file)
+
+    def test_no_warning_no_limit(self):
+        # Arrange
+        # Turn limit off
+        Image.MAX_IMAGE_PIXELS = None
+        self.assertEqual(Image.MAX_IMAGE_PIXELS, None)
+
+        # Act / Assert
+        # Implicit assert: no warning.
+        # A warning would cause a failure.
+        Image.open(test_file)
+
+    def test_warning(self):
+        # Arrange
+        # Set limit to a low, easily testable value
+        Image.MAX_IMAGE_PIXELS = 10
+        self.assertEqual(Image.MAX_IMAGE_PIXELS, 10)
+
+        # Act / Assert
+        self.assert_warning(
+            Image.DecompressionBombWarning,
+            lambda: Image.open(test_file))
+
+if __name__ == '__main__':
+    unittest.main()
+
+# End of file