Fix/winml image dimension overflow

winml/lib/Api.Image/VideoFrameToTensorConverter.cpp

@@ -7,6 +7,7 @@
 #include <TraceLoggingProvider.h>
 #include <TraceloggingConfig.h>
 #include <evntrace.h>
+#include <limits>
 #include <MemoryBuffer.h>
 
 #include "inc/VideoFrameToTensorConverter.h"
@@ -600,9 +601,25 @@ void VideoFrameToTensorConverter::ConvertSoftwareBitmapToGPUTensor(
 
   D3D12_RESOURCE_DESC outputDesc = pOutputResource->GetDesc();
 
-  uint32_t tensorElementSize = tensorDesc.dataType == kImageTensorDataTypeFloat32 ? 4 : 2;
-  uint32_t bufferSize =
-    static_cast<uint32_t>(tensorDesc.sizes[1] * tensorDesc.sizes[2] * tensorDesc.sizes[3] * tensorElementSize);
+  const UINT64 tensorElementSize = tensorDesc.dataType == kImageTensorDataTypeFloat32 ? 4 : 2;
+  UINT64 bufferSize = tensorElementSize;
+  for (UINT index = 1; index < kImageTensorDimensionCountMax; ++index) {
+    WINML_THROW_HR_IF_FALSE_MSG(E_INVALIDARG, tensorDesc.sizes[index] > 0, "Image tensor dimensions must be positive.");
+    WINML_THROW_IF_FAILED(ULongLongMult(bufferSize, static_cast<UINT64>(tensorDesc.sizes[index]), &bufferSize));
+  }
+
+  UINT64 outputOffset = 0;
+  UINT64 outputEnd = 0;
+  WINML_THROW_IF_FAILED(ULongLongMult(bufferSize, batchIdx, &outputOffset));
+  WINML_THROW_IF_FAILED(ULongLongAdd(outputOffset, bufferSize, &outputEnd));
+  WINML_THROW_HR_IF_FALSE_MSG(
+    E_INVALIDARG,
+    outputDesc.Dimension == D3D12_RESOURCE_DIMENSION_BUFFER && outputDesc.Height == 1 && outputDesc.Width >= outputEnd,
+    "Output tensor resource is too small for the requested image batch."
+  );
+  WINML_THROW_HR_IF_FALSE_MSG(
+    E_INVALIDARG, bufferSize <= std::numeric_limits<SIZE_T>::max(), "Image tensor upload size is too large."
+  );
 
   // TODO: Make an allocator for upload heaps
   if (!upload_heap_ || upload_heap_->GetDesc().Width < bufferSize) {
@@ -624,7 +641,7 @@ void VideoFrameToTensorConverter::ConvertSoftwareBitmapToGPUTensor(
   // instead of the initial input bounds
   ConvertSoftwareBitmapToCPUTensor(convertedSoftwareBitmap, tensorDesc, scaledBounds, pCPUTensorBuffer);
 
-  upload_heap_->Unmap(0, unmove_ptr(CD3DX12_RANGE(0, bufferSize)));
+  upload_heap_->Unmap(0, unmove_ptr(CD3DX12_RANGE(0, static_cast<SIZE_T>(bufferSize))));
 
   ResetCommandList(device_cache);
 
@@ -633,7 +650,7 @@ void VideoFrameToTensorConverter::ConvertSoftwareBitmapToGPUTensor(
   );
   command_list_->ResourceBarrier(1, &barrier);
 
-  command_list_->CopyBufferRegion(pOutputResource, bufferSize * batchIdx, upload_heap_.Get(), 0, bufferSize);
+  command_list_->CopyBufferRegion(pOutputResource, outputOffset, upload_heap_.Get(), 0, bufferSize);
 
   WINML_THROW_IF_FAILED(command_list_->Close());
   ID3D12CommandList* ppCommandLists[] = {command_list_.Get()};

winml/lib/Api.Ort/OnnxruntimeDescriptorConverter.cpp

@@ -4,6 +4,7 @@
 #include "lib/Api.Ort/pch.h"
 
 #include <evntrace.h>
+#include <limits>
 
 #include "OnnxruntimeDescriptorConverter.h"
 #include "ImageFeatureDescriptor.h"
@@ -447,6 +448,13 @@ static winml::ILearningModelFeatureDescriptor CreateImageFeatureDescriptor(
   // Should the model metadata be read instead???
   const int c_height_dimension = 2;
   const int c_width_dimension = 3;
+  WINML_THROW_HR_IF_FALSE_MSG(
+    E_INVALIDARG,
+    shape.size() > static_cast<size_t>(c_width_dimension) && shape[c_height_dimension] > 0 &&
+      shape[c_width_dimension] > 0 && shape[c_height_dimension] <= std::numeric_limits<int32_t>::max() &&
+      shape[c_width_dimension] <= std::numeric_limits<int32_t>::max(),
+    "Image tensor height and width must be positive and no greater than INT32_MAX."
+  );
   auto height = static_cast<uint32_t>(shape[c_height_dimension]);
   auto width = static_cast<uint32_t>(shape[c_width_dimension]);
   auto descriptor = winrt::make<winmlp::ImageFeatureDescriptor>(

winml/test/api/LearningModelAPITest.cpp

@@ -5,6 +5,8 @@
 #include "LearningModelAPITest.h"
 #include "APITest.h"
 
+#include <limits>
+
 using namespace winrt;
 using namespace winml;
 using namespace wfc;
@@ -272,6 +274,16 @@ static void CheckLearningModelPixelRange() {
   }
 }
 
+static void RejectOversizedImageDimensions() {
+  WINML_EXPECT_THROW_SPECIFIC(
+    ProtobufHelpers::CreateModel(
+      TensorKind::Float, {1, 3, static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1, 1}, 1, true
+    ),
+    winrt::hresult_error,
+    [](const winrt::hresult_error& e) -> bool { return e.code() == E_INVALIDARG; }
+  );
+}
+
 static void CloseModelCheckEval() {
   LearningModel learningModel = nullptr;
   WINML_EXPECT_NO_THROW(APITest::LoadModel(L"model.onnx", learningModel));
@@ -330,6 +342,7 @@ const LearningModelApiTestsApi& getapi() {
     EnumerateOutputs,
     CloseModelCheckMetadata,
     CheckLearningModelPixelRange,
+    RejectOversizedImageDimensions,
     CloseModelCheckEval,
     CloseModelNoNewSessions,
     CheckMetadataCaseInsensitive,

winml/test/api/LearningModelAPITest.h

@@ -19,6 +19,7 @@ struct LearningModelApiTestsApi {
   VoidTest EnumerateOutputs;
   VoidTest CloseModelCheckMetadata;
   VoidTest CheckLearningModelPixelRange;
+  VoidTest RejectOversizedImageDimensions;
   VoidTest CloseModelCheckEval;
   VoidTest CloseModelNoNewSessions;
   VoidTest CheckMetadataCaseInsensitive;
@@ -44,6 +45,7 @@ WINML_TEST(LearningModelAPITests, EnumerateInputs)
 WINML_TEST(LearningModelAPITests, EnumerateOutputs)
 WINML_TEST(LearningModelAPITests, CloseModelCheckMetadata)
 WINML_TEST(LearningModelAPITests, CheckLearningModelPixelRange)
+WINML_TEST(LearningModelAPITests, RejectOversizedImageDimensions)
 WINML_TEST(LearningModelAPITests, CloseModelNoNewSessions)
 WINML_TEST(LearningModelAPITests, CloseModelCheckEval)
 WINML_TEST(LearningModelAPITests, CheckMetadataCaseInsensitive)

winml/test/common/protobufHelpers.cpp

@@ -209,7 +209,7 @@ TensorFloat16Bit ProtobufHelpers::LoadTensorFloat16FromProtobufFile(const std::w
 }
 
 winml::LearningModel ProtobufHelpers::CreateModel(
-  winml::TensorKind kind, const std::vector<int64_t>& shape, uint32_t num_elements
+  winml::TensorKind kind, const std::vector<int64_t>& shape, uint32_t num_elements, bool image_input
 ) {
   onnx::ModelProto model;
   model.set_ir_version(onnx::Version::IR_VERSION);
@@ -292,6 +292,9 @@ winml::LearningModel ProtobufHelpers::CreateModel(
     onnx::ValueInfoProto& variable = *graph.add_input();
     variable.set_name("input");
     variable.mutable_type()->mutable_tensor_type()->set_elem_type(dataType);
+    if (image_input) {
+      variable.mutable_type()->set_denotation("IMAGE");
+    }
     for (auto dim : shape) {
       if (dim == -1) {
         variable.mutable_type()->mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_param(&dim_param, 1);

winml/test/common/protobufHelpers.h

@@ -21,7 +21,9 @@ winml::ITensor LoadTensorFromProtobufFile(const std::wstring& filePath, bool isF
 // LoadTensorFloat16FromProtobufFile takes a path to a FP16 data file and loads it into a 16bit array
 winml::TensorFloat16Bit LoadTensorFloat16FromProtobufFile(const std::wstring& filePath);
 
-winml::LearningModel CreateModel(winml::TensorKind kind, const std::vector<int64_t>& shape, uint32_t num_elements = 1);
+winml::LearningModel CreateModel(
+  winml::TensorKind kind, const std::vector<int64_t>& shape, uint32_t num_elements = 1, bool image_input = false
+);
 
 // Populates TensorProto with tensor from protobuf file
 bool LoadOnnxTensorFromProtobufFile(onnx::TensorProto& tensor, std::wstring filePath);