Skip to content

fix: prevent integer overflow in tensor byte-size calculations#3588

Open
Ashutosh0x wants to merge 1 commit into
tensorflow:mainfrom
Ashutosh0x:fix/integer-overflow-memory-helpers
Open

fix: prevent integer overflow in tensor byte-size calculations#3588
Ashutosh0x wants to merge 1 commit into
tensorflow:mainfrom
Ashutosh0x:fix/integer-overflow-memory-helpers

Conversation

@Ashutosh0x

@Ashutosh0x Ashutosh0x commented Jun 3, 2026

Copy link
Copy Markdown

Summary

Fix signed integer overflow in BytesRequiredForTensor, TfLiteEvalTensorByteLength, and AllocateOutputDimensionsFromInput that leads to heap out-of-bounds read/write when loading untrusted models.

Vulnerability (#3552)

All three functions compute a running product of tensor dimensions using signed int arithmetic. A model with shape [65536, 65536] over float32 requires 2^34 bytes (17 GiB), but the signed int product wraps to 0. This causes a 0-byte allocation while the tensor metadata holds the original dimensions, leading to heap OOB read/write during kernel execution.

Fix

  • Switch running product from int to size_t
  • Add checked multiplication (SIZE_MAX / x) before each multiply
  • Reject dimensions <= 0
  • Return kTfLiteError on overflow instead of silently wrapping

All three vulnerable functions in memory_helpers.cc are fixed with the same pattern.

BUG=#3552

@Ashutosh0x Ashutosh0x requested a review from a team as a code owner June 3, 2026 07:32
@Ashutosh0x

Copy link
Copy Markdown
Author

Hi @advaitjain — this fixes the integer overflow in tensor byte-size calculations reported in #3552.

All three functions (BytesRequiredForTensor, TfLiteEvalTensorByteLength, AllocateOutputDimensionsFromInput) use signed int for the running product of tensor dimensions, which wraps to 0 on shapes like [65536, 65536]. The allocator then gives a 0-byte buffer while the tensor metadata says it's huge — classic heap OOB write.

The fix uses size_t with checked multiplication (SIZE_MAX / x) before each multiply, and rejects dimensions <= 0 with kTfLiteError. Happy to adjust!

BytesRequiredForTensor, TfLiteEvalTensorByteLength, and
AllocateOutputDimensionsFromInput use signed int for the running
product of tensor dimensions. A model with shape [65536, 65536]
over float32 requires 2^32 * 4 = 17 GiB, but the signed int product
wraps to 0, causing a 0-byte allocation while the tensor metadata
holds the original huge dimensions. This leads to heap OOB
read/write during kernel execution.

Fix: switch to size_t with checked multiplication (SIZE_MAX / x)
before each multiply. Reject dimensions <= 0. Return kTfLiteError
on overflow instead of silently wrapping.

Fixes tensorflow#3552
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant