onnxruntime

55975681 - Pre-check ConstantOfShape output size against input initializer before constant folding (#28751)

Pre-check ConstantOfShape output size against input initializer before constant folding (#28751) A 152-byte ONNX model with a `ConstantOfShape` whose shape initializer encodes huge dims causes `InferenceSession::Initialize()` to materialize the full output tensor (287 MB in the PoC, up to PB-scale with fuzz-mutated dims) via the ConstantFolding optimizer. The existing pre-execution size cap in `EstimateNodeOutputSizeInBytes` relies on shape inference having populated `output_def->Shape()`, which is not guaranteed. ### Description - **`onnxruntime/core/optimizer/constant_folding.cc`** - New `EstimateConstantOfShapeOutputSizeInBytes(node, graph)`: looks up the shape input via `Graph::GetConstantInitializer` (it is constant by the time we reach this node), multiplies its int64 values with `SafeInt<int64_t>` (rejects negative dims, lets overflow propagate as an exception caught upstream), and multiplies by the element size derived from the `value` attribute's tensor type (defaulting to float per ONNX spec). - `EstimateNodeOutputSizeInBytes` now takes `const Graph&` and dispatches to the new estimator for `ConstantOfShape`, falling back to the generic shape-based path if the initializer can't be resolved. - **`onnxruntime/test/optimizer/graph_transform_test.cc`** - `ConstantFoldingConstantOfShapeUsesInputInitializerForSizeCheck`: shape `[100M]` with int64 `value` ⇒ 800 MB derived size; with `kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes=256MB` the node must remain unfolded, proving the size check fires from the initializer alone. - `ConstantFoldingConstantOfShapeBlockedWhenOutputShapeMissing`: same model, but the `pre_graph_checker` (which runs after `Graph::Resolve()` and before the transformer) calls `ClearShape()` on the ConstantOfShape output NodeArg to simulate the documented attack where shape inference has not propagated the output shape. With the inferred shape stripped, the generic shape-based estimator returns -1, so only the new `EstimateConstantOfShapeOutputSizeInBytes` path can derive the 800 MB size and block folding — isolating regressions in the new estimator from the pre-existing shape-inference path. ### Motivation and Context The byte cap added in #28055 only triggers when shape inference has propagated the output shape; a crafted model can bypass it and force unbounded allocation during `Initialize()`. Deriving the size from the (necessarily constant) shape input makes the cap effective for the documented attack vector and tightens the same code path used by the configurable `kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes` setting — no new knob, no behavior change for legitimate models within the existing 1 GB default. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Xavier Dupré <xadupre@microsoft.com>