add more ORT_ENFORCE when building the tree structure of TreeEnsemble operators

[xadupre]

Description

Hardens TreeEnsemble initialization against malformed/unvalidated ONNX models by adding missing bounds checks and fixing existing ones.

Validation additions (tree_ensemble_attribute.h)

• Validate base_values / base_values_as_tensor size against n_targets_or_classes: must be 0 or N (or ≤ 2 for binary classifiers)

Validation additions/fixes (tree_ensemble_common.h)

• Pre-loop: add target_class_treeids.size() == limit check alongside existing target_class_ids, target_class_nodeids, and weights size checks
• nodes_featureids[i]: validate original int64_t value is in [0, INT_MAX] before narrowing cast to int — prevents large values wrapping to positive ints and bypassing the old post-cast >= 0 check; original attribute value included in error message
• Per-row ProcessTreeNodeLeave calls (sections C/D/E, C2/D2/E2): pass x_data + (i+1)*stride (per-row end) instead of x_data + N*C (global tensor end); removes x_data_end from per-row lambda captures
• Error message for target_class_ids range check uses target_class_ids (not the old hard-coded target_ids)

Tests

• Negative tests for out-of-range target_ids, negative nodes_featureids, wrong-sized base_values (regressor and classifier, binary/multi-class)

Motivation and Context

Tree-building code assumed ONNX models had been pre-validated. These changes prevent out-of-bounds memory access and wrap-around bugs when loading unvalidated or adversarially crafted models.

[Copilot]

Pull request overview

This PR hardens TreeEnsemble initialization by adding ORT_ENFORCE bounds checks around target/weight indexing to prevent out-of-bounds access when models have not been pre-validated, and adjusts a couple of ML operator tests accordingly.

Changes:

• Added additional ORT_ENFORCE checks in TreeEnsembleCommon::Init to validate indices and target id ranges before dereferencing.
• Updated TreeRegressor/TreeEnsembler CPU tests to use in-range target_ids / leaf_targetids.

Reviewed changes

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

File	Description
onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h	Adds extra index/range enforcement when mapping leaf weights to targets/classes.
onnxruntime/test/providers/cpu/ml/treeregressor_test.cc	Adjusts test data to keep target_ids within the allowed target range.
onnxruntime/test/providers/cpu/ml/tree_ensembler_test.cc	Adjusts test data to keep leaf_targetids within the allowed target range.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

[Copilot]

@xadupre I've opened a new pull request, #27678, to work on those changes. Once the pull request is ready, I'll request review from you.

[Copilot]

@xadupre I've opened a new pull request, #27679, to work on those changes. Once the pull request is ready, I'll request review from you.

[github-actions]

You can commit the suggested changes from lintrunner.

[yuslepukhin]

There are suggestions to consider and CI to fix.

[xadupre]

There are suggestions to consider and CI to fix.

done

[yuslepukhin]

Please, resolve the ICM ticket accordingly.

revoking review

[yuslepukhin]

Please, resolve comments for tests

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

This PR hardens CPU ML TreeEnsemble tree-building against malformed ONNX model attributes by adding additional runtime validations (primarily via ORT_ENFORCE) and adjusting existing tests to comply with the stricter checks.

Changes:

• Add bounds/range validation for target_ids/weights mapping when building leaf weight structures.
• Enforce non-negative nodes_featureids for non-leaf nodes (while normalizing leaf feature ids).
• Update/add unit tests to reflect and validate the new failure behavior.

Reviewed changes

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

File	Description
onnxruntime/test/providers/cpu/ml/treeregressor_test.cc	Fixes an existing test’s invalid target_ids for n_targets=1 and adds new negative tests covering out-of-range target_ids and negative featureids.
onnxruntime/test/providers/cpu/ml/tree_ensembler_test.cc	Fixes leaf_targetids in n_targets=1 coverage to conform to new validations.
onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h	Adds new ORT_ENFORCE validations during tree construction (target id range, feature id validity) and moves n_targets_or_classes_ initialization earlier.
onnxruntime/core/providers/cpu/ml/tree_ensemble_attribute.h	Adds validation of base_values / base_values_as_tensor sizes vs n_targets_or_classes for regressor/classifier attribute parsing.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[xadupre]

@copilot open a new pull request to apply changes based on the comments in this thread

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[titaiwangms]

Multi-reviewer synthesis (readability + code + critical + deep)

