martilaio/tfjs-layers-dense-units-oom-poc

F-12 — Host OOM in @tensorflow/tfjs-layers via attacker-controlled Dense.units (31,935,321× amplification)

Authorized security research artifact disclosed via huntr.com's TensorFlow.js Model Format Vulnerability program. Source commit 7f5309fef0a47545e34049903dbdae0f97285f7e. All capture data was collected against a synthetic /tmp/victim_host/ CI-runner lab — no real PII present.

Real impact captured (sanitized)

Node process terminated under realistic host RAM caps — exit 134 (SIGABRT)

• prlimit --as=1GB → exit 134 + V8 stack trace inside Heap::PerformGarbageCollection
• prlimit --as=2GB → exit 134 + pthread_create(tf_numa_-1_Eigen) failure
• prlimit --as=4GB → exit 134 + pthread_create(tf_Compute) failure
• prlimit --as=8GB → JS-level Array buffer allocation failed thrown
• Amplification ratio: 31,935,321× (single units int → 4-byte-each weight allocation)

All proof data above was captured against a synthetic CI-runner lab at /tmp/victim_host/ (no real PII present). Full capture: F12_REAL_IMPACT_PROOF_2026-06-11.txt.

---

Summary

A Node.js service that calls tf.loadLayersModel on an attacker-supplied model.json will be OOM-killed by a 540-byte attacker artefact, as @tensorflow/tfjs-layers reads the Dense.units value verbatim from attacker-controlled Keras config and uses it to size the kernel tensor. assertPositiveInteger permits any positive integer up to Number.MAX_SAFE_INTEGER (2⁵³); Dense.build then asks libtensorflow to allocate [inputDim, units] × float32 bytes. A units = 2²⁴ with batch_input_shape = [null, 256] requests 17,179,869,184 bytes (≈ 16 GiB) — measured 16,446 MB RSS in the audit lab. Amplification factor 31,935,321×.

Root Cause

Lines of Code:

• tfjs-layers/src/utils/generic_utils.ts L441-L454 (assertPositiveInteger)
• tfjs-layers/src/layers/core.ts L221 (assertPositiveInteger(this.units, 'units') in Dense.constructor)
• tfjs-layers/src/layers/core.ts L240-L256 (Dense.build)

In tfjs-layers/src/utils/generic_utils.ts:441-454:

export function assertPositiveInteger(value: number|number[], name: string) {
  if (Array.isArray(value)) {
    util.assert(value.length > 0, ...);
    value.forEach((v, i) => assertPositiveInteger(v, `element ${i+1}…`));
  } else {
    util.assert(
        Number.isInteger(value) && value > 0,
        () => `Expected ${name} to be a positive integer, but got …`);
  }
  // NO upper bound check!
}

The only constraints are Number.isInteger(value) && value > 0. There is no upper bound, so units = 2²⁴ (16,777,216) passes — and so would Number.MAX_SAFE_INTEGER.

In tfjs-layers/src/layers/core.ts:240-256 (Dense.build):

public override build(inputShape: Shape|Shape[]): void {
  inputShape = getExactlyOneShape(inputShape);
  const inputLastDim = inputShape[inputShape.length - 1];
  if (this.kernel == null) {
    this.kernel = this.addWeight(
        'kernel', [inputLastDim, this.units], null, this.kernelInitializer,
        this.kernelRegularizer, true, this.kernelConstraint);
    ...
  }
  this.built = true;
}

build runs during loadLayersModel itself (Sequential knows its input shape from batch_input_shape in the saved config), so the unbounded allocation happens at load — no predict() call required.

Why this is NOT a duplicate of the GraphModel allocation finding (F-16): F-16 covers creation_executor.ts (Fill/Ones/Zeros/RandomUniform/ Range) — the GraphModel path. F-12 covers the LayersModel path — Dense.units, Embedding.inputDim, Embedding.outputDim, Conv*.filters, LSTM.units, RNN.units, GRU.units, batchInputShape. Different file (tfjs-layers/src/utils/generic_utils.ts vs tfjs-converter/src/operations/executors/creation_executor.ts), different attacker JSON shape (Keras layer config vs GraphDef Const nodes), different trigger (load-time during Layer.build vs execute-time during model.execute). A maintainer who patches one will ship the other.

Internal Pre-conditions

1. Victim Node.js process or browser dashboard calls tf.loadLayersModel(<url>) (or tf.loadModel, or any LayersModel-based AutoML / model-marketplace ingestion flow).
2. The process uses @tensorflow/tfjs-node, @tensorflow/tfjs-node-gpu, or @tensorflow/tfjs (any package that bundles tfjs-layers).
3. The host's available heap is less than inputDim × units × dtype_bytes for the attacker-chosen units. With units = 2²⁴ and inputDim = 256 that is ~16 GiB — well above any commodity host's free heap.

External Pre-conditions

None. The bug is entirely internal to the loader.

Attack Path

