Fix int overflow in tensor byte-size calculation

tensorflow/lite/micro/memory_helpers.cc

@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <cstddef>
 #include <cstdint>
+#include <limits>
 
 #include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/c/common.h"
@@ -106,19 +107,40 @@ 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;
+  // 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.
   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<size_t>(dim);
+      if (udim != 0 &&
+          element_count > std::numeric_limits<size_t>::max() / udim) {
+        return kTfLiteError;
+      }
+      element_count *= udim;
     }
   }
 
   TfLiteType tf_lite_type;
   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<size_t>::max() / *type_size) {
+    return kTfLiteError;
+  }
   *bytes = element_count * (*type_size);
   return kTfLiteOk;
 }
@@ -127,15 +149,33 @@ TfLiteStatus TfLiteEvalTensorByteLength(const TfLiteEvalTensor* eval_tensor,
                                         size_t* out_bytes) {
   TFLITE_DCHECK(out_bytes != nullptr);
 
-  int element_count = 1;
+  // 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) {
     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<size_t>(dim);
+      if (udim != 0 &&
+          element_count > std::numeric_limits<size_t>::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<size_t>::max() / type_size) {
+    return kTfLiteError;
+  }
   *out_bytes = element_count * type_size;
   return kTfLiteOk;
 }
@@ -151,17 +191,35 @@ 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);
+  // `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++) {
-    size *= input->dims->data[i];
+    const int dim = input->dims->data[i];
+    const size_t udim = static_cast<size_t>(dim);
+    // 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 || bytes <= std::numeric_limits<size_t>::max() / udim));
+    bytes *= udim;
   }
-  output->bytes = size;
+  output->bytes = bytes;
 
   output->dims =
       reinterpret_cast<TfLiteIntArray*>(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];

tensorflow/lite/micro/memory_helpers_test.cc

@@ -34,6 +34,15 @@ void* FakeAllocatePersistentBuffer(TfLiteContext* context, size_t bytes) {
   return reinterpret_cast<void*>(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<void*>(global_persistent_buffer);
+}
+
 }  // namespace
 
 TEST(MemoryHelpersTest, TestAlignPointerUp) {
@@ -187,6 +196,58 @@ TEST(MemoryHelpersTest, TestBytesRequiredForTensor) {
   EXPECT_EQ(static_cast<size_t>(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<size_t>(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};
@@ -223,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<int32_t>(
+      nullptr, tflite::testing::IntArrayFromInts(input1_dims));
+  TfLiteTensor input_tensor2 = tflite::testing::CreateTensor<int32_t>(
+      nullptr, tflite::testing::IntArrayFromInts(input2_dims));
+  TfLiteTensor output_tensor = tflite::testing::CreateTensor<int32_t>(
+      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<size_t>(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