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 #include +#include #include #include #include @@ -58,6 +59,13 @@ Result 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 #include #include +#include #include #include #include @@ -62,6 +63,13 @@ Result 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,