Reject non-control-flow nodes that carry subgraphs during session state finalization

[Bin Miao (miaobin)]

Description

SessionState::FinalizeSessionStateImpl iterates over the subgraphs attached to each node and unconditionally downcasts the node's kernel to controlflow::IControlFlowKernel, then calls SetupSubgraphExecutionInfo on it:

// Downcast is safe, since only control flow nodes have subgraphs
auto& control_flow_kernel = static_cast<controlflow::IControlFlowKernel&>(*p_op_kernel);

The "only control flow nodes have subgraphs" invariant is not enforced anywhere. Node::Init materializes a subgraph for any attribute of type GRAPH, with no schema gate, so an ordinary (non-control-flow) node can carry a subgraph. When it does, the downcast above is invalid: IControlFlowKernel adds a vtable slot that a plain OpKernel does not have, so the call reads an out-of-bounds vtable slot — an undefined-behavior type confusion that crashes (observed as an access violation inside FinalizeSessionStateImpl).

This is reachable by loading a crafted/malformed model whose non-control-flow node has a GRAPH-typed attribute (for example, an op that permits unchecked attributes). ORT should reject such a model with a clear error instead of executing the bad cast.

Fix

Gate the downcast on a virtual predicate:

• Add OpKernel::IsControlFlowKernel() returning false by default.
• Override it to true on IControlFlowKernel, which covers If/Loop/Scan and their CUDA derivatives.
• Check it in FinalizeSessionStateImpl before the cast and return an error otherwise.

A virtual predicate is used rather than dynamic_cast because onnxruntime_DISABLE_RTTI is on by default. There is no behavior change for valid models: only the control-flow kernels inherit IControlFlowKernel, and they always return true.

Adds InferenceSessionTests.SubgraphAttributeOnNonControlFlowNodeIsRejected test case, which builds a model whose non-control-flow node carries a GRAPH attribute and asserts that session initialization fails gracefully instead of triggering the downcast. Existing If/Loop/Scan subgraph tests cover the no-regression path.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

🟢 Approval recommended

The guard prevents the unsafe cast while preserving valid control-flow behavior and includes targeted regression coverage.

Pull request overview

Hardens session finalization against invalid control-flow kernel casts.

Changes:

• Adds a virtual control-flow kernel predicate.
• Rejects non-control-flow kernels carrying subgraphs.
• Adds a regression test for crafted models.

No actionable issues identified.

File summaries

File	Description
include/onnxruntime/core/framework/op_kernel.h	Adds the kernel predicate.
onnxruntime/core/providers/cpu/controlflow/utils.h	Marks control-flow kernels.
onnxruntime/core/framework/session_state.cc	Guards the downcast.
onnxruntime/test/framework/inference_session_test.cc	Tests graceful rejection.

Review details

• Files reviewed: 4/4 changed files
• Comments generated: 0
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟢 Approval recommended

The guard prevents the invalid cast while preserving valid control-flow and plugin EP paths.

Review details

• Files reviewed: 4/4 changed files
• Comments generated: 0 new
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.