Validate SpaceDepth shape arithmetic

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

@@ -1,11 +1,6 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-// TODO: x86 release builds produce warning C4723: potential divide by 0
-#ifdef _MSC_VER
-#pragma warning(disable : 4723)
-#endif
-
 #include "core/providers/cpu/tensor/space_depth_ops.h"
 #include "core/common/eigen_common_wrapper.h"
 #include <array>

onnxruntime/core/providers/cpu/tensor/space_depth_ops.h

@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include <limits>
 #include <string>
 
 #if !defined(SHARED_PROVIDER) && !defined(BUILD_CUDA_EP_AS_PLUGIN)
@@ -52,6 +53,10 @@ inline Status InputValidationsAndOutputDimsCalc(int64_t blocksize,
                            input_shape.NumDimensions());
   }
 
+  if (blocksize <= 0) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "SpaceDepth ops require blocksize to be greater than 0.");
+  }
+
   batch = input_shape[0];
   if constexpr (IsNHWC) {
     input_depth = input_shape[3];
@@ -72,16 +77,26 @@ inline Status InputValidationsAndOutputDimsCalc(int64_t blocksize,
       return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "SpaceToDepth requires input width to be a multiple of block_size");
     }
 
+    const auto int64_max = std::numeric_limits<int64_t>::max();
+    if (input_depth > int64_max / blocksize || input_depth * blocksize > int64_max / blocksize) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "SpaceToDepth output depth exceeds int64_t limits");
+    }
+
     output_depth = input_depth * blocksize * blocksize;
     output_height = input_height / blocksize;
     output_width = input_width / blocksize;
 
   } else {  // DepthToSpace op
-    if ((input_depth % (blocksize * blocksize) != 0)) {
+    if (input_depth % blocksize != 0 || (input_depth / blocksize) % blocksize != 0) {
       return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
                              "DepthToSpace requires input depth to be a multiple of (block_size * block_size)");
     }
 
+    const auto int64_max = std::numeric_limits<int64_t>::max();
+    if (input_height > int64_max / blocksize || input_width > int64_max / blocksize) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "DepthToSpace output dimensions exceed int64_t limits");
+    }
+
     output_depth = input_depth / blocksize / blocksize;
     output_height = input_height * blocksize;
     output_width = input_width * blocksize;

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

@@ -6,6 +6,7 @@
 
 #include "gtest/gtest.h"
 #include "test/providers/provider_test_utils.h"
+#include "core/framework/allocator.h"
 #include "core/providers/cpu/tensor/space_depth_ops.h"
 #include "core/mlas/inc/mlas.h"
 
@@ -19,6 +20,46 @@ class TensorOpTest : public ::testing::Test {
 using TensorOpTestTypes = ::testing::Types<float, MLFloat16, uint8_t, int8_t>;
 TYPED_TEST_SUITE(TensorOpTest, TensorOpTestTypes);
 
+TEST(TensorOpTest, RejectsNonPositiveBlocksizeBeforeShapeArithmetic) {
+  auto allocator = std::make_shared<CPUAllocator>();
+  Tensor input(DataTypeImpl::GetType<float>(), TensorShape({1, 1, 2, 2}), allocator);
+
+  for (const auto& [blocksize, is_space_to_depth] :
+       std::vector<std::pair<int64_t, bool>>{{0, true}, {-1, false}}) {
+    int64_t batch = 0;
+    int64_t input_depth = 0;
+    int64_t input_height = 0;
+    int64_t input_width = 0;
+    int64_t output_depth = 0;
+    int64_t output_height = 0;
+    int64_t output_width = 0;
+    const auto status = space_depth_internal::InputValidationsAndOutputDimsCalc(
+        blocksize, input, batch, input_depth, input_height, input_width,
+        output_depth, output_height, output_width, is_space_to_depth);
+    EXPECT_EQ(status.Code(), common::StatusCode::INVALID_ARGUMENT);
+    EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("blocksize to be greater than 0"));
+  }
+}
+
+TEST(TensorOpTest, RejectsBlocksizeWhoseSquareExceedsInt64) {
+  auto allocator = std::make_shared<CPUAllocator>();
+  Tensor input(DataTypeImpl::GetType<float>(), TensorShape({1, 1, 2, 2}), allocator);
+  int64_t batch = 0;
+  int64_t input_depth = 0;
+  int64_t input_height = 0;
+  int64_t input_width = 0;
+  int64_t output_depth = 0;
+  int64_t output_height = 0;
+  int64_t output_width = 0;
+
+  const auto status = space_depth_internal::InputValidationsAndOutputDimsCalc(
+      int64_t{1} << 32, input, batch, input_depth, input_height, input_width,
+      output_depth, output_height, output_width, false);
+
+  EXPECT_EQ(status.Code(), common::StatusCode::INVALID_ARGUMENT);
+  EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("multiple of (block_size * block_size)"));
+}
+
 TEST(TensorOpTest, SpaceToDepthTest_1) {
   OpTester test("SpaceToDepth");
   constexpr int64_t blocksize = 2;