Skip to content

Commit a87632f

Browse files
JacobSzwejbkafacebook-github-bot
authored andcommitted
Validate kTensorDimensionLimit
Summary: Need to check this limit at deserialization since the ops rely on it Differential Revision: D78675986
1 parent 6aa864e commit a87632f

File tree

3 files changed

+11
-0
lines changed

3 files changed

+11
-0
lines changed

runtime/executor/targets.bzl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,7 @@ def define_common_targets():
133133
],
134134
deps = [
135135
"//executorch/schema:program",
136+
"//executorch/runtime/core/exec_aten/util:tensor_dimension_limit"
136137
],
137138
visibility = [
138139
"//executorch/runtime/executor/...",

runtime/executor/tensor_parser_aten.cpp

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
#include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
1212
#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
13+
#include <executorch/runtime/core/exec_aten/util/tensor_dimension_limit.h>
1314
#include <executorch/runtime/core/named_data_map.h>
1415
#include <executorch/runtime/executor/memory_manager.h>
1516
#include <executorch/runtime/executor/program.h>
@@ -58,6 +59,10 @@ Result<at::Tensor> parseTensor(
5859
s_tensor->sizes() != nullptr, InvalidProgram, "Missing sizes field");
5960
size_t ndim = s_tensor->sizes()->size();
6061

62+
ET_CHECK_OR_RETURN_ERROR(
63+
ndim <= kTensorDimensionLimit, InvalidProgram, "Tensor rank too large %" ET_PRIsize_t " > %zu", ndim, kTensorDimensionLimit
64+
)
65+
6166
ET_CHECK_OR_RETURN_ERROR(
6267
s_tensor->dim_order() != nullptr,
6368
InvalidProgram,

runtime/executor/tensor_parser_portable.cpp

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
#include <executorch/runtime/core/exec_aten/exec_aten.h>
1212
#include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
1313
#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
14+
#include <executorch/runtime/core/exec_aten/util/tensor_dimension_limit.h>
1415
#include <executorch/runtime/core/named_data_map.h>
1516
#include <executorch/runtime/executor/memory_manager.h>
1617
#include <executorch/runtime/executor/program.h>
@@ -62,6 +63,10 @@ Result<Tensor> parseTensor(
6263
const auto serialized_sizes = s_tensor->sizes()->data();
6364
const auto dim = s_tensor->sizes()->size();
6465

66+
ET_CHECK_OR_RETURN_ERROR(
67+
dim <= kTensorDimensionLimit, InvalidProgram, "Tensor rank too large %" PRIu32 " > %zu", dim, kTensorDimensionLimit
68+
)
69+
6570
ET_CHECK_OR_RETURN_ERROR(
6671
s_tensor->dim_order() != nullptr,
6772
InvalidProgram,

0 commit comments

Comments
 (0)