EnigmaConsultant/huntr-poc-keras-embedding-inputdim-alloc-oom

Uncontrolled memory allocation via attacker-controlled Embedding input_dim in a .keras layer config during default load_model (CWE-789 / CWE-400 DoS)

Target

• Project: Keras (keras on PyPI)
• Version verified: keras 3.15.0 (numpy backend), Python 3.13
• Vulnerable file/line: keras/src/layers/core/embedding.py:149
• Load path: default keras.saving.load_model(path) with safe_mode=True, compile=True (the documented defaults)
• Class: CWE-789 Memory Allocation with Excessive Size Value / CWE-400 Uncontrolled Resource Consumption (DoS)

Summary

Loading an untrusted .keras model file with the ordinary, fully-default keras.saving.load_model() call can be driven to an unbounded host-side memory allocation by a ~2.8 KB file. The amplifier is the Embedding layer's own scalar config field input_dim, which is echoed straight into a weight-tensor shape and materialized during model deserialization — before the caller ever runs, trains, or predicts.

safe_mode=True provides no protection here: safe_mode only guards __lambda__ (arbitrary-code) deserialization, not resource sizing.

Root cause

During .keras deserialization, serialization_lib.py:787 unconditionally calls instance.build_from_config(build_config) for any layer that carries a build_config. For an Embedding layer this reaches Embedding.build(), which computes the embeddings table shape purely from the layer's own scalar config fields and ignores input_shape entirely:

# keras/src/layers/core/embedding.py
def build(self, input_shape=None):
    if self.built:
        return
    embeddings_shape = (self.input_dim, self.output_dim)   # line 149 — attacker-controlled input_dim
    ...
    if self.quantization_mode not in ("int8", "int4"):
        self._embeddings = self.add_weight(
            shape=embeddings_shape,
            initializer=self.embeddings_initializer,
            name="embeddings",
            ...
        )

input_dim is read from the JSON config.json inside the .keras zip. It is only validated as a positive integer (> 0, int) — there is no upper bound and no cross-check against the tiny build_config.input_shape. It is a 64-bit JSON integer, so an attacker can set it arbitrarily large in a file that stays a few kilobytes.

add_weight() -> backend.Variable(...) -> _initialize_with_initializer -> the default RandomUniform initializer -> keras/src/backend/numpy/random.py:23:

return rng.uniform(size=shape, low=minval, high=maxval).astype(dtype)

NumPy's Generator.uniform allocates the array in float64 first and only then .astype(dtype) down to float32. That adds an ~8x resident spike over the nominal float32 weight size, making the DoS cheaper for the attacker (for shape (300000000, 8) the float64 buffer alone is 17.9 GiB).

Proof of concept

Three scripts (mk.py, craft.py, load.py) plus the artifacts (benign.keras, evil.keras) are included in this repo.

1. mk.py — builds a benign functional model containing Embedding(input_dim=16, output_dim=8) and saves benign.keras (19962 B).
2. craft.py — opens benign.keras, rewrites only config.json to set the Embedding layer config's input_dim = 300000000, and re-zips as evil.keras (2796 B). Nothing else is changed — in particular the layer's build_config.input_shape stays tiny at [None, 4].
3. load.py — sets a 2 GB RLIMIT_AS cap and calls the fully-default keras.saving.load_model(path) (safe_mode=True, compile=True).

Captured evidence (verbatim, keras 3.15.0 / numpy backend)

########## NEGATIVE CONTROL (benign) ##########
=== default keras.saving.load_model('benign.keras') under 2GB RLIMIT_AS ===
keras 3.15.0 backend numpy
LOADED OK: Functional params: 194

########## EVIL (huge input_dim) ##########
  File ".../keras/src/layers/core/embedding.py", line 157, in build
    self._embeddings = self.add_weight(
  File ".../keras/src/layers/layer.py", line 622, in add_weight
    variable = backend.Variable(
  File ".../keras/src/backend/common/variables.py", line 210, in __init__
    self._initialize_with_initializer(initializer)
  File ".../keras/src/backend/common/variables.py", line 418, in _initialize_with_initializer
    initializer(self._shape, dtype=self._dtype)
  File ".../keras/src/initializers/random_initializers.py", line 187, in __call__
    return random.uniform(
  File ".../keras/src/backend/numpy/random.py", line 23, in uniform
    return rng.uniform(size=shape, low=minval, high=maxval).astype(dtype)
  File "numpy/random/_generator.pyx", line 1102, in numpy.random._generator.Generator.uniform
numpy._core._exceptions._ArrayMemoryError: Unable to allocate 17.9 GiB for an array with shape (300000000, 8) and data type float64
=== default keras.saving.load_model('evil.keras') under 2GB RLIMIT_AS ===
keras 3.15.0 backend numpy
MemoryError: MemoryError((300000000, 8), dtype('float64'))

Proof the amplifier is input_dim, not build_config.input_shape

Embedding layer config input_dim = 300000000
Embedding build_config = {'input_shape': [None, 4]}
evil.keras size = 2796 bytes

The tiny input_shape is untouched; the allocation is driven entirely by the scalar input_dim config field. The negative-control benign file loads normally under the same 2 GB cap.

Impact

A ~2.8 KB .keras file, loaded via the default load_model() API, forces a multi-gigabyte-to-unbounded host memory allocation and crashes the process with MemoryError before any user code runs. Any service or pipeline that deserializes user-supplied Keras models (model hubs, CI validators, inference-serving upload endpoints, notebooks) is exposed to a trivial denial-of-service. The float64-then-astype behavior of the default initializer multiplies the resident spike ~8x, lowering the attacker's cost further.

Suggested remediation

Validate Embedding.input_dim (and, generally, config-derived weight shapes reconstructed during build_from_config) against a sane bound and/or the model's declared input_shape before calling add_weight(); alternatively cap total allocatable weight elements during deserialization. Deserialization should never size a raw allocation directly from an unvalidated file-supplied scalar.

Dedup / novelty note

This is distinct from other Keras allocation findings in the same audit:

• NOT Dense/build_config.input_shape or units (dense.py, "buildconfig-alloc") — here the tiny build_config.input_shape is deliberately left untouched; the amplifier is the Embedding-specific scalar input_dim in embedding.py.
• NOT compile-time metric amplifiers AUC.num_thresholds or IoU.num_classes.
• NOT the config.json zip/JSON size bomb, numpy dtype/subarray, or legacy H5 model-config paths.

No existing HuggingFace PoC repo under this account and no prior flip-list entry covers the Embedding input_dim allocation path. No known CVE for this specific parameter at time of writing.