check for overflow in calculate_nbytes (#11217)

Summary:

Got a fuzzer error around overflow here

Differential Revision: D75544079

Author: JacobSzwejbka

runtime/executor/method_meta.cpp

@@ -55,12 +55,28 @@ Result<Tag> get_tag(
 size_t calculate_nbytes(
     Span<const int32_t> sizes,
     executorch::aten::ScalarType scalar_type) {
-  ssize_t n = 1;
+  size_t n = 1;
+  size_t prev_n = 1;
   for (size_t i = 0; i < sizes.size(); i++) {
+    prev_n = n;
     n *= sizes[i];
+    // Check for overflow
+    ET_CHECK_MSG(
+      sizes[i] == 0 || n / sizes[i] == prev_n,
+      "overflow in calculate_nbytes");
   }
-  // Use the full namespace to disambiguate from c10::elementSize.
-  return n * executorch::runtime::elementSize(scalar_type);
+
+  size_t elem_size = executorch::runtime::elementSize(scalar_type);
+
+  prev_n = n;
+  n = n * elem_size;
+
+  // Check for overflow before multiplying by element size
+  ET_CHECK_MSG(
+    elem_size == 0 || n / elem_size == prev_n,
+    "overflow in calculate_nbytes");
+  
+  return n;
 }
 
 } // namespace

runtime/executor/method_meta.h

@@ -21,6 +21,10 @@ struct ExecutionPlan;
 
 namespace executorch {
 namespace ET_RUNTIME_NAMESPACE {
+namespace testing {
+  // Provides test access to private Program methods.
+  class TensorInfoTestFriend;
+  } // namespace testing
 
 /**
  * Metadata about a specific tensor of an ExecuTorch Program.
@@ -71,6 +75,7 @@ class TensorInfo final {
  private:
   // Let MethodMeta create TensorInfo.
   friend class MethodMeta;
+  friend class testing::TensorInfoTestFriend;
 
   TensorInfo(
       Span<const int32_t> sizes,

runtime/executor/test/method_meta_test.cpp

@@ -9,21 +9,48 @@
 #include <executorch/runtime/executor/method_meta.h>
 
 #include <cstdlib>
-#include <filesystem>
+#include <limits>
+#include <vector>
 
 #include <executorch/extension/data_loader/file_data_loader.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/executor/program.h>
+#include <executorch/test/utils/DeathTest.h>
 #include <gtest/gtest.h>
 
 using namespace ::testing;
 using executorch::runtime::Error;
 using executorch::runtime::MethodMeta;
 using executorch::runtime::Program;
 using executorch::runtime::Result;
+using executorch::runtime::Span;
 using executorch::runtime::TensorInfo;
 using torch::executor::util::FileDataLoader;
 
+namespace executorch {
+namespace runtime {
+namespace testing {
+// Provides access to private TensorInfo methods.
+class TensorInfoTestFriend final {
+  public:
+  ET_NODISCARD static TensorInfo get(
+    Span<const int32_t> sizes,
+    Span<const uint8_t> dim_order,
+    executorch::aten::ScalarType scalar_type,
+    const bool is_memory_planned,
+    executorch::aten::string_view name) {
+    return TensorInfo(
+      Span<const int32_t>(sizes.data(), sizes.size()),
+      Span<const uint8_t>(dim_order.data(), dim_order.size()),
+      scalar_type,
+      is_memory_planned,
+      name);
+  }
+};
+} // namespace testing
+} // namespace runtime
+} // namespace executorch
+
 class MethodMetaTest : public ::testing::Test {
  protected:
   void load_program(const char* path, const char* module_name) {
@@ -163,3 +190,25 @@ TEST_F(MethodMetaTest, MethodMetaAttribute) {
   auto bad_access = method_meta->attribute_tensor_meta(1);
   ASSERT_EQ(bad_access.error(), Error::InvalidArgument);
 }
+
+TEST_F(MethodMetaTest, TensorInfoSizeOverflow) {
+  // Create sizes that will cause overflow when multiplied
+  std::vector<int32_t> overflow_sizes = {
+      std::numeric_limits<int32_t>::max(),
+      std::numeric_limits<int32_t>::max(),
+      std::numeric_limits<int32_t>::max(),
+      std::numeric_limits<int32_t>::max(),
+  };
+  
+  // Create a minimal dim_order
+  std::vector<uint8_t> dim_order = {0, 1, 2, 3};
+  
+  // Create a TensorInfo with the overflow sizes and expect it to fail.
+  ET_EXPECT_DEATH(executorch::runtime::testing::TensorInfoTestFriend::get(
+    Span<const int32_t>(overflow_sizes.data(), overflow_sizes.size()),
+    Span<const uint8_t>(dim_order.data(), dim_order.size()),
+    executorch::aten::ScalarType::Float,
+    false, // is_memory_planned
+    executorch::aten::string_view{nullptr, 0}),"");
+
+}