Fix int overflow in tensor byte-size calculation

[Forbiddem]

BUG=#3552

Re-submission of #3551. That PR was auto-closed by the stale bot on 2026-06-22 before any review took place — the only outstanding gate was the maintainer-approved CI matrix (approval-gate / call-*), which was never triggered for an external contributor. The PR could not be reopened (the closed-PR head ref had become detached and both API reopen and push-to-branch were refused), so I'm re-submitting the identical commit. CLA is signed and the local tests pass.

Summary

BytesRequiredForTensor, TfLiteEvalTensorByteLength, and
AllocateOutputDimensionsFromInput in
tensorflow/lite/micro/memory_helpers.cc all multiply tensor shape dimensions
into a signed 32-bit int element_count running product before assigning the
result to a size_t byte count. Shape dimensions and the element type size
come from the (untrusted) flatbuffer model, so a model with a shape such as
[65536, 65536] over a 4-byte element type produces a product of 2^32, which
wraps to 0 in the original code and yields *bytes = 0. Subsequent buffer
arithmetic in the allocator then operates on a tensor whose advertised size
is 0 while its dimensions imply ~17 GiB of storage, opening a memory
corruption surface when a model is loaded from an untrusted source.

This PR:

• Performs the running product in size_t, matching the result type, and
  rejects negative shape dimensions with kTfLiteError.
• Guards every multiplication (shape * shape and elements * type_size)
  against overflow of size_t and returns kTfLiteError when the next
  multiplication would wrap.
• Propagates the previously discarded status of TfLiteTypeSizeOf in
  AllocateOutputDimensionsFromInput.

Test plan

• make -f tensorflow/lite/micro/tools/make/Makefile memory_helpers_test
  builds clean.
• gen/linux_x86_64_default_gcc/bin/memory_helpers_test — all 10 cases
  pass, including the 3 new regression tests:
  • TfLiteEvalTensorByteLengthDoesNotTruncateAcrossInt32Boundary — shape
    [65536, 65536] float32 now reports 17,179,869,184 bytes on 64-bit
    size_t platforms (previously: 0), and kTfLiteError on 32-bit
    size_t platforms.
  • TfLiteEvalTensorByteLengthRejectsSizeTOverflow — a shape that
    overflows even a 64-bit size_t is rejected with kTfLiteError.
  • TfLiteEvalTensorByteLengthRejectsNegativeDimension — a negative
    shape dimension is rejected with kTfLiteError.

Notes for review

• The overflow checks use std::numeric_limits<size_t>::max() / x, which is
  portable and MSVC-safe — deliberately not __builtin_mul_overflow, which
  would break the call-windows build.
• @veblush — since this is a memory-safety fix for untrusted-model parsing,
  would you be able to approve the workflow run and take a look when you have a
  moment? Happy to switch the manual checks to any form the tree prefers and to
  address any other feedback.

[veblush]

Thank you for the PR. Please read the Error Handling & Defensive Programming Guide and modify this PR accordingly.

[Forbiddem]

Thanks @veblush — I've gone through the Error Handling & Defensive Programming Guide and reworked the PR to match its macro/phase model. My reasoning per touched function:

Function	Phase it runs in	TfLiteContext*?	Macro (per guide)
BytesRequiredForTensor	Setup (MicroAllocator)	no	low-cost return kTfLiteError;
TfLiteEvalTensorByteLength	Setup planning and Eval (squeeze, kernel_util)	no	low-cost return kTfLiteError;
AllocateOutputDimensionsFromInput	Setup/Prepare (kernel helper)	yes	TF_LITE_ENSURE

I'm treating an oversized/negative shape as an Invalid Model Topology (semantic malformation) — "wrong tensor shapes" that the guide says the runtime should defend against in the Setup phase — rather than a structural corruption of the FlatBuffer itself (the shape vector and its offsets are intact; only the product overflows size_t, silently yielding *bytes = 0 for a tensor whose dims imply ~17 GiB). That places the fix squarely in the "Preventing Buffer Overflows" bucket.

Changes in this revision:

• AllocateOutputDimensionsFromInput is the only edited function that receives a TfLiteContext, so it keeps TF_LITE_ENSURE (guide default for Setup/Prepare). Following "The Hidden Cost of TF_LITE_ENSURE", I collapsed the two per-iteration checks (dim >= 0 and the overflow guard) into a single TF_LITE_ENSURE so only one error string lands in .rodata per call site. dim >= 0 is the first && term, so the division guard is never reached with a wrapped-around udim.
• BytesRequiredForTensor and TfLiteEvalTensorByteLength have no TfLiteContext, so per the cheat sheet I use the low-cost if (!cond) return kTfLiteError; branch (no embedded __FILE__/__LINE__/#cond strings). TfLiteEvalTensorByteLength is additionally reachable from Eval, where the guide prescribes exactly this raw-return form for runtime memory-corruption guards.
• Kept TF_LITE_ENSURE_STATUS for propagating TfLiteTypeSizeOf's status.
• Regression tests use GoogleTest EXPECT_* macros, matching the guide's unit-test guidance.

One judgement call I'd like your read on: if you'd instead classify the shape-product overflow as a structural FlatBuffer concern (i.e. the app layer owns validating untrusted/OTA models before handing them to TFLM), then per the guide the correct form is a zero-cost TFLITE_DCHECK developer aid rather than production returns. I'm happy to convert all three sites to TFLITE_DCHECK if that's the bucket you prefer — just let me know and I'll push that variant.