Skip to content

Validate initializer data length and guard element-count computation in contrib Range shape inference#29265

Open
titaiwangms wants to merge 4 commits into
microsoft:mainfrom
titaiwangms:fix/contrib-range-input-validation
Open

Validate initializer data length and guard element-count computation in contrib Range shape inference#29265
titaiwangms wants to merge 4 commits into
microsoft:mainfrom
titaiwangms:fix/contrib-range-input-validation

Conversation

@titaiwangms

Copy link
Copy Markdown
Contributor

Summary

The com.microsoft Range operator's shape-inference helper read a fixed number of bytes from an initializer's raw_data without first checking the buffer length, and the element-count computation could pass non-finite or out-of-range values to an int64 cast. This change adds the missing validations and aligns the shape-inference and CPU kernel paths.

Changes

  • GetFirstElement now checks raw_data length is at least the element size before reading, and reads via std::memcpy into an aligned local.
  • CalcRangeDim and the CPU kernel ComputeRange now reject non-finite computed counts, handle non-positive counts before the int64 cast, and reject counts that are not representable as int64 (>= 2^63). Both paths use identical messages and semantics.
  • The output dimension for empty/backward ranges is clamped to 0 in shape inference to match the kernel.

Tests

Added contrib Range model-load regression tests in range_test.cc covering:

  • truncated raw_data for start, limit, and delta (double, plus float and int64 element types),
  • zero delta,
  • a finite-but-too-large element count,
  • an exact-size success boundary,
  • backward-range zero-dimension inference.

Tests assert on Status (safe for no-exception builds) and are guarded by #ifndef DISABLE_CONTRIB_OPS (throwing cases additionally by !defined(ORT_NO_EXCEPTIONS)). 21/21 RangeTest cases pass locally.

Follow-up

int16 inputs are currently supported only via raw_data (there is no get_data<int16_t> specialization for the non-raw-data path); this is left as a separate follow-up.

titaiwangms and others added 3 commits June 25, 2026 19:04
…e inference

GetFirstElement<T> read sizeof(T) bytes from a graph initializer's raw_data guarded only by the presence bit. Add a length check before the cast and fail shape inference cleanly when raw_data is shorter than the element size. Also clamp the computed element count to >= 0 in CalcRangeDim so shape inference matches the CPU kernel for empty/backward ranges. Add model-load regression tests in range_test.cc (guarded by #ifndef DISABLE_CONTRIB_OPS, Status-only assertions for no-exception builds).

Agent-signed-off: Developer (b2fe149b) [claude-opus-4.8 via copilot]
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- GetFirstElement<T>: read raw_data via std::memcpy into an aligned local after the
  length check, avoiding an unaligned access on strict-alignment targets.
- CalcRangeDim<T> and the CPU kernel ComputeRange: guard the element-count cast with
  std::isfinite so a non-finite computed count fails cleanly instead of an out-of-range
  cast (fail_shape_inference in the schema, a non-OK Status in the kernel).
- range_test.cc: generalize the model-building helper over element type, dimensions and
  per-input bytes. Truncated raw_data is now modeled as dims=[0] with empty raw_data so
  the cases are tight regressions for the length check; add coverage for start/limit/delta
  positions, float and int64 element types, a zero-delta failure, and an exact-size success
  boundary. Throwing tests are additionally guarded by !defined(ORT_NO_EXCEPTIONS).

Agent-signed-off: Developer (b2fe149b) [claude-opus-4.8 via copilot]
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…cast

The std::isfinite guard catches inf/NaN but not a finite count larger than the
int64 range, where static_cast<int64_t>(count) is out of range. Handle the
non-positive case before the cast (so a large-magnitude negative count cannot
overflow it) and reject counts at or above 2^63 with a neutral message. Applied
identically in the schema CalcRangeDim (fail_shape_inference) and the CPU kernel
ComputeRange (non-OK Status), keeping the two sites consistent.

Add a guarded regression test (start=0, limit=1e19, delta=1) asserting a clean
non-OK status; verified it fails without the new guard.

Agent-signed-off: Developer (b2fe149b) [claude-opus-4.8 via copilot]
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR hardens the com.microsoft Range operator’s shape inference and the CPU kernel against malformed initializer raw_data and invalid/overflowing element-count computations, ensuring model load and runtime execution fail cleanly instead of reading out-of-bounds or hitting undefined conversions.

Changes:

  • Added raw_data length validation + aligned reads (via memcpy) in contrib Range shape inference.
  • Added non-finite and out-of-int64_t-range guards for the computed element count in both shape inference and the CPU kernel; clamped empty/backward ranges to dimension 0.
  • Added contrib Range model-load regression tests covering truncated raw_data, zero delta, too-large counts, and backward-range inferred zero-dim.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File Description
onnxruntime/core/graph/contrib_ops/range_schema_defs.cc Adds raw_data length checks/aligned reads and validates/clamps computed output length during shape inference.
onnxruntime/core/providers/cpu/generator/range.cc Adds runtime validation for computed element counts (finite + representable in int64_t) and clamps empty/backward ranges.
onnxruntime/test/providers/cpu/generator/range_test.cc Adds model-load regression tests for contrib Range shape inference failure/success boundary cases.

Comment thread onnxruntime/core/providers/cpu/generator/range.cc Outdated
Comment thread onnxruntime/test/providers/cpu/generator/range_test.cc
…e kernel guards

Address automated review feedback on PR microsoft#29265:
- range.cc ComputeRange now computes the element count as
  ceil((double(limit) - double(start)) / double(delta)), byte-identical to the
  shape-inference path in CalcRangeDim, so integral inputs are promoted to double
  before the subtraction and cannot overflow in T.
- range_schema_defs.cc CalcRangeDim uses the same expression; comments on both
  sites note they must stay identical.
- Added two CPU-EP-pinned kernel-execution tests with runtime (non-constant)
  inputs so the ComputeRange element-count guards are exercised at execution
  time: a count that exceeds the int64 range and a non-finite count.

Agent-signed-off: Developer (b2fe149b) [claude-opus-4.8 via copilot]
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants