diff --git a/PIL/Image.py b/PIL/Image.py index e064ed9ef..b0b90a20b 100644 --- a/PIL/Image.py +++ b/PIL/Image.py @@ -31,11 +31,18 @@ from PIL import VERSION, PILLOW_VERSION, _plugins 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 @@ _fromarray_typemap[((1, 1), _ENDIAN + "i4")] = ("I", "I") _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() diff --git a/Tests/test_decompression_bomb.py b/Tests/test_decompression_bomb.py new file mode 100644 index 000000000..b83f33f46 --- /dev/null +++ b/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