Harden OneHot operator input validation and output size computation

onnxruntime/core/providers/cpu/tensor/onehot.cc

@@ -15,7 +15,11 @@ limitations under the License.
 /* Modifications Copyright (c) Microsoft. */
 
 #include "core/providers/cpu/tensor/onehot.h"
+
+#include <limits>
+
 #include "core/common/eigen_common_wrapper.h"
+#include "core/common/safeint.h"
 #include "core/platform/env.h"
 #include "core/providers/common.h"
 
@@ -91,6 +95,13 @@ Status PrepareOutputShape(const Tensor* indices, const int64_t depth_val, const
   const auto& indices_shape = indices->Shape();
   const auto indices_dims = indices_shape.GetDims();
   const auto indices_num_dims = indices_shape.NumDimensions();
+
+  // ONNX spec requires indices to have rank >= 1.
+  if (indices_num_dims == 0) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "OneHot: indices tensor must have rank >= 1.");
+  }
+
   output_shape = indices_shape.AsShapeVector();
 
   // output rank is always 1 more than the input rank as a new dimension is added to the input shape
@@ -100,11 +111,31 @@ Status PrepareOutputShape(const Tensor* indices, const int64_t depth_val, const
 
   output_shape.insert(output_shape.begin() + true_axis, depth_val);
 
-  prefix_dim_size = 1;
+  // Validate that the total output tensor element count does not overflow int64.
+  {
+    int64_t total_elements = 1;
+    for (auto dim : output_shape) {
+      if (dim > 0 && total_elements > std::numeric_limits<int64_t>::max() / dim) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "OneHot: output tensor size would overflow for the given indices shape "
+                               "and depth value (",
+                               depth_val, ").");
+      }
+      total_elements *= dim;
+    }
+  }
+
+  // Use SafeInt for prefix_dim_size computation to guard against overflow.
+  // SafeInt is defensive here -- the total-element overflow check above already covers this case,
+  // so a SafeIntException should never fire in practice.
+  SafeInt<int64_t> safe_prefix = 1;
   for (int64_t i = 0; i < true_axis; ++i) {
-    prefix_dim_size *= indices_dims[onnxruntime::narrow<size_t>(i)];
+    safe_prefix *= indices_dims[onnxruntime::narrow<size_t>(i)];
   }
-  suffix_dim_size = indices_shape.Size() / prefix_dim_size;
+  prefix_dim_size = safe_prefix;
+
+  // Guard against division by zero when indices have a zero-sized dimension before the axis.
+  suffix_dim_size = (prefix_dim_size > 0) ? (indices_shape.Size() / prefix_dim_size) : 0;
 
   return Status::OK();
 }
@@ -166,6 +197,7 @@ Status OneHotOp<in_type, out_type, depth_type>::Compute(OpKernelContext* p_op_ke
   // allocate output
   const auto* values_data = values->Data<out_type>();
   Tensor* output = p_op_kernel_context->Output(0, TensorShape(output_shape));
+  ORT_RETURN_IF_NOT(output, "OneHot: failed to allocate output tensor. Output shape may be too large.");
 
   // edge case where we have a dim with a value of 0
   if (output->Shape().Size() == 0)

onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h

@@ -15,8 +15,11 @@
 
 #pragma once
 
+#include <limits>
+
 #include "core/common/status.h"
 #include "core/common/narrow.h"
+#include "core/common/safeint.h"
 #include "core/common/float16.h"
 #include "core/common/float8.h"
 #include "core/framework/float4.h"
@@ -769,15 +772,43 @@ inline Status PrepareOutputShape(const Tensor* indices, const int64_t depth_val,
   const auto& indices_shape = indices->Shape();
   const auto indices_dims = indices_shape.GetDims();
   const auto indices_num_dims = indices_shape.NumDimensions();
+
+  // ONNX spec requires indices to have rank >= 1.
+  if (indices_num_dims == 0) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "OneHot: indices tensor must have rank >= 1.");
+  }
+
   output_shape = indices_shape.AsShapeVector();
   const auto output_rank = static_cast<int64_t>(indices_num_dims) + 1;
   auto true_axis = HandleNegativeAxis(axis, output_rank);
   output_shape.insert(output_shape.begin() + true_axis, depth_val);
