Skip to content

Commit 5e8a0df

Browse files
lucylqGithub Executorch
andauthored
MethodMeta validation (#18724)
- Reject negative memory planning values - Reject negative tensor dimension sizes in calculate_nbytes() - Add null-pointer checks before TensorInfo::create to catch invalid tensor metadata in the flatbuffer. This PR was authored with the assistance of Claude. Co-authored-by: Github Executorch <github_executorch@arm.com>
1 parent 6217c75 commit 5e8a0df

File tree

1 file changed

+44
-5
lines changed

1 file changed

+44
-5
lines changed

runtime/executor/method_meta.cpp

Lines changed: 44 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
* LICENSE file in the root directory of this source tree.
77
*/
88

9+
#include <cinttypes> // @donotremove
10+
911
#include <c10/util/safe_numerics.h>
1012
#include <executorch/runtime/core/error.h>
1113
#include <executorch/runtime/core/exec_aten/exec_aten.h>
@@ -58,13 +60,20 @@ Result<size_t> calculate_nbytes(
5860
executorch::aten::ScalarType scalar_type) {
5961
size_t n = 1;
6062
for (size_t i = 0; i < sizes.size(); i++) {
63+
ET_CHECK_OR_RETURN_ERROR(
64+
sizes[i] >= 0,
65+
InvalidProgram,
66+
"Invalid size[%zu]: %" PRId32 ". Size must not be negative",
67+
i,
68+
sizes[i]);
6169
size_t next_n;
6270
bool overflow =
6371
c10::mul_overflows(n, static_cast<size_t>(sizes[i]), &next_n);
6472
ET_CHECK_OR_RETURN_ERROR(
6573
!overflow,
6674
InvalidArgument,
67-
"Invalid size[%zu]: %d. Potentially overflowed, expect to be 0 or n: %zu",
75+
"Invalid size[%zu]: %" PRId32
76+
". Potentially overflowed, expect to be 0 or n: %zu",
6877
i,
6978
sizes[i],
7079
n);
@@ -186,6 +195,12 @@ Result<TensorInfo> MethodMeta::input_tensor_meta(size_t index) const {
186195
auto input_index = s_plan_->inputs()->Get(index);
187196
// input_index was already validated by input_tag().
188197
auto tensor_value = s_plan_->values()->Get(input_index)->val_as_Tensor();
198+
ET_CHECK_OR_RETURN_ERROR(
199+
tensor_value != nullptr && tensor_value->sizes() != nullptr &&
200+
tensor_value->dim_order() != nullptr,
201+
InvalidProgram,
202+
"Null tensor metadata for input %zu",
203+
index);
189204
return TensorInfo::create(
190205
Span<const int32_t>(
191206
tensor_value->sizes()->data(), tensor_value->sizes()->size()),
@@ -237,7 +252,12 @@ Result<TensorInfo> MethodMeta::output_tensor_meta(size_t index) const {
237252
auto output_index = s_plan_->outputs()->Get(index);
238253
// output_index was already validated by output_tag().
239254
auto tensor_value = s_plan_->values()->Get(output_index)->val_as_Tensor();
240-
255+
ET_CHECK_OR_RETURN_ERROR(
256+
tensor_value != nullptr && tensor_value->sizes() != nullptr &&
257+
tensor_value->dim_order() != nullptr,
258+
InvalidProgram,
259+
"Null tensor metadata for output %zu",
260+
index);
241261
return TensorInfo::create(
242262
Span<const int32_t>(
243263
tensor_value->sizes()->data(), tensor_value->sizes()->size()),
@@ -257,7 +277,10 @@ size_t MethodMeta::num_attributes() const {
257277
auto value = values->Get(i);
258278
if (value->val_type() == executorch_flatbuffer::KernelTypes::Tensor) {
259279
auto tensor_value = value->val_as_Tensor();
260-
if (tensor_value->extra_tensor_info() != nullptr &&
280+
if (tensor_value != nullptr &&
281+
tensor_value->extra_tensor_info() != nullptr &&
282+
tensor_value->extra_tensor_info()->fully_qualified_name() !=
283+
nullptr &&
261284
tensor_value->extra_tensor_info()->fully_qualified_name()->c_str() !=
262285
nullptr) {
263286
++counter;
@@ -274,10 +297,19 @@ Result<TensorInfo> MethodMeta::attribute_tensor_meta(size_t index) const {
274297
auto value = values->Get(i);
275298
if (value->val_type() == executorch_flatbuffer::KernelTypes::Tensor) {
276299
auto tensor_value = value->val_as_Tensor();
277-
if (tensor_value->extra_tensor_info() != nullptr &&
300+
if (tensor_value != nullptr &&
301+
tensor_value->extra_tensor_info() != nullptr &&
302+
tensor_value->extra_tensor_info()->fully_qualified_name() !=
303+
nullptr &&
278304
tensor_value->extra_tensor_info()->fully_qualified_name()->c_str() !=
279305
nullptr) {
280306
if (counter == index) {
307+
ET_CHECK_OR_RETURN_ERROR(
308+
tensor_value->sizes() != nullptr &&
309+
tensor_value->dim_order() != nullptr,
310+
InvalidProgram,
311+
"Null tensor metadata for attribute %zu",
312+
index);
281313
auto t_name =
282314
tensor_value->extra_tensor_info()->fully_qualified_name();
283315
// Count constant returns as memory planned
@@ -322,7 +354,14 @@ Result<int64_t> MethodMeta::memory_planned_buffer_size(size_t index) const {
322354
num_buffers);
323355
// Index zero is reserved internally, and we hide it from users. Adjust the
324356
// provided index to point to one of the actual buffers.
325-
return s_plan_->non_const_buffer_sizes()->Get(index + 1);
357+
int64_t size = s_plan_->non_const_buffer_sizes()->Get(index + 1);
358+
ET_CHECK_OR_RETURN_ERROR(
359+
size >= 0,
360+
InvalidProgram,
361+
"memory_planned_buffer_size(%zu) has invalid negative size: %" PRId64,
362+
index,
363+
size);
364+
return size;
326365
}
327366

328367
bool MethodMeta::uses_backend(const char* backend_name) const {

0 commit comments

Comments
 (0)