Skip to content

Commit 59adaf3

Browse files
author
Github Executorch
committed
Fix TOB-EXECUTORCH-28: validate buffer sizes and tensor metadata in MethodMeta
memory_planned_buffer_size() returned buffer sizes from the PTE FlatBuffer without validation. Negative int64_t values would become huge size_t values when used by callers for allocation, causing OOM/DoS. - Reject negative values from memory_planned_buffer_size() with Error::InvalidProgram - Reject negative tensor dimension sizes in calculate_nbytes() before the implicit cast to size_t, preventing wraparound to huge values - Add null-pointer checks on tensor metadata (tensor_value, sizes, dim_order) in input_tensor_meta() and output_tensor_meta() to guard against malformed PTE files This PR was authored with the assistance of Claude.
1 parent 3d2c853 commit 59adaf3

1 file changed

Lines changed: 59 additions & 7 deletions

File tree

runtime/executor/method_meta.cpp

Lines changed: 59 additions & 7 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,19 @@ 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 ". Potentially overflowed, expect to be 0 or n: %zu",
6876
i,
6977
sizes[i],
7078
n);
@@ -185,7 +193,20 @@ Result<TensorInfo> MethodMeta::input_tensor_meta(size_t index) const {
185193
index);
186194
auto input_index = s_plan_->inputs()->Get(index);
187195
// input_index was already validated by input_tag().
188-
auto tensor_value = s_plan_->values()->Get(input_index)->val_as_Tensor();
196+
auto serialization_value = s_plan_->values()->Get(input_index);
197+
ET_CHECK_OR_RETURN_ERROR(
198+
serialization_value != nullptr,
199+
InvalidProgram,
200+
"Null value for input %zu at index %d",
201+
index,
202+
input_index);
203+
auto tensor_value = serialization_value->val_as_Tensor();
204+
ET_CHECK_OR_RETURN_ERROR(
205+
tensor_value != nullptr && tensor_value->sizes() != nullptr &&
206+
tensor_value->dim_order() != nullptr,
207+
InvalidProgram,
208+
"Null tensor metadata for input %zu",
209+
index);
189210
return TensorInfo::create(
190211
Span<const int32_t>(
191212
tensor_value->sizes()->data(), tensor_value->sizes()->size()),
@@ -236,8 +257,20 @@ Result<TensorInfo> MethodMeta::output_tensor_meta(size_t index) const {
236257
index);
237258
auto output_index = s_plan_->outputs()->Get(index);
238259
// output_index was already validated by output_tag().
239-
auto tensor_value = s_plan_->values()->Get(output_index)->val_as_Tensor();
240-
260+
auto serialization_value = s_plan_->values()->Get(output_index);
261+
ET_CHECK_OR_RETURN_ERROR(
262+
serialization_value != nullptr,
263+
InvalidProgram,
264+
"Null value for output %zu at index %d",
265+
index,
266+
output_index);
267+
auto tensor_value = serialization_value->val_as_Tensor();
268+
ET_CHECK_OR_RETURN_ERROR(
269+
tensor_value != nullptr && tensor_value->sizes() != nullptr &&
270+
tensor_value->dim_order() != nullptr,
271+
InvalidProgram,
272+
"Null tensor metadata for output %zu",
273+
index);
241274
return TensorInfo::create(
242275
Span<const int32_t>(
243276
tensor_value->sizes()->data(), tensor_value->sizes()->size()),
@@ -257,7 +290,10 @@ size_t MethodMeta::num_attributes() const {
257290
auto value = values->Get(i);
258291
if (value->val_type() == executorch_flatbuffer::KernelTypes::Tensor) {
259292
auto tensor_value = value->val_as_Tensor();
260-
if (tensor_value->extra_tensor_info() != nullptr &&
293+
if (tensor_value != nullptr &&
294+
tensor_value->extra_tensor_info() != nullptr &&
295+
tensor_value->extra_tensor_info()->fully_qualified_name() !=
296+
nullptr &&
261297
tensor_value->extra_tensor_info()->fully_qualified_name()->c_str() !=
262298
nullptr) {
263299
++counter;
@@ -274,10 +310,19 @@ Result<TensorInfo> MethodMeta::attribute_tensor_meta(size_t index) const {
274310
auto value = values->Get(i);
275311
if (value->val_type() == executorch_flatbuffer::KernelTypes::Tensor) {
276312
auto tensor_value = value->val_as_Tensor();
277-
if (tensor_value->extra_tensor_info() != nullptr &&
313+
if (tensor_value != nullptr &&
314+
tensor_value->extra_tensor_info() != nullptr &&
315+
tensor_value->extra_tensor_info()->fully_qualified_name() !=
316+
nullptr &&
278317
tensor_value->extra_tensor_info()->fully_qualified_name()->c_str() !=
279318
nullptr) {
280319
if (counter == index) {
320+
ET_CHECK_OR_RETURN_ERROR(
321+
tensor_value->sizes() != nullptr &&
322+
tensor_value->dim_order() != nullptr,
323+
InvalidProgram,
324+
"Null tensor metadata for attribute %zu",
325+
index);
281326
auto t_name =
282327
tensor_value->extra_tensor_info()->fully_qualified_name();
283328
// Count constant returns as memory planned
@@ -322,7 +367,14 @@ Result<int64_t> MethodMeta::memory_planned_buffer_size(size_t index) const {
322367
num_buffers);
323368
// Index zero is reserved internally, and we hide it from users. Adjust the
324369
// provided index to point to one of the actual buffers.
325-
return s_plan_->non_const_buffer_sizes()->Get(index + 1);
370+
int64_t size = s_plan_->non_const_buffer_sizes()->Get(index + 1);
371+
ET_CHECK_OR_RETURN_ERROR(
372+
size >= 0,
373+
InvalidProgram,
374+
"memory_planned_buffer_size(%zu) has invalid negative size: %" PRId64,
375+
index,
376+
size);
377+
return size;
326378
}
327379

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

0 commit comments

Comments
 (0)