120 lines
5.0 KiB
Python
120 lines
5.0 KiB
Python
"""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)
|