From e8f2fca7cc4370a07da18cd2f4c7a7e5de276b6a Mon Sep 17 00:00:00 2001
From: Github Executorch <github_executorch@arm.com>
Date: Thu, 23 Apr 2026 10:38:47 -0700
Subject: [PATCH] Use compute_numel_overflow() in tensor_ptr builders

make_tensor_ptr() and related extension builders return TensorPtr (a shared_ptr) and cannot propagate Error, so the existing compute_numel()*elementSize() size checks silently wrapped on overflow and let an undersized buffer back an oversized tensor (heap OOB). Migrate these callers to compute_numel_overflow() so overflow aborts, and add explicit c10::mul_overflows() guards around the subsequent numel*elementSize multiplication. Touched: make_tensor_ptr(vector<uint8_t>), clone_tensor_ptr, empty_strided, and the two tensor_ptr.h template overloads.

Authored with Claude.

Differential Revision: [D102082922](https://our.internmc.facebook.com/intern/diff/D102082922/)

[ghstack-poisoned]
---
 extension/tensor/tensor_ptr.cpp       | 30 ++++++++++++++++++++++-----
 extension/tensor/tensor_ptr.h         |  5 +++--
 extension/tensor/tensor_ptr_maker.cpp | 15 ++++++++++++--
 3 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/extension/tensor/tensor_ptr.cpp b/extension/tensor/tensor_ptr.cpp
index bb76311bd67..956e406108a 100644
--- a/extension/tensor/tensor_ptr.cpp
+++ b/extension/tensor/tensor_ptr.cpp
@@ -10,6 +10,8 @@
 
 #include <numeric>
 
+#include <c10/util/safe_numerics.h>
+
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
 
 namespace executorch {
@@ -147,11 +149,22 @@ TensorPtr make_tensor_ptr(
     std::vector<executorch::aten::StridesType> strides,
     executorch::aten::ScalarType type,
     executorch::aten::TensorShapeDynamism dynamism) {
+  const ssize_t numel =
+      executorch::aten::compute_numel_overflow(sizes.data(), sizes.size());
+  size_t nbytes;
+  ET_CHECK_MSG(
+      !c10::mul_overflows(
+          static_cast<size_t>(numel),
+          executorch::aten::elementSize(type),
+          &nbytes),
+      "Overflow computing nbytes: numel=%zd element_size=%zu",
+      numel,
+      executorch::aten::elementSize(type));
   ET_CHECK_MSG(
-      data.size() ==
-          executorch::aten::compute_numel(sizes.data(), sizes.size()) *
-              executorch::aten::elementSize(type),
-      "Data size does not match tensor size.");
+      data.size() == nbytes,
+      "Data size (%zu) does not match tensor size (%zu).",
+      data.size(),
+      nbytes);
   auto data_ptr = data.data();
   return make_tensor_ptr(
       std::move(sizes),
@@ -205,7 +218,14 @@ TensorPtr clone_tensor_ptr(
       runtime::canCast(tensor_type, type),
       "Cannot cast tensor type to desired type.");
   const auto tensor_numel = static_cast<size_t>(tensor.numel());
-  std::vector<uint8_t> data(tensor_numel * aten::elementSize(type));
+  size_t clone_nbytes;
+  ET_CHECK_MSG(
+      !c10::mul_overflows(
+          tensor_numel, aten::elementSize(type), &clone_nbytes),
+      "Overflow computing clone nbytes: numel=%zu element_size=%zu",
+      tensor_numel,
+      aten::elementSize(type));
+  std::vector<uint8_t> data(clone_nbytes);
 
   // Create a minimal context for error handling in ET_SWITCH
   struct {
diff --git a/extension/tensor/tensor_ptr.h b/extension/tensor/tensor_ptr.h
index 47124bdeca6..4afa1c7be45 100644
--- a/extension/tensor/tensor_ptr.h
+++ b/extension/tensor/tensor_ptr.h
@@ -111,7 +111,8 @@ inline TensorPtr make_tensor_ptr(
         executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND) {
   ET_CHECK_MSG(
       data.size() ==
-          executorch::aten::compute_numel(sizes.data(), sizes.size()),
+          static_cast<size_t>(executorch::aten::compute_numel_overflow(
+              sizes.data(), sizes.size())),
       "Data size does not match tensor size.");
   if (type != deduced_type) {
     ET_CHECK_MSG(
@@ -359,7 +360,7 @@ inline TensorPtr make_tensor_ptr(
   const auto same_shape = same_rank &&
       std::equal(sizes.begin(), sizes.end(), tensor.sizes().begin());
   const auto element_count =
-      executorch::aten::compute_numel(sizes.data(), sizes.size());
+      executorch::aten::compute_numel_overflow(sizes.data(), sizes.size());
   const auto parent_element_count = tensor.numel();
   ET_CHECK_MSG(
       element_count <= parent_element_count,
diff --git a/extension/tensor/tensor_ptr_maker.cpp b/extension/tensor/tensor_ptr_maker.cpp
index b71dfab8eeb..c361e2f7d2f 100644
--- a/extension/tensor/tensor_ptr_maker.cpp
+++ b/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,18 @@ 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()) *
+  const ssize_t numel =
+      executorch::aten::compute_numel_overflow(sizes.data(), sizes.size());
+  size_t nbytes;
+  ET_CHECK_MSG(
+      !c10::mul_overflows(
+          static_cast<size_t>(numel),
+          executorch::aten::elementSize(type),
+          &nbytes),
+      "Overflow computing nbytes: numel=%zd element_size=%zu",
+      numel,
       executorch::aten::elementSize(type));
+  std::vector<uint8_t> data(nbytes);
   return make_tensor_ptr(
       std::move(sizes),
       std::move(data),