From eba56444f5ea367a024aeea57a479b65a2f9a69b Mon Sep 17 00:00:00 2001
From: Github Executorch <github_executorch@arm.com>
Date: Wed, 4 Mar 2026 23:02:05 -0800
Subject: [PATCH] Fix TOB-EXECUTORCH-28: validate buffer sizes and tensor
 metadata in MethodMeta

memory_planned_buffer_size() returned buffer sizes from the PTE FlatBuffer
without validation. Negative int64_t values would become huge size_t values
when used by callers for allocation, causing OOM/DoS.

- Reject negative values from memory_planned_buffer_size() with
  Error::InvalidProgram
- Reject negative tensor dimension sizes in calculate_nbytes() before
  the implicit cast to size_t, preventing wraparound to huge values
- Add null-pointer checks on tensor metadata (tensor_value, sizes,
  dim_order) in input_tensor_meta() and output_tensor_meta() to guard
  against malformed PTE files

This PR was authored with the assistance of Claude.
---
 runtime/executor/method_meta.cpp | 49 ++++++++++++++++++++++++++++----
 1 file changed, 44 insertions(+), 5 deletions(-)

diff --git a/runtime/executor/method_meta.cpp b/runtime/executor/method_meta.cpp
index 75dadfd893a..ba73828378a 100644
--- a/runtime/executor/method_meta.cpp
+++ b/runtime/executor/method_meta.cpp
@@ -6,6 +6,8 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <cinttypes> // @donotremove
+
 #include <c10/util/safe_numerics.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
@@ -58,13 +60,20 @@ Result<size_t> calculate_nbytes(
     executorch::aten::ScalarType scalar_type) {
   size_t n = 1;
   for (size_t i = 0; i < sizes.size(); i++) {
+    ET_CHECK_OR_RETURN_ERROR(
+        sizes[i] >= 0,
+        InvalidProgram,
+        "Invalid size[%zu]: %" PRId32 ". Size must not be negative",
+        i,
+        sizes[i]);
     size_t next_n;
     bool overflow =
         c10::mul_overflows(n, static_cast<size_t>(sizes[i]), &next_n);
     ET_CHECK_OR_RETURN_ERROR(
         !overflow,
         InvalidArgument,
-        "Invalid size[%zu]: %d. Potentially overflowed, expect to be 0 or n: %zu",
+        "Invalid size[%zu]: %" PRId32
+        ". Potentially overflowed, expect to be 0 or n: %zu",
         i,
         sizes[i],
         n);
@@ -186,6 +195,12 @@ Result<TensorInfo> MethodMeta::input_tensor_meta(size_t index) const {
   auto input_index = s_plan_->inputs()->Get(index);
   // input_index was already validated by input_tag().
   auto tensor_value = s_plan_->values()->Get(input_index)->val_as_Tensor();
+  ET_CHECK_OR_RETURN_ERROR(
+      tensor_value != nullptr && tensor_value->sizes() != nullptr &&
+          tensor_value->dim_order() != nullptr,
+      InvalidProgram,
+      "Null tensor metadata for input %zu",
+      index);
   return TensorInfo::create(
       Span<const int32_t>(
           tensor_value->sizes()->data(), tensor_value->sizes()->size()),
@@ -237,7 +252,12 @@ Result<TensorInfo> MethodMeta::output_tensor_meta(size_t index) const {
   auto output_index = s_plan_->outputs()->Get(index);
   // output_index was already validated by output_tag().
   auto tensor_value = s_plan_->values()->Get(output_index)->val_as_Tensor();
-
+  ET_CHECK_OR_RETURN_ERROR(
+      tensor_value != nullptr && tensor_value->sizes() != nullptr &&
+          tensor_value->dim_order() != nullptr,
+      InvalidProgram,
+      "Null tensor metadata for output %zu",
+      index);
   return TensorInfo::create(
       Span<const int32_t>(
           tensor_value->sizes()->data(), tensor_value->sizes()->size()),
@@ -257,7 +277,10 @@ size_t MethodMeta::num_attributes() const {
     auto value = values->Get(i);
     if (value->val_type() == executorch_flatbuffer::KernelTypes::Tensor) {
       auto tensor_value = value->val_as_Tensor();
-      if (tensor_value->extra_tensor_info() != nullptr &&
+      if (tensor_value != nullptr &&
+          tensor_value->extra_tensor_info() != nullptr &&
+          tensor_value->extra_tensor_info()->fully_qualified_name() !=
+              nullptr &&
           tensor_value->extra_tensor_info()->fully_qualified_name()->c_str() !=
               nullptr) {
         ++counter;
@@ -274,10 +297,19 @@ Result<TensorInfo> MethodMeta::attribute_tensor_meta(size_t index) const {
     auto value = values->Get(i);
     if (value->val_type() == executorch_flatbuffer::KernelTypes::Tensor) {
       auto tensor_value = value->val_as_Tensor();
-      if (tensor_value->extra_tensor_info() != nullptr &&
+      if (tensor_value != nullptr &&
+          tensor_value->extra_tensor_info() != nullptr &&
+          tensor_value->extra_tensor_info()->fully_qualified_name() !=
+              nullptr &&
           tensor_value->extra_tensor_info()->fully_qualified_name()->c_str() !=
               nullptr) {
         if (counter == index) {
+          ET_CHECK_OR_RETURN_ERROR(
+              tensor_value->sizes() != nullptr &&
+                  tensor_value->dim_order() != nullptr,
+              InvalidProgram,
+              "Null tensor metadata for attribute %zu",
+              index);
           auto t_name =
               tensor_value->extra_tensor_info()->fully_qualified_name();
           // Count constant returns as memory planned
@@ -322,7 +354,14 @@ Result<int64_t> MethodMeta::memory_planned_buffer_size(size_t index) const {
       num_buffers);
   // Index zero is reserved internally, and we hide it from users. Adjust the
   // provided index to point to one of the actual buffers.
-  return s_plan_->non_const_buffer_sizes()->Get(index + 1);
+  int64_t size = s_plan_->non_const_buffer_sizes()->Get(index + 1);
+  ET_CHECK_OR_RETURN_ERROR(
+      size >= 0,
+      InvalidProgram,
+      "memory_planned_buffer_size(%zu) has invalid negative size: %" PRId64,
+      index,
+      size);
+  return size;
 }
 
 bool MethodMeta::uses_backend(const char* backend_name) const {