Validate tensor slicing, sparse CSR shapes, and opset versions

onnxruntime/core/framework/ort_value_tensor_slicer.cc

@@ -21,11 +21,14 @@ OrtValueTensorSlicer<T> OrtValueTensorSlicer<T>::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<int64_t>(tensor_shape.NumDimensions()) >= slice_dimension,
+  ORT_ENFORCE(slice_dimension >= 0 &&
+                  gsl::narrow_cast<int64_t>(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};
 };

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;

onnxruntime/core/graph/model.cc

@@ -282,6 +282,8 @@ Model::Model(ModelProto&& model_proto, const PathString& model_path,
   std::unordered_map<std::string, int> 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<int>::max(),
+                "Invalid opset version for domain ", domain);
     const auto version = gsl::narrow_cast<int>(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 {

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> 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();

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<int32_t>(), TensorShape{1, 1},
+                   cpu_allocator);
+  const std::vector<int32_t> values = {1};
+  const std::vector<int64_t> inner = {0};
+  const std::vector<int64_t> 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<TensorShape&>(src.DenseShape()) =
+      TensorShape{std::numeric_limits<int64_t>::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];

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<float>(), TensorShape({2, 3}),
+                       allocator, value);
+
+  EXPECT_NO_THROW({
+    auto slicer = OrtValueTensorSlicer<OrtValue>::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<float>(), TensorShape({2, 3}),
+                       allocator, value);
+
+  EXPECT_THROW(OrtValueTensorSlicer<OrtValue>::Create(
+                   value, /*slice_dimension=*/-1, /*dim0_offset=*/0),
+               OnnxRuntimeException);
+  EXPECT_THROW(OrtValueTensorSlicer<OrtValue>::Create(
+                   value, /*slice_dimension=*/2, /*dim0_offset=*/0),
+               OnnxRuntimeException);
+  EXPECT_THROW(OrtValueTensorSlicer<OrtValue>::Create(
+                   value, /*slice_dimension=*/1, /*dim0_offset=*/-1),
+               OnnxRuntimeException);
+  EXPECT_THROW(OrtValueTensorSlicer<OrtValue>::Create(
+                   value, /*slice_dimension=*/1, /*dim0_offset=*/2),
+               OnnxRuntimeException);
+  EXPECT_THROW(OrtValueTensorSlicer<OrtValue>::Create(
+                   value, /*slice_dimension=*/0, /*dim0_offset=*/1),
+               OnnxRuntimeException);
+}
+
 TEST(TensorTest, StringTensorTest) {
 // add scope to explicitly delete tensor
 #ifdef _MSC_VER

onnxruntime/test/ir/onnx_model_test.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include <google/protobuf/io/zero_copy_stream_impl.h>
+#include <limits>
 #include <memory>
 #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<int64_t>(std::numeric_limits<int>::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<const ORTCHAR_T*> {
   // You can implement all the usual fixture class members here.
   // To access the test parameter, call GetParam() from class