The hardening direction is right and several individual checks are well-placed. The cross-reviewer team surfaced one backward-compat regression, two structurally-incomplete coverage gaps (other narrowing casts; V5 path bypass), and one placement issue where new checks land after the function that can OOB on the data they validate. No Critical issues — but the Major findings together are substantive.

Major

1. base_values carve-out misses the n_targets_or_classes == 2 && size == 1 case → backward-compat regression.
tree_ensemble_attribute.h only relaxes the size check when n_targets_or_classes == 1. For n_targets_or_classes == 2 it falls into the strict size == 0 || size == N branch. But the existing aggregator at tree_ensemble_aggregator.h:593–619 explicitly accepts base_values.size() == 1 in the binary branch (n_targets_or_classes_ <= 2), with the comment: "ONNX is vague about two classes and only one base_values."

Models with classlabels_int64s = [0, 1] (so n_targets_or_classes == 2) and base_values of length 1 — historically accepted and computed sensibly — will now hard-fail at session Init.

Fix: change the outer guard from if (n_targets_or_classes == 1) to if (n_targets_or_classes <= 2) (and update messages to "≤ 2 for binary classifiers"). Add a regression test for n_classes == 2, base_values.size() == 1 to pin the behavior.

2. Other narrowing casts in the same init path are still unguarded.
The PR adds a pre-cast [0, INT_MAX] check on nodes_featureids[i], but the same pattern still applies to:

• nodes_treeids[i] → static_cast<int> at tree_ensemble_common.h:225
• nodes_nodeids[i] → static_cast<int> at :225
• nodes_truenodeids[i] → static_cast<int> at :241
• nodes_falsenodeids[i] → static_cast<int> at :253

