Validate tensor dimensions in xnnpack flatbuffer#18896
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/18896
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 1 New Failure, 2 Unrelated FailuresAs of commit 459b8f5 with merge base 5e8a0df ( NEW FAILURE - The following job has failed:
BROKEN TRUNK - The following jobs failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
bot
added
the
CLA Signed
|
@claude review |
This PR needs a
|
|
Claude finished @lucylq's task in 1m 33s —— View job Review of PR #18896
This PR adds input validation to prevent heap buffer overflows (TOB-EXECUTORCH-39, -42) from malicious flatbuffer payloads. The approach is sound — validating at deserialization time before passing untrusted dimension values to XNNPACK APIs. Here's a detailed review: Looks GoodNull check in
Combined null check in Error propagation at all call sites: All callers ( Issues1.
2. constexpr uint32_t kMaxDimensionValue = 1 << 24; // 16M per dimensionThe literal 3. No total-elements overflow check (suggestion, not a blocker) The per-dimension cap of 16M prevents any single dimension from being absurdly large, but a tensor with dimensions SummaryThe PR is a clear security improvement. It correctly validates deserialized flatbuffer tensor data before it reaches XNNPACK APIs, addressing heap buffer overflow risks from crafted inputs. The |
|
addressed claude's comments |
There was a problem hiding this comment.
Pull request overview
This PR hardens the XNNPACK runtime flatbuffer deserialization path by validating shape vectors before passing raw pointers into XNNPACK APIs, aiming to prevent heap buffer overflows and unbounded allocations from malformed or malicious programs.
Changes:
- Change
flatbufferDimsToVectorto returnResult<std::vector<T>>, add NULL-checking, and validate each dimension value against a max limit. - Add validation in
defineTensorto ensure dims exist andnum_dimsmatches the flatbuffer vector length, and propagateResulterrors to callers. - Update several node-definition helpers (transpose/reshape/slice/pad) to handle
ResultfromflatbufferDimsToVector.
Comments suppressed due to low confidence (3)
backends/xnnpack/runtime/XNNCompiler.cpp:1160
xnn_define_static_reshapeis called withgraph_node->num_dims()butnew_shape()->size()is not validated againstnum_dims. A malformed flatbuffer could cause XNNPACK to read pastdims_data.
Please addnew_shape != nullptrandnew_shape->size() == num_dimsvalidation (and a reasonablenum_dimscap) before the call.
auto graph_node = node->xnode_union_as_XNNStaticReshape();
// Get tensor dims, we need to convert the uint32_t* to size_t*
auto dims_result = flatbufferDimsToVector(graph_node->new_shape());
if (!dims_result.ok()) {
return dims_result.error();
}
std::vector<size_t> dims_data = std::move(dims_result.get());
xnn_status status = xnn_define_static_reshape(
subgraph_ptr,
graph_node->num_dims(),
dims_data.data(),
remapped_ids.at(graph_node->input_id()),
backends/xnnpack/runtime/XNNCompiler.cpp:1465
xnn_define_static_sliceis called withgraph_node->num_dims()but there is no validation thatoffsets()->size()andsizes()->size()both equalnum_dims. If either vector is shorter, XNNPACK will read pastoffsets/sizes.
Add checks for non-null vectors, matching sizes, and a reasonablenum_dimscap before calling XNNPACK.
auto graph_node = node->xnode_union_as_XNNStaticSlice();
auto offsets_result = flatbufferDimsToVector(graph_node->offsets());
if (!offsets_result.ok()) {
return offsets_result.error();
}
std::vector<size_t> offsets = std::move(offsets_result.get());
auto sizes_result = flatbufferDimsToVector(graph_node->sizes());
if (!sizes_result.ok()) {
return sizes_result.error();
}
std::vector<size_t> sizes = std::move(sizes_result.get());
xnn_status status = xnn_define_static_slice(
subgraph_ptr,
graph_node->num_dims(),
offsets.data(),
sizes.data(),
remapped_ids.at(graph_node->input_id()),
backends/xnnpack/runtime/XNNCompiler.cpp:1004
xnn_define_static_transposeis called withgraph_node->num_dims()but there is no validation thatperm()->size()matchesnum_dims. Ifnum_dimsis larger than the perm vector length, XNNPACK will read pastdims_data.
Add a check thatperm != nullptrandperm->size() == num_dims(and ideallynum_dims <= XNN_MAX_TENSOR_DIMS) before calling into XNNPACK.
auto graph_node = node->xnode_union_as_XNNStaticTranspose();
// Get tensor dims, we need to convert the uint32_t* to size_t*
auto dims_result = flatbufferDimsToVector(graph_node->perm());
if (!dims_result.ok()) {
return dims_result.error();
}
std::vector<size_t> dims_data = std::move(dims_result.get());
xnn_status status = xnn_define_static_transpose(
subgraph_ptr,
graph_node->num_dims(),
dims_data.data(),
remapped_ids.at(graph_node->input_id()),
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| ET_CHECK_OR_RETURN_ERROR( | ||
| tensor_value != nullptr, | ||
| Internal, | ||
| "Deserialized Tensor is Null, this should never happen"); | ||
| tensor_value != nullptr && tensor_value->dims() != nullptr, | ||
| InvalidProgram, | ||
| "Deserialized tensor is null, or tensor dims is null"); | ||
|
|
||
| ET_CHECK_OR_RETURN_ERROR( | ||
| tensor_value->num_dims() == tensor_value->dims()->size(), | ||
| InvalidProgram, | ||
| "Tensor num_dims %u does not match dims array size %u", | ||
| tensor_value->num_dims(), | ||
| tensor_value->dims()->size()); |
| Converts dims from uint32 to size_t. Takes in a flatbuffer vector | ||
| of uint32_t and returns a std::vector of size_t. XNNPACK takes in | ||
| dims of size_t* but tensor shape is serialized in flatbuffer as | ||
| int32_t. As a result, we need to static cast the shapes to size_t | ||
| int32_t. As a result, we need to static cast the shapes to size_t. | ||
| Individual dimension values are validated to prevent unbounded memory | ||
| allocation from malicious inputs. |
| ET_CHECK_OR_RETURN_ERROR( | ||
| fb_dims != nullptr, | ||
| InvalidProgram, | ||
| "flatbufferDimsToVector: dims vector is null"); | ||
| std::vector<T> dims_data; | ||
| dims_data.reserve(fb_dims->size()); | ||
| for (auto fb_dim : *fb_dims) { |
| ET_CHECK_OR_RETURN_ERROR( | ||
| tensor_value != nullptr, | ||
| Internal, | ||
| "Deserialized Tensor is Null, this should never happen"); | ||
| tensor_value != nullptr && tensor_value->dims() != nullptr, | ||
| InvalidProgram, | ||
| "Deserialized tensor is null, or tensor dims is null"); | ||
|
|
| const fb_xnnpack::XNNStaticConstantPad* graph_node = | ||
| node->xnode_union_as_XNNStaticConstantPad(); | ||
|
|
||
| std::vector<size_t> pre_paddings_dims = | ||
| flatbufferDimsToVector(graph_node->pre_paddings()); | ||
| std::vector<size_t> post_paddings_dims = | ||
| flatbufferDimsToVector(graph_node->post_paddings()); | ||
| auto pre_result = flatbufferDimsToVector(graph_node->pre_paddings()); | ||
| if (!pre_result.ok()) { | ||
| return pre_result.error(); | ||
| } | ||
| std::vector<size_t> pre_paddings_dims = std::move(pre_result.get()); | ||
| auto post_result = flatbufferDimsToVector(graph_node->post_paddings()); | ||
| if (!post_result.ok()) { | ||
| return post_result.error(); | ||
| } | ||
| std::vector<size_t> post_paddings_dims = std::move(post_result.get()); |
| Individual dimension values are validated to prevent unbounded memory | ||
| allocation from malicious inputs. | ||
| */ | ||
| constexpr uint32_t kMaxDimensionValue = 1 << 24; // 16M per dimension |
|
Clamping max elements at 16 million makes me a little bit nervous. Flattening large tensors isn't uncommon, and if a graph break happened in between the flatten call and whatever uses it, it could break existing PTEs. How critical is this check? |
…piler 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
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