-  prefix_dim_size = 1;
+
+  // Validate that the total output tensor element count does not overflow int64.
+  {
+    int64_t total_elements = 1;
+    for (auto dim : output_shape) {
+      if (dim > 0 && total_elements > std::numeric_limits<int64_t>::max() / dim) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "OneHot: output tensor size would overflow for the given indices shape "
+                               "and depth value (",
+                               depth_val, ").");
+      }
+      total_elements *= dim;
+    }
+  }
+
+  // Use SafeInt for prefix_dim_size to guard against overflow.
+  // SafeInt is defensive here -- the total-element overflow check above already covers this case,
+  // so a SafeIntException should never fire in practice.
+  SafeInt<int64_t> safe_prefix = 1;
   for (int64_t i = 0; i < true_axis; ++i) {
-    prefix_dim_size *= indices_dims[narrow<size_t>(i)];
+    safe_prefix *= indices_dims[narrow<size_t>(i)];
   }
-  suffix_dim_size = indices_shape.Size() / prefix_dim_size;
+  prefix_dim_size = safe_prefix;
+
+  // Guard against division by zero when indices have a zero-sized dimension before the axis.
+  suffix_dim_size = (prefix_dim_size > 0) ? (indices_shape.Size() / prefix_dim_size) : 0;
   return Status::OK();
 }
 

onnxruntime/core/providers/cuda/tensor/onehot.cc

@@ -3,6 +3,9 @@
 
 #include "core/providers/cuda/tensor/onehot.h"
 
