From 20b8aaf2642714f45ef8075ba0db9cc42a986cb3 Mon Sep 17 00:00:00 2001 From: Forbiddem Date: Fri, 8 May 2026 10:25:01 +0000 Subject: [PATCH 1/3] Fix int overflow in tensor byte-size calculation `BytesRequiredForTensor`, `TfLiteEvalTensorByteLength`, and `AllocateOutputDimensionsFromInput` in `tensorflow/lite/micro/memory_helpers.cc` all multiplied 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 are sourced from the (untrusted) flatbuffer model, so 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 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. This change: * 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`. Adds three regression tests in `memory_helpers_test.cc`: * Shape [65536, 65536] float32 reports the correct 17,179,869,184 byte length on platforms with a 64-bit `size_t` and `kTfLiteError` on platforms with a 32-bit `size_t` (instead of the pre-fix value of 0). * A shape that overflows even a 64-bit `size_t` is rejected with `kTfLiteError`. * A negative shape dimension is rejected with `kTfLiteError`. All existing `memory_helpers_test` cases continue to pass. --- tensorflow/lite/micro/memory_helpers.cc | 49 +++++++++++++++--- tensorflow/lite/micro/memory_helpers_test.cc | 52 ++++++++++++++++++++ 2 files changed, 95 insertions(+), 6 deletions(-) diff --git a/tensorflow/lite/micro/memory_helpers.cc b/tensorflow/lite/micro/memory_helpers.cc index d78e34d4d96..59568b832cc 100644 --- a/tensorflow/lite/micro/memory_helpers.cc +++ b/tensorflow/lite/micro/memory_helpers.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include "flatbuffers/flatbuffers.h" // from @flatbuffers #include "tensorflow/lite/c/common.h" @@ -106,12 +107,25 @@ TfLiteStatus TfLiteTypeSizeOf(TfLiteType type, size_t* size) { TfLiteStatus BytesRequiredForTensor(const tflite::Tensor& flatbuffer_tensor, size_t* bytes, size_t* type_size) { - int element_count = 1; + // Use size_t for the running product so the byte computation is performed + // in the same unsigned domain as the result. Each multiplication is guarded + // against overflow because shape dimensions and the element type size all + // come from the (untrusted) flatbuffer model. + size_t element_count = 1; // If flatbuffer_tensor.shape == nullptr, then flatbuffer_tensor is a scalar // so has 1 element. if (flatbuffer_tensor.shape() != nullptr) { for (size_t n = 0; n < flatbuffer_tensor.shape()->size(); ++n) { - element_count *= flatbuffer_tensor.shape()->Get(n); + const int32_t dim = flatbuffer_tensor.shape()->Get(n); + if (dim < 0) { + return kTfLiteError; + } + const size_t udim = static_cast(dim); + if (udim != 0 && + element_count > std::numeric_limits::max() / udim) { + return kTfLiteError; + } + element_count *= udim; } } @@ -119,6 +133,10 @@ TfLiteStatus BytesRequiredForTensor(const tflite::Tensor& flatbuffer_tensor, TF_LITE_ENSURE_STATUS( ConvertTensorType(flatbuffer_tensor.type(), &tf_lite_type)); TF_LITE_ENSURE_STATUS(TfLiteTypeSizeOf(tf_lite_type, type_size)); + if (*type_size != 0 && + element_count > std::numeric_limits::max() / *type_size) { + return kTfLiteError; + } *bytes = element_count * (*type_size); return kTfLiteOk; } @@ -127,15 +145,28 @@ TfLiteStatus TfLiteEvalTensorByteLength(const TfLiteEvalTensor* eval_tensor, size_t* out_bytes) { TFLITE_DCHECK(out_bytes != nullptr); - int element_count = 1; + size_t element_count = 1; // If eval_tensor->dims == nullptr, then tensor is a scalar so has 1 element. if (eval_tensor->dims != nullptr) { for (int n = 0; n < eval_tensor->dims->size; ++n) { - element_count *= eval_tensor->dims->data[n]; + const int dim = eval_tensor->dims->data[n]; + if (dim < 0) { + return kTfLiteError; + } + const size_t udim = static_cast(dim); + if (udim != 0 && + element_count > std::numeric_limits::max() / udim) { + return kTfLiteError; + } + element_count *= udim; } } size_t type_size; TF_LITE_ENSURE_STATUS(TfLiteTypeSizeOf(eval_tensor->type, &type_size)); + if (type_size != 0 && + element_count > std::numeric_limits::max() / type_size) { + return kTfLiteError; + } *out_bytes = element_count * type_size; return kTfLiteOk; } @@ -152,10 +183,16 @@ TfLiteStatus AllocateOutputDimensionsFromInput(TfLiteContext* context, input = input1->dims->size > input2->dims->size ? input1 : input2; TF_LITE_ENSURE(context, output->type == input->type); size_t size = 0; - TfLiteTypeSizeOf(input->type, &size); + TF_LITE_ENSURE_STATUS(TfLiteTypeSizeOf(input->type, &size)); const int dimensions_count = tflite::GetTensorShape(input).DimensionsCount(); for (int i = 0; i < dimensions_count; i++) { - size *= input->dims->data[i]; + const int dim = input->dims->data[i]; + TF_LITE_ENSURE(context, dim >= 0); + const size_t udim = static_cast(dim); + TF_LITE_ENSURE(context, + udim == 0 || + size <= std::numeric_limits::max() / udim); + size *= udim; } output->bytes = size; diff --git a/tensorflow/lite/micro/memory_helpers_test.cc b/tensorflow/lite/micro/memory_helpers_test.cc index a004bb61bcd..7e03d9faeb7 100644 --- a/tensorflow/lite/micro/memory_helpers_test.cc +++ b/tensorflow/lite/micro/memory_helpers_test.cc @@ -187,6 +187,58 @@ TEST(MemoryHelpersTest, TestBytesRequiredForTensor) { EXPECT_EQ(static_cast(4), type_size); } +TEST(MemoryHelpersTest, + TfLiteEvalTensorByteLengthDoesNotTruncateAcrossInt32Boundary) { + // Shape [65536, 65536] with float32: 65536 * 65536 * 4 = 17179869184 bytes. + // The original implementation used a signed 32-bit running product, so the + // element count wrapped to 0 and *out_bytes was reported as 0 even though + // any subsequent allocation would address ~17 GiB. The hardened + // implementation performs the arithmetic in size_t and must report the + // mathematically correct value (where size_t is wide enough), or refuse the + // request via kTfLiteError when it would otherwise overflow size_t. + int dims[] = {2, 65536, 65536}; + TfLiteEvalTensor eval_tensor = {}; + eval_tensor.dims = tflite::testing::IntArrayFromInts(dims); + eval_tensor.type = kTfLiteFloat32; + + size_t out_bytes = 0; + const TfLiteStatus status = + tflite::TfLiteEvalTensorByteLength(&eval_tensor, &out_bytes); + + if (sizeof(size_t) >= 8) { + EXPECT_EQ(kTfLiteOk, status); + EXPECT_EQ(static_cast(17179869184ULL), out_bytes); + } else { + // 32-bit size_t cannot represent the result; the hardened code must + // refuse rather than silently truncate. + EXPECT_EQ(kTfLiteError, status); + } +} + +TEST(MemoryHelpersTest, TfLiteEvalTensorByteLengthRejectsSizeTOverflow) { + // A shape whose product overflows size_t even on 64-bit platforms must be + // rejected. INT32_MAX^4 * 4 vastly exceeds 2^64. + int dims[] = {4, 0x7fffffff, 0x7fffffff, 0x7fffffff, 0x7fffffff}; + TfLiteEvalTensor eval_tensor = {}; + eval_tensor.dims = tflite::testing::IntArrayFromInts(dims); + eval_tensor.type = kTfLiteFloat32; + + size_t out_bytes = 0; + EXPECT_EQ(kTfLiteError, + tflite::TfLiteEvalTensorByteLength(&eval_tensor, &out_bytes)); +} + +TEST(MemoryHelpersTest, TfLiteEvalTensorByteLengthRejectsNegativeDimension) { + int dims[] = {2, -1, 4}; + TfLiteEvalTensor eval_tensor = {}; + eval_tensor.dims = tflite::testing::IntArrayFromInts(dims); + eval_tensor.type = kTfLiteFloat32; + + size_t out_bytes = 0; + EXPECT_EQ(kTfLiteError, + tflite::TfLiteEvalTensorByteLength(&eval_tensor, &out_bytes)); +} + TEST(MemoryHelpersTest, TestAllocateOutputDimensionsFromInput) { constexpr int kDimsLen = 4; int input1_dims[] = {1, 1}; From 6131b79a4a5f1c1a8fcc05b2e0870092ef7f1644 Mon Sep 17 00:00:00 2001 From: Forbiddem Date: Fri, 3 Jul 2026 04:56:43 +0000 Subject: [PATCH 2/3] Align overflow guards with Error Handling & Defensive Programming Guide Addresses review feedback on #3606: map each edited helper to its execution phase and pick the guide-prescribed macro. - AllocateOutputDimensionsFromInput is the only edited helper that receives a TfLiteContext and runs in Setup/Prepare, so it keeps TF_LITE_ENSURE. Per "The Hidden Cost of TF_LITE_ENSURE", the two per-iteration preconditions (non-negative dimension and no size_t overflow of the running product) are combined into a single TF_LITE_ENSURE so only one error string is emitted into .rodata per call site. The dim >= 0 term is evaluated first, so the division guard is never reached with a wrapped-around udim. - BytesRequiredForTensor and TfLiteEvalTensorByteLength have no TfLiteContext (and the latter is reachable from Eval), so they retain the guide's low-cost `return kTfLiteError;` branch, which carries no embedded __FILE__/__LINE__/#cond strings. - Comments now state the phase/trust-boundary rationale (oversized shape = invalid model topology, validated in the Setup phase). No behavioral change to the tests; the 3 regression cases still pass. --- tensorflow/lite/micro/memory_helpers.cc | 33 +++++++++++++++++++------ 1 file changed, 25 insertions(+), 8 deletions(-) diff --git a/tensorflow/lite/micro/memory_helpers.cc b/tensorflow/lite/micro/memory_helpers.cc index 59568b832cc..cb43f1eaf73 100644 --- a/tensorflow/lite/micro/memory_helpers.cc +++ b/tensorflow/lite/micro/memory_helpers.cc @@ -107,10 +107,14 @@ TfLiteStatus TfLiteTypeSizeOf(TfLiteType type, size_t* size) { TfLiteStatus BytesRequiredForTensor(const tflite::Tensor& flatbuffer_tensor, size_t* bytes, size_t* type_size) { - // Use size_t for the running product so the byte computation is performed - // in the same unsigned domain as the result. Each multiplication is guarded - // against overflow because shape dimensions and the element type size all - // come from the (untrusted) flatbuffer model. + // Shape dimensions and the element type size come from the partially-trusted + // flatbuffer model, so an oversized shape (e.g. [65536, 65536]) is an invalid + // model topology that -- per the Error Handling & Defensive Programming Guide + // -- must be rejected during the Setup phase (this runs under MicroAllocator). + // The running product is done in size_t (matching the result type) and every + // multiplication is guarded against wrap-around. This helper is not given a + // TfLiteContext, so the guide's low-cost `return kTfLiteError;` branch is used + // rather than TF_LITE_ENSURE. size_t element_count = 1; // If flatbuffer_tensor.shape == nullptr, then flatbuffer_tensor is a scalar // so has 1 element. @@ -145,6 +149,11 @@ TfLiteStatus TfLiteEvalTensorByteLength(const TfLiteEvalTensor* eval_tensor, size_t* out_bytes) { TFLITE_DCHECK(out_bytes != nullptr); + // Called from both the Setup planning path and from Eval helpers (e.g. + // squeeze, kernel_util). The running product is computed in size_t and each + // multiplication is guarded against wrap-around. Because there is no + // TfLiteContext here and the guard protects against runtime memory + // corruption, the guide's low-cost `return kTfLiteError;` branch is used. size_t element_count = 1; // If eval_tensor->dims == nullptr, then tensor is a scalar so has 1 element. if (eval_tensor->dims != nullptr) { @@ -187,11 +196,19 @@ TfLiteStatus AllocateOutputDimensionsFromInput(TfLiteContext* context, const int dimensions_count = tflite::GetTensorShape(input).DimensionsCount(); for (int i = 0; i < dimensions_count; i++) { const int dim = input->dims->data[i]; - TF_LITE_ENSURE(context, dim >= 0); const size_t udim = static_cast(dim); - TF_LITE_ENSURE(context, - udim == 0 || - size <= std::numeric_limits::max() / udim); + // This helper runs in the Setup/Prepare phase, so per the Error Handling & + // Defensive Programming Guide a bad (partially-trusted model) shape is + // rejected with TF_LITE_ENSURE. The two preconditions -- a non-negative + // dimension and a running product that does not overflow size_t -- are + // combined into a single TF_LITE_ENSURE so only one error string is emitted + // into .rodata per call site (see "The Hidden Cost of TF_LITE_ENSURE"). + // The `dim >= 0` term is evaluated first, so the division guard is never + // reached with a wrapped-around `udim`. + TF_LITE_ENSURE( + context, + dim >= 0 && + (udim == 0 || size <= std::numeric_limits::max() / udim)); size *= udim; } output->bytes = size; From 154950a12f94d7d0e5306f02226a0131da55a9e6 Mon Sep 17 00:00:00 2001 From: Forbiddem Date: Fri, 3 Jul 2026 05:05:16 +0000 Subject: [PATCH 3/3] Fix dim/byte coupling in AllocateOutputDimensionsFromInput + regression test Follow-up to 6131b79 that fixes the byte-count/dimension-count coupling @ZX41R flagged on #3552, in the existing inline-guard style (no new helper). AllocateOutputDimensionsFromInput accumulated the tensor byte count and then passed that same value to TfLiteIntArrayGetSizeInBytes(), which expects the *number of dimension entries*. For an int32 [5, 5, 5, 5] tensor that requested storage for 2500 int entries to hold a 4-entry dims array. It never crashed the existing test only because the fake allocator ignores the requested size. - Rename the byte accumulator to `bytes` (stored in output->bytes) and pass `dimensions_count` -- already computed one line above -- to TfLiteIntArrayGetSizeInBytes(). No behavioral change to the overflow guard. - Add AllocateOutputDimensionsSizesDimsByDimensionCount, which records the size argument handed to AllocatePersistentBuffer and asserts it equals TfLiteIntArrayGetSizeInBytes(dimensions_count), locking in the fix. BUG=#3552 --- tensorflow/lite/micro/memory_helpers.cc | 16 +++++--- tensorflow/lite/micro/memory_helpers_test.cc | 39 ++++++++++++++++++++ 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/tensorflow/lite/micro/memory_helpers.cc b/tensorflow/lite/micro/memory_helpers.cc index cb43f1eaf73..82689697bd9 100644 --- a/tensorflow/lite/micro/memory_helpers.cc +++ b/tensorflow/lite/micro/memory_helpers.cc @@ -191,8 +191,12 @@ TfLiteStatus AllocateOutputDimensionsFromInput(TfLiteContext* context, input = input1->dims->size > input2->dims->size ? input1 : input2; TF_LITE_ENSURE(context, output->type == input->type); - size_t size = 0; - TF_LITE_ENSURE_STATUS(TfLiteTypeSizeOf(input->type, &size)); + // `bytes` is the tensor storage size (element size * shape product) that is + // stored in output->bytes. Keep it distinct from `dimensions_count`: the + // output->dims array is sized from the number of dimension entries, not from + // the byte count. + size_t bytes = 0; + TF_LITE_ENSURE_STATUS(TfLiteTypeSizeOf(input->type, &bytes)); const int dimensions_count = tflite::GetTensorShape(input).DimensionsCount(); for (int i = 0; i < dimensions_count; i++) { const int dim = input->dims->data[i]; @@ -208,14 +212,14 @@ TfLiteStatus AllocateOutputDimensionsFromInput(TfLiteContext* context, TF_LITE_ENSURE( context, dim >= 0 && - (udim == 0 || size <= std::numeric_limits::max() / udim)); - size *= udim; + (udim == 0 || bytes <= std::numeric_limits::max() / udim)); + bytes *= udim; } - output->bytes = size; + output->bytes = bytes; output->dims = reinterpret_cast(context->AllocatePersistentBuffer( - context, TfLiteIntArrayGetSizeInBytes(size))); + context, TfLiteIntArrayGetSizeInBytes(dimensions_count))); output->dims->size = input->dims->size; for (int i = 0; i < dimensions_count; i++) { output->dims->data[i] = input->dims->data[i]; diff --git a/tensorflow/lite/micro/memory_helpers_test.cc b/tensorflow/lite/micro/memory_helpers_test.cc index 7e03d9faeb7..c3f6b7213aa 100644 --- a/tensorflow/lite/micro/memory_helpers_test.cc +++ b/tensorflow/lite/micro/memory_helpers_test.cc @@ -34,6 +34,15 @@ void* FakeAllocatePersistentBuffer(TfLiteContext* context, size_t bytes) { return reinterpret_cast(global_persistent_buffer); } +// Records the size argument of the most recent AllocatePersistentBuffer call so +// a test can assert that output->dims is sized from the dimension count rather +// than from the (much larger) tensor byte count. +size_t g_last_requested_bytes = 0; +void* RecordingAllocatePersistentBuffer(TfLiteContext* context, size_t bytes) { + g_last_requested_bytes = bytes; + return reinterpret_cast(global_persistent_buffer); +} + } // namespace TEST(MemoryHelpersTest, TestAlignPointerUp) { @@ -275,4 +284,34 @@ TEST(MemoryHelpersTest, TestAllocateOutputDimensionsFromInput) { } EXPECT_EQ(output_tensor.bytes, input_tensor2.bytes); } + +TEST(MemoryHelpersTest, AllocateOutputDimensionsSizesDimsByDimensionCount) { + // Regression test for the coupling between the byte count and the dimension + // count: output->dims must be allocated from the number of dimension entries, + // not from the tensor byte count. int32 [5, 5, 5, 5] is 2500 bytes but only 4 + // dimensions, so a correct implementation requests + // TfLiteIntArrayGetSizeInBytes(4), never TfLiteIntArrayGetSizeInBytes(2500). + constexpr int kDimsLen = 4; + int input1_dims[] = {1, 1}; + int input2_dims[] = {kDimsLen, 5, 5, 5, 5}; + int output_dims[] = {0, 0, 0, 0, 0}; + TfLiteTensor input_tensor1 = tflite::testing::CreateTensor( + nullptr, tflite::testing::IntArrayFromInts(input1_dims)); + TfLiteTensor input_tensor2 = tflite::testing::CreateTensor( + nullptr, tflite::testing::IntArrayFromInts(input2_dims)); + TfLiteTensor output_tensor = tflite::testing::CreateTensor( + nullptr, tflite::testing::IntArrayFromInts(output_dims)); + TfLiteContext context; + context.AllocatePersistentBuffer = RecordingAllocatePersistentBuffer; + + g_last_requested_bytes = 0; + EXPECT_EQ(kTfLiteOk, + tflite::AllocateOutputDimensionsFromInput( + &context, &input_tensor1, &input_tensor2, &output_tensor)); + + EXPECT_EQ(static_cast(TfLiteIntArrayGetSizeInBytes(kDimsLen)), + g_last_requested_bytes); + EXPECT_EQ(kDimsLen, output_tensor.dims->size); + EXPECT_EQ(input_tensor2.bytes, output_tensor.bytes); +} TF_LITE_MICRO_TESTS_MAIN