Contrib Op Shape Inference Memory Safety

Contrib-Op Shape-Inference Output-Index Safety

Reusable method for finding and fixing the bug class where an operator's TypeAndShapeInference function guards an output write with getNumOutputs() > N but then writes an output index greater than N. For a node that declares fewer outputs, the written index is past the end of the inference context's output vector.

Scope: schema-level shape inference in onnxruntime/core/graph/contrib_ops/*.cc and shape_inference_functions.cc. This runs once during Graph::Resolve (model-load time), EP-agnostic - there is no per-EP (CPU/CUDA/ROCm) kernel duplicate of this code to chase. Op kernels allocate outputs via the bounds-safe OpKernelContext::Output(index) and are a separate concern.

1. The pattern

// onnxruntime/core/graph/contrib_ops/bert_defs.cc  (before)
propagateElemTypeFromInputToOutput(ctx, 0, 0);
if (ctx.getNumOutputs() > 1) {                       // guard says "> 1"
  propagateElemTypeFromInputToOutput(ctx, 0, 1);
  propagateElemTypeFromInputToOutput(ctx, 0, 2);     // but writes index 2
}

The guard getNumOutputs() > 1 admits a node with exactly 2 outputs (indices 0, 1), yet the body writes index 2. The implication "> 1 ⇒ index 2 exists" is false: > 1 only guarantees indices 0 and 1.

Why a smaller output count is valid

Trailing outputs declared OpSchema::Optional lower min_output. ONNX derives min_output = number of required outputs, max_output = total declared. The model checker (checker::check_node) only enforces min_output <= N <= max_output.

So a node with output=['out','present_key'] is schema-valid, passes the checker, and then reaches the index-2 write. A passing checker is not a guarantee the index is in range.

2. The sink (why the write is not caught)

// onnxruntime/core/graph/graph.cc  -  InferenceContextImpl
const TypeProto* getInputType(size_t index) const override {
  return node_.InputDefs().at(index)->TypeAsProto();   // .at()  -> bounds-checked
}
TypeProto* getOutputType(size_t index) override {
  return &node_output_types_[index];                   // operator[]  -> NOT bounds-checked
}

• node_output_types_ is sized to node.OutputDefs().size() in the InferenceContextImpl ctor, so for a 2-output node it has 2 elements; getOutputType(2) returns one past the end.
• getInputType uses .at() (would throw on a bad index); getOutputType uses raw operator[] (no check) - the asymmetry is the root cause.
• The call runs at Graph::Resolve → InferAndVerifyTypeMatch → RunInferencing. The surrounding ORT_TRY/ORT_CATCH(const std::exception&) only catches thrown fail_shape_inference; a raw out-of-range operator[] does not throw, so the catch does not help.
• Because this is schema-level inference, it is EP-independent - no CUDA/ROCm copy.

3. Audit technique — always sweep siblings

Do not stop at the reported function. Grep every shape-inference guard and compare its threshold against the highest output index written before the next guard.

git grep -n 'getNumOutputs' -- \
  onnxruntime/core/graph/contrib_ops/*.cc \
  onnxruntime/core/graph/contrib_ops/shape_inference_functions.cc

For each if (ctx.getNumOutputs() > N) block, find the largest index passed to propagateElemTypeFromInputToOutput(ctx, _, index) / updateOutputShape(ctx, index, _) / getOutputType(index) inside it. Rule: the guard must require strictly more outputs than the highest index written (write index k ⇒ guard must ensure getNumOutputs() > k).

Correct exemplars already in the tree to copy:

Gotcha — conditional writes can hide a vacuous audit. A write may sit behind an inner condition (e.g. hasInputShape(past_key_index) before writing index 2). The site is still a bug, but you can only observe it when that inner condition is also satisfied. Keep this in mind both for the audit and for tests (§5).

4. Fix patterns

Point fix (required): raise the guard to cover the highest index written.

// before
if (ctx.getNumOutputs() > 1) { ... writes idx 2 ... }
// after
if (ctx.getNumOutputs() > 2) {  // both present_key (idx 1) AND present_value (idx 2)
  ...
}

Justify the threshold with the op's output semantics. For these attention ops the two trailing outputs - present_key (idx 1) and present_value (idx 2) for MultiHeadAttention, new_key_cache / new_value_cache for DecoderAttention (see the §1 table for each op's exact output names) - are a both-or-neither pair: there is no valid configuration that emits one without the other, so requiring all three outputs before populating indices 1 and 2 is behavior-preserving. (PagedAttention encodes the same invariant via its nested != 3 check.)

Defense-in-depth (recommended): bound the sink so a future author cannot reintroduce the class.

// onnxruntime/core/graph/graph.cc  -  InferenceContextImpl::getOutputType
TypeProto* getOutputType(size_t index) override {
  if (index >= node_output_types_.size()) {
    fail_type_inference("output index ", index, " is out of range; node has ",
                        node_output_types_.size(), " outputs");
  }
  return &node_output_types_[index];
}

This mirrors getInputType's .at() and the existing bounds checks in the sibling DataPropagationContextImpl. Placing it at the base layer transitively protects the NHWC and quantization wrapper contexts. After the point fix this branch is unreachable through a normal model (the guard already prevents the out-of-range index), so it is pure defense-in-depth. Its failure mode is build-dependent: with exceptions enabled, fail_type_inference raises InferenceError (a std::exception), caught by the existing ORT_CATCH(const std::exception&) around RunInferencing and surfaced as a clean load-time error; under ORT_NO_EXCEPTIONS it is not compiled out - ONNX's no-exceptions path prints the message to std::cerr and calls abort(), a deterministic fail-fast (consistent with getInputType's .at(), which likewise terminates under no-exceptions). Either way the result is a controlled failure rather than an out-of-range write.

5. Test recipe

Tests live in onnxruntime/test/contrib_ops/*.cc and are auto-globbed into the onnxruntime_provider_test target by cmake/onnxruntime_unittests.cmake (test/contrib_ops/*.cc pattern) - no cmake edit needed for a new file. See the ort-test skill for the executable taxonomy (onnxruntime_provider_test vs onnxruntime_test_all).

Rules that make the regression test actually guard the fix:

1. Drive through Model + Graph::Resolve, not ONNX's standalone TestShapeInference. Only the full resolve path constructs the real InferenceContextImpl and hits the getOutputType sink described in §2. A standalone ONNX shape-inference helper uses a different context and bypasses the sink, so it cannot reproduce the bug.
2. Negative tests must be NON-VACUOUS - they must actually enter the write branch on pre-fix source. If a write is gated by an inner condition (§3 gotcha), satisfy it: e.g. for MultiHeadAttention/DecoderMaskedMultiHeadAttention, supply a shaped past_key (and past_sequence_length / past_present_share_buffer as the op requires) so the index-2 block runs. A negative test that only supplies query skips the block and passes even on pre-fix source - regression-proof in name only.
3. Add positive (all-outputs) cases: a node with every output present must still infer the trailing output types - proves the tightened guard did not over-restrict.
4. Keep tests throw-free post-fix so they are valid under ORT_NO_EXCEPTIONS. Any case that is expected to fail_shape_inference (throws) must be excluded with #ifndef ORT_NO_EXCEPTIONS. The "2 outputs must not go out of range" case is throw-free after the point fix and is safe in all builds.

Verify the negative test is non-vacuous (sanitizer A/B) - the most reliable way to prove a negative test enters the previously-out-of-range branch: build the test at the pre-fix commit with AddressSanitizer and confirm it flags the out-of-range output access; then confirm it is clean after the fix.

# Functional run (any Debug build):
cmake --build build/Linux/Debug --target onnxruntime_provider_test -j"$(nproc)"
./build/Linux/Debug/onnxruntime_provider_test \
  --gtest_filter='AttentionOptionalOutputsShapeInferenceTest.*'

# A/B proof (isolated worktree at the pre-fix commit, CPU-only Debug + sanitizer):
git worktree add --detach ../ort-prefix-check <fix_commit>~1
# copy the new test file in, then:
python3 tools/ci_build/build.py --build_dir build/asan --config Debug --parallel \
  --skip_tests --enable_address_sanitizer --skip_submodule_sync \
  --cmake_generator Ninja --target onnxruntime_provider_test
# Pre-fix: the negative tests fail (the sanitizer flags the out-of-range output access).
# Post-fix (cherry-pick the guard fix): all tests pass, no sanitizer report.

6. Process / wording conventions

• Run lintrunner -a before pushing so the CLANGFORMAT / Python-format gate passes. See the ort-lint skill.
• Use correctness/robustness framing in code, comments, commit messages, and the PR body
  • describe the change as fixing an optional-output guard, not as a security fix. This matches repo convention (compare python-kwargs-setattr-security) and keeps the PR neutral.

7. Audit checklist (per-operator review)

When reviewing or hardening any operator implementation or its shape inference:

• Read the op's spec - ONNX standard op page, or for a contrib op its OpSchema registration (.Input/.Output/.Attr, and Optional/Variadic markers). A local ONNX checkout has the standard-op spec pages; contrib ops are defined only in ORT.
• Enumerate all inputs, attributes, and outputs, noting which are optional and the resulting min/max input and output counts.
• Validate every input/attribute before indexing into it, to avoid out-of-range reads (which can cascade into worse failures). Match each output-index write to a guard that guarantees the index is in range (§3 rule).
• Prefer ORT_RETURN_IF / ORT_RETURN_IF_NOT for validation; use ORT_ENFORCE in constructors. In shape inference use fail_shape_inference / fail_type_inference.
• Use SafeInt<> / narrow<>() for index and size arithmetic and casts to avoid overflow or truncation that yields a wrong index. See core/common/safeint.h and docs/Coding_Conventions_and_Standards.md.
• Ensure tests build and pass under no-exceptions builds; #ifndef ORT_NO_EXCEPTIONS around any case expected to throw.
• Exclude EPs known not to support the op, with a comment explaining why.
• Check whether other EPs (notably CUDA/ROCm) implement the same op and whether the same issue exists there. (For shape inference specifically, the logic is EP-agnostic and single-source - confirm there is no kernel-side analogue.)

References

• PR #28176 - "Fix ... in EmbedLayerNormalizationShapeInference": the precedent that fixed the identical > 1 → > 2 primitive in one site; the sibling attention sites were missed, motivating the sweep in §3.
• PR #29268 - this fix: guards corrected in DecoderAttention / MultiHeadAttention / DecoderMaskedMultiHeadAttention shape inference, plus the getOutputType bounds check and non-vacuous regression tests.
• Sibling skill: ort-test (test executables, --gtest_filter, contrib-op test layout); ort-lint (lintrunner -a); ort-build (build flags, ASan).