check for overflow in calculate_nbytes (#11217)

Summary:
Got a fuzzer error around overflow here

Pull Request resolved: #11217

Test Plan:
Imported from GitHub, without a `Test Plan:` line.

Rollback Plan:

Reviewed By: lucylq

Differential Revision: D75544079

Pulled By: JacobSzwejbka

Author: JacobSzwejbka

.github/workflows/pull.yml

@@ -371,7 +371,7 @@ jobs:
         size=${arr[4]}
         # threshold=48120 on devserver with gcc11.4
         # todo(lfq): update once binary size is below 50kb.
-        threshold="51408"
+        threshold="55504"
         if [[ "$size" -le "$threshold" ]]; then
           echo "Success $size <= $threshold"
         else
@@ -406,7 +406,7 @@ jobs:
         output=$(ls -la cmake-out/test/size_test)
         arr=($output)
         size=${arr[4]}
-        threshold="47560"
+        threshold="51656"
         if [[ "$size" -le "$threshold" ]]; then
           echo "Success $size <= $threshold"
         else

runtime/executor/method_meta.cpp

@@ -55,12 +55,24 @@ 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(sizes[i] == 0 || n / sizes[i] == prev_n);
   }
-  // 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
+  ET_CHECK(elem_size == 0 || n / elem_size == prev_n);
+
+  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,26 @@ 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}),
+      "");
+}