From 76e483bd23c9b398acfb4e06e0aa0abbe35689f3 Mon Sep 17 00:00:00 2001
From: Github Executorch <github_executorch@arm.com>
Date: Tue, 3 Feb 2026 22:29:32 -0800
Subject: [PATCH] Add data buffer validation to impl_like() for portable/lean
 builds

This addresses the security concern (TOB-EXECUTORCH-9) where the non-ATen
code path in impl_like() was missing validation that existed in
tensor_like().

Followup fix to suggestions: https://github.com/pytorch/executorch/pull/17163/
---
 devtools/bundled_program/bundled_program.cpp | 43 +++++++++++++++++++-
 1 file changed, 41 insertions(+), 2 deletions(-)

diff --git a/devtools/bundled_program/bundled_program.cpp b/devtools/bundled_program/bundled_program.cpp
index 7fe8bfd5f5d..10e857dc2ed 100644
--- a/devtools/bundled_program/bundled_program.cpp
+++ b/devtools/bundled_program/bundled_program.cpp
@@ -20,6 +20,7 @@
 #include <executorch/devtools/bundled_program/schema/bundled_program_schema_generated.h>
 #include <executorch/runtime/core/event_tracer_hooks.h>
 #include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
+#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
 #include <executorch/runtime/core/memory_allocator.h>
 #include <executorch/runtime/executor/method.h>
 #include <executorch/runtime/platform/log.h>
@@ -56,7 +57,9 @@ at::Tensor tensor_like(bundled_program_flatbuffer::Tensor* bundled_tensor) {
       at::dtype(static_cast<ScalarType>(bundled_tensor->scalar_type())));
 
   // Validate data buffer exists and has sufficient size
-  ET_CHECK(bundled_tensor->data() != nullptr);
+  ET_CHECK(
+      bundled_tensor->data() != nullptr,
+      "Tensor flatbuffer is missing its data field");
   ET_CHECK_MSG(
       bundled_tensor->data()->size() >= ret_tensor.nbytes(),
       "Tensor data buffer too small: got %zu bytes, need %zu bytes",
@@ -77,11 +80,47 @@ TensorImpl impl_like(bundled_program_flatbuffer::Tensor* bundled_tensor) {
   ScalarType scalar_type =
       static_cast<ScalarType>(bundled_tensor->scalar_type());
   ssize_t dim = bundled_tensor->sizes()->size();
+
+  // Validate dimension count
+  ET_CHECK(
+      dim <= static_cast<ssize_t>(kMaxDim),
+      "Tensor rank too large: got %zd, max allowed %zu",
+      dim,
+      kMaxDim);
+
   executorch::aten::SizesType* sizes = bundled_tensor->mutable_sizes()->data();
-  void* data = bundled_tensor->mutable_data()->data();
   executorch::aten::DimOrderType* dim_order =
       bundled_tensor->mutable_dim_order()->data();
 
+  // Calculate expected tensor size in bytes
+  size_t numel = 1;
+  for (ssize_t i = 0; i < dim; i++) {
+    ET_CHECK_MSG(
+        sizes[i] >= 0,
+        "Tensor has negative size at dimension %zd: %d",
+        static_cast<ssize_t>(i),
+        static_cast<int>(sizes[i]));
+    size_t new_numel;
+    ET_CHECK_MSG(
+        !c10::mul_overflows(numel, static_cast<size_t>(sizes[i]), &new_numel),
+        "Integer overflow calculating tensor numel at dim %zd",
+        static_cast<size_t>(i));
+    numel = new_numel;
+  }
+  size_t expected_bytes = numel * executorch::runtime::elementSize(scalar_type);
+
+  // Validate data buffer exists and has sufficient size
+  ET_CHECK(
+      bundled_tensor->data() != nullptr,
+      "Tensor flatbuffer is missing its data field");
+  ET_CHECK_MSG(
+      bundled_tensor->data()->size() >= expected_bytes,
+      "Tensor data buffer too small: got %zu bytes, need %zu bytes",
+      static_cast<size_t>(bundled_tensor->data()->size()),
+      static_cast<size_t>(expected_bytes));
+
+  void* data = bundled_tensor->mutable_data()->data();
+
   // The strides of created tensorimpl will only be actually used when
   // comparsion (`tensor_are_close` below). To eliminate the usage of memory
   // allocator, here we set the initial strides as null and reconstruct the