Numel and Nbytes Validation (#19121) (#19121)

JacobSzwejbka · web-flow · commit 56da9646ccb2 · 2026-04-24T15:25:51.000-07:00
Summary:

Attempt 3 to check numel and nbytes overflow. This time we defer
checking dynamic sized inputs until their size is realized.

Reviewed By: lucylq

Differential Revision: D98148157
diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp
@@ -9,6 +9,7 @@
 #include <executorch/runtime/executor/method.h>
 
 #include <c10/util/irange.h>
+#include <c10/util/safe_numerics.h>
 #include <array>
 #include <cinttypes> // @donotremove
 #include <cstdint>
@@ -1194,6 +1195,33 @@ Method::set_input(const EValue& input_evalue, size_t input_idx) {
         input_idx,
         executorch::runtime::toString(t_dst.scalar_type()),
         executorch::runtime::toString(t_src.scalar_type()));
+
+    ssize_t numel = 1;
+    for (ssize_t i = 0; i < t_src.dim(); i++) {
+      bool overflow = c10::mul_overflows(
+          numel, static_cast<ssize_t>(t_src.size(i)), &numel);
+      ET_CHECK_OR_RETURN_ERROR(
+          !overflow,
+          InvalidArgument,
+          "Input %" ET_PRIsize_t
+          ": numel overflowed at dimension %zd with size %zd",
+          input_idx,
+          (size_t)i,
+          (size_t)t_src.size(i));
+    }
+    size_t nbytes;
+    bool nbytes_overflow = c10::mul_overflows(
+        static_cast<size_t>(numel),
+        executorch::runtime::elementSize(t_src.scalar_type()),
+        &nbytes);
+    ET_CHECK_OR_RETURN_ERROR(
+        !nbytes_overflow,
+        InvalidArgument,
+        "Input %" ET_PRIsize_t
+        ": nbytes overflowed: numel %zd with element size %zu",
+        input_idx,
+        numel,
+        executorch::runtime::elementSize(t_src.scalar_type()));
     // Reset the shape for the Method's input as the size of forwarded input
     // tensor for shape dynamism. Also is a safety check if need memcpy.
     ET_CHECK_OK_OR_RETURN_ERROR(
diff --git a/runtime/executor/program_validation.cpp b/runtime/executor/program_validation.cpp
@@ -14,7 +14,7 @@
 #include <executorch/runtime/platform/log.h>
 #include <executorch/schema/program_generated.h>
 
-// #include <c10/util/safe_numerics.h>
+#include <c10/util/safe_numerics.h>
 
 namespace executorch {
 namespace runtime {
@@ -32,7 +32,8 @@ validate_tensor(const executorch_flatbuffer::Tensor* tensor) {
     return Error::InvalidProgram;
   }
 
-  // ssize_t numel = 1;
+  ssize_t numel = 1;
+  bool numel_overflowed = false;
   for (flatbuffers::uoffset_t i = 0; i < sizes->size(); i++) {
     int32_t size = sizes->Get(i);
 
@@ -45,16 +46,10 @@ validate_tensor(const executorch_flatbuffer::Tensor* tensor) {
       return Error::InvalidProgram;
     }
 
-    // bool overflow =
-    //     c10::mul_overflows(numel, static_cast<ssize_t>(size), &numel);
-    // if (overflow) {
-    //   ET_LOG(
-    //       Error,
-    //       "numel overflowed at dimension %u with size %d",
-    //       static_cast<unsigned>(i),
-    //       size);
-    //   return Error::InvalidProgram;
-    // }
+    if (!numel_overflowed) {
+      numel_overflowed =
+          c10::mul_overflows(numel, static_cast<ssize_t>(size), &numel);
+    }
   }
 
   auto scalar_type =
@@ -64,19 +59,18 @@ validate_tensor(const executorch_flatbuffer::Tensor* tensor) {
     return Error::InvalidProgram;
   }
 
-  // size_t nbytes;
-  // bool nbytes_overflow = c10::mul_overflows(
-  //     static_cast<size_t>(numel),
-  //     executorch::runtime::elementSize(scalar_type),
-  //     &nbytes);
-  // if (nbytes_overflow) {
-  //   ET_LOG(
-  //       Error,
-  //       "nbytes overflowed: numel %zd with element size %zu",
-  //       numel,
-  //       executorch::runtime::elementSize(scalar_type));
-  //   return Error::InvalidProgram;
-  // }
+  if (numel_overflowed) {
+    return Error::InvalidProgram;
+  }
+
+  size_t nbytes;
+  bool nbytes_overflow = c10::mul_overflows(
+      static_cast<size_t>(numel),
+      executorch::runtime::elementSize(scalar_type),
+      &nbytes);
+  if (nbytes_overflow) {
+    return Error::InvalidProgram;
+  }
 
   return Error::Ok;
 }
@@ -114,6 +108,27 @@ validate_program(const executorch_flatbuffer::Program* program) {
       return Error::InvalidProgram;
     }
 
+    const auto* inputs = plan->inputs();
+    auto is_dynamic_input = [&](flatbuffers::uoffset_t idx) -> bool {
+      if (inputs == nullptr) {
+        return false;
+      }
+      for (flatbuffers::uoffset_t i = 0; i < inputs->size(); i++) {
+        if (inputs->Get(i) == static_cast<int32_t>(idx)) {
+          const auto* value = values->Get(idx);
+          if (value == nullptr) {
+            return false;
+          }
+          const auto* tensor =
+              static_cast<const executorch_flatbuffer::Tensor*>(value->val());
+          return tensor != nullptr &&
+              tensor->shape_dynamism() !=
+              executorch_flatbuffer::TensorShapeDynamism::STATIC;
+        }
+      }
+      return false;
+    };
+
     for (flatbuffers::uoffset_t value_idx = 0; value_idx < values->size();
          value_idx++) {
       const auto* value = values->Get(value_idx);
@@ -128,12 +143,25 @@ validate_program(const executorch_flatbuffer::Program* program) {
 
         Error err = validate_tensor(tensor);
         if (err != Error::Ok) {
-          ET_LOG(
-              Error,
-              "Tensor validation failed for value %u in execution plan %u",
-              static_cast<unsigned>(value_idx),
-              static_cast<unsigned>(plan_idx));
-          return err;
+          // Dynamic input tensors may have upper-bound sizes serialized for
+          // 64-bit machines that would overflow on 32-bit. Since their actual
+          // sizes are provided at set_input time, we defer overflow checks
+          // for those to Method::set_input.
+          if (is_dynamic_input(value_idx)) {
+            ET_LOG(
+                Info,
+                "Skipping validation failure for dynamic input tensor "
+                "at value %u in execution plan %u",
+                static_cast<unsigned>(value_idx),
+                static_cast<unsigned>(plan_idx));
+          } else {
+            ET_LOG(
+                Error,
+                "Tensor validation failed for value %u in execution plan %u",
+                static_cast<unsigned>(value_idx),
+                static_cast<unsigned>(plan_idx));
+            return err;
+          }
         }
       }
 
diff --git a/runtime/executor/program_validation.h b/runtime/executor/program_validation.h
@@ -22,13 +22,12 @@ namespace executorch {
 namespace runtime {
 
 /**
- * Validates that computing numel (number of elements) from the tensor's sizes
- * will not overflow. This check should be performed before creating TensorImpl
- * objects to prevent undefined behavior from integer overflow.
+ * Validates that a tensor's metadata is semantically valid: sizes are
+ * non-negative, scalar type is valid, and computing numel/nbytes will not
+ * overflow.
  *
  * @param[in] tensor The flatbuffer Tensor to validate.
- * @return Error::Ok if the numel calculation is safe, Error::InvalidProgram
- *     if computing numel would overflow.
+ * @return Error::Ok if validation passes, Error::InvalidProgram otherwise.
  */
 ET_NODISCARD Error validate_tensor(const executorch_flatbuffer::Tensor* tensor);
 
diff --git a/runtime/executor/test/method_test.cpp b/runtime/executor/test/method_test.cpp
@@ -310,6 +310,31 @@ TEST_F(MethodTest, AliasedIOTest) {
   }
 }
 
+TEST_F(MethodTest, SetInputRejectsOverflowingSizes) {
+  // The "cat" model (ModuleDynamicCatUnallocatedIO) has a 2D input.
+  // set_input validates numel/nbytes overflow before resize_tensor.
+  ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
+  Result<Method> method = programs_["cat"]->load_method("forward", &mmm.get());
+  ASSERT_EQ(method.error(), Error::Ok);
+
+  // Create a 2D tensor with enormous sizes. On 32-bit platforms the numel
+  // multiplication overflows; on 64-bit the resize bounds check rejects it.
+  int32_t sizes[2] = {2000000000, 2000000000};
+  uint8_t dim_order[2] = {0, 1};
+  int32_t strides[2] = {1, 1};
+  executorch::aten::TensorImpl impl(
+      executorch::aten::ScalarType::Float,
+      2,
+      sizes,
+      nullptr,
+      dim_order,
+      strides);
+
+  auto input_err =
+      method->set_input(EValue(executorch::aten::Tensor(&impl)), 0);
+  EXPECT_NE(input_err, Error::Ok);
+}
+
 TEST_F(MethodTest, ConstantSegmentTest) {
   // Execute model with constants stored in segment.
   ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
diff --git a/runtime/executor/test/program_validation_test.cpp b/runtime/executor/test/program_validation_test.cpp
@@ -64,12 +64,15 @@ struct EValueConfig {
   EValueType type;
   std::vector<int32_t> tensor_sizes; // For Tensor type.
   std::vector<int32_t> tensor_list_items; // For TensorList type (indices).
+  bool is_dynamic = false; // For Tensor type: if true, uses DYNAMIC_BOUND.
 };
 
 // Unified helper to create a minimal valid PTE flatbuffer with configurable
-// evalues. Returns a buffer containing the flatbuffer data.
+// evalues. Returns a buffer containing the flatbuffer data. input_indices
+// specifies which value indices appear in the execution plan's inputs list.
 std::vector<uint8_t> CreateTestProgram(
-    const std::vector<EValueConfig>& configs) {
+    const std::vector<EValueConfig>& configs,
+    const std::vector<int32_t>& input_indices = {}) {
   flatbuffers::FlatBufferBuilder builder(1024);
 
   std::vector<flatbuffers::Offset<executorch_flatbuffer::EValue>> evalues;
@@ -93,7 +96,9 @@ std::vector<uint8_t> CreateTestProgram(
             /*data_buffer_idx=*/0,
             /*allocation_info=*/0,
             /*layout=*/0,
-            executorch_flatbuffer::TensorShapeDynamism::STATIC,
+            config.is_dynamic
+                ? executorch_flatbuffer::TensorShapeDynamism::DYNAMIC_BOUND
+                : executorch_flatbuffer::TensorShapeDynamism::STATIC,
             /*extra_tensor_info=*/0);
         evalues.push_back(executorch_flatbuffer::CreateEValue(
             builder,
@@ -124,6 +129,7 @@ std::vector<uint8_t> CreateTestProgram(
 
   auto values_vec = builder.CreateVector(evalues);
   auto plan_name = builder.CreateString("forward");
+  auto inputs_vec = builder.CreateVector(input_indices);
   auto empty_int_vec = builder.CreateVector(std::vector<int32_t>{});
   auto empty_int64_vec = builder.CreateVector(std::vector<int64_t>{0});
   auto empty_chain_vec = builder.CreateVector(
@@ -139,8 +145,8 @@ std::vector<uint8_t> CreateTestProgram(
       plan_name,
       /*container_meta_type=*/0,
       values_vec,
-      empty_int_vec,
-      empty_int_vec,
+      /*inputs=*/inputs_vec,
+      /*outputs=*/empty_int_vec,
       empty_chain_vec,
       empty_operators_vec,
       empty_delegates_vec,
@@ -206,29 +212,93 @@ TEST_F(ProgramValidationTest, InternalConsistencyDetectsTruncatedData) {
   ASSERT_EQ(program.error(), Error::InvalidProgram);
 }
 
-// TEST_F(ProgramValidationTest, TensorNumelOverflowDetected) {
-//   std::vector<EValueConfig> configs = {
-//       {EValueType::Tensor, {2000000000, 2000000000, 2000000000}, {}}};
-//
-//   AlignedBuffer buf(CreateTestProgram(configs));
-//   auto loader = buf.loader();
-//
-//   Result<Program> program =
-//       Program::load(&loader, Program::Verification::InternalConsistency);
-//   EXPECT_EQ(program.error(), Error::InvalidProgram);
-// }
-
-// TEST_F(ProgramValidationTest, TensorNumelOverflowNotDetectedWithMinimal) {
-//   std::vector<EValueConfig> configs = {
-//       {EValueType::Tensor, {2000000000, 2000000000, 2000000000}, {}}};
-//
-//   AlignedBuffer buf(CreateTestProgram(configs));
-//   auto loader = buf.loader();
-//
-//   // Minimal verification doesn't run program validation.
-//   Result<Program> program =
-//       Program::load(&loader, Program::Verification::Minimal);
-// }
+TEST_F(ProgramValidationTest, TensorNumelOverflowDetectedForStaticTensor) {
+  // Static tensors always have their overflow checked at validation time.
+  std::vector<EValueConfig> configs = {
+      {EValueType::Tensor,
+       {2000000000, 2000000000, 2000000000},
+       {},
+       /*is_dynamic=*/false}};
+
+  AlignedBuffer buf(CreateTestProgram(configs));
+  auto loader = buf.loader();
+
+  Result<Program> program =
+      Program::load(&loader, Program::Verification::InternalConsistency);
+  EXPECT_EQ(program.error(), Error::InvalidProgram);
+}
+
+TEST_F(ProgramValidationTest, TensorNumelOverflowNotDetectedWithMinimal) {
+  std::vector<EValueConfig> configs = {
+      {EValueType::Tensor,
+       {2000000000, 2000000000, 2000000000},
+       {},
+       /*is_dynamic=*/false}};
+
+  AlignedBuffer buf(CreateTestProgram(configs));
+  auto loader = buf.loader();
+
+  // Minimal verification doesn't run program validation.
+  Result<Program> program =
+      Program::load(&loader, Program::Verification::Minimal);
+  EXPECT_EQ(program.error(), Error::Ok);
+}
+
+TEST_F(ProgramValidationTest, TensorNumelOverflowSkippedForDynamicInput) {
+  // Dynamic input tensors skip overflow checks at validation time; the check
+  // is deferred to set_input where actual sizes are known.
+  std::vector<EValueConfig> configs = {
+      {EValueType::Tensor,
+       {2000000000, 2000000000, 2000000000},
+       {},
+       /*is_dynamic=*/true}};
+
+  // Mark value index 0 as a plan input.
+  AlignedBuffer buf(CreateTestProgram(configs, /*input_indices=*/{0}));
+  auto loader = buf.loader();
+
+  Result<Program> program =
+      Program::load(&loader, Program::Verification::InternalConsistency);
+  EXPECT_EQ(program.error(), Error::Ok);
+}
+
+TEST_F(
+    ProgramValidationTest,
+    TensorNumelOverflowDetectedForDynamicNonInputTensor) {
+  // A dynamic tensor that is NOT in the inputs list should still have its
+  // overflow checked at validation time.
+  std::vector<EValueConfig> configs = {
+      {EValueType::Tensor,
+       {2000000000, 2000000000, 2000000000},
+       {},
+       /*is_dynamic=*/true}};
+
+  // No input indices — the tensor is not a plan input.
+  AlignedBuffer buf(CreateTestProgram(configs));
+  auto loader = buf.loader();
+
+  Result<Program> program =
+      Program::load(&loader, Program::Verification::InternalConsistency);
+  EXPECT_EQ(program.error(), Error::InvalidProgram);
+}
+
+TEST_F(ProgramValidationTest, TensorNumelOverflowDetectedForStaticInputTensor) {
+  // A static input tensor should still have its overflow checked at
+  // validation time since its sizes cannot change.
+  std::vector<EValueConfig> configs = {
+      {EValueType::Tensor,
+       {2000000000, 2000000000, 2000000000},
+       {},
+       /*is_dynamic=*/false}};
+
+  // Mark value index 0 as a plan input.
+  AlignedBuffer buf(CreateTestProgram(configs, /*input_indices=*/{0}));
+  auto loader = buf.loader();
+
+  Result<Program> program =
+      Program::load(&loader, Program::Verification::InternalConsistency);
+  EXPECT_EQ(program.error(), Error::InvalidProgram);
+}
 
 TEST_F(ProgramValidationTest, NegativeSizeDetected) {
   std::vector<EValueConfig> configs = {{EValueType::Tensor, {10, -5, 10}, {}}};
