Patched code sample

import zlib
import sys

# This constant is the core of the fix, mirroring the one added to python-jose.
# It sets a hard limit on the size of compressed data to be processed.
# The original library uses 10 * 1024 * 1024 (10MB). We use a smaller value
# for a more practical demonstration.
_DECOMPRESSION_BOMB_PREVENTION_LIMIT = 2 * 1024 * 1024  # 2MB

def fixed_jwe_decompress(compressed_data: bytes):
    """
    Simulates the fixed JWE decompression logic from python-jose >= 3.3.0.post1.

    The fix involves checking the size of the compressed data against a predefined
    limit *before* attempting decompression. This prevents the allocation of
    massive amounts of memory from a small but highly compressed payload (a "bomb").
    """
    print(f"\nAttempting to decompress data of size: {len(compressed_data) / 1024:.2f} KB")

    # --- THE FIX ---
    # This check is the security patch for CVE-2024-29370.
    # If the compressed input is larger than the allowed limit, an error is
    # raised immediately, preventing the resource-intensive decompression call.
    if len(compressed_data) > _DECOMPRESSION_BOMB_PREVENTION_LIMIT:
        raise ValueError(
            f"Decompression bomb suspected. Compressed data size ({len(compressed_data)}) "
            f"exceeds limit ({_DECOMPRESSION_BOMB_PREVENTION_LIMIT})."
        )
    # --- END OF FIX ---

    # This line is only reached if the data is within the safe limit.
    print("Compressed data is within the safe limit. Proceeding with decompression.")
    decompressed_data = zlib.decompress(compressed_data)
    return decompressed_data

if __name__ == "__main__":
    # 1. Create a "compression bomb" payload.
    # A large sequence of repeating bytes compresses very effectively.
    # Here, we create 50MB of null bytes.
    uncompressed_bomb_size_mb = 50
    malicious_uncompressed_payload = b'\x00' * (uncompressed_bomb_size_mb * 1024 * 1024)

    # Compress the payload. It will become very small.
    compressed_bomb = zlib.compress(malicious_uncompressed_payload, level=9)

    print("--- Demonstration of a Malicious Payload (Compression Bomb) ---")
    print(f"Original size of malicious payload: {len(malicious_uncompressed_payload) / (1024*1024):.2f} MB")
    print(f"Compressed size of malicious payload: {len(compressed_bomb) / 1024:.2f} KB")
    print("-" * 60)

    # 2. Demonstrate the fix with the malicious payload.
    # The 'fixed_jwe_decompress' function will detect the large compressed size
    # and raise an error before memory exhaustion occurs.
    try:
        # Without the fix, this line would consume ~50MB of RAM and significant CPU.
        # With the fix, it raises a ValueError immediately.
        fixed_jwe_decompress(compressed_bomb)
    except ValueError as e:
        print(f"\nSUCCESS: The vulnerability was mitigated.")
        print(f"Caught expected error: {e}")

    print("\n" + "=" * 60 + "\n")

    # 3. Create a normal, non-malicious payload.
    safe_payload = b"This is a legitimate and reasonably sized payload for a JWE token."
    compressed_safe_payload = zlib.compress(safe_payload)

    print("--- Demonstration of a Benign Payload ---")
    print(f"Original size of safe payload: {len(safe_payload)} bytes")
    print(f"Compressed size of safe payload: {len(compressed_safe_payload)} bytes")
    print("-" * 60)

    # 4. Show that the fix does not affect legitimate payloads.
    try:
        decompressed_data = fixed_jwe_decompress(compressed_safe_payload)
        print("\nSUCCESS: Benign payload was decompressed correctly.")
        print(f"Decompressed data matches original: {decompressed_data == safe_payload}")
    except ValueError as e:
        print(f"\nFAILURE: The fix incorrectly blocked a valid payload.")
        print(f"Caught unexpected error: {e}")