Fix missing bounds checks in subgraph API functions

ynnpack/subgraph/broadcast.cc

@@ -37,7 +37,10 @@ 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) {
-    axes_set[axis_to_slinky_dim(input.rank(), axes[i])] = true;
+    const int axis = axis_to_slinky_dim(input.rank(), axes[i]);
+    if (axis < input.rank()) {
+      axes_set[axis] = true;
+    }
   }
 
   ynn_node node;

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;

ynnpack/xnnpack/subgraph.cc

@@ -773,6 +773,11 @@ 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) {
+    YNN_LOG_ERROR() << "num_new_axes " << num_new_axes
+                    << " exceeds XNN_MAX_TENSOR_DIMS " << 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 +809,11 @@ 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) {
+    YNN_LOG_ERROR() << "num_reduction_axes " << num_reduction_axes
+                    << " exceeds XNN_MAX_TENSOR_DIMS " << 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];