TheAuroraAi/tf-savedmodel-oom-poc

TF SavedModel OOM PoC — CWE-789

A 62-byte TensorFlow SavedModel file that forces 4 GB memory allocation when loaded.

Summary

CWE-789: Uncontrolled Memory Allocation via Const float_val broadcasting.

A single float_val in a Const node TensorProto is replicated to fill shape [2^30], forcing TensorFlow to allocate 4,294 MB from a 62-byte file.

• Amplification: 62 bytes → 4,294 MB (69 million to 1)
• Affected: TensorFlow 2.x (tested 2.21.0, latest)
• Attack: Load untrusted SavedModel files

Reproduction

import tensorflow as tf
import resource
# Limit to 512 MB to safely observe the OOM
resource.setrlimit(resource.RLIMIT_AS, (512*1024*1024, 512*1024*1024))
try:
    tf.saved_model.load('.')  # Will attempt 4 GB allocation, OOM
except MemoryError:
    print("OOM confirmed")

Root Cause

tensorflow/core/kernels/constant_op.cc:83 calls Tensor::FromProto() which broadcasts float_val[0] to fill all N elements. No size limit on N.