Validate kTensorDimensionLimit

Differential Revision: D78675986

Pull Request resolved: #12684

Author: JacobSzwejbka

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/...",

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,

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,