Validate input rank/size for BifurcationDetector before indexing - #31701
Open
Ti-Tai Wang (titaiwangms) wants to merge 1 commit into
Open
Validate input rank/size for BifurcationDetector before indexing#31701Ti-Tai Wang (titaiwangms) wants to merge 1 commit into
Ti-Tai Wang (titaiwangms) wants to merge 1 commit into
Conversation
BifurcationDetector unconditionally reads and writes element [0] of several tensors without first checking they contain enough elements to do so safely: - src_tokens/cur_tokens/pred_tokens are treated as flat 1-D sequences via Shape().GetDims()[0] and linear DataRaw() indexing, but a 0-D (scalar) input has no dimensions, making GetDims()[0] an out-of-range access. - prev_suffix_match_idx (and the output tensor that mirrors its shape) is read and written at index [0] regardless of its element count; a 0-element input backs those accesses with an empty/null buffer. Add explicit rank and size checks for these inputs before they are used, returning a failure status instead. Adds regression tests for each of these cases, plus a case with no predicted tokens. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 77dcaf1b-748a-4379-94a7-478f7a924d73
Contributor
There was a problem hiding this comment.
Pull request overview
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Adds explicit rank/size validation to BifurcationDetector inputs to prevent out-of-bounds shape access and unsafe buffer indexing, and introduces regression tests that assert the operator fails fast with clear errors on invalid inputs.
Changes:
- Validate
src_tokens,cur_tokens, and optionalpred_tokensare rank-1 before usingGetDims()[0]/ linear indexing. - Validate
prev_suffix_match_idxcontains exactly one element (matching unconditional[0]read/write semantics). - Add regression tests for scalar token tensors and 0-element
prev_suffix_match_idx(with/withoutpred_tokens).
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| onnxruntime/contrib_ops/cpu/bert/bifurcation_detector.h | Adds early input validation for rank/size to prevent unsafe indexing. |
| onnxruntime/test/contrib_ops/bifurcation_detector_op_test.cc | Adds tests covering scalar token tensors and empty prev_suffix_match_idx failure cases. |
Comment on lines
+34
to
+42
| ORT_RETURN_IF_NOT(src_tokens->Shape().NumDimensions() == 1, "src_tokens must be a 1-D tensor"); | ||
| ORT_RETURN_IF_NOT(cur_tokens->Shape().NumDimensions() == 1, "cur_tokens must be a 1-D tensor"); | ||
| // prev_suffix_match_idx (and, by construction below, the suffix_match_idx output that mirrors its | ||
| // shape) holds a single index value; element [0] is read and/or written unconditionally. | ||
| ORT_RETURN_IF_NOT(prev_suffix_match_idx->Shape().Size() == 1, | ||
| "prev_suffix_match_idx must contain exactly one element"); | ||
| if (pred_tokens != nullptr) { | ||
| ORT_RETURN_IF_NOT(pred_tokens->Shape().NumDimensions() == 1, "pred_tokens must be a 1-D tensor"); | ||
| } |
Comment on lines
+381
to
+395
| TEST(BifurcationDetectorTest, ScalarSrcTokensRejected) { | ||
| OpTester tester("BifurcationDetector", 1, onnxruntime::kMSDomain); | ||
|
|
||
| tester.AddInput<int64_t>("src_tokens", {}, {1}); | ||
| tester.AddInput<int64_t>("cur_tokens", {1}, {2}); | ||
| tester.AddInput<int64_t>("prev_suffix_match_idx", {}, {0}); | ||
| tester.AddOutput<int64_t>("tokens", {1}, {0}); | ||
| tester.AddOutput<int64_t>("suffix_match_idx", {}, {0}); | ||
|
|
||
| std::vector<std::unique_ptr<IExecutionProvider>> execution_providers; | ||
| execution_providers.push_back(DefaultCpuExecutionProvider()); | ||
| tester.Run(OpTester::ExpectResult::kExpectFailure, | ||
| "src_tokens must be a 1-D tensor", | ||
| {}, nullptr, &execution_providers); | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
BifurcationDetectorunconditionally reads and/or writes element[0]of several tensors without first checking that they contain enough elements to do so safely:src_tokens,cur_tokens, andpred_tokensare treated as flat 1-D sequences, both viaShape().GetDims()[0](their reported "length") and via linearDataRaw()indexing. A 0-D (scalar) input has no dimensions, soGetDims()[0]is an out-of-range access on an empty span.prev_suffix_match_idxis read at index[0], and the output tensor that mirrors its shape is written at index[0], regardless of how many elements it actually has. A 0-element input backs both of those accesses with an empty/null buffer.This adds explicit validation for these inputs before they are used, returning a failure
Status(viaORT_RETURN_IF_NOT) rather than proceeding to index into a buffer that may not have enough elements:src_tokens,cur_tokens, andpred_tokens(when present) must be 1-D tensors.prev_suffix_match_idxmust contain exactly one element (consistent with how it, and the output that mirrors its shape, are used elsewhere in the kernel).Testing
Added regression tests to
onnxruntime/test/contrib_ops/bifurcation_detector_op_test.cccovering:prev_suffix_match_idx, both with and withoutpred_tokenspresent.src_tokens,cur_tokens, andpred_tokens, each rejected individually.All existing tests continue to pass unchanged.
BifurcationDetectoris CPU-only; no other execution provider implements this operator.