ericblackgachara/tf-dos-poc

TensorFlow DoS PoC — Unbounded Tensor Dimension in MakeNdarray()

Vulnerability

The tensorflow Python package (v2.21.0) fails to validate tensor dimensions before allocating memory in tf.make_ndarray() / MakeNdarray(). A crafted 70-byte .pb GraphDef file containing a Const node with tensor_shape.dim=[5,000,000,000] but only 1 actual float value causes np.pad() to request ~18.6 GiB of memory — crashing the process instantly.

• CWE: CWE-789 — Memory Allocation with Excessive Size Value
• CVSS: 7.5 High (AV:N/AC:L/PR:N/UI:N/S:U/C:N/I:N/A:H)
• Affected: tensorflow Python package, tested on v2.21.0

Affected Code Paths

Three separate allocation paths in MakeNdarray() are vulnerable — all use attacker-controlled tensor.tensor_shape.dim with no bounds check:

Vector	Data Type	Serialized Size	Vulnerable Call	Result
np.pad()	DT_FLOAT (1 value, shape 5B)	18 bytes	np.pad(values, (0, num_elements - values.size), "edge")	MemoryError in 0.00s
np.zeros()	DT_BFLOAT16 (0 values, shape 10B)	12 bytes	np.zeros(shape, dtype)	MemoryError in 0.00s
String padding	DT_STRING (1 value, shape 1B)	15 bytes	values.extend([last] * padding)	MemoryError in 0.00s

Real-World Attack Surface

MakeNdarray() is called internally during multiple TensorFlow workflows:

• SavedModel loading — saved_model_utils.py:_deserialize_from_proto() calls MakeNdarray() on Const node tensors
• Model optimization — optimize_for_inference_lib.py during constant folding passes
• Constant evaluation — tensor_util.py:_ConstantValue() during graph analysis
• Debug/inspection tools — debug_data.py and debug_events_reader.py
• Public API — tf.make_ndarray() documented for TensorProto-to-numpy conversion

Files

File	Description
craft_malicious_tf.py	Generates the 70-byte malicious .pb GraphDef file
demo_dos_tf.py	Runs the PoC with a 512 MB memory safety cap
malicious_tf_dos.pb	The 70-byte malicious GraphDef file

Reproduction

pip install tensorflow  # tested on 2.21.0

# Step 1: Generate malicious file (or use the included one)
python craft_malicious_tf.py

# Step 2: Trigger the crash
python demo_dos_tf.py

Expected Output

============================================================
TensorFlow DoS PoC - Unbounded Tensor Dimension Allocation
============================================================

TensorFlow version: 2.21.0

[1] Crafted TensorProto:
    Serialized size: 18 bytes
    Declared shape: [5,000,000,000]
    Actual float_val count: 1
    Expected allocation: ~18.6 GB (5B * 4 bytes/float32)

[2] Memory limit set to 512 MB (safety cap)

[3] Calling tf.make_ndarray(tensor)...

[CONFIRMED] MemoryError triggered after 0.00 seconds
    Error: Unable to allocate 18.6 GiB for an array with shape
    (5000000000,) and data type float32

    DoS vulnerability confirmed - 18-byte protobuf caused OOM crash

Impact

TensorFlow has 200M+ monthly pip downloads and is the most widely deployed ML framework. The SavedModel .pb format is the standard serialization format. A malicious .pb uploaded to any model repository would crash automated inspection pipelines, production inference servers (TF Serving), model converters (TF-to-TFLite, TF-to-ONNX), cloud ML platforms, and any user loading the model's constants.

Disclosure

Submitted to huntr.com — TensorFlow SavedModel (.pb) – Google target.