1. Attacker authors a model.json containing a Sequential model with one InputLayer (batch_input_shape: [null, 256]) and one Dense layer with units: 16777216 and kernel_initializer: { class_name: 'Zeros' }. The whole file is 540 bytes.
2. Attacker delivers the file to the victim — a user upload field, a model registry URL, a browser-side tfjs-vis "browse before download" UI, a CI step.
3. Victim calls tf.loadLayersModel('file:///path/to/uploaded/model.json').
4. Sequential.fromConfig instantiates the layers. Dense.build runs because the input shape is known statically.
5. addWeight('kernel', [256, 16777216], …) reaches the libtensorflow allocator, which emits: tensorflow/core/framework/cpu_allocator_impl.cc:82] Allocation of 17179869184 exceeds 10% of free system memory. and then completes the allocation if the host has the space, or aborts with RESOURCE_EXHAUSTED if it does not.
6. On any host with less than 16 GiB free heap, the process is OOM-killed before loadLayersModel returns to the caller.

Impact

Quantified from the captured PoC run on @tensorflow/tfjs-node@4.22.0:

Metric	Value
Attacker model.json size	540 bytes
Server RSS allocated during load	16,446 MB
Amplification factor	31,935,321×
libtensorflow allocator warning	Allocation of 17179869184 exceeds 10% of free system memory
Result on a 4 GiB host	OOM-kill of the inference process
Result on a 16 GiB host	Sustained 16 GiB allocation, near-certain OOM-kill under any concurrent load

Service-level impact:

• Any tfjs-node service that loads model.json from sources not under full operator control (AutoML pipelines, model marketplaces, file-upload features, browser pages loading models from attacker-controlled CDN paths) can be DoS'd by a single 540-byte file. The bug is purely on the loader path — no inference call required.

Extended Impact — same-root-cause manifestations

The same assertPositiveInteger accepts unbounded values for every other shape-axis field read from the Keras config. The same one-line fix closes all of them at once:

Layer / param	Lookup site	Allocation cost
Dense.units	core.ts:221 (this finding)	inputDim × units × 4 B
Embedding.inputDim	embeddings.ts	inputDim × outputDim × 4 B
Embedding.outputDim	embeddings.ts	same
Conv*.filters	convolutional.ts	prod(kernelShape) × filters × 4 B
LSTM.units / GRU.units / RNN.units	recurrent.ts	units² × 4 B (recurrent kernel)
batchInputShape axes	engine/topology.ts:488	prod(axes) × 4 B for the input tensor

A single guard in assertPositiveInteger (or a wrapper assertReasonableShapeDim) that caps each axis at e.g. 2²² = 4,194,304 elements covers every one.

PoC

The repository ships a package.json so install is one step. Tested on Node 22 + @tensorflow/tfjs-node@4.22.0.

git clone https://huggingface.co/martilaio/tfjs-layers-dense-units-oom-poc
cd tfjs-layers-dense-units-oom-poc
npm install                  # pulls every dep from package.json
node reproduce.js            # minimal canary PoC — primitive proven
bash reproduce_real_impact.sh

Captured signal lands in F12_REAL_IMPACT_PROOF_2026-06-11.txt (sanitized; collected against the synthetic /tmp/victim_host/ CI-runner lab).

Mitigation

In tfjs-layers/src/utils/generic_utils.ts:441-454, add an upper-bound check for the shape-dimension call sites:

// Hard ceiling per shape axis. 2**22 ≈ 4 M elements per axis is well above
// any real model and far below the smallest VM's heap budget.
export const MAX_LAYER_SHAPE_DIM = 1 << 22;

export function assertPositiveInteger(value: number|number[], name: string) {
  if (Array.isArray(value)) {
    util.assert(value.length > 0, …);
    value.forEach((v, i) => assertPositiveInteger(v, `element ${i+1}…`));
  } else {
    util.assert(
      Number.isInteger(value) && value > 0 && value <= MAX_LAYER_SHAPE_DIM,
      () => `Expected ${name} to be a positive integer ≤ ${MAX_LAYER_SHAPE_DIM},
             but got ${formatAsFriendlyString(value)}.`);
  }
}

Apply at every layer-config site that reads a shape dimension from the model file (Dense.units, Embedding.inputDim, Embedding.outputDim, Conv*.filters, LSTM.units, RNN.units, GRU.units, batchInputShape).

CVSS

CVSS 3.1 7.5 / High — AV:N/AC:L/PR:N/UI:R/S:U/C:N/I:N/A:H.

A:H — process termination. UI:R — victim must invoke loadLayersModel. PR:N, AV:N — artifact arrives over the network.

Bug classification

• CWE-1284 (Improper Validation of Specified Quantity in Input)
• CWE-770 (Allocation of Resources Without Limits or Throttling)

Affected versions

@tensorflow/tfjs-layers ≤ 4.22.0 (bundled in @tensorflow/tfjs-node, @tensorflow/tfjs-node-gpu, @tensorflow/tfjs).

Files in this repository

File	Purpose
README.md	this disclosure
reproduce.js	minimal PoC — Dense.units = 1e9 triggers Array buffer allocation failed
reproduce_real_impact.sh	host-OOM emulation — runs the PoC under prlimit --as=1/2/4/8 GB to capture exit signal
F12_REAL_IMPACT_PROOF_2026-06-11.txt	captured exit-code 134 (SIGABRT) at 1/2/4 GB host emulation; V8 / pthread_create stack traces