Skip to content

Commit 010f8ca

Browse files
JacobSzwejbkaJacobSzwejbka
authored
Skeleton out more robust program validation
Differential Revision: D92787645 Pull Request resolved: #17491
1 parent 20f9719 commit 010f8ca

8 files changed

Lines changed: 564 additions & 2 deletions

File tree

‎runtime/executor/program.cpp‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
#include <executorch/runtime/core/event_tracer_hooks.h>
1515
#include <executorch/runtime/executor/memory_manager.h>
1616
#include <executorch/runtime/executor/method.h>
17+
#include <executorch/runtime/executor/program_validation.h>
1718
#include <executorch/runtime/platform/profiler.h>
1819
#include <executorch/schema/extended_header.h>
1920
#include <executorch/schema/program_generated.h>
@@ -150,6 +151,13 @@ Result<executorch_flatbuffer::ExecutionPlan*> get_execution_plan(
150151
ok,
151152
InvalidProgram,
152153
"Verification failed; data may be truncated or corrupt");
154+
const executorch_flatbuffer::Program* flatbuffer_program =
155+
executorch_flatbuffer::GetProgram(program_data->data());
156+
Error err = validate_program(flatbuffer_program);
157+
ET_CHECK_OR_RETURN_ERROR(
158+
err == Error::Ok,
159+
InvalidProgram,
160+
"Program validation failed: likely a corrupt file");
153161
#else
154162
ET_LOG(
155163
Info, "InternalConsistency verification requested but not available");

‎runtime/executor/program.h‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88

99
#pragma once
1010

11-
#include <cinttypes>
1211
#include <cstdint>
1312
#include <optional>
1413

