SIGN IN SIGN UP

fix(security): add SafeInt overflow protection in Expand and constant folding output size limit (#28055)

### Description

Harden the constant folding optimizer and the `Expand` CPU kernel
against integer overflow attacks from crafted ONNX models.

**Problem:** The `Expand::Compute()` kernel performs cumulative
dimension multiplications (`input_count *= input_dim`, `output_count *=
output_dim`) using raw `int64_t` arithmetic. When triggered during
constant folding at `CreateSession()` time via a crafted model with
extreme shape values, signed integer overflow can produce corrupted
values used for buffer offset calculations and `memcpy` lengths,
creating a potential out-of-bounds write. The downstream SafeInt check
in the allocator catches overflow only when the total byte count wraps,
but carefully chosen dimensions can make the overflowed value appear
valid.

Additionally, the constant folding optimizer has no output size budget —
any deterministic node with constant inputs is eligible for constant
folding regardless of output size, enabling memory exhaustion attacks at
model load time.

### Key Changes

**1. SafeInt-protected arithmetic in `expand.cc`**

Wraps all dimension accumulation and offset/length calculations with
`SafeInt<int64_t>` or `SafeInt<size_t>` to catch overflow before it can
corrupt buffer arithmetic:

| Location | Before | After |
|---|---|---|
| Accumulator loop (L97-98) | `input_count *= input_dim` |
`SafeInt<int64_t>(input_count) * input_dim` |
| Accumulator loop (L109) | `last_dim_size *= expand_dim_size[...]` |
`SafeInt<int64_t>(last_dim_size) * ...` |
| copy_byte (L116) | `copy_len * sizeof(T)` | `SafeInt<size_t>(copy_len)
* sizeof(T)` |
| input_offset (L122) | `i * copy_len` | `SafeInt<int64_t>(i) *
copy_len` |
| output_offset (L126) | `output_offset += current_count * ...` |
`SafeInt<int64_t>(output_offset) + SafeInt<int64_t>(current_count) *
...` |

**2. Constant folding output size limit in `constant_folding.cc`**

- **Pre-execution check**: `EstimateNodeOutputSizeInBytes()` uses shape
inference results with SafeInt-protected arithmetic to estimate total
output bytes. Nodes exceeding the limit are skipped.
- **Post-execution check**: After `kernel->Compute()`, actual output
`SizeInBytes()` is verified against the limit (catches cases where shape
inference couldn't determine output size).
- **Exception isolation**: `kernel->Compute()` is wrapped in `try/catch`
so that SafeInt overflow exceptions from individual nodes skip the node
rather than aborting the entire optimization pass.
- **Configurable limit**: New session option
`optimization.constant_folding_max_output_size_in_bytes` (default: 1 GB,
`"0"` to disable).

**3. Session option**

New key `kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes` in
`onnxruntime_session_options_config_keys.h`.

### Motivation and Context

This addresses a security vulnerability where a malicious ONNX model can
cause signed integer overflow in the Expand kernel during constant
folding at model load time (`CreateSession()`), potentially leading to
out-of-bounds memory writes. The constant folding size limit provides
defense-in-depth against memory exhaustion attacks from untrusted
models.

### Testing

- `ConstantFoldingOutputSizeLimit` — Verifies 4 MB Expand is blocked at
1 MB limit, allowed at 8 MB limit.
- `ConstantFoldingDefaultLimitBlocksLargeExpand` — Verifies 1 GB
ConstantOfShape is blocked at 512 MB limit.
- `ConstantFoldingSmallOutputAllowed` — Verifies small Expand (64 bytes)
is still folded normally.
- `ConstantFoldingExpandOverflowDimsSkipped` — Verifies Expand with
`[2^32, 2^32]` dimensions (int64 overflow) is gracefully skipped during
constant folding.
T
Tianlei Wu committed
04f744032cf21e182b2d03391d64fa2a970a3aea
Parent: 2979bab
Committed by GitHub <noreply@github.com> on 5/14/2026, 9:08:25 PM