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 ```cpp // 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`. | Op | Output decls | min / max | A 2-output node? | |---|---|---|---| | `DecoderAttention` | out (req), new_key_cache (Opt), new_value_cache (Opt) | 1 / 3 | passes checker | | `MultiHeadAttention` | out (req), present_key (Opt), present_value (Opt), qk (Opt) | 1 / 4 | passes checker | | `DecoderMaskedMultiHeadAttention` | out (req) + 3 Optional | 1 / 4 | passes checker | 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) ```cpp // 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**. ```bash 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: | Exemplar | Pattern | Why it is safe | |---|---|---| | `BaseGroupQueryAttention...` | `if (getNumOutputs() >= 3)` then writes idx 2 | guard covers highest index | | `PagedAttention...` | nested `> 1` + inner `if (getNumOutputs() != 3) fail_shape_inference` | fails before any write | | `EmbedLayerNormalizationShapeInference` | `> 2` then writes idx 2 | fixed by PR #28176 (precedent) | | `SkipLayerNormalizationShapeInference` | each idx `k` guarded by `> k` | per-index guard | > **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.** ```cpp // 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. ```cpp // 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. ```bash # 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).