Integer overflow in tensor byte-size calculation in memory_helpers.cc

[Forbiddem]

Summary

tensorflow/lite/micro/memory_helpers.cc computes the byte length of tensors using a 32-bit signed int running product before assigning the result to a size_t. Because tensor shape dimensions and element type sizes come from the (untrusted) flatbuffer model, a model with a shape such as [65536, 65536] over a 4-byte element type produces an element-count product of 2^32, which wraps to 0 in the existing code and causes *bytes = 0.

The same pattern appears in three places:

• BytesRequiredForTensor (called during initial allocation from the flatbuffer model)
• TfLiteEvalTensorByteLength
• AllocateOutputDimensionsFromInput

A downstream allocator that trusts *bytes then operates on a tensor whose advertised size is 0 while its dimensions imply ~17 GiB of storage, which is a memory corruption surface when models are loaded from an untrusted source.

Reproducer

The arithmetic from BytesRequiredForTensor extracted into a self-contained C++ program:

size_t VulnerableBytesRequiredForTensor(const int32_t* shape, size_t rank,
                                        size_t type_size) {
  int element_count = 1;
  for (size_t n = 0; n < rank; ++n) {
    element_count *= shape[n];
  }
  return element_count * type_size;
}

int main() {
  const int32_t shape[] = {65536, 65536};
  const size_t type_size = 4;  // float32
  // Mathematically: 17,179,869,184 bytes.
  // Observed:       0 (signed 32-bit wraparound).
  std::cout << VulnerableBytesRequiredForTensor(shape, 2, type_size) << "\n";
}

Proposed fix

See PR — running product moved to size_t, every multiplication guarded against size_t overflow, negative dimensions rejected, and three regression tests added in memory_helpers_test.cc.

[ZX41R]

One detail to keep separate in the patch: AllocateOutputDimensionsFromInput() currently uses the same size variable for byte-count accumulation and then passes it to TfLiteIntArrayGetSizeInBytes(size). That helper wants the number of dimension entries, not the tensor byte length.

So while fixing the overflow, it would be safer to split the state into something like element_count, bytes, and dimensions_count, then allocate output->dims with TfLiteIntArrayGetSizeInBytes(dimensions_count). Otherwise a correct byte-size fix can still leave this function harder to audit, and tests may miss the accidental coupling because over-allocation usually does not fail.

For the arithmetic itself, size_t is the right destination type, but TFLM still runs on 32-bit targets. The regression set should include a case that overflows 32-bit size_t as well as the current signed-int wrap case, so the guard is real on MCU-class builds too.

[Forbiddem]

Good catches, both folding into the next push.

Confirmed the size reuse — it ends up as the tensor byte count and then gets handed to TfLiteIntArrayGetSizeInBytes, which wants the dim slot count. Splitting into
type_size / element_count / bytes and allocating output->dims with dimensions_count. Adding the missing final-multiply guard while I'm in there.

For 32-bit coverage, adding regression cases that overflow at the per-dim multiply and the final type-size multiply on uint32_t, alongside the existing signed-wrap
case.

[ZX41R]

Nice. One small caveat on the 32-bit tests: if those run as normal host unit tests on a 64-bit CI worker, a shape that would overflow uint32_t still will not overflow size_t, so that path may not actually exercise the MCU-sized failure mode.

If there is no 32-bit test job for this file, I would either factor the checked multiply into a tiny helper that can be tested with an explicit max/cap, or keep a dedicated uint32_t accumulator test alongside the real size_t path. Otherwise the test proves the old signed-int wrap is gone, but the 32-bit size_t guard is only covered by inspection.

[Forbiddem]

Fair point — you're right that a uint32_t-overflowing shape pushed through a size_t accumulator on 64-bit CI doesn't exercise the guard at all, it just multiplies
cleanly.

Going with option one: pulling the checked multiply into a tiny templated helper, CheckedMulInPlace(T cap, T& acc, T factor), called from the production path with
cap = std::numeric_limits<size_t>::max() and from tests with cap = UINT32_MAX (and a small cap for boundary fuzz). That way the 32-bit guard is exercised on any host,
and the production call site stays a one-liner.

[ZX41R]

Helper shape looks right. Two thoughts:

1. With cap = std::numeric_limits<size_t>::max() the production-path branch should optimize out at -O2; worth a quick objdump -d on the built object to confirm the guard isn't paying for itself in MCU code. If it isn't elided, a __builtin_expect on the overflow branch or a template <size_t Cap> non-type parameter (so the cap is a compile-time constant) handles it.

