From b0f4660bc6314c844cb0f5ae7cb98e616756c94b Mon Sep 17 00:00:00 2001 From: m00dy Date: Mon, 26 Oct 2026 18:53:43 -0500 Subject: [PATCH] Validate tensor slicing, sparse CSR shapes, and opset versions ## Summary Reject invalid tensor-slicer indices, overflow-prone sparse CSR shape metadata, and ONNX opset versions that cannot be represented by the runtime. Also propagate model-loading failures instead of continuing after a failed status. ## Changes - Validate the slice dimension and dimension-zero offset before creating tensor views. - Express the CSR row invariant without overflowing `rows + 1`. - Validate imported opset versions before narrowing them to `int`. - Return the model-loading status immediately on failure. - Preserve transformer coverage by removing a duplicate zero-version import from one legacy fixture in memory before that test loads the model. - Add valid and invalid slicer tests, a sparse CSR overflow-boundary test, and invalid opset version tests. ## Testing - Built the full onnxruntime_test_all target in Release mode on Linux. - Test-only baseline: 2,504 passed, 2 added regression tests failed, and 26 existing environment-dependent tests skipped. - Patched full suite: 2,506 passed and the same 26 existing tests skipped. - Pre-run and post-run hashes match within each execution, and the red and green executable hashes differ. --- .../core/framework/ort_value_tensor_slicer.cc | 7 +++- onnxruntime/core/framework/sparse_utils.cc | 2 +- onnxruntime/core/graph/model.cc | 3 ++ .../framework/insert_cast_transformer_test.cc | 14 ++++++- .../test/framework/sparse_kernels_test.cc | 24 ++++++++++++ onnxruntime/test/framework/tensor_test.cc | 38 +++++++++++++++++++ onnxruntime/test/ir/onnx_model_test.cc | 21 ++++++++++ 7 files changed, 105 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/framework/ort_value_tensor_slicer.cc b/onnxruntime/core/framework/ort_value_tensor_slicer.cc index 2acc1e301d6e0..a909338e6786c 100644 --- a/onnxruntime/core/framework/ort_value_tensor_slicer.cc +++ b/onnxruntime/core/framework/ort_value_tensor_slicer.cc @@ -21,11 +21,14 @@ OrtValueTensorSlicer OrtValueTensorSlicer::Create(T& ort_value, int64_t sl ORT_ENFORCE(!prim_type->HasSubElems(), "Can't slice a tensor with a subbyte element type"); } auto& tensor_shape = tensor.Shape(); - ORT_ENFORCE(gsl::narrow_cast(tensor_shape.NumDimensions()) >= slice_dimension, + ORT_ENFORCE(slice_dimension >= 0 && + gsl::narrow_cast(tensor_shape.NumDimensions()) > slice_dimension, "Insufficient dimensions to slice on ", slice_dimension, ". Shape:", tensor_shape); auto dim0_size = tensor_shape[0]; - ORT_ENFORCE(dim0_offset < dim0_size, "Invalid dim0_offset of ", dim0_offset, ". Dimension 0 is ", dim0_size); + ORT_ENFORCE(dim0_offset >= 0 && dim0_offset < dim0_size && + (slice_dimension != 0 || dim0_offset == 0), + "Invalid dim0_offset of ", dim0_offset, ". Dimension 0 is ", dim0_size); return OrtValueTensorSlicer{ort_value, slice_dimension, dim0_offset}; }; diff --git a/onnxruntime/core/framework/sparse_utils.cc b/onnxruntime/core/framework/sparse_utils.cc index 48b28db9f9028..5de30ef20a3ba 100644 --- a/onnxruntime/core/framework/sparse_utils.cc +++ b/onnxruntime/core/framework/sparse_utils.cc @@ -213,7 +213,7 @@ Status SparseCsrToDenseTensor(const DataTransferManager& data_manager, const Spa const auto inner_num = csr_view.Inner().Shape().Size(); const auto outer_num = csr_view.Outer().Shape().Size(); ORT_ENFORCE(inner_num == src.Values().Shape().Size(), "Expecting inner indices to be same as nnz. Got: ", inner_num); - ORT_ENFORCE(outer_num == (rows + 1), "Outer indices must be M + 1. Got: ", outer_num); + ORT_ENFORCE(outer_num > 0 && rows == outer_num - 1, "Outer indices must be M + 1. Got: ", outer_num); } CopyElementFunc copy_func; diff --git a/onnxruntime/core/graph/model.cc b/onnxruntime/core/graph/model.cc index 4b15bac58d239..732aa223ff8ca 100644 --- a/onnxruntime/core/graph/model.cc +++ b/onnxruntime/core/graph/model.cc @@ -282,6 +282,8 @@ Model::Model(ModelProto&& model_proto, const PathString& model_path, std::unordered_map domain_to_version; for (auto& opSet : model_proto_.opset_import()) { const auto& domain = opSet.domain(); + ORT_ENFORCE(opSet.version() > 0 && opSet.version() <= std::numeric_limits::max(), + "Invalid opset version for domain ", domain); const auto version = gsl::narrow_cast(opSet.version()); // empty domain and 'ai.onnx' are equivalent if ((domain.empty() || domain == kOnnxDomainAlias) && version < 7) { @@ -662,6 +664,7 @@ static Status LoadModelHelper(const T& file_path, Loader loader) { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "system error number ", status.Code()); } } + return status; } ORT_TRY { diff --git a/onnxruntime/test/framework/insert_cast_transformer_test.cc b/onnxruntime/test/framework/insert_cast_transformer_test.cc index 21bd7ba8ff622..588a43880e489 100644 --- a/onnxruntime/test/framework/insert_cast_transformer_test.cc +++ b/onnxruntime/test/framework/insert_cast_transformer_test.cc @@ -294,8 +294,20 @@ TEST(TransformerTest, RandomNormalLikeWithFloat16Inputs) { // as is (int32) TEST(TransformerTest, MultinomialWithFloat16Input) { auto model_uri = MODEL_FOLDER ORT_TSTR("multinomial_float16.onnx"); + ModelProto model_proto; + auto status = Model::Load(model_uri, model_proto); + ASSERT_TRUE(status.IsOK()) << status; + + // This legacy fixture has a duplicate, empty-domain opset import whose + // version is zero. Remove it before validating and resolving the model. + ASSERT_EQ(model_proto.opset_import_size(), 2); + ASSERT_TRUE(model_proto.opset_import(0).domain().empty()); + ASSERT_EQ(model_proto.opset_import(0).version(), 0); + model_proto.mutable_opset_import()->DeleteSubrange(0, 1); + std::shared_ptr model; - auto status = Model::Load(model_uri, model, nullptr, DefaultLoggingManager().DefaultLogger()); + status = Model::Load(std::move(model_proto), model_uri, model, nullptr, + DefaultLoggingManager().DefaultLogger()); ASSERT_TRUE(status.IsOK()) << status; Graph& graph = model->MainGraph(); diff --git a/onnxruntime/test/framework/sparse_kernels_test.cc b/onnxruntime/test/framework/sparse_kernels_test.cc index b37ef94c0315c..1d3947057b291 100644 --- a/onnxruntime/test/framework/sparse_kernels_test.cc +++ b/onnxruntime/test/framework/sparse_kernels_test.cc @@ -1554,6 +1554,30 @@ TEST(SparseTensorConversionTests, CsrConversion) { #endif } +TEST(SparseTensorConversionTests, SparseCsrToDenseRejectsMaxRowsWithoutOverflow) { + auto* cpu_provider = TestCPUExecutionProvider(); + auto cpu_allocator = cpu_provider->CreatePreferredAllocators()[0]; + + DataTransferManager dtm; + ASSERT_STATUS_OK(dtm.RegisterDataTransfer(cpu_provider->GetDataTransfer())); + + SparseTensor src(DataTypeImpl::GetType(), TensorShape{1, 1}, + cpu_allocator); + const std::vector values = {1}; + const std::vector inner = {0}; + const std::vector outer = {0, 1}; + ASSERT_STATUS_OK(src.MakeCsrData( + *cpu_provider->GetDataTransfer(), cpu_allocator->Info(), values.size(), + values.data(), gsl::make_span(inner), gsl::make_span(outer))); + + const_cast(src.DenseShape()) = + TensorShape{std::numeric_limits::max(), 0}; + Tensor dense_dst; + EXPECT_THROW(sparse_utils::SparseCsrToDenseTensor( + dtm, src, cpu_allocator, cpu_allocator, dense_dst), + OnnxRuntimeException); +} + TEST(SparseTensorConversionTests, CooConversion) { auto* cpu_provider = TestCPUExecutionProvider(); auto cpu_allocator = cpu_provider->CreatePreferredAllocators()[0]; diff --git a/onnxruntime/test/framework/tensor_test.cc b/onnxruntime/test/framework/tensor_test.cc index 2125de0a36e4f..eb20bcf6781b7 100644 --- a/onnxruntime/test/framework/tensor_test.cc +++ b/onnxruntime/test/framework/tensor_test.cc @@ -3,6 +3,7 @@ #include "core/framework/tensor.h" #include "core/framework/allocator_utils.h" +#include "core/framework/ort_value_tensor_slicer.h" #include "test/unittest_util/framework_test_utils.h" #include "gmock/gmock.h" @@ -144,6 +145,43 @@ TEST(TensorTest, EmptyTensorTest) { EXPECT_EQ(location.alloc_type, expected_allocator_type); } +// Positive control: valid slicer parameters create the expected tensor views. +TEST(TensorTest, OrtValueTensorSlicerValidParameters) { + auto allocator = TestCPUExecutionProvider()->CreatePreferredAllocators()[0]; + OrtValue value; + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape({2, 3}), + allocator, value); + + EXPECT_NO_THROW({ + auto slicer = OrtValueTensorSlicer::Create( + value, /*slice_dimension=*/1, /*dim0_offset=*/1); + EXPECT_NE(slicer.begin(), slicer.end()); + }); +} + +TEST(TensorTest, OrtValueTensorSlicerRejectsInvalidParameters) { + auto allocator = TestCPUExecutionProvider()->CreatePreferredAllocators()[0]; + OrtValue value; + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape({2, 3}), + allocator, value); + + EXPECT_THROW(OrtValueTensorSlicer::Create( + value, /*slice_dimension=*/-1, /*dim0_offset=*/0), + OnnxRuntimeException); + EXPECT_THROW(OrtValueTensorSlicer::Create( + value, /*slice_dimension=*/2, /*dim0_offset=*/0), + OnnxRuntimeException); + EXPECT_THROW(OrtValueTensorSlicer::Create( + value, /*slice_dimension=*/1, /*dim0_offset=*/-1), + OnnxRuntimeException); + EXPECT_THROW(OrtValueTensorSlicer::Create( + value, /*slice_dimension=*/1, /*dim0_offset=*/2), + OnnxRuntimeException); + EXPECT_THROW(OrtValueTensorSlicer::Create( + value, /*slice_dimension=*/0, /*dim0_offset=*/1), + OnnxRuntimeException); +} + TEST(TensorTest, StringTensorTest) { // add scope to explicitly delete tensor #ifdef _MSC_VER diff --git a/onnxruntime/test/ir/onnx_model_test.cc b/onnxruntime/test/ir/onnx_model_test.cc index 55fc4f42bec64..ac8ba28820f43 100644 --- a/onnxruntime/test/ir/onnx_model_test.cc +++ b/onnxruntime/test/ir/onnx_model_test.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include +#include #include #include "core/platform/env.h" #include "core/graph/graph_viewer.h" @@ -10,6 +11,7 @@ #include "core/session/onnxruntime_c_api.h" #include "test/providers/provider_test_utils.h" //For ASSERT_STATUS_OK #include "test/test_environment.h" +#include "gmock/gmock.h" #include "gtest/gtest.h" #include "onnx/defs/function.h" #include "onnx/defs/parser.h" @@ -42,6 +44,25 @@ TEST_F(ONNXModelsTest, future_opset) { ASSERT_FALSE(st.IsOK()); } +TEST_F(ONNXModelsTest, InvalidOpsetVersionsRejectedBeforeNarrowing) { + for (const int64_t version : + {int64_t{0}, static_cast(std::numeric_limits::max()) + 1}) { + ModelProto model_proto; + model_proto.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); + model_proto.mutable_graph()->set_name("invalid_opset_version"); + auto* opset = model_proto.add_opset_import(); + opset->set_domain(onnxruntime::kOnnxDomain); + opset->set_version(version); + + try { + Model model(std::move(model_proto), nullptr, *logger_); + FAIL() << "Expected invalid opset version " << version << " to be rejected"; + } catch (const OnnxRuntimeException& ex) { + EXPECT_THAT(ex.what(), testing::HasSubstr("Invalid opset version")); + } + } +} + class ONNXModelsTest1 : public ::testing::TestWithParam { // You can implement all the usual fixture class members here. // To access the test parameter, call GetParam() from class