Integer overflow in tensor size calculations leads to heap buffer overflow

tensorflow/lite/kernels/internal/runtime_shape.h

@@ -93,11 +93,11 @@ class RuntimeShape {
 
   // Returns the total count of elements, that is the size when flattened into a
   // vector.
-  int FlatSize() const {
-    int buffer_size = 1;
+  size_t FlatSize() const {
+    size_t buffer_size = 1;
     const int* dims_data = reinterpret_cast<const int*>(DimsData());
     for (int i = 0; i < size_; i++) {
-      buffer_size *= dims_data[i];
+      buffer_size *= static_cast<size_t>(dims_data[i]);
     }
     return buffer_size;
   }

tensorflow/lite/micro/memory_helpers.cc

@@ -106,12 +106,12 @@ 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;
+  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);
+      element_count *= static_cast<size_t>(flatbuffer_tensor.shape()->Get(n));
     }
   }
 
@@ -127,11 +127,11 @@ 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];
+      element_count *= static_cast<size_t>(eval_tensor->dims->data[n]);
     }
   }
   size_t type_size;

tensorflow/lite/micro/micro_utils.cc

@@ -27,10 +27,10 @@ limitations under the License.
 
 namespace tflite {
 
-int ElementCount(const TfLiteIntArray& dims) {
-  int result = 1;
+size_t ElementCount(const TfLiteIntArray& dims) {
+  size_t result = 1;
   for (int i = 0; i < dims.size; ++i) {
-    result *= dims.data[i];
+    result *= static_cast<size_t>(dims.data[i]);
   }
   return result;
 }

tensorflow/lite/micro/micro_utils.h

@@ -27,7 +27,7 @@ namespace tflite {
 
 // Returns number of elements in the shape array.
 
-int ElementCount(const TfLiteIntArray& dims);
+size_t ElementCount(const TfLiteIntArray& dims);
 
 size_t EvalTensorBytes(const TfLiteEvalTensor* tensor);