+#include <algorithm>
+#include <limits>
+
 using namespace onnxruntime::common;
 
 namespace onnxruntime {
@@ -55,11 +58,19 @@ Status OneHotOp<in_type, out_type, depth_type>::ComputeInternal(OpKernelContext*
   // allocate output
   const auto* values_data = reinterpret_cast<const CudaT_Out*>(values->Data<out_type>());
   Tensor* output = ctx->Output(0, TensorShape(output_shape));
+  ORT_RETURN_IF_NOT(output, "OneHot: failed to allocate output tensor. Output shape may be too large.");
 
   // edge case where we have a dim with a value of 0
   if (output->Shape().Size() == 0)
     return Status::OK();
 
+  // Validate that suffix_dim_size fits in int32 range. fast_divmod requires int32 operands
+  // and fdm_suffix is constructed on every code path below.
+  constexpr int64_t kInt32Max = std::numeric_limits<int>::max();
+  ORT_RETURN_IF_NOT(suffix_dim_size <= kInt32Max,
+                    "OneHot: suffix dimension size (", suffix_dim_size,
+                    ") exceeds int32 range supported by the CUDA kernel.");
+
   const fast_divmod fdm_suffix(gsl::narrow_cast<int>(suffix_dim_size));
   const auto* indices_data = indices->Data<in_type>();
   auto* output_data = reinterpret_cast<CudaT_Out*>(output->MutableData<out_type>());
@@ -76,6 +87,10 @@ Status OneHotOp<in_type, out_type, depth_type>::ComputeInternal(OpKernelContext*
     return Status::OK();
   }
 
+  // depth * suffix is only needed for fdm_depth_suffix on the non-zero-off-value path.
+  ORT_RETURN_IF_NOT(depth_val <= kInt32Max / std::max(suffix_dim_size, int64_t{1}),
+                    "OneHot: depth (", depth_val, ") * suffix dimension size (", suffix_dim_size,
+                    ") exceeds int32 range supported by the CUDA kernel.");
   const fast_divmod fdm_depth_suffix(gsl::narrow_cast<int>(depth_val * suffix_dim_size));
   OneHotImpl(Stream(ctx),
              indices_data, fdm_depth_suffix, fdm_suffix, depth_val,

onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include <limits>
+
 #include "gtest/gtest.h"
 #include "test/providers/provider_test_utils.h"
 #include "test/common/trt_op_test_utils.h"
@@ -499,6 +501,89 @@ TEST(OneHotOpTest, DimWithZero) {
   test.Run();
 }
 
+// Test that extremely large depth values that would cause output tensor size overflow are rejected.
+TEST(OneHotOpTest, DepthTooLarge_OutputSizeOverflow) {
+  OpTester test("OneHot", 11);
+  // indices shape [2, 3] with depth = INT64_MAX causes output shape [2, 3, INT64_MAX]
+  // which would overflow when computing total element count.
+  test.AddInput<int64_t>("indices", {2, 3}, {1, 2, 3, 4, 5, 6});
+  test.AddInput<int64_t>("depth", {1}, {std::numeric_limits<int64_t>::max()});
+  test.AddInput<int64_t>("values", {2}, {0, 1});
+  test.AddOutput<int64_t>("output", {2, 3, 1}, {0, 0, 0, 0, 0, 0});
+  // Exclude TensorRT and DML EPs: they fail internally on INT64_MAX depth before our kernel's
+  // validation runs, producing a different error message.
+  test.Run(OpTester::ExpectResult::kExpectFailure, "output tensor size would overflow",
+           {kTensorrtExecutionProvider, kDmlExecutionProvider});
+}
+
+// Test that a very large depth value that overflows with multi-dimensional indices is rejected.
+TEST(OneHotOpTest, DepthTooLarge_OutputSizeOverflow_LargeIndices) {
+  OpTester test("OneHot", 11);
+  // indices shape [1000] with depth = INT64_MAX / 500 causes overflow in element count.
+  const int64_t large_depth = std::numeric_limits<int64_t>::max() / 500;
+  std::vector<int64_t> indices(1000, 0);
+  std::vector<int64_t> dummy_output(1000, 0);
+  test.AddInput<int64_t>("indices", {1000}, indices);
+  test.AddInput<int64_t>("depth", {1}, {large_depth});
+  test.AddInput<int64_t>("values", {2}, {0, 1});
+  test.AddOutput<int64_t>("output", {1000, 1}, dummy_output);
+  // Exclude TensorRT and DML EPs: they fail internally on overflow-inducing depth before our
+  // kernel's validation runs.
+  test.Run(OpTester::ExpectResult::kExpectFailure, "output tensor size would overflow",
+           {kTensorrtExecutionProvider, kDmlExecutionProvider});
+}
+
+// Test that a negative depth value is rejected.
+TEST(OneHotOpTest, NegativeDepth) {
+  OpTester test("OneHot", 11);
+  test.AddInput<int64_t>("indices", {2, 3}, {1, 2, 3, 4, 5, 6});
+  test.AddInput<int64_t>("depth", {1}, {-5});
+  test.AddInput<int64_t>("values", {2}, {0, 1});
+  test.AddOutput<int64_t>("output", {2, 3, 1}, {0, 0, 0, 0, 0, 0});
+  // Exclude TensorRT and DML EPs: they reject negative depth with their own error messages rather
+  // than ours.
+  test.Run(OpTester::ExpectResult::kExpectFailure, "Depth is negative",
+           {kTensorrtExecutionProvider, kDmlExecutionProvider});
+}
+
+// Test minimum valid depth value of 1.
+TEST(OneHotOpTest, DepthOne) {
+  OpTester test("OneHot", 11);
+  test.AddInput<int64_t>("indices", {3}, {0, 0, 0});
+  test.AddInput<int64_t>("depth", {1}, {1});
+  test.AddInput<int64_t>("values", {2}, {0, 1});
+  test.AddOutput<int64_t>("output", {3, 1}, {1, 1, 1});
+  test.Run();
+}
+
+// Test scalar (rank-0) indices are rejected per ONNX spec (indices must have rank >= 1).
+TEST(OneHotOpTest, ScalarIndicesRejected) {
+  OpTester test("OneHot", 11);
+  test.AddInput<int64_t>("indices", {}, {2});
+  test.AddInput<int64_t>("depth", {1}, {5});
+  test.AddInput<int64_t>("values", {2}, {0, 1});
+  test.AddOutput<int64_t>("output", {5}, {0, 0, 1, 0, 0});
+  // Match either the ONNX shape-inference error ("Indices tensor must have rank >= 1") or the
+  // explicit kernel-level rejection ("OneHot: indices tensor must have rank >= 1.").
+  test.Run(OpTester::ExpectResult::kExpectFailure, "ndices tensor must have rank >= 1");
+}
+
+// Test with opset 9.
+TEST(OneHotOpTest, DefaultAxis_Opset9) {
+  OpTester test("OneHot", 9);
+  test.AddInput<int64_t>("indices", {2, 3}, {1, 9, 8, 2, 4, 6});
+  test.AddInput<int64_t>("depth", {1}, {10});
+  test.AddInput<int64_t>("values", {2}, {0, 1});
+  test.AddOutput<int64_t>("output", {2, 3, 10},
+                          {0, 1, 0, 0, 0, 0, 0, 0, 0, 0,
+                           0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
+                           0, 0, 0, 0, 0, 0, 0, 0, 1, 0,
+                           0, 0, 1, 0, 0, 0, 0, 0, 0, 0,
+                           0, 0, 0, 0, 1, 0, 0, 0, 0, 0,
+                           0, 0, 0, 0, 0, 0, 1, 0, 0, 0});
+  test.Run();
+}
+
 #ifdef USE_CUDA
 
 TEST(OneHotOpTest, DefaultAxis_int64_MLFloat16_int64 /*indices, output, depth*/) {