For true/false node ids, the bounds check happens after the narrowing, so INT64_MAX-style values can wrap into valid-looking small ids before being checked. TreeNodeElementId actually stores int64_t, so several of these casts may be unnecessary altogether; otherwise, apply the same pre-cast [0, INT_MAX] guard. Also: missing test for nodes_featureids > INT_MAX (the wrap-to-positive case the PR's narrowing fix actually targets).

3. Pre-loop target_class_* size checks are placed AFTER AddNodes, which can OOB on the data being validated.
The new size checks (target_class_treeids/ids/weights[_as_tensor].size() == limit) land just before the second target_class_nodeids loop at tree_ensemble_common.h:295–299, but AddNodes runs at line 287 — and AddNodes recursively calls CheckIfSubtreesAreEqual, which dereferences target_class_weights[left_target_node] / target_class_weights_as_tensor[left_target_node] at :402. A malformed model where class_weights* is shorter than target_class_nodeids can therefore still OOB during categorical-folding before the new check fires.

Also, the first parallel-array consumer is the loop at line 271 (indices.emplace_back({target_class_treeids[i], target_class_nodeids[i]}, i)), which dereferences target_class_treeids[i] before the new size check on target_class_treeids.

Fix: move all four target_class_* size checks to before indices.reserve(...) at line 270, so they protect both the indices loop and AddNodes.

4. V5 → V3 conversion path bypasses the V3 attribute ctor.
TreeEnsembleAttributesV5::convert_to_v3 populates V3 attributes directly without invoking the V3 constructor, so:

• The new base_values size check (added in the V3 ctor) does not run for ai.onnx.ml v5 ops.
• The pre-existing min_ids/max_ids range check on target_class_ids does not run on the V5 path either.
• V5 conversion itself indexes parallel arrays (leaf_targetids, leaf_weights, nodes_featureids, nodes_missing_value_tracks_true, nodes_splits, nodes_falseleafs, nodes_trueleafs, tree_roots) without size validation.

The pre-loop target_class_* size checks added by this PR do cover V5 (they're in Init) — good. But base_values and target_class_ids value-range coverage on V5 remains a gap.

Fix: factor the V3 ctor's validations into a reusable method and run it on converted V5 attributes before common init.

Minor

5. Wrong comment + typo: "// AddNodes already maded sure that w.i >= 0 && w.i < n_targets_or_classes_." Both:

   • maded → made.
   • Mis-attributed: AddNodes does NOT check w.i. The check lives in TreeEnsembleAttributesV3::check() (and is bypassed on V5). Suggested replacement: "Validated in TreeEnsembleAttributesV3 ctor for V3 ops; not enforced on the V5→V3 path."

6. base_values error messages omit the actual offending size. All four new ORT_ENFORCE in tree_ensemble_attribute.h state the expected count but not the actual one — unlike the well-formed pre-existing message just below ("… (", min_ids, ")"). Compare:

   // Current
   ORT_ENFORCE(base_values.size() <= 2, "base_values should have 0, 1, or 2 values.");
   // Better
   ORT_ENFORCE(base_values.size() <= 2,
               "base_values must have 0, 1, or 2 elements for a binary classifier, got ",
               base_values.size(), ".");

7. TREE_FIND_VALUE macro is duplicated wholesale for debug vs. release — only the ORT_ENFORCE lines differ. Any future change to traversal logic must edit both copies. Move the guard inside a single definition (e.g. #ifndef NDEBUG inline, or an ORT_ENFORCE_DEBUG_ONLY helper).

8. ProcessTreeNodeLeave parameter conditionally named with #ifndef NDEBUG / x_data_end / #endif — legal C++ but jarring. Use [[maybe_unused]] to silence the release-only warning while keeping the name.

9. Per-row x_data_end guard is debug-only; release builds still rely entirely on init-time validation. Also, comparing x_data + feature_id < x_data_end forms an out-of-range pointer before comparing — safer to compare integer offsets (feature_id < row_feature_count).

10. Error-message string in TREE_FIND_VALUE macro reads as a comparison: "root->feature_id=", id, ">", x_data_end - x_data, " outside of boundaries" produces text like "feature_id=5>3 outside of boundaries" — the > character looks like a relational operator, not a separator. Same pattern repeats across multiple sites.

11. Negative tests duplicate ~30 lines of minimal-tree boilerplate 5 times. A small helper (MakeMinimalTreeClassifier(n_classes, featureids, classlabels)) would reduce each test to its delta and document the canonical minimal-tree shape (e.g., the featureids = {0, -2, -2} magic — -2 is never explained).

12. nodes_missing_value_tracks_true not fully validated. A shorter array silently disables tracking for later nodes; values other than 0/1 aren't rejected (has_missing_tracks_ treats nonzero as enabled, but AddNodes only sets the flag when value is exactly 1, leading to inconsistent semantics). Worth a follow-up.

13. TreeEnsembleRegressorTargetIdsOutsideBoundary test exercises a pre-existing check, not a check added by this PR. Worth either renaming for clarity or noting in the PR description.

Nits

• // if debug build comment after #ifndef NDEBUG is redundant noise — appears 5+ times.
• Commented-out dead code in TreeEnsembleNegativeFeatureIds: // std::vector<int64_t> classes = {0, 1, 2, 3}; — looks like leftover.
• Magic sentinel featureids = {0, -2, -2} for leaves never explained — single comment on first occurrence would help.
• Test naming inconsistency: pre-existing TreeRegressorOutsideBoundaryTargetIds (attribute-first) vs new TreeEnsembleRegressorTargetIdsOutsideBoundary (condition-first). Pick one.
• The reorder of n_targets_or_classes_ = attributes.n_targets_or_classes; (moved before base_values_ population) is unjustified by the PR — either revert the churn or add the check that motivates it.

Verifications (these hold up)

• Per-row end pointer derivation for the multi-row sections is correct: x_data + (i+1)*stride is the exclusive end of row i. Single-row sections use x_data + N*C which equals stride when N == 1. The dereferenced expression is x_data[root->feature_id], so the strict < (exclusive) is the correct operator.
• [0, INT_MAX] bound on nodes_featureids is the precise upper bound to prevent ill-defined int64_t → int narrowing. The actual feature-count check (max_feature_id_ >= C) at :564 runs at compute-time and catches the tighter constraint, so INT_MAX is "safe but loose" by intent.
• Leaf-node bypass + feature_id = -1 is correct: leaves don't index x_data, and the is_not_leaf() loop predicate prevents leaf feature_id from being read.
• Pre-loop parallel-array size checks: the schema (cmake/external/onnx/onnx/defs/traditionalml/defs.cc:836–881 for classifier, :947–992 for regressor) groups class_treeids/class_nodeids/class_ids/class_weights as a parallel-array contract that must all be of equal length. The PR's new checks correctly cover all four (with the _as_tensor xor branch).
• Refactor of while (1) { case LEAF: return root; } to while (root->is_not_leaf()) { ... } is behaviorally equivalent including the "root is already a leaf at entry" case.
• All six negative tests trigger their intended check (deep-reviewer walked each one).

Praise

• Pre-cast nodes_featureids validation is the correct pattern and message correctly identifies the offending value.
• Lifting parallel-array size checks into Init covers the V5 path that previously had no validation at all (modulo finding update HighLevelDesign.md #3 about placement).
• The while (root->is_not_leaf()) refactor is a clarity improvement that eliminates the implicit LEAF case inside the switch.
• Tests use kExpectFailure with substring matching — proper regression discipline.

Recommendation

Request changes on finding #1 (the binary-classifier regression — straightforward fix) and finding #3 (move size checks before line 270 so they're actually defensive). Findings #2 and #4 are larger but high-value — at minimum, file follow-up issues if they won't fit this PR. Minors and nits are nice-to-have.

---

Reviewed by a cross-family reviewer team (Claude readability + GPT code & critical + Claude-high deep). Schema citations grounded in cmake/external/onnx/onnx/defs/traditionalml/defs.cc. Pre-loop check placement and the V5 bypass independently verified by the synthesizer.

[tianleiwu]

Review Summary

Solid security-hardening PR. The validation additions are correct and close real vulnerabilities:

• base_values / base_values_as_tensor size validation prevents OOB when model specifies wrong count
• target_class_treeids size check completes the consistency validation set
• nodes_featureids range check before narrowing cast closes a wrap-around vulnerability
• Per-row bounds tightening (x_data + (i+1)*stride instead of global end) is genuine defense-in-depth
• Debug-only bounds checks in ProcessTreeNodeLeave are reasonable given the existing release-mode validations (feature_id ∈ [0, INT_MAX] at init, max_feature_id_ < C at inference)
• while(1) → while(root->is_not_leaf()) in the mixed-mode branch is cleaner
• Test coverage is adequate: negative tests for binary/multi-class base_values overflow, negative feature IDs, and out-of-range target IDs

Two minor suggestions below; neither blocks merge.

[github-actions]

You can commit the suggested changes from lintrunner.

[tianleiwu]

Review Summary

Well-scoped security-hardening PR. Both concerns from my prior review round (macro duplication and comment typo) have been addressed. The validation additions are correct and close real OOB memory-access vulnerabilities.

One remaining inconsistency:

Missing TREE_FIND_VALUE_CHECK in BRANCH_LEQ same_mode else branch — In ProcessTreeNodeLeave, TREE_FIND_VALUE_CHECK() was added to the BRANCH_LEQ has_missing_tracks_ true branch but not the else branch (no missing tracks). Every other branch type has the check in both code paths: TREE_FIND_VALUE macro (LT/GTE/GT/EQ/NEQ), BRANCH_MEMBER (both branches), BRANCH_MEMBER_BIGSET (both branches), and the mixed-mode loop. BRANCH_LEQ without missing tracks is the most common path for gradient-boosted tree ensembles.

Suggested fix (around line 831 in the new file):

} else {
  while (root->is_not_leaf()) {
    TREE_FIND_VALUE_CHECK();  // <-- add this
    val = x_data[root->feature_id];
    root = val <= root->value_or_unique_weight ? root->truenode_or_weight.ptr : root + 1;
  }
}

This is low-severity since the init-time checks (feature_id ∈ [0, INT_MAX] + max_feature_id_ < C) provide the release-build guard, but the debug check should be consistent across all branch modes.

[xadupre]

root = val <= root->value_or_unique_weight ? root->truenode_or_weight.ptr : root + 1;

Added.

[tianleiwu]

Review Summary

All concerns from prior review rounds have been addressed:

1. Macro duplication (round 1): Resolved — TREE_FIND_VALUE_CHECK() is now a separate conditional macro, with a single TREE_FIND_VALUE definition.
2. Comment typo (round 1): Fixed — "already made sure" with correct attribution.
3. Missing check in BRANCH_LEQ else branch (round 2): Fixed — TREE_FIND_VALUE_CHECK() is now consistently applied to all traversal paths.

The validation additions are correct and comprehensive:

• base_values size validation prevents OOB when model attributes are malformed.
• target_class_treeids size check completes the consistency validation set.
• nodes_featureids range check validates int64_t before narrowing to int, closing the wrap-around vulnerability.
• Per-row end pointer (x_data + (i+1)*stride) tightens debug bounds checking.
• while(1) → while(root->is_not_leaf()) avoids dereferencing x_data[-1] for leaf nodes.

Test coverage is good: negative tests for binary/multi-class base_values, negative feature IDs, out-of-range target IDs, and a backward-compat test for binary classifiers with 1 base value.

Ready to merge.