ericblackgachara/tflite-oom-poc

TFLite CWE-789: Uncapped Tensor Size Allocation

Status: CONFIRMED — READY TO SUBMIT

Severity: Medium-High (P2) — DoS

Target

• Repo: tensorflow/tensorflow (tensorflow/lite/)
• Platform: huntr.com
• Format: .tflite (FlatBuffers)

What Exists vs What's Missing

Protection	Present?	Notes
Integer overflow guard (MultiplyAndCheckOverflow)	✅ Yes	Prevents wrapping (e.g., INT64_MAX * 2)
Per-dimension value cap	❌ No	dim=32768 accepted without limit
Maximum tensor byte cap (kMaxTensorBytes)	❌ No	4 GB allocation reaches aligned_alloc unchecked

Vulnerable Call Chain

.tflite file → FlatBuffers parse
  interpreter_builder.cc:598   FlatBufferIntArrayToVector(tensor->shape())  ← no per-dim cap
  util.cc:220-256               BytesRequired() → MultiplyAndCheckOverflow   ← overflow only, no size cap
  arena_planner.cc              arena_.Allocate(tensor.bytes)                ← uncapped
  simple_memory_arena.cc        aligned_alloc(allocation_size)               ← 4 GB, no ceiling

Exploit

Shape [1, 32768, 32768] float32 → 32768 × 32768 × 4 = 4,294,967,296 bytes (4 GB). BytesRequired() returns kTfLiteOk — no integer overflow. No size ceiling checked. aligned_alloc(4 GB) fires; on Linux with default overcommit, malloc succeeds and OOM-killer fires.

PoC Files

• poc_tflite_oom.py — builds evil.tflite and triggers allocate_tensors()
• evil.tflite — generated by the script

Reproduction

pip install flatbuffers tflite-runtime
python poc_tflite_oom.py

Fix

Add a size cap in BytesRequired() (tensorflow/lite/util.cc):

constexpr size_t kMaxTensorBytes = 1ULL << 30;  // 1 GB
if (*bytes > kMaxTensorBytes) return kTfLiteError;