hacnho/tensorflow-savedmodel-variadic-n-dos-poc

TensorFlow SavedModel variadic N attribute load-time DoS PoC

This repository contains a proof of concept for a TensorFlow SavedModel loader denial of service in tensorflow-cpu==2.21.0.

Summary

tf.saved_model.load(model_dir) validates NodeDef entries inside saved_model.pb. For variadic ops such as Pack and StringJoin, TensorFlow trusts the serialized N attribute while building the expected input type signature. A malicious SavedModel can set N=250000000 while keeping only the original small number of inputs. During load, TensorFlow expands the expected signature before rejecting the malformed node, causing multi-GB memory use and CPU burn from an approximately 8 KB saved_model.pb.

Contents

• models/control_saved_model/: valid control SavedModel.
• models/pack_n_250m_saved_model/: malicious SavedModel with Pack.N=250000000.
• models/string_join_n_250m_saved_model/: malicious SavedModel with StringJoin.N=250000000.
• repro_savedmodel_variadic_n_dos.py: bounded reproduction script.
• evidence/n-sweep-time-rss-20260623.json: measured N sweep showing RSS growth.
• evidence/structured-load-v2-20260623.json: structured mutant loader results.

Reproduction

Use Linux with GNU time and timeout.

python3 -m venv .venv
. .venv/bin/activate
pip install -r requirements.txt
python repro_savedmodel_variadic_n_dos.py --exploit pack --timeout 12 --json-out reproduce-pack.json
python repro_savedmodel_variadic_n_dos.py --exploit string_join --timeout 12 --json-out reproduce-string-join.json

Expected result:

• The control SavedModel loads successfully.
• The malicious SavedModel does not complete within the timeout.
• /usr/bin/time reports multi-GB maximum RSS for the malicious load.

On the reporter machine, the packaged reproduction produced:

• Pack.N=250000000: timeout after 12 seconds at 6.39 GiB RSS.
• StringJoin.N=250000000: timeout after 12 seconds at 5.43 GiB RSS.

Root Cause

The root cause is in TensorFlow node_def_util.cc.

In AddArgToSig(), the number_attr value is read from the untrusted serialized NodeDef. If it fits in int32, TensorFlow pushes one dtype per repeat into DataTypeVector:

for (int i = 0; i < repeats; ++i) {
  sig->push_back(dtype);
}

ValidateNodeDef() then compares the expanded expected input vector with the actual serialized input count. On mismatch, it formats the entire expanded vector with DataTypeVectorString(inputs) for the exception message.

This allows a tiny SavedModel file to force allocation and formatting work linear in attacker-controlled N.

Environment

• TensorFlow: tensorflow-cpu==2.21.0
• Trigger: tf.saved_model.load(model_dir)
• Affected file: saved_model.pb
• Impact: local denial of service when a victim loads or validates an untrusted SavedModel.