onnxruntime
b95cf781 - Reject non-control-flow nodes that carry subgraphs during session state finalization (#32641)

Commit
8 days ago
Reject non-control-flow nodes that carry subgraphs during session state finalization (#32641) ### 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: ```cpp // 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.
Author
Parents
Loading