From d92b2eabd972630a480947faf539d440f978d642 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sat, 31 Oct 2026 15:18:02 -0700 Subject: [PATCH] test(images): cover the text budget, XMP and layout guard of the PNG text walk (#523) --- tests/test_png_text.py | 30 ++++++++++++++++++++++++++++-- utils/png_text.py | 10 +++++++++- 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/tests/test_png_text.py b/tests/test_png_text.py index 9dc5577b..9633e506 100644 --- a/tests/test_png_text.py +++ b/tests/test_png_text.py @@ -148,6 +148,13 @@ def curated(): itxt_chunk("workflow", WORKFLOW, compressed=True), ) ), + "xmp-after-pixels": assemble( + before_iend( + base_png(), + itxt_chunk("XML:com.adobe.xmp", ""), + text_chunk("prompt", PROMPT), + ) + ), "itxt-bad-utf8": assemble( before_iend( base_png(), @@ -288,9 +295,9 @@ class PngTextParityTests(_TempFiles): self.assertEqual(walked_text(path), legacy_text(path)) self.assertEqual(walked_text(path)["workflow"], WORKFLOW) - def test_truncated_pixels_still_fail_when_truncated_images_load(self): + def test_truncated_pixels_read_alike_when_truncated_images_load(self): """With LOAD_TRUNCATED_IMAGES set (a node pack may set it) both readers - accept a short file.""" + accept a short file alike.""" path = self.write("t.png", curated()["truncated-in-pixels"]) with patch.object(ImageFile, "LOAD_TRUNCATED_IMAGES", True): self.assertEqual(walked_text(path), legacy_text(path)) @@ -318,6 +325,25 @@ class PngTextParityTests(_TempFiles): self.assertEqual(walked_text(path), old) self.assertEqual(old, {"prompt": PROMPT}) + def test_the_text_memory_budget_counts_text_before_the_pixels(self): + """Pillow's MAX_TEXT_MEMORY spans the whole file: text before and + after the pixels, each under it, fail together in both readers.""" + path = self.write( + "m.png", + assemble( + before_iend( + base_png(text={"prompt": PROMPT}), text_chunk("workflow", WORKFLOW) + ) + ), + ) + budget = len(PROMPT) + len(WORKFLOW) - 1 + with patch.object(PngImagePlugin, "MAX_TEXT_MEMORY", budget): + self.assertEqual(legacy_text(path), {}) + self.assertEqual(walked_text(path), {}) + with patch.object(PngImagePlugin, "MAX_TEXT_MEMORY", budget + 1): + self.assertEqual(walked_text(path), legacy_text(path)) + self.assertEqual(len(walked_text(path)), 2) + def test_complete_but_undecodable_pixels_read_more_than_legacy(self): path = self.write("c.png", corrupt_pixels()) self.assertEqual(legacy_text(path), {}) diff --git a/utils/png_text.py b/utils/png_text.py index 90a78331..36f43e25 100644 --- a/utils/png_text.py +++ b/utils/png_text.py @@ -28,6 +28,8 @@ import struct from PIL import ImageFile _PIXEL_CHUNKS = (b"IDAT", b"fdAT") +# Signature, IHDR chunk and IDAT chunk header: the earliest pixel data offset. +_FIRST_DATA_OFFSET = 8 + 25 + 8 def png_text(image) -> dict: @@ -57,6 +59,9 @@ def _walkable(image) -> bool: and isinstance(tile, list) and len(tile) == 1 and tile[0][0] == "zip" + # The first IDAT's data offset; anything else is not a layout we know. + and isinstance(tile[0][2], int) + and tile[0][2] >= _FIRST_DATA_OFFSET ) @@ -65,7 +70,10 @@ def _walk_after_pixels(stream, fp, offset: int) -> dict: ``stream`` is the image's own ``PngStream``: its text so far and its text memory budget carry on exactly as a ``load()`` would carry them. The text - and info dictionaries are copied first, so the image's state is unchanged. + and info dictionaries are copied first, so ``image.text`` and + ``image.info`` are unchanged. Like ``load_end``, a handler for a chunk + misplaced after the pixels (a second IHDR or PLTE) may still update the + stream's other header fields; Pillow ignores them once the image is open. """ end = fp.seek(0, os.SEEK_END) saved = stream.im_text, stream.im_info, stream.text_memory