From 6b223ba7c2abbe4b08df53bf78d55b5e13d0d482 Mon Sep 17 00:00:00 2001 From: Michele Dolfi Date: Thu, 22 Oct 2026 17:44:55 +0200 Subject: [PATCH] fix: decode image pages on demand and limit decoded image pixels The image backend decoded every frame of a multi-frame image when it was opened, before the page count was checked against max_num_pages, and kept all frames in memory. It now reads only the frame count when opened and decodes each frame in load_page, so page limits apply before any pixel data is decoded and memory holds only the pages in flight. A new max_image_pixels option (default 89,478,485, the default of Pillow's Image.MAX_IMAGE_PIXELS) limits the width x height of one decoded image: - the image backend checks each frame before it is decoded; an oversized first frame rejects the document, an oversized later frame fails that page; - the shared image resource loader (HTML, Markdown, AsciiDoc, JATS, ODF) and the MHTML resources check the image header before decoding, and skip an oversized image, or one Pillow refuses to open, with a warning instead of failing the document. The loader also checks the decoded size of a base64 data URI before it decodes it. Signed-off-by: Michele Dolfi --- .../skills/docling/references/python-sdk.md | 3 +- docling/backend/html_backend.py | 11 +- docling/backend/image_backend.py | 80 +++++++------ docling/backend/md_backend.py | 1 + .../backend/utils/image_resource_loader.py | 20 +++- docling/datamodel/backend_options.py | 27 +++++ docling/utils/image_utils.py | 14 +++ docs/usage/advanced_options.md | 25 +++- tests/test_backend_html.py | 30 ++++- tests/test_backend_image_native.py | 108 +++++++++++------- tests/test_backend_markdown_images.py | 16 +++ tests/test_image_resource_loader.py | 51 +++++++++ 12 files changed, 304 insertions(+), 82 deletions(-) diff --git a/docling/.agents/skills/docling/references/python-sdk.md b/docling/.agents/skills/docling/references/python-sdk.md index be17238dad..e44ae5fd95 100644 --- a/docling/.agents/skills/docling/references/python-sdk.md +++ b/docling/.agents/skills/docling/references/python-sdk.md @@ -465,7 +465,8 @@ Or save with `ImageRefMode.EMBEDDED` so the JSON carries its images. HTML image downloads are off by default. Enable them on the HTML backend options; downloads only connect to public addresses, re-check every redirect -and stop at `max_remote_image_bytes`. +and stop at `max_remote_image_bytes`. An image larger than `max_image_pixels` +(width x height, default 89,478,485) is skipped with a warning. ```python from docling.datamodel.backend_options import HTMLBackendOptions diff --git a/docling/backend/html_backend.py b/docling/backend/html_backend.py index da645055bc..9794a8718f 100644 --- a/docling/backend/html_backend.py +++ b/docling/backend/html_backend.py @@ -79,6 +79,7 @@ detect_code_language, normalize_code_language, ) +from docling.utils.image_utils import check_image_pixels _BS4_AVAILABLE: bool = False _BS4_IMPORT_ERROR: ImportError | None = None @@ -536,6 +537,7 @@ def __init__( enable_remote_fetch=options.enable_remote_fetch, max_image_data_base64_bytes=options.max_image_data_base64_bytes, max_remote_image_bytes=options.max_remote_image_bytes, + max_image_pixels=options.max_image_pixels, max_redirects=options.max_redirects, headers=options.headers, header_origins=self._get_header_origins(options, configured_base_path), @@ -5375,11 +5377,18 @@ def _create_image_ref(self, src_url: str) -> Optional[ImageRef]: return None try: image = Image.open(BytesIO(image_data)) + check_image_pixels(image.size, self.options.max_image_pixels) image.load() return ImageRef.from_pil( image, dpi=int(image.info.get("dpi", (72,))[0]) ) - except (UnidentifiedImageError, OSError, TypeError, ValueError) as exc: + except ( + Image.DecompressionBombError, + UnidentifiedImageError, + OSError, + TypeError, + ValueError, + ) as exc: warnings.warn( f"Could not process an embedded MHTML image from {src_url}: " f"{exc}" diff --git a/docling/backend/image_backend.py b/docling/backend/image_backend.py index 90da683f93..864691613d 100644 --- a/docling/backend/image_backend.py +++ b/docling/backend/image_backend.py @@ -3,6 +3,7 @@ import logging import math +import threading from collections.abc import Iterable from io import BytesIO from pathlib import Path @@ -24,7 +25,7 @@ from docling.datamodel.base_models import InputFormat from docling.datamodel.document import InputDocument from docling.exceptions import DocumentLoadError -from docling.utils.image_utils import to_rgb +from docling.utils.image_utils import check_image_pixels, to_rgb _log = logging.getLogger(__name__) @@ -200,8 +201,10 @@ class ImageDocumentBackend(PdfDocumentBackend): - Subclasses PdfDocumentBackend to satisfy pipeline type checks. - Intentionally avoids calling PdfDocumentBackend.__init__ to skip the image→PDF conversion and any pypdfium2 usage. - - Handles multi-page TIFF by extracting frames eagerly to separate - Image objects to keep thread-safety when pages process in parallel. + - Reads only the frame count when opened, so the page limits are + checked before any pixel data is decoded. Each page decodes its + frame on ``load_page`` into a separate Image object, which keeps + pages thread-safe when they process in parallel. """ def __init__( @@ -221,47 +224,52 @@ def __init__( f"Incompatible file format {self.input_format} was passed to ImageDocumentBackend." ) - # Load frames eagerly for thread-safety across pages - self._frames: list[Image.Image] = [] - self._frame_dpi: list[tuple[float, float]] = [] + # The source stays open until unload(); the lock serializes the seeks + # on it, because pages can load from several threads. + self._lock = threading.Lock() + self._source: Image.Image | None = None + self._frame_count = 0 try: - with Image.open(self.path_or_stream) as img: # type: ignore[arg-type] - # Handle multi-frame and single-frame images - # - multiframe formats: TIFF, GIF, ICO - # - singleframe formats: JPEG (.jpg, .jpeg), PNG (.png), BMP, WEBP (unless animated), HEIC - frame_count = getattr(img, "n_frames", 1) - - if frame_count > 1: - for i in range(frame_count): - img.seek(i) - frame, dpi = _oriented_frame(img) - self._frame_dpi.append(dpi) - self._frames.append(to_rgb(frame)) - else: - frame, dpi = _oriented_frame(img) - self._frame_dpi.append(dpi) - self._frames.append(to_rgb(frame)) + self._source = Image.open(self.path_or_stream) # type: ignore[arg-type] + # Multi-frame formats: TIFF, GIF, ICO, animated WEBP/PNG. Other + # formats have no n_frames attribute and hold one frame. + self._frame_count = getattr(self._source, "n_frames", 1) + # The first frame is the current one after open (and after + # n_frames, which seeks back): reject a bad first frame here, + # without decoding it. + check_image_pixels(self._source.size, self.options.max_image_pixels) + _get_frame_dpi(self._source) except Exception as e: - for frame in self._frames: - frame.close() - self._frames = [] - self._frame_dpi = [] + self._close_source() raise DocumentLoadError( - f"Could not load image for document {self.file}" + f"Could not load image for document {self.file}: {e}" ) from e + def _close_source(self) -> None: + if self._source is not None: + self._source.close() + self._source = None + self._frame_count = 0 + def is_valid(self) -> bool: - return len(self._frames) > 0 + return self._frame_count > 0 def page_count(self) -> int: - return len(self._frames) + return self._frame_count def load_page(self, page_no: int) -> _ImagePageBackend: - if not (0 <= page_no < len(self._frames)): + if not (0 <= page_no < self._frame_count): raise IndexError(f"Page index out of range: {page_no}") - return _ImagePageBackend( - self._frames[page_no], page_no, self._frame_dpi[page_no] - ) + with self._lock: + assert self._source is not None + if self._source.tell() != page_no: + self._source.seek(page_no) + # Frames of a multi-page TIFF have their own size, so each one is + # checked before it is decoded. + check_image_pixels(self._source.size, self.options.max_image_pixels) + frame, dpi = _oriented_frame(self._source) + image = to_rgb(frame) + return _ImagePageBackend(image, page_no, dpi) @classmethod def supported_formats(cls) -> set[InputFormat]: @@ -273,8 +281,6 @@ def supports_pagination(cls) -> bool: return True def unload(self): - for frame in self._frames: - frame.close() - self._frames = [] - self._frame_dpi = [] + with self._lock: + self._close_source() super().unload() diff --git a/docling/backend/md_backend.py b/docling/backend/md_backend.py index 99cbcaf23a..7e77cb8054 100644 --- a/docling/backend/md_backend.py +++ b/docling/backend/md_backend.py @@ -1112,6 +1112,7 @@ def _get_image_loader(self) -> ImageResourceLoader: enable_local_fetch=md_options.enable_local_fetch, enable_remote_fetch=md_options.enable_remote_fetch, max_image_data_base64_bytes=md_options.max_image_data_base64_bytes, + max_image_pixels=md_options.max_image_pixels, ) return self._image_loader diff --git a/docling/backend/utils/image_resource_loader.py b/docling/backend/utils/image_resource_loader.py index 5d498761a9..e1e1a99715 100644 --- a/docling/backend/utils/image_resource_loader.py +++ b/docling/backend/utils/image_resource_loader.py @@ -5,8 +5,8 @@ Turns an image source (a URI, a local file, or a remote URL) into a DoclingDocument `ImageRef`, enforcing the relevant safety limits: a path-traversal guard when resolving -relative paths, the `enable_local_fetch` / `enable_remote_fetch` toggles, and the -base64 and remote download size caps. +relative paths, the `enable_local_fetch` / `enable_remote_fetch` toggles, the +base64 and remote download size caps, and a limit on the decoded image pixels. """ import base64 @@ -32,7 +32,9 @@ from pydantic import ValidationError from requests.utils import get_environ_proxies, select_proxy +from docling.datamodel.backend_options import DEFAULT_MAX_IMAGE_PIXELS from docling.exceptions import OperationNotAllowed +from docling.utils.image_utils import check_image_pixels _log = logging.getLogger(__name__) @@ -267,6 +269,7 @@ def __init__( enable_remote_fetch: bool = False, max_image_data_base64_bytes: int = 20 * 1024 * 1024, max_remote_image_bytes: int = 20 * 1024 * 1024, + max_image_pixels: int = DEFAULT_MAX_IMAGE_PIXELS, max_redirects: int = 5, headers: Optional[dict[str, str]] = None, header_origins: Iterable[str] = (), @@ -275,6 +278,7 @@ def __init__( self.enable_remote_fetch = enable_remote_fetch self.max_image_data_base64_bytes = max_image_data_base64_bytes self.max_remote_image_bytes = max_remote_image_bytes + self.max_image_pixels = max_image_pixels self.max_redirects = max_redirects self.headers = headers # Only requests to these origins carry the configured `headers`. @@ -348,8 +352,12 @@ def create_image_ref( img_data = self.load_image_data(src_url, base_path) if img_data: img = Image.open(BytesIO(img_data)) + # Image.open reads only the header: check the size before + # ImageRef.from_pil decodes the pixels. + check_image_pixels(img.size, self.max_image_pixels) return ImageRef.from_pil(img, dpi=int(img.info.get("dpi", (72,))[0])) except ( + Image.DecompressionBombError, requests.RequestException, urllib3.exceptions.HTTPError, ValidationError, @@ -454,6 +462,14 @@ def load_image_data( return self.fetch_remote(src_loc).content elif src_loc.startswith("data:"): encoded_data = re.sub(r"^data:image/.+;base64,", "", src_loc) + # Check the size before decoding: 4 base64 characters hold 3 + # bytes, less one byte for each "=" padding character. + padding = len(encoded_data[-2:]) - len(encoded_data[-2:].rstrip("=")) + decoded_size = len(encoded_data) * 3 // 4 - padding + if decoded_size > self.max_image_data_base64_bytes: + raise ValueError( + f"Decoded image exceeds size limit of {self.max_image_data_base64_bytes} bytes." + ) decoded_data = base64.b64decode(encoded_data) if len(decoded_data) > self.max_image_data_base64_bytes: diff --git a/docling/datamodel/backend_options.py b/docling/datamodel/backend_options.py index 0e1a84e73a..adeb33789b 100644 --- a/docling/datamodel/backend_options.py +++ b/docling/datamodel/backend_options.py @@ -20,6 +20,11 @@ model_validator, ) +# Default limit on the pixels of one decoded image or image frame. It is the +# default of Pillow's ``Image.MAX_IMAGE_PIXELS`` (about 268 MB once decoded to +# RGB), fixed here so that the limit does not depend on global Pillow state. +DEFAULT_MAX_IMAGE_PIXELS = 89_478_485 + class BaseBackendOptions(BaseModel): """Common options for all declarative document backends.""" @@ -196,6 +201,13 @@ class HTMLBackendOptions(BaseBackendOptions): 20 * 1024 * 1024, # 20 MB description="The maximum number of bytes for remote image downloads.", ) + max_image_pixels: PositiveInt = Field( + DEFAULT_MAX_IMAGE_PIXELS, + description=( + "The maximum number of pixels (width x height) of an embedded or " + "fetched image. Larger images are skipped with a warning." + ), + ) max_redirects: Annotated[int, Field(ge=0)] = Field( 5, description="Maximum number of HTTP redirects to follow when fetching remote resources. Set to 0 to disable redirects.", @@ -236,6 +248,13 @@ class MarkdownBackendOptions(TextBackendOptions): 20 * 1024 * 1024, # 20 MB description="The maximum number of base64 data bytes that the backend will accept.", ) + max_image_pixels: PositiveInt = Field( + DEFAULT_MAX_IMAGE_PIXELS, + description=( + "The maximum number of pixels (width x height) of an embedded or " + "fetched image. Larger images are skipped with a warning." + ), + ) class JatsBackendOptions(BaseBackendOptions): @@ -306,6 +325,14 @@ class PdfBackendOptions(BaseBackendOptions): "decoding costs time and memory, so it is off by default." ), ) + max_image_pixels: PositiveInt = Field( + DEFAULT_MAX_IMAGE_PIXELS, + description=( + "The maximum number of pixels (width x height) of one page of an " + "image input, such as one frame of a multi-page TIFF. Used by the " + "image backend; an image with a larger page is rejected." + ), + ) class ThreadedDoclingParseBackendOptions(PdfBackendOptions): diff --git a/docling/utils/image_utils.py b/docling/utils/image_utils.py index a952422552..06ece1058f 100644 --- a/docling/utils/image_utils.py +++ b/docling/utils/image_utils.py @@ -4,6 +4,20 @@ from PIL import Image +def check_image_pixels(size: tuple[int, int], max_pixels: int) -> None: + """Raise ``ValueError`` if an image of ``size`` exceeds ``max_pixels``. + + Pillow reads the dimensions from the image header, so call this after + ``Image.open`` (or ``seek``) and before the pixel data is decoded. + """ + width, height = size + if width * height > max_pixels: + raise ValueError( + f"Image of {width}x{height} pixels exceeds the limit of " + f"{max_pixels} pixels." + ) + + def to_rgb(image: Image.Image) -> Image.Image: """Convert an image to RGB, flattening any transparency onto white. diff --git a/docs/usage/advanced_options.md b/docs/usage/advanced_options.md index b548926d63..5275dcf2a3 100644 --- a/docs/usage/advanced_options.md +++ b/docs/usage/advanced_options.md @@ -568,7 +568,9 @@ The HTML backend only downloads images referenced by a page when you opt in with `--html-image-fetch remote`). Downloads connect only to public, globally routable addresses: every address of a host is checked, every redirect is checked again before it is followed (up to `max_redirects`), and downloads stop -at `max_remote_image_bytes`. When `render_page=True`, the browser requests +at `max_remote_image_bytes`. An image larger than `max_image_pixels` (width x +height, default 89,478,485) is skipped with a warning, and the rest of the +document converts; the Markdown backend has the same option. When `render_page=True`, the browser requests remote resources through the same download path, and navigating the page away from the source document is refused. @@ -620,6 +622,27 @@ converter = DocumentConverter() result = converter.convert(source, max_num_pages=100, max_file_size=20971520) ``` +For image inputs, the page count is checked against `max_num_pages` before any +page is decoded, and each page (for example each frame of a multi-page TIFF) is +decoded only when it is processed. The `max_image_pixels` backend option limits +the pixels (width x height) of one page; the default, 89,478,485, is the default +of Pillow's `Image.MAX_IMAGE_PIXELS`. A larger first page rejects the document, +and a larger later page fails that page: + +```python +from docling.datamodel.backend_options import PdfBackendOptions +from docling.datamodel.base_models import InputFormat +from docling.document_converter import DocumentConverter, ImageFormatOption + +converter = DocumentConverter( + format_options={ + InputFormat.IMAGE: ImageFormatOption( + backend_options=PdfBackendOptions(max_image_pixels=50_000_000) + ) + } +) +``` + ## Convert from binary PDF streams You can convert PDFs from a binary stream instead of from the filesystem as follows: diff --git a/tests/test_backend_html.py b/tests/test_backend_html.py index a1818b2cb5..f991ef316a 100644 --- a/tests/test_backend_html.py +++ b/tests/test_backend_html.py @@ -20,6 +20,7 @@ RichTableCell, ) from docling_core.types.doc.document import ContentLayer +from PIL import Image from pydantic import AnyUrl, ValidationError from docling.backend.html_backend import ( @@ -31,7 +32,7 @@ validate_url_safety as _validate_url_safety, ) from docling.datamodel.backend_options import HTMLBackendOptions -from docling.datamodel.base_models import ConversionStatus, InputFormat +from docling.datamodel.base_models import ConversionStatus, DocumentStream, InputFormat from docling.datamodel.document import ( ConversionResult, DoclingDocument, @@ -2010,6 +2011,33 @@ def test_load_image_data_enforces_data_uri_size_limit(): backend._load_image_data(data_uri) +def test_image_over_pixel_limit_is_skipped(): + def png_data_uri(width: int, height: int) -> str: + buffer = BytesIO() + Image.new("RGB", (width, height), (255, 0, 0)).save(buffer, format="PNG") + return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() + + html = ( + "

Before.

" + f'' + "

After.

" + ) + converter = _create_html_converter( + HTMLBackendOptions(fetch_images=True, max_image_pixels=63) + ) + with pytest.warns(UserWarning, match="exceeds the limit of 63 pixels"): + result = converter.convert( + DocumentStream(name="test.html", stream=BytesIO(html.encode())) + ) + doc = result.document + + assert result.status == ConversionStatus.SUCCESS + assert len(doc.pictures) == 2 + assert doc.pictures[0].image is None + assert doc.pictures[1].get_image(doc).size == (7, 5) + assert [text.text for text in doc.texts] == ["Before.", "After."] + + def test_anchor_fragment_links_with_source_uri(): """Fragment-only hrefs must not be mangled when source_uri is set.""" html_path = Path("tests/data/html/sources/hyperlink_06.html") diff --git a/tests/test_backend_image_native.py b/tests/test_backend_image_native.py index 664d84246c..a8a6c230dd 100644 --- a/tests/test_backend_image_native.py +++ b/tests/test_backend_image_native.py @@ -2,13 +2,14 @@ # SPDX-License-Identifier: MIT from io import BytesIO -from unittest.mock import MagicMock import pytest from docling_core.types.doc import BoundingBox, CoordOrigin from PIL import Image, TiffImagePlugin +from docling.backend import image_backend from docling.backend.image_backend import ImageDocumentBackend, _ImagePageBackend +from docling.datamodel.backend_options import PdfBackendOptions from docling.datamodel.base_models import DocumentStream, InputFormat from docling.datamodel.document import ( InputDocument, @@ -16,6 +17,7 @@ _DummyBackend, get_input_rejection_cause, ) +from docling.datamodel.settings import DocumentLimits from docling.document_converter import DocumentConverter, ImageFormatOption from docling.document_extractor import DocumentExtractor from docling.exceptions import DocumentLoadError @@ -419,36 +421,17 @@ def test_invalid_explicit_dpi_rejects_document(): assert isinstance(get_input_rejection_cause(in_doc), DocumentLoadError) -def test_source_image_is_closed_after_backend_init(tmp_path, monkeypatch): +def test_source_image_is_closed_on_unload(tmp_path, monkeypatch): image_path = tmp_path / "test.png" Image.new("RGB", (32, 32), (10, 20, 30)).save(image_path) opened_images = [] original_open = Image.open - class TrackingImage: - def __init__(self, image): - self._image = image - self.closed = False - - def __getattr__(self, attr): - return getattr(self._image, attr) - - def close(self): - self.closed = True - return self._image.close() - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - self.close() - return False - def tracking_open(*args, **kwargs): - tracked_image = TrackingImage(original_open(*args, **kwargs)) - opened_images.append(tracked_image) - return tracked_image + image = original_open(*args, **kwargs) + opened_images.append(image) + return image input_doc = InputDocument( path_or_stream=image_path, @@ -462,25 +445,72 @@ def tracking_open(*args, **kwargs): in_doc=input_doc, path_or_stream=image_path, ) + page_image = backend.load_page(0).get_page_image() + backend.unload() assert len(opened_images) == 1 - assert opened_images[0].closed is True - backend.unload() + assert opened_images[0].fp is None + # The decoded page stays usable after the source is closed. + assert page_image.getpixel((0, 0)) == (10, 20, 30) -def test_unload_closes_cached_frames(): - stream = _make_multipage_tiff_stream(num_pages=3, size=(32, 32)) - doc_backend = _get_backend_from_stream(stream) +def test_frame_count_over_max_num_pages_is_rejected_before_decoding(monkeypatch): + decoded = [] + original_to_rgb = image_backend.to_rgb - tracked_closers = [] - for frame in doc_backend._frames: - closer = MagicMock(wraps=frame.close) - frame.close = closer - tracked_closers.append(closer) + def counting_to_rgb(image): + decoded.append(image.size) + return original_to_rgb(image) + + monkeypatch.setattr(image_backend, "to_rgb", counting_to_rgb) + stream = _make_multipage_tiff_stream(num_pages=5, size=(16, 16)) + in_doc = InputDocument( + path_or_stream=stream.stream, + format=InputFormat.IMAGE, + backend=ImageDocumentBackend, + filename=stream.name, + limits=DocumentLimits(max_num_pages=2), + ) + + assert in_doc.valid is False + assert in_doc.page_count == 5 + assert "max_num_pages limit of 2" in in_doc._rejection.message + assert decoded == [] + in_doc._backend.unload() - doc_backend.unload() - assert doc_backend._frames == [] - assert doc_backend._frame_dpi == [] - for closer in tracked_closers: - closer.assert_called_once() +def test_image_over_pixel_limit_is_rejected(): + stream = _make_png_stream(width=64, height=48) + in_doc = InputDocument( + path_or_stream=stream.stream, + format=InputFormat.IMAGE, + backend=ImageDocumentBackend, + backend_options=PdfBackendOptions(max_image_pixels=64 * 48 - 1), + filename=stream.name, + ) + + assert in_doc.valid is False + cause = get_input_rejection_cause(in_doc) + assert isinstance(cause, DocumentLoadError) + assert "exceeds the limit of 3071 pixels" in str(cause) + + +def test_tiff_frame_over_pixel_limit_is_rejected_on_load(): + frames = [Image.new("RGB", (16, 16)), Image.new("RGB", (64, 64))] + buf = BytesIO() + frames[0].save(buf, format="TIFF", save_all=True, append_images=frames[1:]) + buf.seek(0) + in_doc = InputDocument( + path_or_stream=buf, + format=InputFormat.IMAGE, + backend=ImageDocumentBackend, + backend_options=PdfBackendOptions(max_image_pixels=1000), + filename="test.tiff", + ) + doc_backend = in_doc._backend + + assert in_doc.valid is True + assert doc_backend.load_page(0).get_size().width == 16 + with pytest.raises(ValueError, match="64x64 pixels exceeds the limit"): + doc_backend.load_page(1) + doc_backend.unload() diff --git a/tests/test_backend_markdown_images.py b/tests/test_backend_markdown_images.py index 57c9c2d6bf..e843575035 100644 --- a/tests/test_backend_markdown_images.py +++ b/tests/test_backend_markdown_images.py @@ -87,6 +87,22 @@ def test_mixed_html_keeps_size_rejected_image_placeholder(html_block: str) -> No assert doc.pictures[0].caption_text(doc) == "caption" +def test_image_over_pixel_limit_is_skipped() -> None: + markdown = ( + f"Before.\n\n![big]({_png_data_uri(8, 8)})\n\n" + f"![small]({_png_data_uri(7, 5)})\n\nAfter.\n" + ) + with pytest.warns(UserWarning, match="exceeds the limit of 63 pixels"): + doc = _convert_markdown( + markdown, MarkdownBackendOptions(fetch_images=True, max_image_pixels=63) + ) + + assert len(doc.pictures) == 2 + assert doc.pictures[0].image is None + assert doc.pictures[1].get_image(doc).size == (7, 5) + assert [text.text for text in doc.texts] == ["Before.", "big", "small", "After."] + + @pytest.mark.parametrize("fetch_images", [False, True]) def test_original_html_images_keep_the_normal_loader(fetch_images: bool) -> None: markdown = ( diff --git a/tests/test_image_resource_loader.py b/tests/test_image_resource_loader.py index 37ba0d6859..04dd602b82 100644 --- a/tests/test_image_resource_loader.py +++ b/tests/test_image_resource_loader.py @@ -3,11 +3,16 @@ """Unit tests for the shared image-resource loader and its safety limits.""" +import base64 import socket +import struct import threading +import zlib +from io import BytesIO from urllib.parse import quote import pytest +from PIL import Image from docling.backend.utils.image_resource_loader import ( ImageResourceLoader, @@ -242,3 +247,49 @@ def test_load_image_data_local_requires_base_path(): loader = ImageResourceLoader(enable_local_fetch=True) with pytest.raises(OperationNotAllowed, match="requires base_path"): loader.load_image_data("/some/where/image.png", None) + + +def _png_header_data_uri(width: int, height: int) -> str: + """A PNG that declares ``width`` x ``height`` pixels but has no pixel data.""" + + def chunk(kind: bytes, data: bytes) -> bytes: + body = kind + data + return struct.pack(">I", len(data)) + body + struct.pack(">I", zlib.crc32(body)) + + png = ( + b"\x89PNG\r\n\x1a\n" + + chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 0, 0, 0, 0)) + + chunk(b"IDAT", zlib.compress(b"")) + + chunk(b"IEND", b"") + ) + return "data:image/png;base64," + base64.b64encode(png).decode() + + +def _png_data_uri(width: int, height: int) -> str: + buffer = BytesIO() + Image.new("RGB", (width, height)).save(buffer, format="PNG") + return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() + + +def test_image_over_pixel_limit_is_skipped(): + loader = ImageResourceLoader(max_image_pixels=63) + + assert loader.create_image_ref(_png_data_uri(7, 9), None) is not None + with pytest.warns(UserWarning, match="8x8 pixels exceeds the limit of 63"): + assert loader.create_image_ref(_png_data_uri(8, 8), None) is None + + +def test_image_over_pillow_pixel_limit_is_skipped(): + # Above twice Image.MAX_IMAGE_PIXELS, Pillow refuses to open the image. + loader = ImageResourceLoader(max_image_pixels=10**9) + + with pytest.warns(UserWarning, match="400000000 pixels"): + assert loader.create_image_ref(_png_header_data_uri(20000, 20000), None) is None + + +def test_data_uri_over_size_limit_is_rejected(): + loader = ImageResourceLoader(max_image_data_base64_bytes=8) + + assert loader.load_image_data("data:image/png;base64,AAAAAAAAAAA=", None) + with pytest.raises(ValueError, match="exceeds size limit of 8 bytes"): + loader.load_image_data("data:image/png;base64,AAAAAAAAAAAAAAAA", None)