From f4872e6fe2a1a711de93be86bac7b63f771b7a84 Mon Sep 17 00:00:00 2001
From: Jacob Szwejbka <jakeszwe@meta.com>
Date: Tue, 22 Jul 2025 09:25:59 -0700
Subject: [PATCH] Validate kTensorDimensionLimit (#12684)

Summary:

Need to check this limit at deserialization since the ops rely on it

Reviewed By: larryliu0820

Differential Revision: D78675986
---
 runtime/executor/targets.bzl                | 1 +
 runtime/executor/tensor_parser_aten.cpp     | 8 ++++++++
 runtime/executor/tensor_parser_portable.cpp | 8 ++++++++
 3 files changed, 17 insertions(+)

diff --git a/runtime/executor/targets.bzl b/runtime/executor/targets.bzl
index 424cc3e147b..ec0fb19ff96 100644
--- a/runtime/executor/targets.bzl
+++ b/runtime/executor/targets.bzl
@@ -133,6 +133,7 @@ def define_common_targets():
             ],
             deps = [
                 "//executorch/schema:program",
+                "//executorch/runtime/core/exec_aten/util:tensor_dimension_limit"
             ],
             visibility = [
                 "//executorch/runtime/executor/...",
diff --git a/runtime/executor/tensor_parser_aten.cpp b/runtime/executor/tensor_parser_aten.cpp
index 2d454d15be5..ad980177cf1 100644
--- a/runtime/executor/tensor_parser_aten.cpp
+++ b/runtime/executor/tensor_parser_aten.cpp
@@ -10,6 +10,7 @@
 
 #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/exec_aten/util/tensor_dimension_limit.h>
 #include <executorch/runtime/core/named_data_map.h>
 #include <executorch/runtime/executor/memory_manager.h>
 #include <executorch/runtime/executor/program.h>
@@ -58,6 +59,13 @@ Result<at::Tensor> parseTensor(
       s_tensor->sizes() != nullptr, InvalidProgram, "Missing sizes field");
   size_t ndim = s_tensor->sizes()->size();
 
+  ET_CHECK_OR_RETURN_ERROR(
+      ndim <= kTensorDimensionLimit,
+      InvalidProgram,
+      "Tensor rank too large %" ET_PRIsize_t " > %zu",
+      ndim,
+      kTensorDimensionLimit)
+
   ET_CHECK_OR_RETURN_ERROR(
       s_tensor->dim_order() != nullptr,
       InvalidProgram,
diff --git a/runtime/executor/tensor_parser_portable.cpp b/runtime/executor/tensor_parser_portable.cpp
index e1f09d557ac..02cb019a1da 100644
--- a/runtime/executor/tensor_parser_portable.cpp
+++ b/runtime/executor/tensor_parser_portable.cpp
@@ -11,6 +11,7 @@
 #include <executorch/runtime/core/exec_aten/exec_aten.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/exec_aten/util/tensor_dimension_limit.h>
 #include <executorch/runtime/core/named_data_map.h>
 #include <executorch/runtime/executor/memory_manager.h>
 #include <executorch/runtime/executor/program.h>
@@ -62,6 +63,13 @@ Result<Tensor> parseTensor(
   const auto serialized_sizes = s_tensor->sizes()->data();
   const auto dim = s_tensor->sizes()->size();
 
+  ET_CHECK_OR_RETURN_ERROR(
+      dim <= kTensorDimensionLimit,
+      InvalidProgram,
+      "Tensor rank too large %" PRIu32 " > %zu",
+      dim,
+      kTensorDimensionLimit)
+
   ET_CHECK_OR_RETURN_ERROR(
       s_tensor->dim_order() != nullptr,
       InvalidProgram,