Update PIllow's image size limits through an environment variable

arkindex_worker/image.py

@@ -2,6 +2,8 @@
 Helper methods to download and open IIIF images, and manage polygons.
 """
 
+import functools
+import os
 import re
 import tempfile
 from collections import namedtuple
@@ -47,6 +49,53 @@ IIIF_MAX = "max"
 IMAGE_RATIO = [1, 0.9, 0.85, 0.80, 0.75, 0.70, 0.60, 0.50, 0.40, 0.30]
 
 
+def update_pillow_image_size_limit(func):
+    """
+    Update Pillow Image size limit
+    """
+
+    @functools.wraps(func)
+    def wrapper(
+        *args,
+        max_image_pixels: str | int | None = os.getenv("ARKINDEX_MAX_IMAGE_PIXELS"),
+        **kwargs,
+    ):
+        """
+        Wrapper to update Pillow Image size limit and restore it at the end of the function.
+
+        :param *args: Positional arguments passed to the function.
+        :param max_image_pixels: Pillow Image size limit to use.
+        :param **kwargs: Keyword arguments passed to the function.
+        """
+        MAX_IMAGE_PIXELS = Image.MAX_IMAGE_PIXELS
+
+        # Override Pillow Image size limit
+        if max_image_pixels is not None:
+            max_image_pixels = int(max_image_pixels)
+            # Override Pillow limit for detecting decompression bombs, disabled if set to 0
+            if max_image_pixels == 0:
+                logger.warning(
+                    "Pillow Image size limit is completely disabled, make sure you trust the image source."
+                )
+                Image.MAX_IMAGE_PIXELS = None
+            else:
+                Image.MAX_IMAGE_PIXELS = max_image_pixels
+
+        try:
+            results = func(*args, **kwargs)
+        except:
+            # Restore initial Pillow Image size limit
+            Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
+            raise
+
+        # Restore initial Pillow Image size limit
+        Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
+        return results
+
+    return wrapper
+
+
+@update_pillow_image_size_limit
 def open_image(
     path: str,
     mode: str | None = "RGB",
@@ -155,16 +204,20 @@ def upload_image(image: Image, url: str) -> requests.Response:
 
 
 def resized_images(
+    *args,
     element: "Element",
     max_pixels: int | None = None,
     max_bytes: int | None = None,
+    **kwargs,
 ) -> Iterator[Generator[tempfile.NamedTemporaryFile, None, None]]:
     """
     Build resized images according to the pixel and byte limits.
 
+    :param *args: Positional arguments passed to [arkindex_worker.models.Element.open_image_tempfile][].
     :param element: Element whose image needs to be resized.
     :param max_pixels: Maximum pixel size of the resized images.
     :param max_bytes: Maximum byte size of the resized images.
+    :param **kwargs: Keyword arguments passed to [arkindex_worker.models.Element.open_image_tempfile][].
     :returns: An iterator of the temporary file of the resized image.
     """
     _, _, element_width, element_height = polygon_bounding_box(element.polygon)
@@ -189,7 +242,9 @@ def resized_images(
         ),
         reverse=True,
     ):
-        with element.open_image_tempfile(**{param: resized_pixel}) as image:
+        with element.open_image_tempfile(
+            *args, **{**kwargs, param: resized_pixel}
+        ) as image:
             pillow_image = Image.open(image)
             if (
                 pillow_image.width != element_width

tests/test_image.py

@@ -22,6 +22,7 @@ from arkindex_worker.image import (
     resized_images,
     revert_orientation,
     trim_polygon,
+    update_pillow_image_size_limit,
     upload_image,
 )
 from arkindex_worker.models import Element
@@ -76,6 +77,30 @@ def _root_mean_square(img_a, img_b):
     )
 
 
+@pytest.mark.parametrize(
+    ("max_image_pixels", "expected_image_pixels"),
+    [
+        # Pillow Image size limit not updated
+        (None, Image.MAX_IMAGE_PIXELS),
+        # Pillow Image size limit set to None
+        ("0", None),
+        (0, None),
+        # Update Pillow Image size limit
+        ("1", 1),
+        (1, 1),
+    ],
+)
+def test_update_pillow_image_size_limit(max_image_pixels, expected_image_pixels):
+    MAX_IMAGE_PIXELS = Image.MAX_IMAGE_PIXELS
+
+    @update_pillow_image_size_limit
+    def function() -> int | None:
+        return Image.MAX_IMAGE_PIXELS
+
+    assert function(max_image_pixels=max_image_pixels) == expected_image_pixels
+    assert Image.MAX_IMAGE_PIXELS == MAX_IMAGE_PIXELS
+
+
 def test_download_tiles(responses):
     expected = Image.open(FULL_IMAGE).convert("RGB")
     tile_bytes = TILE.read_bytes()
@@ -792,7 +817,9 @@ def test_resized_images(
 
     assert [
         Image.open(image).size
-        for image in resized_images(mock_page, max_pixels, max_bytes)
+        for image in resized_images(
+            element=mock_page, max_pixels=max_pixels, max_bytes=max_bytes
+        )
     ] == expected_sizes
     assert [
         (record.levelno, record.message) for record in caplog.records