# 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:

```python
# 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`:

```python
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.