From 50772ec6fe74c5cbede8dbfa87d20057ae9c5b58 Mon Sep 17 00:00:00 2001 From: mohammadmseet-hue Date: Wed, 1 Apr 2026 10:46:47 +0200 Subject: [PATCH] Fix missing bounds checks in multiple subgraph API functions Add bounds validation for user-controlled parameters that were used as loop bounds or array indices without checking: - xnn_define_static_expand_dims: num_new_axes was used to write into int32_t ynn_axes[XNN_MAX_TENSOR_DIMS] without bounds check. Add check that num_new_axes <= XNN_MAX_TENSOR_DIMS. - xnn_define_static_reduce: num_reduction_axes was used to write into int64_t signed_reduction_axes[XNN_MAX_TENSOR_DIMS] without bounds check. Add check that num_reduction_axes <= XNN_MAX_TENSOR_DIMS. - ynn_define_broadcast: axes were passed to axis_to_slinky_dim without validate_axis, producing out-of-range bitset indices. Add validate_axis call. - ynn_define_broadcast_like: same missing validate_axis. Add it. - ynn_define_reduce: same missing validate_axis. Add it. - ynn_define_fuse_dims: axes[i]+1 was passed to axis_to_slinky_dim but only axes[i] was validated. When axes[i]=rank-1, axes[i]+1=rank produces an invalid index. Add validate_axis for axes[i]+1. --- ynnpack/subgraph/broadcast.cc | 2 + ynnpack/subgraph/broadcast_like.cc | 2 + ynnpack/subgraph/copy.cc | 3 + ynnpack/subgraph/reduce.cc | 2 + ynnpack/subgraph/test/BUILD | 8 ++ ynnpack/subgraph/test/oob_write_poc.cc | 166 +++++++++++++++++++++++++ ynnpack/xnnpack/subgraph.cc | 6 + 7 files changed, 189 insertions(+) create mode 100644 ynnpack/subgraph/test/oob_write_poc.cc diff --git a/ynnpack/subgraph/broadcast.cc b/ynnpack/subgraph/broadcast.cc index ed5e65c7c98..7c34305cec9 100644 --- a/ynnpack/subgraph/broadcast.cc +++ b/ynnpack/subgraph/broadcast.cc @@ -37,6 +37,8 @@ ynn_status ynn_define_broadcast(ynn_subgraph_t subgraph, size_t num_axes, ynn::axes_set axes_set; for (size_t i = 0; i < num_axes; ++i) { + YNN_RETURN_IF_ERROR( + validate_axis("broadcast", "input", input.rank(), axes[i])); axes_set[axis_to_slinky_dim(input.rank(), axes[i])] = true; } diff --git a/ynnpack/subgraph/broadcast_like.cc b/ynnpack/subgraph/broadcast_like.cc index 0721ef2fda9..91a348d0799 100644 --- a/ynnpack/subgraph/broadcast_like.cc +++ b/ynnpack/subgraph/broadcast_like.cc @@ -40,6 +40,8 @@ ynn_status ynn_define_broadcast_like(ynn_subgraph_t subgraph, size_t num_axes, const ynn_value& template_value = subgraph->value(template_id); ynn::axes_set axes_set; for (size_t i = 0; i < num_axes; ++i) { + YNN_RETURN_IF_ERROR( + validate_axis("broadcast_like", "input", input.rank(), axes[i])); const int axis = axis_to_slinky_dim(input.rank(), axes[i]); if (axis < template_value.rank()) { axes_set[axis] = true; diff --git a/ynnpack/subgraph/copy.cc b/ynnpack/subgraph/copy.cc index 05b84aa2715..6acab8b8086 100644 --- a/ynnpack/subgraph/copy.cc +++ b/ynnpack/subgraph/copy.cc @@ -554,6 +554,9 @@ ynn_status ynn_define_fuse_dims(ynn_subgraph_t subgraph, size_t num_axes, for (size_t i = 0; i < num_axes; ++i) { YNN_RETURN_IF_ERROR( validate_axis("fuse_dims", "input", input.rank(), axes[i])); + // Fusing requires a next dimension, so axes[i]+1 must also be valid. + YNN_RETURN_IF_ERROR( + validate_axis("fuse_dims", "input", input.rank(), axes[i] + 1)); // Since we are reversing the axes, the first dimension to fuse is actually // the next dimension. op.axes[axis_to_slinky_dim(input.rank(), axes[i] + 1)] = true; diff --git a/ynnpack/subgraph/reduce.cc b/ynnpack/subgraph/reduce.cc index e9b453b9914..a5784da53be 100644 --- a/ynnpack/subgraph/reduce.cc +++ b/ynnpack/subgraph/reduce.cc @@ -454,6 +454,8 @@ ynn_status ynn_define_reduce(ynn_subgraph_t subgraph, ynn::axes_set k_dims; for (size_t i = 0; i < num_axes; ++i) { + YNN_RETURN_IF_ERROR( + validate_axis("reduce", "input_a", a.rank(), axes[i])); const int axis = axis_to_slinky_dim(a.rank(), axes[i]); if (axis < a.rank()) { k_dims[axis] = true; diff --git a/ynnpack/subgraph/test/BUILD b/ynnpack/subgraph/test/BUILD index 8b8200cbbc2..4a4ec6edda4 100644 --- a/ynnpack/subgraph/test/BUILD +++ b/ynnpack/subgraph/test/BUILD @@ -198,6 +198,14 @@ cc_test( "stencil_copy", ]] +cc_test( + name = "oob_write_poc", + srcs = ["oob_write_poc.cc"], + linkopts = ynn_binary_linkopts(), + malloc = ynn_binary_malloc(), + deps = TEST_DEPS + ["//ynnpack/xnnpack"], +) + cc_test( name = "reduce_window", srcs = ["reduce_window.cc"], diff --git a/ynnpack/subgraph/test/oob_write_poc.cc b/ynnpack/subgraph/test/oob_write_poc.cc new file mode 100644 index 00000000000..e5b43a58c79 --- /dev/null +++ b/ynnpack/subgraph/test/oob_write_poc.cc @@ -0,0 +1,166 @@ +// Copyright 2025 Google LLC +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +// PoC: Stack buffer overflows in XNNPACK/YNNPACK public API functions. +// +// Bug class: User-controlled count parameters used as loop bounds to write +// into fixed-size stack arrays without bounds checking. +// +// xnn_define_static_expand_dims: int32_t ynn_axes[XNN_MAX_TENSOR_DIMS=6] +// written with num_new_axes as loop bound. No bounds check. +// +// xnn_define_static_reduce: int64_t signed_reduction_axes[XNN_MAX_TENSOR_DIMS=6] +// written with num_reduction_axes as loop bound. No bounds check. + +#include +#include + +#include +#include "include/xnnpack.h" +#include "ynnpack/include/ynnpack.h" + +namespace { + +// Bug 1: xnn_define_static_expand_dims stack buffer overflow. +// ynnpack/xnnpack/subgraph.cc:774-776: +// int32_t ynn_axes[XNN_MAX_TENSOR_DIMS]; // size 6 +// for (size_t i = 0; i < num_new_axes; ++i) { +// ynn_axes[i] = new_axes[i]; // OOB when num_new_axes > 6 +// } +TEST(OobWrite, StaticExpandDimsStackOverflow) { + xnn_subgraph_t subgraph = nullptr; + ASSERT_EQ(xnn_create_subgraph(2, 0, &subgraph), xnn_status_success); + + size_t dims[] = {4, 4}; + uint32_t input_id = XNN_INVALID_VALUE_ID; + ASSERT_EQ(xnn_define_tensor_value(subgraph, xnn_datatype_fp32, 2, dims, + nullptr, 0, XNN_VALUE_FLAG_EXTERNAL_INPUT, + &input_id), + xnn_status_success); + uint32_t output_id = XNN_INVALID_VALUE_ID; + ASSERT_EQ(xnn_define_tensor_value(subgraph, xnn_datatype_fp32, 2, dims, + nullptr, 1, XNN_VALUE_FLAG_EXTERNAL_OUTPUT, + &output_id), + xnn_status_success); + + // num_new_axes=10 > XNN_MAX_TENSOR_DIMS=6. + // The loop writes 10 int32_t elements into a 6-element stack array. + // ASAN detects: stack-buffer-overflow WRITE of size 4. + size_t new_axes[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}; + xnn_status status = xnn_define_static_expand_dims(subgraph, 10, new_axes, + input_id, output_id, 0); + // Should return error, not overflow. + EXPECT_NE(status, xnn_status_success); + + xnn_delete_subgraph(subgraph); +} + +// Bug 2: xnn_define_static_reduce stack buffer overflow. +// ynnpack/xnnpack/subgraph.cc:805-807: +// int64_t signed_reduction_axes[XNN_MAX_TENSOR_DIMS]; // size 6 +// for (int i = 0; i < num_reduction_axes; i++) { +// signed_reduction_axes[i] = reduction_axes[i]; // OOB when > 6 +// } +TEST(OobWrite, StaticReduceStackOverflow) { + xnn_subgraph_t subgraph = nullptr; + ASSERT_EQ(xnn_create_subgraph(2, 0, &subgraph), xnn_status_success); + + size_t dims[] = {2, 3, 4, 5}; + uint32_t input_id = XNN_INVALID_VALUE_ID; + ASSERT_EQ(xnn_define_tensor_value(subgraph, xnn_datatype_fp32, 4, dims, + nullptr, 0, XNN_VALUE_FLAG_EXTERNAL_INPUT, + &input_id), + xnn_status_success); + uint32_t output_id = XNN_INVALID_VALUE_ID; + size_t out_dims[] = {1}; + ASSERT_EQ(xnn_define_tensor_value(subgraph, xnn_datatype_fp32, 1, out_dims, + nullptr, 1, XNN_VALUE_FLAG_EXTERNAL_OUTPUT, + &output_id), + xnn_status_success); + + // num_reduction_axes=10 > XNN_MAX_TENSOR_DIMS=6. + // The loop writes 10 int64_t elements into a 6-element stack array. + // ASAN detects: stack-buffer-overflow WRITE of size 8. + size_t reduction_axes[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}; + xnn_status status = xnn_define_static_reduce(subgraph, xnn_reduce_sum, 10, + reduction_axes, input_id, + output_id, 0); + EXPECT_NE(status, xnn_status_success); + + xnn_delete_subgraph(subgraph); +} + +// Bug 3: Missing axis validation in ynn_define_broadcast. +// broadcast.cc:38-41: axes_set[axis_to_slinky_dim(rank, axes[i])] = true +// No validate_axis call. Out-of-range axes produce invalid bitset indices. +TEST(OobWrite, BroadcastMissingAxisValidation) { + ynn_subgraph_t sg = nullptr; + ynn_create_subgraph(2, 0, &sg); + size_t dims[] = {4, 4}; + uint32_t input_id = YNN_INVALID_VALUE_ID; + ynn_define_tensor(sg, ynn_type_fp32, 2, dims, nullptr, + YNN_VALUE_FLAG_EXTERNAL_INPUT, &input_id); + uint32_t output_id = YNN_INVALID_VALUE_ID; + ynn_define_tensor(sg, ynn_type_fp32, 2, dims, nullptr, + YNN_VALUE_FLAG_EXTERNAL_OUTPUT, &output_id); + + // axis=100 on rank-2 tensor. axis_to_slinky_dim(2, 100) = -99. + // This is UB: writes to bitset at index SIZE_MAX-98. + int32_t axes[] = {100}; + ynn_status status = + ynn_define_broadcast(sg, 1, axes, input_id, &output_id, 0); + EXPECT_NE(status, ynn_status_success); + ynn_delete_subgraph(sg); +} + +// Bug 4: Off-by-one in ynn_define_fuse_dims. +// copy.cc:560: op.axes[axis_to_slinky_dim(rank, axes[i]+1)] = true +// axes[i]=rank-1 passes validate_axis, but axes[i]+1 goes to +// axis_to_slinky_dim which returns -1 → bitset[SIZE_MAX]. +TEST(OobWrite, FuseDimsOffByOneOverflow) { + ynn_subgraph_t sg = nullptr; + ynn_create_subgraph(2, 0, &sg); + size_t dims[] = {2, 3, 4, 5}; + uint32_t input_id = YNN_INVALID_VALUE_ID; + ynn_define_tensor(sg, ynn_type_fp32, 4, dims, nullptr, + YNN_VALUE_FLAG_EXTERNAL_INPUT, &input_id); + uint32_t output_id = YNN_INVALID_VALUE_ID; + size_t out_dims[] = {2, 3, 20}; + ynn_define_tensor(sg, ynn_type_fp32, 3, out_dims, nullptr, + YNN_VALUE_FLAG_EXTERNAL_OUTPUT, &output_id); + + // axes[0]=3 (last valid axis for rank-4). + // axis_to_slinky_dim(4, 4) = -1 → bitset[SIZE_MAX] = UB/OOB write. + int32_t axes[] = {3}; + ynn_status status = + ynn_define_fuse_dims(sg, 1, axes, input_id, &output_id, 0); + EXPECT_NE(status, ynn_status_success); + ynn_delete_subgraph(sg); +} + +// Bug 5: Missing axis validation in ynn_define_reduce. +// reduce.cc:456-462: k_dims[axis_to_slinky_dim(rank, axes[i])] = true +// No validate_axis call. +TEST(OobWrite, ReduceMissingAxisValidation) { + ynn_subgraph_t sg = nullptr; + ynn_create_subgraph(2, 0, &sg); + size_t dims[] = {2, 3, 4, 5}; + uint32_t input_id = YNN_INVALID_VALUE_ID; + ynn_define_tensor(sg, ynn_type_fp32, 4, dims, nullptr, + YNN_VALUE_FLAG_EXTERNAL_INPUT, &input_id); + uint32_t output_id = YNN_INVALID_VALUE_ID; + size_t out_dims[] = {2}; + ynn_define_tensor(sg, ynn_type_fp32, 1, out_dims, nullptr, + YNN_VALUE_FLAG_EXTERNAL_OUTPUT, &output_id); + + int32_t axes[] = {100}; + ynn_status status = + ynn_define_reduce(sg, ynn_reduce_sum, 1, axes, input_id, + YNN_INVALID_VALUE_ID, &output_id, 0); + EXPECT_NE(status, ynn_status_success); + ynn_delete_subgraph(sg); +} + +} // namespace diff --git a/ynnpack/xnnpack/subgraph.cc b/ynnpack/xnnpack/subgraph.cc index 153ab6df30f..a14ab70e5c5 100644 --- a/ynnpack/xnnpack/subgraph.cc +++ b/ynnpack/xnnpack/subgraph.cc @@ -773,6 +773,9 @@ xnn_status xnn_define_static_expand_dims(xnn_subgraph_t subgraph, const size_t* new_axes, uint32_t input_id, uint32_t output_id, uint32_t flags) { + if (num_new_axes > XNN_MAX_TENSOR_DIMS) { + return xnn_status_invalid_parameter; + } int32_t ynn_axes[XNN_MAX_TENSOR_DIMS]; for (size_t i = 0; i < num_new_axes; ++i) { ynn_axes[i] = new_axes[i]; @@ -804,6 +807,9 @@ xnn_status xnn_define_static_reduce(xnn_subgraph_t subgraph, const size_t* reduction_axes, uint32_t input_id, uint32_t output_id, uint32_t flags) { + if (num_reduction_axes > XNN_MAX_TENSOR_DIMS) { + return xnn_status_invalid_parameter; + } int64_t signed_reduction_axes[XNN_MAX_TENSOR_DIMS]; for (int i = 0; i < num_reduction_axes; i++) { signed_reduction_axes[i] = reduction_axes[i];