check for overflow in calculate_nbytes (#11217)

Summary: Got a fuzzer error around overflow here Differential Revision: D75544079

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_tcalculate_nbytes(
  Span<constint32_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;
 
 namespaceexecutorch{
 namespaceET_RUNTIME_NAMESPACE{
+namespacetesting{
+// Provides test access to private Program methods.
+classTensorInfoTestFriend;
+} // namespace testing
 
 /**
  * Metadata about a specific tensor of an ExecuTorch Program.
@@ -71,6 +75,7 @@ class TensorInfo final{
 private:
 // Let MethodMeta create TensorInfo.
 friendclassMethodMeta;
+friendclasstesting::TensorInfoTestFriend;
 
 TensorInfo(
  Span<constint32_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>
 
 usingnamespace ::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;
 
+namespaceexecutorch{
+namespaceruntime{
+namespacetesting{
+// Provides access to private TensorInfo methods.
+classTensorInfoTestFriendfinal{
+public:
+ ET_NODISCARD static TensorInfo get(
+ Span<constint32_t> sizes,
+ Span<constuint8_t> dim_order,
+ executorch::aten::ScalarType scalar_type,
+constbool is_memory_planned,
+ executorch::aten::string_view name){
+returnTensorInfo(
+ Span<constint32_t>(sizes.data(), sizes.size()),
+ Span<constuint8_t>(dim_order.data(), dim_order.size()),
+ scalar_type,
+ is_memory_planned,
+ name);
+ }
+};
+} // namespace testing
+} // namespace runtime
+} // namespace executorch
+
 classMethodMetaTest : public ::testing::Test{
 protected:
 voidload_program(constchar* path, constchar* 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<constint32_t>(overflow_sizes.data(), overflow_sizes.size()),
+ Span<constuint8_t>(dim_order.data(), dim_order.size()),
+ executorch::aten::ScalarType::Float,
+false, // is_memory_planned
+ executorch::aten::string_view{nullptr, 0}),
+"");
+}