Fix integer overflows in tensor byte-size computations (pytorch#19055)

Three tensor-byte-size multiplications had no overflow check, letting a
malicious PTE trigger wrap-to-small size_t values while kernels iterate
on the un-wrapped element count, producing heap buffer overflows.

Fixed here:
- extension/tensor/tensor_ptr.h: data.size() * elementSize(type) in
make_tensor_ptr cast path.
- extension/tensor/tensor_ptr_maker.cpp: compute_numel(...) *
elementSize(type) in empty_strided.
- runtime/core/tensor_layout.cpp: dim-product loop and final *
elementSize(scalar_type) in calculate_nbytes; now returns
Error::InvalidArgument on overflow since the function already returns
Result<size_t>.

All guards use c10::mul_overflows, matching the existing pattern in
MethodMeta::calculate_nbytes, the data loaders, and
PlatformMemoryAllocator.

runtime/core/portable_type/tensor_impl.cpp is intentionally left alone
in this branch; guarding the nbytes() / compute_numel multiplications
there breaks internal callers and will be handled separately.

Authored with Claude.

Author: lucylq

extension/tensor/tensor_ptr.h

@@ -14,6 +14,7 @@
 #include <vector>
 
 #include <c10/macros/Macros.h>
+#include <c10/util/safe_numerics.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
@@ -117,7 +118,16 @@ inline TensorPtr make_tensor_ptr(
     ET_CHECK_MSG(
         runtime::canCast(deduced_type, type),
         "Cannot cast deduced type to specified type.");
-    std::vector<uint8_t> casted_data(data.size() * aten::elementSize(type));
+    size_t casted_bytes = 0;
+    ET_CHECK_MSG(
+        !c10::mul_overflows(
+            data.size(),
+            static_cast<size_t>(aten::elementSize(type)),
+            &casted_bytes),
+        "casted_data size overflow: %zu elements * %zu bytes/element",
+        data.size(),
+        static_cast<size_t>(aten::elementSize(type)));
+    std::vector<uint8_t> casted_data(casted_bytes);
 
     // Create a minimal context for error handling in ET_SWITCH
     struct {

extension/tensor/tensor_ptr_maker.cpp

@@ -10,6 +10,8 @@
 
 #include <random>
 
+#include <c10/util/safe_numerics.h>
+
 namespace executorch {
 namespace extension {
 namespace {
@@ -111,9 +113,17 @@ TensorPtr empty_strided(
     std::vector<executorch::aten::StridesType> strides,
     executorch::aten::ScalarType type,
     executorch::aten::TensorShapeDynamism dynamism) {
-  std::vector<uint8_t> data(
-      executorch::aten::compute_numel(sizes.data(), sizes.size()) *
-      executorch::aten::elementSize(type));
+  const auto numel = static_cast<size_t>(
+      executorch::aten::compute_numel(sizes.data(), sizes.size()));
+  const auto elem_size =
+      static_cast<size_t>(executorch::aten::elementSize(type));
+  size_t nbytes = 0;
+  ET_CHECK_MSG(
+      !c10::mul_overflows(numel, elem_size, &nbytes),
+      "empty_strided size overflow: numel %zu * element size %zu",
+      numel,
+      elem_size);
+  std::vector<uint8_t> data(nbytes);
   return make_tensor_ptr(
       std::move(sizes),
       std::move(data),

runtime/core/tensor_layout.cpp

@@ -7,6 +7,7 @@
  */
 
 #include <c10/util/irange.h>
+#include <c10/util/safe_numerics.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
 #include <executorch/runtime/core/span.h>
@@ -19,15 +20,25 @@ namespace {
 Result<size_t> calculate_nbytes(
     const Span<const int32_t>& sizes,
     const executorch::aten::ScalarType& scalar_type) {
-  ssize_t n = 1;
+  size_t n = 1;
   for (const auto i : c10::irange(sizes.size())) {
     if (sizes[i] < 0) {
       return Error::InvalidArgument;
     }
-    n *= sizes[i];
+    size_t next = 0;
+    if (c10::mul_overflows(n, static_cast<size_t>(sizes[i]), &next)) {
+      return Error::InvalidArgument;
+    }
+    n = next;
   }
   // Use the full namespace to disambiguate from c10::elementSize.
-  return n * executorch::runtime::elementSize(scalar_type);
+  const size_t elem_size =
+      static_cast<size_t>(executorch::runtime::elementSize(scalar_type));
+  size_t total = 0;
+  if (c10::mul_overflows(n, elem_size, &total)) {
+    return Error::InvalidArgument;
+  }
+  return total;
 }
 } // namespace