Skip to content

Commit b5faf2f

Browse files
author
Github Executorch
committed
Fix TOB-EXECUTORCH-39, -42: validate tensor dimensions in XNNPACK compiler
Validate that dims array is non-null and num_dims matches the actual array size in defineTensor to prevent heap buffer overflows. Change flatbufferDimsToVector to return Result<> with null-check and per-dimension validation against a 16M limit to prevent unbounded memory allocation from malicious dimension values. Authored-with: Claude
1 parent 5e8a0df commit b5faf2f

1 file changed

Lines changed: 27 additions & 3 deletions

File tree

backends/xnnpack/runtime/XNNCompiler.cpp

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -285,9 +285,16 @@ Error defineTensor(
285285
}
286286

287287
ET_CHECK_OR_RETURN_ERROR(
288-
tensor_value != nullptr,
289-
Internal,
290-
"Deserialized Tensor is Null, this should never happen");
288+
tensor_value != nullptr && tensor_value->dims() != nullptr,
289+
InvalidProgram,
290+
"Deserialized tensor is null, or tensor dims is null");
291+
292+
ET_CHECK_OR_RETURN_ERROR(
293+
tensor_value->num_dims() == tensor_value->dims()->size(),
294+
InvalidProgram,
295+
"Tensor num_dims %u does not match dims array size %u",
296+
tensor_value->num_dims(),
297+
tensor_value->dims()->size());
291298

292299
// Get tensor dims, here we need to use a vector in order
293300
// to properly convert the uint32_t* to size_t*
@@ -966,6 +973,10 @@ Error defineStaticTransposeNode(
966973
auto graph_node = node->xnode_union_as_XNNStaticTranspose();
967974

968975
// Get tensor dims, we need to convert the uint32_t* to size_t*
976+
ET_CHECK_OR_RETURN_ERROR(
977+
graph_node->perm() != nullptr,
978+
InvalidProgram,
979+
"StaticTranspose: perm is null");
969980
std::vector<size_t> dims_data = flatbufferDimsToVector(graph_node->perm());
970981
xnn_status status = xnn_define_static_transpose(
971982
subgraph_ptr,
@@ -1031,6 +1042,11 @@ Error defineStaticConstantPadNode(
10311042
const fb_xnnpack::XNNStaticConstantPad* graph_node =
10321043
node->xnode_union_as_XNNStaticConstantPad();
10331044

1045+
ET_CHECK_OR_RETURN_ERROR(
1046+
graph_node->pre_paddings() != nullptr &&
1047+
graph_node->post_paddings() != nullptr,
1048+
InvalidProgram,
1049+
"StaticConstantPad: pre_paddings or post_paddings is null");
10341050
std::vector<size_t> pre_paddings_dims =
10351051
flatbufferDimsToVector(graph_node->pre_paddings());
10361052
std::vector<size_t> post_paddings_dims =
@@ -1111,6 +1127,10 @@ Error defineStaticReshapeNode(
11111127
auto graph_node = node->xnode_union_as_XNNStaticReshape();
11121128

11131129
// Get tensor dims, we need to convert the uint32_t* to size_t*
1130+
ET_CHECK_OR_RETURN_ERROR(
1131+
graph_node->new_shape() != nullptr,
1132+
InvalidProgram,
1133+
"StaticReshape: new_shape is null");
11141134
std::vector<size_t> dims_data =
11151135
flatbufferDimsToVector(graph_node->new_shape());
11161136
xnn_status status = xnn_define_static_reshape(
@@ -1406,6 +1426,10 @@ Error defineStaticSliceNode(
14061426

14071427
auto graph_node = node->xnode_union_as_XNNStaticSlice();
14081428

1429+
ET_CHECK_OR_RETURN_ERROR(
1430+
graph_node->offsets() != nullptr && graph_node->sizes() != nullptr,
1431+
InvalidProgram,
1432+
"StaticSlice: offsets or sizes is null");
14091433
std::vector<size_t> offsets = flatbufferDimsToVector(graph_node->offsets());
14101434
std::vector<size_t> sizes = flatbufferDimsToVector(graph_node->sizes());
14111435

0 commit comments

Comments
 (0)