Integer overflow in tensor size calculations leads to heap buffer overflow

[chupaohong]

Vulnerability

BytesRequiredForTensor, TfLiteEvalTensorByteLength, ElementCount, and RuntimeShape::FlatSize use int to accumulate element counts. A malformed model with large tensor dimensions (e.g. shape [1024, 1024, 4096]) overflows int and wraps to zero. Kernel ops then write past the undersized buffer, causing out of bounds write and memory corruption.

Same bug class as #3516, but in the allocation helpers rather than a single kernel, so all operators are affected.

Attack scenario

Any application that loads a .tflite model from an external source (OTA model update, SD card, model marketplace, user upload) is vulnerable. The attacker only needs to modify the tensor shape field in the FlatBuffer; no special configuration or flags are required on the target.

tflite-micro runs primarily on microcontrollers where ASLR and heap guards are usually absent, so going from heap corruption to code execution is not difficult. The attacker controls both the allocation size (via tensor dimensions) and the data that gets written (via model weights).

On Android or Linux hosts the same overflow causes heap corruption and likely a crash. Code execution depends on heap layout.

Root cause

• int element_count in BytesRequiredForTensor (memory_helpers.cc:109) and TfLiteEvalTensorByteLength (:130)
• int result in ElementCount (micro_utils.cc:30)
• int buffer_size in RuntimeShape::FlatSize (runtime_shape.h:97)

All four overflow on large positive dimensions and the result feeds into arena allocation sizes.

Fix

Change int to size_t for the accumulator in each location and cast individual dims to size_t before multiplying.

Verification

Before: BytesRequiredForTensor returns 0 for shape [1024, 1024, 4096] uint8.
After: returns 4294967296.

BUG=None

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[chupaohong]

Hello,

For detailed information about the exploit code (poc), please review the ticket via link: https://issuetracker.google.com/issues/501064564

Thanks