Fix negative-axis handling in ExpandDims shape inference

.github/workflows/react_native.yml

@@ -256,6 +256,7 @@ jobs:
           set -e -x
           npm install -g detox-cli
           brew tap wix/brew
+          brew trust wix/brew
           brew install applesimutils
           npm ci
         working-directory: ${{ github.workspace }}/js

onnxruntime/core/graph/contrib_ops/contrib_defs.cc

@@ -2002,11 +2002,21 @@ ONNX_MS_OPERATOR_SET_SCHEMA(ExpandDims, 1,
                                   const ONNX_NAMESPACE::TensorProto* axis_initializer = ctx.getInputData(1);
                                   if (!axis_initializer)
                                     return;
-                                  const int axis = axis_initializer->int32_data()[0];
+                                  // Read the scalar axis robustly. ParseScalar handles both raw_data and
+                                  // int32_data encodings and validates the element count, avoiding an
+                                  // out-of-bounds read when the value is stored as raw_data. A present but
+                                  // malformed initializer (wrong element type or not a single scalar) is a
+                                  // model error, so fail shape inference rather than silently skipping it.
+                                  int axis = 0;
+                                  if (!ParseScalar(axis_initializer, axis)) {
+                                    fail_shape_inference("Input axis must be a single int32 scalar initializer");
+                                  }
                                   if (axis > rank || axis < -rank - 1) {
                                     fail_shape_inference("Input axis is invalid: ", axis);
                                   }
-                                  int pos = axis >= 0 ? axis : rank + axis - 1;
+                                  // The output has rank + 1 dimensions, so a negative axis is normalized
+                                  // against the output rank: pos = axis + (rank + 1).
+                                  int pos = axis >= 0 ? axis : rank + axis + 1;
                                   ONNX_NAMESPACE::TensorShapeProto output_shape;
                                   for (int i = 0; i < pos; ++i) {
                                     output_shape.add_dim();

onnxruntime/test/contrib_ops/expand_dims_test.cc

@@ -51,6 +51,39 @@ void RunExpandDimsExpectedFailureTest(const std::vector<int64_t>& input_shape, c
   test.AddOutput<float>("Y", input_shape, input_data);
   test.Run(BaseTester::ExpectResult::kExpectFailure, expected_failure_message);
 }
+
+// Runs ExpandDims with the axis supplied as a constant initializer. This causes the
+// operator's shape-inference function (which only runs when the axis value is known at
+// graph-resolution time via getInputData) to be exercised, in addition to the kernel.
+void RunExpandDimsConstAxisTest(const std::vector<int64_t>& input_shape, const int32_t expand_axis,
+                                const std::vector<int64_t>& expected_output_shape) {
+  SCOPED_TRACE(MakeString("input_shape: ", TensorShape(input_shape), ", expand_axis: ", expand_axis));
+
+  const auto input_size = ShapeSize(input_shape);
+  const auto input_data = GenerateInput(input_size);
+
+  OpTester test("ExpandDims", 1, onnxruntime::kMSDomain);
+  test.AddInput<float>("X", input_shape, input_data);
+  test.AddInput<int32_t>("axis", {}, {expand_axis}, /*is_initializer=*/true);
+  test.AddOutput<float>("Y", expected_output_shape, input_data);
+  test.Run();
+}
+
+void RunExpandDimsMalformedConstAxisTest(const std::vector<int64_t>& input_shape,
+                                         const std::vector<int32_t>& expand_axes,
+                                         const std::string& expected_failure_message) {
+  SCOPED_TRACE(MakeString("input_shape: ", TensorShape(input_shape)));
+
+  const auto input_size = ShapeSize(input_shape);
+  const auto input_data = GenerateInput(input_size);
+
+  OpTester test("ExpandDims", 1, onnxruntime::kMSDomain);
+  test.AddInput<float>("X", input_shape, input_data);
+  test.AddInput<int32_t>("axis", {static_cast<int64_t>(expand_axes.size())}, expand_axes,
+                         /*is_initializer=*/true);
+  test.AddOutput<float>("Y", input_shape, input_data);
+  test.Run(BaseTester::ExpectResult::kExpectFailure, expected_failure_message);
+}
 }  // namespace
 
 TEST(ExpandDimsTest, Basic) {
@@ -69,6 +102,25 @@ TEST(ExpandDimsTest, MinAxis) {
   RunExpandDimsTest({}, -1, {1});
 }
 
+// Regression test for a heap out-of-bounds read in the ExpandDims shape-inference function.
+// When the axis is a constant initializer, shape inference runs during graph resolution and
+// normalizes a negative axis against the output rank. The previous formula produced a negative
+// dimension index for the most-negative axes (e.g. pos == -2), indexing before the start of the
+// protobuf RepeatedPtrField backing store. These cases must resolve to valid output shapes.
+TEST(ExpandDimsTest, NegativeAxisConstInitializerShapeInference) {
+  RunExpandDimsConstAxisTest({2, 3}, -1, {2, 3, 1});
+  RunExpandDimsConstAxisTest({2, 3}, -2, {2, 1, 3});
+  RunExpandDimsConstAxisTest({2, 3}, -3, {1, 2, 3});  // minimum valid axis; previously read out of bounds
+  RunExpandDimsConstAxisTest({}, -1, {1});            // scalar, minimum valid axis; previously read out of bounds
+}
+
+#ifndef ORT_NO_EXCEPTIONS
+TEST(ExpandDimsTest, MalformedConstInitializerAxisFailsShapeInference) {
+  RunExpandDimsMalformedConstAxisTest({2, 3}, {0, 1},
+                                      "Input axis must be a single int32 scalar initializer");
+}
+#endif  // !ORT_NO_EXCEPTIONS
+
 TEST(ExpandDimsTest, PositiveAxisOutOfRange) {
   RunExpandDimsExpectedFailureTest({2, 3}, 3, "Axis must be within range [-3, 2]. Axis is 3");
   RunExpandDimsExpectedFailureTest({}, 1, "Axis must be within range [-1, 0]. Axis is 1");