2. For the boundary fuzz, a dedicated small cap (e.g. 1024) is more discriminating than UINT32_MAX — UINT32_MAX mostly re-exercises the original wrap case, while a small cap forces the multiply-then-check ordering to be right (acc * factor checked before assignment, not after). Both are fine to keep; just calling out that the small-cap test catches a different mistake.

[Forbiddem]

Both fair, taking both into the next push.

On the codegen: running objdump -d on the -O2 build of memory_helpers.o to confirm the acc > cap / factor branch is elided when cap ==
std::numeric_limits<size_t>::max() (since SIZE_MAX / factor should constant-fold and the compare-against-SIZE_MAX goes away after constant propagation). If it isn't
elided cleanly across the toolchains TFLM cares about, going with the template <size_t Cap> non-type parameter rather than __builtin_expect — __builtin_expect is only a
branch-prediction hint, the compare itself still emits, while a compile-time Cap lets the whole guard fold away deterministically. Production call site becomes
CheckedMulInPlace<std::numeric_limits<size_t>::max()>(acc, factor), test call site instantiates with UINT32_MAX or 1024. Will post the disassembly diff in the PR for
the record.

On the boundary fuzz: good distinction, you're right that UINT32_MAX only catches "did the type get wider," not "is the check on the right side of the assignment."
Adding a dedicated Cap = 1024 test with acc = 64, factor = 32 (product = 2048, no native wrap, but exceeds cap) so a buggy acc *= factor; if (acc > cap) ordering would
still report success on the assigned value and fail the test. Keeping the UINT32_MAX case too since it documents the original 32-bit MCU failure mode — the two tests
are checking different invariants.

Will tag you on the updated PR.

[ReplikanteK]

Additional affected functions found (same overflow pattern):

1. ElementCount() in micro_utils.cc:30-36 — used by EvalTensorBytes() and many kernels
2. EvalTensorBytes() in micro_utils.cc:38-43 — uses ElementCount() under the hood
3. AllocateOutputDimensionsFromInput() in memory_helpers.cc:143 — reuses same size variable for byte counting and dim allocation

Same root cause: int element_count overflow with large tensor dimensions.

We independently confirmed these via code review and built an ASAN fuzzer harness (FuzzOpResolver accepting any operator, 2000+ models tested). Happy to share if helpful.

[Forbiddem]

Thanks for the cross-check.

Triage on each:

• AllocateOutputDimensionsFromInput() (memory_helpers.cc:143) — already in the next push. @ZX41R flagged the size variable reuse on day one; the current change splits it into type_size / bytes / dimensions_count, routes the multiply through CheckedMulInPlace<size_t>, and feeds dimensions_count (not bytes) into TfLiteIntArrayGetSizeInBytes. The accumulator there is already size_t so the int wrap isn't the issue, but the unchecked size_t *= dim is still 32-bit-vulnerable on MCU targets and is now guarded.

• ElementCount() (micro_utils.cc:30-36) and EvalTensorBytes() (38-43) — confirmed, same int result = 1; result *= dims.data[i] pattern, and EvalTensorBytes then does ElementCount(...) * bytes_per_element with no guard on the final multiply either. Folding into the same patch series rather than a follow-up:

  • ElementCount accumulator promoted to size_t, per-dim multiply routed through CheckedMulInPlace<size_t>(SIZE_MAX, acc, factor).
  • EvalTensorBytes becomes the canonical bytes = element_count; CheckedMulInPlace(bytes, bytes_per_element); return bytes;.
  • Caller signature for ElementCount becomes bool ElementCount(const TfLiteIntArray&, size_t* out) returning overflow status — the fan-out into kernels (add, mul, conv, softmax, …) means silently truncating to int on success would just re-introduce the bug at every call site. Each kernel site gets the same TF_LITE_ENSURE_* treatment as the rest of TFLM.

One PR rather than a follow-up because partial coverage is worse than the current uniform state — kernel A taking the safe path while kernel B still trusts a wrapped int is harder to reason about and to test cleanly than fixing the whole fan-out at once.

Yes please on the fuzzer — FuzzOpResolver accepting any operator with a 2000-model corpus is exactly the regression evidence the patch needs. If you can drop a gist or branch link, I'll thread it into the PR description as the validation reference and credit the additional sites to you in the changelog.