US07-03: Harden Media and Metadata Edge Cases

photo_pipeline/imaging.py

@@ -0,0 +1,119 @@
+"""Bounded, defensive image decoding — one door for every pixel this app reads.
+
+A photo library contains files nobody planned for: truncated downloads, zero-byte
+placeholders, a PNG whose header claims 200000×200000, a TIFF with a broken ICC
+profile, an extension that lies about its content. None of them may take down a
+request or a worker, and none may decode more pixels than the caller allowed
+(concept §17: decoded pixels, not file size, are what exhausts memory).
+
+``open_image`` is that single door:
+
+* the declared dimensions are checked **before** a pixel is decoded;
+* Pillow's decompression-bomb *warning* is promoted to an error, because the
+  warning band (between Pillow's limit and twice it) still decodes the image;
+* every decoder failure — at open time or during the caller's decode — becomes one
+  of two typed errors, so callers map them to their own item state instead of
+  catching ``Exception``;
+* error text names no path: it reaches API responses, and the full reason goes to
+  the server log instead (US07-02).
+
+``to_srgb`` and ``draft`` are the other two bounded-decode helpers: colour-manage a
+profile-bearing image into sRGB, and let JPEG decode straight to a size near the
+requested one rather than at full resolution.
+"""
+
+from __future__ import annotations
+
+import io
+import logging
+import warnings
+from contextlib import contextmanager
+from pathlib import Path
+
+from PIL import Image, ImageCms, ImageFile, UnidentifiedImageError
+
+log = logging.getLogger(__name__)
+
+# Matches ``Config.thumbnail_max_pixels``; used where no configuration is at hand
+# (hashing runs inside discovery, which takes no config).
+DEFAULT_MAX_PIXELS = 100_000_000
+
+
+class MediaError(Exception):
+    """A file could not be turned into pixels safely."""
+
+
+class UndecodableImage(MediaError):
+    """Corrupt, truncated, empty, or not an image at all."""
+
+
+class ImageTooLarge(MediaError):
+    """More pixels than this operation is allowed to decode."""
+
+
+@contextmanager
+def open_image(path: Path | str, *, max_pixels: int = DEFAULT_MAX_PIXELS):
+    """Yield an open :class:`PIL.Image.Image`, bounded and with typed failures.
+
+    Decoder errors raised inside the ``with`` body are translated too — a truncated
+    JPEG only fails when its pixels are actually pulled, which is the caller's line,
+    not this one.
+    """
+    # Pillow's truncation tolerance is a process-global switch that any library in
+    # the process can flip (the donor CLI did). This door decides the policy for its
+    # own callers: half a file is not a picture.
+    tolerated = ImageFile.LOAD_TRUNCATED_IMAGES
+    ImageFile.LOAD_TRUNCATED_IMAGES = False
+    with warnings.catch_warnings():
+        # The warning band is not a warning for us: it means Pillow was willing to
+        # decode an image large enough to be a denial-of-service.
+        warnings.simplefilter("error", Image.DecompressionBombWarning)
+        try:
+            with Image.open(path) as image:
+                width, height = image.size
+                if width * height > max_pixels:
+                    raise ImageTooLarge(f"{width}x{height} exceeds the {max_pixels} pixel limit")
+                yield image
+        except MediaError:
+            raise
+        except (Image.DecompressionBombError, Image.DecompressionBombWarning) as error:
+            log.info("refused oversized image %s: %s", path, error)
+            raise ImageTooLarge("image exceeds the decompression-bomb limit") from error
+        except (UnidentifiedImageError, OSError, ValueError, SyntaxError, MemoryError) as error:
+            log.info("cannot decode %s: %s", path, error)
+            raise UndecodableImage(f"cannot decode image ({type(error).__name__})") from error
+        finally:
+            ImageFile.LOAD_TRUNCATED_IMAGES = tolerated
+
+
+def draft(image: Image.Image, size: int) -> None:
+    """Ask the decoder for a smaller image where the format allows it (JPEG).
+
+    This is the difference between decoding a 40-megapixel JPEG and decoding the
+    roughly 1-megapixel version a 1280px preview needs.
+    """
+    try:
+        image.draft(None, (size, size))
+    except (AttributeError, ValueError, OSError):  # not a draft-capable format
+        pass
+
+
+def to_srgb(image: Image.Image, *, mode: str) -> Image.Image:
+    """Convert into ``mode``, colour-managing through an embedded ICC profile.
+
+    Without this a wide-gamut original renders with visibly wrong colours, because
+    its numbers are interpreted as sRGB. A broken or unreadable profile is not a
+    reason to fail a preview — the plain conversion is still a correct picture.
+    """
+    profile = image.info.get("icc_profile")
+    if profile:
+        try:
+            return ImageCms.profileToProfile(
+                image,
+                ImageCms.ImageCmsProfile(io.BytesIO(profile)),
+                ImageCms.createProfile("sRGB"),
+                outputMode=mode,
+            )
+        except Exception as error:  # noqa: BLE001 - any ICC failure falls back
+            log.info("ignoring unusable ICC profile on %s: %s", getattr(image, "filename", "?"), error)
+    return image.convert(mode)