@@ -58,10 +57,15 @@ class Program final {
5857
*/
5958
Minimal,
6059
/**
60+
* When ET_ENABLE_PROGRAM_VERIFICATION is enabled,
6161
* Do full verification of the data, ensuring that internal pointers are
6262
* self-consistent and that the data has not been truncated or obviously
6363
* corrupted. May not catch all types of corruption, but should guard
6464
* against illegal memory operations during parsing.
65+
* Also performs additional semantic validation such as:
66+
* - Tensor numel overflow checks (ensuring size calculations don't
67+
* overflow)
68+
* - List element type validation
6569
*
6670
* Will have higher runtime overhead, scaling with the complexity of the
6771
* proram data.

‎runtime/executor/program_validation.cpp‎

Lines changed: 210 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,210 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#include <executorch/runtime/executor/program_validation.h>
10+
11+
#include <cstdint>
12+
13+
#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
14+
#include <executorch/runtime/platform/log.h>
15+
#include <executorch/schema/program_generated.h>
16+
17+
#include <c10/util/safe_numerics.h>
18+
19+
namespace executorch {
20+
namespace runtime {
21+
22+
ET_NODISCARD Error
23+
validate_tensor(const executorch_flatbuffer::Tensor* tensor) {
24+
if (tensor == nullptr) {
25+
return Error::InvalidProgram;
26+
}
27+
28+
const auto* sizes = tensor->sizes();
29+
if (sizes == nullptr) {
30+
return Error::InvalidProgram;
31+
}
32+
33+
ssize_t numel = 1;
34+
for (flatbuffers::uoffset_t i = 0; i < sizes->size(); i++) {
35+
int32_t size = sizes->Get(i);
36+
37+
if (size < 0) {
38+
ET_LOG(
39+
Error,
40+
"Size must be non-negative, got %d at dimension %u",
41+
size,
42+
static_cast<unsigned>(i));
43+
return Error::InvalidProgram;
44+
}
45+
46+
bool overflow =
47+
c10::mul_overflows(numel, static_cast<ssize_t>(size), &numel);
48+
if (overflow) {
49+
ET_LOG(
50+
Error,
51+
"numel overflowed at dimension %u with size %d",
52+
static_cast<unsigned>(i),
53+
size);
54+
return Error::InvalidProgram;
55+
}
56+
}
57+
58+
auto scalar_type =
59+
static_cast<executorch::aten::ScalarType>(tensor->scalar_type());
60+
if (!executorch::runtime::isValid(scalar_type)) {
61+
return Error::InvalidProgram;
62+
}
63+
64+
size_t nbytes;
65+
bool nbytes_overflow = c10::mul_overflows(
66+
static_cast<size_t>(numel),
67+
executorch::runtime::elementSize(scalar_type),
68+
&nbytes);
69+
if (nbytes_overflow) {
70+
ET_LOG(
71+
Error,
72+
"nbytes overflowed: numel %zd with element size %zu",
73+
numel,
74+
executorch::runtime::elementSize(scalar_type));
75+
return Error::InvalidProgram;
76+
}
77+
78+
return Error::Ok;
79+
}
80+
81+
ET_NODISCARD Error
82+
validate_program(const executorch_flatbuffer::Program* program) {
83+
if (program == nullptr) {
84+
return Error::InvalidProgram;
85+
}
86+
87+
// Validate all execution plans.
88+
const auto* execution_plans = program->execution_plan();
89+
if (execution_plans == nullptr) {
90+
ET_LOG(Error, "Program has null execution_plan");
91+
return Error::InvalidProgram;
92+
}
93+
94+
for (flatbuffers::uoffset_t plan_idx = 0; plan_idx < execution_plans->size();
95+
plan_idx++) {
96+
const auto* plan = execution_plans->Get(plan_idx);
97+
if (plan == nullptr) {
98+
ET_LOG(
99+
Error, "Execution plan %u is null", static_cast<unsigned>(plan_idx));
100+
return Error::InvalidProgram;
101+
}
102+
103+
// Validate all values in the plan.
104+
const auto* values = plan->values();
105+
if (values == nullptr) {
106+
ET_LOG(
107+
Error,
108+
"Execution plan %u has null values table",
109+
static_cast<unsigned>(plan_idx));
110+
return Error::InvalidProgram;
111+
}
112+
113+
for (flatbuffers::uoffset_t value_idx = 0; value_idx < values->size();
114+
value_idx++) {
115+
const auto* value = values->Get(value_idx);
116+
if (value == nullptr) {
117+
continue;
118+
}
119+
120+
// Check if this value is a tensor.
121+
if (value->val_type() == executorch_flatbuffer::KernelTypes::Tensor) {
122+
const auto* tensor =
123+
static_cast<const executorch_flatbuffer::Tensor*>(value->val());
124+
125+
Error err = validate_tensor(tensor);
126+
if (err != Error::Ok) {
127+
ET_LOG(
128+
Error,
129+
"Tensor validation failed for value %u in execution plan %u",
130+
static_cast<unsigned>(value_idx),
131+
static_cast<unsigned>(plan_idx));
132+
return err;
133+
}
134+
}
135+
136+
// Check if this value is a TensorList.
137+
if (value->val_type() == executorch_flatbuffer::KernelTypes::TensorList) {
138+
const auto* tensor_list =
139+
static_cast<const executorch_flatbuffer::TensorList*>(value->val());
140+
141+
if (tensor_list == nullptr) {
142+
ET_LOG(
143+
Error,
144+
"TensorList is null for value %u in execution plan %u",
145+
static_cast<unsigned>(value_idx),
146+
static_cast<unsigned>(plan_idx));
147+
return Error::InvalidProgram;
148+
}
149+
150+
const auto* items = tensor_list->items();
151+
if (items == nullptr) {
152+
ET_LOG(Error, "TensorList items is null");
153+
return Error::InvalidProgram;
154+
}
155+
156+
// Validate that each item index points to a Tensor evalue.
157+
for (flatbuffers::uoffset_t item_idx = 0; item_idx < items->size();
158+
item_idx++) {
159+
int32_t evalue_index = items->Get(item_idx);
160+
161+
// Check bounds.
162+
if (evalue_index < 0 ||
163+
static_cast<flatbuffers::uoffset_t>(evalue_index) >=
164+
values->size()) {
165+
ET_LOG(
166+
Error,
167+
"TensorList item %u has out-of-bounds index %d (values size "
168+
"%u) in execution plan %u",
169+
static_cast<unsigned>(item_idx),
170+
evalue_index,
171+
static_cast<unsigned>(values->size()),
172+
static_cast<unsigned>(plan_idx));
173+
return Error::InvalidProgram;
174+
}
175+
176+
// Check that the referenced evalue is actually a Tensor.
177+
const auto* referenced_value = values->Get(evalue_index);
178+
if (referenced_value == nullptr) {
179+
ET_LOG(
180+
Error,
181+
"TensorList item %u references null evalue at index %d in "
182+
"execution plan %u",
183+
static_cast<unsigned>(item_idx),
184+
evalue_index,
185+
static_cast<unsigned>(plan_idx));
186+
return Error::InvalidProgram;
187+
}
188+
189+
if (referenced_value->val_type() !=
190+
executorch_flatbuffer::KernelTypes::Tensor) {
191+
ET_LOG(
192+
Error,
193+
"TensorList item %u references non-Tensor evalue (type %d) at "
194+
"index %d in execution plan %u",
195+
static_cast<unsigned>(item_idx),
196+
static_cast<int>(referenced_value->val_type()),
197+
evalue_index,
198+
static_cast<unsigned>(plan_idx));
199+
return Error::InvalidProgram;
200+
}
201+
}
202+
}
203+
}
204+
}
205+
206+
return Error::Ok;
207+
}
208+
209+
} // namespace runtime
210+
} // namespace executorch

‎runtime/executor/program_validation.h‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#pragma once
10+
11+
#include <executorch/runtime/core/error.h>
12+
13+
// Forward declare flatbuffer types.
14+
namespace executorch_flatbuffer {
15+
struct ExecutionPlan;
16+
struct Program;
17+
struct Tensor;
18+
struct TensorList;
19+
} // namespace executorch_flatbuffer
20+
21+
namespace executorch {
22+
namespace runtime {
23+
24+
/**
25+
* Validates that computing numel (number of elements) from the tensor's sizes
26+
* will not overflow. This check should be performed before creating TensorImpl
27+
* objects to prevent undefined behavior from integer overflow.
28+
*
29+
* @param[in] tensor The flatbuffer Tensor to validate.
30+
* @return Error::Ok if the numel calculation is safe, Error::InvalidProgram
31+
* if computing numel would overflow.
32+
*/
33+
ET_NODISCARD Error validate_tensor(const executorch_flatbuffer::Tensor* tensor);
34+
35+
/**
36+
* Performs validation of all tensors and lists in the program, checking that
37+
* their metadata is semantically valid and will not cause issues during
38+
* execution.
39+
*
40+
* Currently validates:
41+
* - Tensor numel overflow (all tensors)
42+
* - TensorList element types (all TensorLists)
43+
*
44+
* @param[in] program The flatbuffer Program to validate.
45+
* @return Error::Ok if validation passes, Error::InvalidProgram if any
46+
* validation check fails.
47+
*/
48+
ET_NODISCARD Error
49+
validate_program(const executorch_flatbuffer::Program* program);
50+
51+
} // namespace runtime
52+
} // namespace executorch

‎runtime/executor/targets.bzl‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ def define_common_targets():
9090
],
9191
headers = [
9292
"platform_memory_allocator.h",
93+
"program_validation.h",
9394
],
9495
exported_headers = [
9596
"method.h",
@@ -121,7 +122,8 @@ def define_common_targets():
121122
],
122123
deps = [
123124
"//executorch/schema:program",
124-
"//executorch/runtime/core/exec_aten/util:tensor_dimension_limit"
125+
"//executorch/runtime/core/exec_aten/util:tensor_dimension_limit",
126+
"//executorch/runtime/core/portable_type/c10/c10:c10",
125127
],
126128
visibility = ["PUBLIC"],
127129
)

0 commit comments

Comments
 (0)