apps: fix integer-overflow and OOB write issues in image loaders

CHANGELOG.md

@@ -3,6 +3,7 @@ Version History
 
 -   Added Intel BMG-G31, Wildcat Lake, Nova Lake, and Crescent Island GPU support
 -   Added AMD RDNA 3.5 GPU support, extended RDNA 2 support
+-   Fixed integer overflow and out-of-bounds write issues in image loaders (only affects the `oidnDenoise` example application)
 
 ### Changes in v2.3.3:
 

apps/utils/image_buffer.cpp

@@ -16,7 +16,7 @@ OIDN_NAMESPACE_BEGIN
       dataType(DataType::Void),
       format(Format::Undefined) {}
 
-  ImageBuffer::ImageBuffer(const DeviceRef& device, int width, int height, int numChannels,
+  ImageBuffer::ImageBuffer(const DeviceRef& device, size_t width, size_t height, size_t numChannels,
                            DataType dataType, Storage storage, bool forceHostCopy)
     : device(device),
       numValues(size_t(width) * height * numChannels),
@@ -26,6 +26,10 @@ OIDN_NAMESPACE_BEGIN
       dataType(dataType),
       format(makeFormat(dataType, numChannels))
   {
+
+    if (width > maxDim || height > maxDim || width * height * getC() > std::numeric_limits<int>::max())
+      throw std::runtime_error("image size is too large");
+
     const size_t valueByteSize = getDataTypeSize(dataType);
     byteSize = std::max(numValues * valueByteSize, size_t(1)); // avoid zero-sized buffer
     buffer = device.newBuffer(byteSize, storage);

apps/utils/image_buffer.h

@@ -15,17 +15,19 @@ OIDN_NAMESPACE_BEGIN
   class ImageBuffer
   {
   public:
+    static constexpr size_t maxDim = 65536;
+
     ImageBuffer();
-    ImageBuffer(const DeviceRef& device, int width, int height, int numChannels,
+    ImageBuffer(const DeviceRef& device, size_t width, size_t height, size_t numChannels,
                 DataType dataType = DataType::Float32,
                 Storage storage = Storage::Undefined,
                 bool forceHostCopy = false);
     ~ImageBuffer();
 
-    oidn_inline int getW() const { return width; }
-    oidn_inline int getH() const { return height; }
-    oidn_inline int getC() const { return numChannels; }
-    std::array<int, 3> getDims() const { return {width, height, numChannels}; }
+    oidn_inline size_t getW() const { return width; }
+    oidn_inline size_t getH() const { return height; }
+    oidn_inline size_t getC() const { return numChannels; }
+    std::array<size_t, 3> getDims() const { return {width, height, numChannels}; }
 
     oidn_inline DataType getDataType() const { return dataType; }
     oidn_inline Format getFormat() const { return format; }
@@ -98,9 +100,9 @@ OIDN_NAMESPACE_BEGIN
     char* hostPtr;
     size_t byteSize;
     size_t numValues;
-    int width;
-    int height;
-    int numChannels;
+    size_t width;
+    size_t height;
+    size_t numChannels;
     DataType dataType;
     Format format;
   };

apps/utils/image_io.cpp

@@ -60,7 +60,7 @@ OIDN_NAMESPACE_BEGIN
       // Read the header
       std::string id;
       file >> id;
-      int C;
+      size_t C;
       if (id == "PF")
         C = 3;
       else if (id == "Pf")
@@ -73,7 +73,7 @@ OIDN_NAMESPACE_BEGIN
       if (dataType == DataType::Void)
         dataType = DataType::Float32;
 
-      int H, W;
+      size_t H, W;
       file >> W >> H;
 
       float scale;
@@ -91,23 +91,22 @@ OIDN_NAMESPACE_BEGIN
       // Read the pixels
       auto image = std::make_shared<ImageBuffer>(device, W, H, C, dataType, storage);
 
-      for (int h = 0; h < H; ++h)
+      for (size_t h = 0; h < H; ++h)
       {
-        for (int w = 0; w < W; ++w)
+        for (size_t w = 0; w < W; ++w)
         {
-          for (int c = 0; c < C; ++c)
+          for (size_t c = 0; c < C; ++c)
           {
             float x;
-            file.read((char*)&x, sizeof(float));
-            if (c < C)
-              image->set((size_t(H-1-h)*W + w) * C + c, x * scale);
+            file.read(reinterpret_cast<char*>(&x), sizeof(float));
+            if (file.fail()) {
+                throw std::runtime_error("invalid PFM image: error reading pixel data");
+            }
+            image->set((size_t(H-1-h)*W + w) * C + c, x * scale);
           }
         }
       }
 
-      if (file.fail())
-        throw std::runtime_error("invalid PFM image");
-
       return image;
     }
 

common/platform.h

@@ -265,6 +265,8 @@ OIDN_NAMESPACE_BEGIN
   void* alignedMalloc(size_t size, size_t alignment = memoryAlignment);
   void alignedFree(void* ptr);
 
+  static_assert(sizeof(size_t) == 8, "size_t is not 64-bit!");
+
   // -----------------------------------------------------------------------------------------------
   // String functions
   // -----------------------------------------------------------------------------------------------