Skip to content

Commit c8339db

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 c8339db

1 file changed

Lines changed: 30 additions & 2 deletions

File tree

runtime/executor/method_meta.cpp

Lines changed: 30 additions & 2 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,6 +60,12 @@ 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]: %d. 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);
@@ -92,6 +100,10 @@ Result<size_t> calculate_nbytes(
92100
executorch::aten::ScalarType scalar_type,
93101
const bool is_memory_planned,
94102
std::string_view name) {
103+
ET_CHECK_OR_RETURN_ERROR(
104+
sizes.data() != nullptr && dim_order.data() != nullptr,
105+
InvalidProgram,
106+
"Null sizes or dim_order in tensor metadata");
95107
auto nbytes = calculate_nbytes(sizes, scalar_type);
96108
ET_CHECK_OR_RETURN_ERROR(
97109
nbytes.ok(),
@@ -186,6 +198,11 @@ Result<TensorInfo> MethodMeta::input_tensor_meta(size_t index) const {
186198
auto input_index = s_plan_->inputs()->Get(index);
187199
// input_index was already validated by input_tag().
188200
auto tensor_value = s_plan_->values()->Get(input_index)->val_as_Tensor();
201+
ET_CHECK_OR_RETURN_ERROR(
202+
tensor_value != nullptr,
203+
InvalidProgram,
204+
"Null tensor value for input %zu",
205+
index);
189206
return TensorInfo::create(
190207
Span<const int32_t>(
191208
tensor_value->sizes()->data(), tensor_value->sizes()->size()),
@@ -237,7 +254,11 @@ Result<TensorInfo> MethodMeta::output_tensor_meta(size_t index) const {
237254
auto output_index = s_plan_->outputs()->Get(index);
238255
// output_index was already validated by output_tag().
239256
auto tensor_value = s_plan_->values()->Get(output_index)->val_as_Tensor();
240-
257+
ET_CHECK_OR_RETURN_ERROR(
258+
tensor_value != nullptr,
259+
InvalidProgram,
260+
"Null tensor value for output %zu",
261+
index);
241262
return TensorInfo::create(
242263
Span<const int32_t>(
243264
tensor_value->sizes()->data(), tensor_value->sizes()->size()),
@@ -322,7 +343,14 @@ Result<int64_t> MethodMeta::memory_planned_buffer_size(size_t index) const {
322343
num_buffers);
323344
// Index zero is reserved internally, and we hide it from users. Adjust the
324345
// provided index to point to one of the actual buffers.
325-
return s_plan_->non_const_buffer_sizes()->Get(index + 1);
346+
int64_t size = s_plan_->non_const_buffer_sizes()->Get(index + 1);
347+
ET_CHECK_OR_RETURN_ERROR(
348+
size >= 0,
349+
InvalidProgram,
350+
"memory_planned_buffer_size(%zu) has invalid negative size: %" PRId64,
351+
index,
352+
size);
353+
return size;
326354
}
327355

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

0 commit comments

Comments
 (0)