PI: Decode the PNG Up predictor a row at a time

pypdf/filters.py

@@ -278,8 +278,11 @@ def _decode_png_prediction(data: bytes, columns: int, row_length: int) -> bytes:
             data += b"\x00" * (row_length - remainder)
             assert len(data) % row_length == 0
         output = bytearray()
-        previous_row_data = bytes(row_length)
+        previous_row_data = bytearray(row_length)
         bpp = (row_length - 1) // columns  # recomputed locally to not change params
+        # Masks to add all bytes of two rows at once, without carrying over into the next byte.
+        low_mask = int.from_bytes(b"\x7f" * row_length, "big")
+        high_mask = int.from_bytes(b"\x80" * row_length, "big")
         for row in range(0, len(data), row_length):
             row_data = bytearray(data[row : row + row_length])
             filter_byte = row_data[0]
@@ -293,8 +296,13 @@ def _decode_png_prediction(data: bytes, columns: int, row_length: int) -> bytes:
                     row_data[i] = (row_data[i] + row_data[i - bpp]) % 256
             elif filter_byte == 2:
                 # PNG Up Predictor
-                for i in range(1, row_length):
-                    row_data[i] = (row_data[i] + previous_row_data[i]) % 256
+                current = int.from_bytes(row_data, "big")
+                previous = int.from_bytes(previous_row_data, "big")
+                # The low 7 bits of each byte are added without carrying into the next byte,
+                # the high bit of each byte is added with XOR, dropping its carry.
+                low_sum = (current & low_mask) + (previous & low_mask)
+                high_bits = (current ^ previous) & high_mask
+                row_data = bytearray((low_sum ^ high_bits).to_bytes(row_length, "big"))
             elif filter_byte == 3:
                 # PNG Average Predictor
                 for i in range(1, bpp + 1):
@@ -313,10 +321,17 @@ def _decode_png_prediction(data: bytes, columns: int, row_length: int) -> bytes:
                     up = previous_row_data[i]
                     up_left = previous_row_data[i - bpp]
 
-                    p = left + up - up_left
-                    dist_left = abs(p - left)
-                    dist_up = abs(p - up)
-                    dist_up_left = abs(p - up_left)
+                    # The distances of p = left + up - up_left to left, up and up_left.
+                    # Avoid calling abs() in this innermost loop.
+                    dist_left = up - up_left
+                    dist_up = left - up_left
+                    dist_up_left = dist_left + dist_up
+                    if dist_left < 0:
+                        dist_left = -dist_left
+                    if dist_up < 0:
+                        dist_up = -dist_up
+                    if dist_up_left < 0:
+                        dist_up_left = -dist_up_left
 
                     if dist_left <= dist_up and dist_left <= dist_up_left:
                         paeth = left
@@ -330,7 +345,7 @@ def _decode_png_prediction(data: bytes, columns: int, row_length: int) -> bytes:
                 raise PdfReadError(
                     f"Unsupported PNG filter {filter_byte!r}"
                 )  # pragma: no cover
-            previous_row_data = bytes(row_data)
+            previous_row_data = row_data
             output += row_data[1:]
         return bytes(output)
 

tests/test_filters.py

@@ -1,6 +1,8 @@
 """Test the pypdf.filters module."""
 import os
+import random
 import string
+import struct
 import subprocess
 import sys
 import zlib
@@ -1335,6 +1337,39 @@ def test_flatedecode__decode_png_prediction__matches_image_without_predictor(pdf
     assert decoded == expected
 
 
+@pytest.mark.parametrize(
+    ("color_type", "bytes_per_pixel"),
+    [(0, 1), (4, 2), (2, 3), (6, 4)],  # Grayscale, grayscale with alpha, RGB, RGBA with 8 bits per sample
+)
+def test_flatedecode__decode_png_prediction__matches_png_decoder(color_type, bytes_per_pixel):
+    """Rows using all PNG filter types are decoded like the PNG decoder of Pillow does."""
+    columns = 9
+    rows = 40
+    row_length = columns * bytes_per_pixel + 1
+    generator = random.Random(color_type)  # noqa: S311
+    data = bytearray()
+    for row in range(rows):
+        data.append(row % 5)  # None, Sub, Up, Average and Paeth filter
+        # Small values lead to ties in the Paeth predictor, large ones to sums exceeding a byte.
+        limit = 4 if row % 10 < 5 else 256
+        data.extend(generator.randrange(limit) for _ in range(row_length - 1))
+
+    def chunk(chunk_type: bytes, chunk_data: bytes) -> bytes:
+        checksum = zlib.crc32(chunk_type + chunk_data)
+        return struct.pack(">I", len(chunk_data)) + chunk_type + chunk_data + struct.pack(">I", checksum)
+
+    png = (
+        b"\x89PNG\r\n\x1a\n"
+        + chunk(b"IHDR", struct.pack(">IIBBBBB", columns, rows, 8, color_type, 0, 0, 0))
+        + chunk(b"IDAT", zlib.compress(data))
+        + chunk(b"IEND", b"")
+    )
+    with Image.open(BytesIO(png)) as image:
+        expected = image.tobytes()
+
+    assert FlateDecode._decode_png_prediction(data=bytes(data), columns=columns, row_length=row_length) == expected
+
+
 _FLATE_IMAGE_DATA = (
     b"\x00\xff\xff\xff\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
     b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xff\xff\xff\x00\x00\x00\x00\x00\x00\x00\x00"