test(images): cover the text budget, XMP and layout guard of the PNG text walk (#523)

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", "<x:xmpmeta/>"),
+                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), {})

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