fix(security): guard ONNX/shape overflow and external data paths - #5168
Open
causten wants to merge 1 commit into
Open
fix(security): guard ONNX/shape overflow and external data paths#5168causten wants to merge 1 commit into
causten wants to merge 1 commit into
Conversation
Addresses ROCM-26638, ROCM-26616, ROCM-26615, ROCM-26617. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR hardens MIGraphX’s ONNX ingestion path by adding checked integer arithmetic to shape byte-size computations and tightening ONNX external tensor loading to prevent overflow and unsafe file path resolution.
Changes:
- Introduces
checked_mul/checked_addhelpers and uses them to guard shape element/byte calculations against integer overflow. - Updates ONNX tensor literal creation to use checked multiplication for element-count computation.
- Adds external tensor size parsing + path resolution logic and enforces raw/external data size matching.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| src/shape.cpp | Uses checked arithmetic for shape::bytes() and size_t element product computation. |
| src/onnx/onnx_parser.cpp | Adds external-data size/path validation and enforces size checks for raw/external tensor buffers. |
| src/include/migraphx/checked_ops.hpp | Adds checked integer add/mul helpers used by shape and ONNX parsing codepaths. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+810
to
+820
| static std::size_t parse_external_size(const std::string& value, const char* field) | ||
| { | ||
| try | ||
| { | ||
| return std::stoull(value); | ||
| } | ||
| catch(const std::exception&) | ||
| { | ||
| MIGRAPHX_THROW(std::string("Invalid ONNX external data ") + field + ": " + value); | ||
| } | ||
| } |
Comment on lines
+822
to
+835
| static fs::path resolve_external_data_path(const fs::path& base_dir, const fs::path& relative) | ||
| { | ||
| const fs::path base = fs::weakly_canonical(base_dir); | ||
| const fs::path resolved = fs::weakly_canonical(base_dir / relative); | ||
| const auto base_str = base.string(); | ||
| const auto resolved_str = resolved.string(); | ||
| if(resolved_str.size() < base_str.size() or | ||
| resolved_str.compare(0, base_str.size(), base_str) != 0 or | ||
| (resolved_str.size() > base_str.size() and resolved_str[base_str.size()] != '/')) | ||
| { | ||
| MIGRAPHX_THROW("ONNX external data path escapes model directory: " + relative.string()); | ||
| } | ||
| return resolved; | ||
| } |
Comment on lines
+34
to
+50
| template <class T, MIGRAPHX_REQUIRES(std::is_integral<T>{})> | ||
| constexpr T checked_mul(T a, T b) | ||
| { | ||
| T c{}; | ||
| if(__builtin_mul_overflow(a, b, &c)) | ||
| MIGRAPHX_THROW("Integer overflow in multiplication"); | ||
| return c; | ||
| } | ||
|
|
||
| template <class T, MIGRAPHX_REQUIRES(std::is_integral<T>{})> | ||
| constexpr T checked_add(T a, T b) | ||
| { | ||
| T c{}; | ||
| if(__builtin_add_overflow(a, b, &c)) | ||
| MIGRAPHX_THROW("Integer overflow in addition"); | ||
| return c; | ||
| } |
Comment on lines
822
to
866
| @@ -820,20 +850,19 @@ literal onnx_parser::parse_tensor(const onnx::TensorProto& t) const | |||
|
|
|||
| if(num_data_fields > 1) // if offset field is present | |||
| { | |||
| offset = std::stoull(t.external_data().at(1).value()); | |||
| offset = parse_external_size(t.external_data().at(1).value(), "offset"); | |||
| } | |||
| if(num_data_fields > 2) // if nbytes field is present | |||
| { | |||
| nbytes = std::stoull(t.external_data().at(2).value()); | |||
| nbytes = parse_external_size(t.external_data().at(2).value(), "length"); | |||
| } | |||
| std::vector<char> raw_buffer; | |||
| if(not external_data_path.empty()) | |||
| const fs::path base_dir = | |||
| external_data_path.empty() ? path : fs::path{external_data_path}; | |||
| const fs::path data_path = resolve_external_data_path(base_dir, data_file); | |||
| std::vector<char> raw_buffer = read_buffer(data_path, offset, nbytes); | |||
| if(raw_buffer.size() != tensor_shape.bytes()) | |||
| { | |||
| raw_buffer = read_buffer(fs::path{external_data_path} / data_file, offset, nbytes); | |||
| } | |||
| else | |||
| { | |||
| raw_buffer = read_buffer(path / data_file, offset, nbytes); | |||
| MIGRAPHX_THROW("ONNX external tensor data size mismatch"); | |||
| } | |||
Regressions detected 🔴 * No develop baseline was found for this PR's branch point; compared against the latest available develop run instead. |
|
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.
Summary
JIRA
Test plan
Made with Cursor