Patched code sample

import zlib
import sys

# This code demonstrates the fix for CVE-2024-39373 in jwcrypto.
# The vulnerability allowed a small, compressed JWE token (a "zip bomb")
# to decompress to a very large size, causing memory exhaustion.
#
# The fix, introduced in jwcrypto 1.5.7, involves limiting the size of the
# decompressed output. This is achieved by using `zlib.decompressobj()`
# with a `max_length` parameter, which prevents the allocation of an
# oversized memory buffer.

# The constant for the maximum allowed size for the decompressed data.
# In the actual patch, this is 100 * 1024 * 1024 (100MB).
# We use a smaller value here for a quicker demonstration.
MAX_DECOMPRESSED_SIZE = 10 * 1024 * 1024  # 10 MB

class JWEException(Exception):
    """A custom exception class to mimic jwcrypto's JWException."""
    pass

def safe_decompress_jwe_payload(compressed_payload: bytes) -> bytes:
    """
    Demonstrates the fixed decompression logic inspired by jwcrypto 1.5.7.

    This function safely decompresses a payload by enforcing a size limit on the
    output, thus preventing a "zip bomb" from causing memory exhaustion.

    Args:
        compressed_payload: The JWE payload after decryption but before
                            decompression.

    Returns:
        The decompressed plaintext if it is within the size limit.

    Raises:
        JWEException: If the decompressed data would exceed the predefined
                      MAX_DECOMPRESSED_SIZE.
    """
    # The core of the fix is to use a decompress object which allows for
    # controlled, streaming decompression.
    decompressor = zlib.decompressobj()

    try:
        # The `max_length` argument is the critical part of the fix.
        # It tells zlib to stop decompressing if the output buffer would
        # exceed this size. We add 1 to the limit to reliably detect if
        # the output is larger than allowed. If it is exactly the limit,
        # this will succeed, but the final length check will catch it.
        # If it is over, zlib itself will raise an error.
        plaintext = decompressor.decompress(compressed_payload, MAX_DECOMPRESSED_SIZE + 1)

        # The patch also checks if there is unconsumed data left, which
        # would indicate the original stream was larger than the limit.
        if decompressor.unconsumed_tail or decompressor.unused_data:
            raise JWEException("Decompression failed: output is too large.")

    except zlib.error as e:
        # A zlib.error is raised if `max_length` is exceeded during decompression.
        # We catch it and wrap it in the application-specific exception.
        raise JWEException(f"Decompression failed: {e}") from e

    # Finally, a redundant but explicit check on the final length of the
    # decompressed data ensures nothing slips through.
    if len(plaintext) > MAX_DECOMPRESSED_SIZE:
        raise JWEException("Decompression failed: output exceeded max size.")

    return plaintext


if __name__ == "__main__":
    # --- Step 1: Create a malicious "zip bomb" payload ---
    # This payload is small when compressed but would expand to a size greater
    # than our defined MAX_DECOMPRESSED_SIZE.
    # We will make it expand to 11 MB, which is over our 10 MB limit.
    target_decompressed_size = 11 * 1024 * 1024  # 11 MB
    large_uncompressed_data = b'\x00' * target_decompressed_size
    zip_bomb_payload = zlib.compress(large_uncompressed_data, level=9)

    print("--- CVE-2024-39373 Fix Demonstration ---")
    print(f"Crafted a malicious JWE payload (a 'zip bomb'):")
    print(f"  - Compressed size: {len(zip_bomb_payload) / 1024:.2f} KB")
    print(f"  - Intended decompressed size: {target_decompressed_size / (1024*1024):.0f} MB")
    print(f"Set security limit (MAX_DECOMPRESSED_SIZE): {MAX_DECOMPRESSED_SIZE / (1024*1024):.0f} MB")
    print("-" * 40)

    # --- Step 2: Demonstrate the fix in action ---
    # The `safe_decompress_jwe_payload` function should detect that the
    # payload is a zip bomb and raise an exception, preventing the
    # allocation of 11 MB of memory.

    # This represents a vulnerable system without the fix.
    # It would attempt to allocate 11MB, potentially leading to a crash.
    # print("Simulating VULNERABLE behavior (for context, will try to allocate memory)...")
    # try:
    #     # On a memory-constrained system, this line could cause a crash.
    #     vulnerable_result = zlib.decompress(zip_bomb_payload)
    #     print(f"  [VULNERABLE] Decompression succeeded. Allocated {len(vulnerable_result)/(1024*1024):.0f} MB.")
    # except MemoryError:
    #     print("  [VULNERABLE] Caught MemoryError. System is under duress.")
    # print("-" * 40)

    print("Attempting to decompress the payload with the FIXED logic...")
    try:
        safe_decompress_jwe_payload(zip_bomb_payload)
        # If we reach this line, the fix has failed.
        print("\n[RESULT: FAILED] The fix did not work. Decompression was successful.")
        sys.exit(1)
    except JWEException as e:
        # This is the expected outcome. The fix worked.
        print(f"\n[RESULT: SUCCESS] The fixed logic correctly prevented the attack.")
        print(f"Caught the expected exception, stopping the decompression:")
        print(f"  -> {e}")
        sys.exit(0)
    except Exception as e:
        print(f"\n[RESULT: FAILED] An unexpected error occurred: {e}")
        sys.exit(1)