Reject Oversized Images in DecodeImage

operators/vision/image_decoder.hpp

@@ -15,6 +15,11 @@
 #include "ext_status.h"
 
 namespace ort_extensions::internal {
+
+// Maximum image dimension (width or height) and total pixel count to prevent decompression bombs.
+static constexpr uint64_t kMaxImageDimension = 16384;
+static constexpr uint64_t kMaxPixelCount = 100'000'000;  // 100 megapixels
+
 struct DecodeImage {
   OrtxStatus OnInit() { return {}; }
 
@@ -83,6 +88,14 @@ struct DecodeImage {
 
     png_read_update_info(png, info);
 
+    // Dimension limit to prevent decompression bombs
+    if (width > kMaxImageDimension || height > kMaxImageDimension ||
+        static_cast<uint64_t>(width) * height > kMaxPixelCount) {
+      png_destroy_read_struct(&png, &info, nullptr);
+      return {kOrtxErrorInvalidArgument,
+              "[ImageDecoder]: PNG dimensions exceed maximum allowed size."};
+    }
+
     std::vector<int64_t> output_dimensions{height, width, 3};
     uint8_t* output_data = output.Allocate(output_dimensions);
     // Read the image row by row
@@ -146,6 +159,15 @@ struct DecodeImage {
       // Start decompression
       jpeg_start_decompress(&cinfo);
 
+      // Dimension limit to prevent decompression bombs
+      if (cinfo.output_width > kMaxImageDimension ||
+          cinfo.output_height > kMaxImageDimension ||
+          static_cast<uint64_t>(cinfo.output_width) * cinfo.output_height > kMaxPixelCount) {
+        jpeg_destroy_decompress(&cinfo);
+        return {kOrtxErrorInvalidArgument,
+                "[ImageDecoder]: JPEG dimensions exceed maximum allowed size."};
+      }
+
       // Safety net: verify 3-channel output after decompression.
       if (cinfo.output_components != 3) {
         jpeg_destroy_decompress(&cinfo);

operators/vision/image_decoder_darwin.hpp

@@ -10,6 +10,10 @@
 
 namespace ort_extensions::internal {
 
+// Maximum image dimension (width or height) and total pixel count to prevent decompression bombs.
+static constexpr uint64_t kMaxImageDimension = 16384;
+static constexpr uint64_t kMaxPixelCount = 100'000'000;  // 100 megapixels
+
 struct DecodeImage {
   OrtxStatus OnInit() {
     CFStringRef optionKeys[2];
@@ -78,8 +82,20 @@ struct DecodeImage {
       }
     }
 
-    const int64_t width = static_cast<int64_t>(CGImageGetWidth(image));
-    const int64_t height = static_cast<int64_t>(CGImageGetHeight(image));
+    // Validate dimensions using unsigned types before any narrowing cast
+    const size_t rawWidth = CGImageGetWidth(image);
+    const size_t rawHeight = CGImageGetHeight(image);
+
+    // Dimension limit to prevent decompression bombs
+    if (rawWidth > kMaxImageDimension || rawHeight > kMaxImageDimension ||
+        static_cast<uint64_t>(rawWidth) * rawHeight > kMaxPixelCount) {
+      CGImageRelease(image);
+      return {kOrtxErrorInvalidArgument,
+              "[ImageDecoder]: Image dimensions exceed maximum allowed size."};
+    }
+
+    const int64_t width = static_cast<int64_t>(rawWidth);
+    const int64_t height = static_cast<int64_t>(rawHeight);
     const int64_t channels = 3;
 
     std::vector<int64_t> output_dimensions{height, width, channels};

operators/vision/image_decoder_win32.hpp

@@ -13,6 +13,11 @@
 #include "ext_status.h"
 
 namespace ort_extensions::internal {
+
+// Maximum image dimension (width or height) and total pixel count to prevent decompression bombs.
+static constexpr uint64_t kMaxImageDimension = 16384;
+static constexpr uint64_t kMaxPixelCount = 100'000'000;  // 100 megapixels
+
 struct DecodeImage {
   OrtxStatus OnInit() {
     HRESULT hr = CoInitializeEx(NULL, COINIT_MULTITHREADED | COINIT_DISABLE_OLE1DDE);
@@ -107,6 +112,13 @@ struct DecodeImage {
       return {kOrtxErrorInternal, "[ImageDecoder]: pIDecoderFrame->GetSize."};
     }
 
+    // Dimension limit to prevent decompression bombs (validate before narrowing cast)
+    if (uiWidth > kMaxImageDimension || uiHeight > kMaxImageDimension ||
+        static_cast<uint64_t>(uiWidth) * uiHeight > kMaxPixelCount) {
+      return {kOrtxErrorInvalidArgument,
+              "[ImageDecoder]: Image dimensions exceed maximum allowed size."};
+    }
+
     const int height = static_cast<int>(uiHeight);
     const int width = static_cast<int>(uiWidth);
     const int channels = 3;  // Asks for RGB

test/pp_api_test/test_imgcodec.cc

@@ -198,3 +198,142 @@ TEST(ImageDecoderTest, TestTiffDecoder) {
 }
 #endif
 #endif
+// Security: verify that oversized PNG images are rejected (decompression bomb mitigation).
+// Crafts a minimal valid PNG with an IHDR claiming 20000x20000 dimensions.
+TEST(ImageDecoderTest, TestPngOversizeDimensionsRejected) {
+  ort_extensions::DecodeImage image_decoder;
+  image_decoder.Init(std::unordered_map<std::string, std::variant<std::string>>());
+
+  // Minimal PNG: signature + IHDR (20000x20000) + IDAT (minimal zlib) + IEND.
+  // png_read_info reads all chunks up to the first IDAT, so we need IDAT present
+  // for libpng to successfully parse the header and reach our dimension check.
+  std::vector<uint8_t> png_oversize = {
+    // PNG signature
+    0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A,
+
+    // IHDR chunk: length=13
+    0x00, 0x00, 0x00, 0x0D,
+    // "IHDR"
+    0x49, 0x48, 0x44, 0x52,
+    // Width = 20000 (0x00004E20)
+    0x00, 0x00, 0x4E, 0x20,
+    // Height = 20000 (0x00004E20)
+    0x00, 0x00, 0x4E, 0x20,
+    // Bit depth = 8, Color type = 2 (RGB), Compression = 0, Filter = 0, Interlace = 0
+    0x08, 0x02, 0x00, 0x00, 0x00,
+    // CRC32 of IHDR
+    0x6C, 0x12, 0xD1, 0x6E,
+
+    // IDAT chunk: length=11 (minimal valid zlib stream)
+    0x00, 0x00, 0x00, 0x0B,
+    // "IDAT"
+    0x49, 0x44, 0x41, 0x54,
+    // Minimal zlib: header(78 01) + stored block BFINAL=1 LEN=0 NLEN=FFFF + Adler32(00000001)
+    0x78, 0x01, 0x01, 0x00, 0x00, 0xFF, 0xFF, 0x00, 0x00, 0x00, 0x01,
+    // CRC32 of IDAT type + data
+    0x89, 0xD6, 0xAE, 0x5F,
+
+    // IEND chunk: length=0
+    0x00, 0x00, 0x00, 0x00,
+    // "IEND"
+    0x49, 0x45, 0x4E, 0x44,
+    // CRC32 of IEND
+    0xAE, 0x42, 0x60, 0x82
+  };
+
+  ortc::Tensor<uint8_t> png_tensor({static_cast<int64_t>(png_oversize.size())}, png_oversize.data());
+  ortc::Tensor<uint8_t> out_tensor{&CppAllocator::Instance()};
+  auto status = image_decoder.Compute(png_tensor, out_tensor);
+
+  // Must be rejected — on libjpeg/libpng platforms this hits our dimension check;
+  // on macOS (CoreGraphics) and Windows (WIC) the platform decoder may reject the
+  // synthetic data before our check runs. Either way, the image must not be accepted.
+  std::cout << "[Expected rejection] PNG 20000x20000: " << status.ToString() << std::endl;
+  ASSERT_FALSE(status.IsOk()) << "Oversized PNG (20000x20000) should have been rejected but was accepted.";
+}
+
+// Security: verify that oversized JPEG images are rejected (decompression bomb mitigation).
+// Crafts a minimal JPEG with SOF0 claiming 17000x17000 dimensions.
+TEST(ImageDecoderTest, TestJpegOversizeDimensionsRejected) {
+  ort_extensions::DecodeImage image_decoder;
+  image_decoder.Init(std::unordered_map<std::string, std::variant<std::string>>());
+
+  // Minimal JPEG: SOI + SOF0 (with oversized dimensions) + EOI
+  // This should be enough for jpeg_read_header + jpeg_start_decompress to parse dimensions.
+  // Craft a minimal but structurally valid JPEG so libjpeg can parse the header
+  // and reach jpeg_start_decompress where our dimension check fires.
+  // Structure: SOI + DQT + SOF0 (oversized dims) + SOS + EOI
+  std::vector<uint8_t> jpeg_oversize = {
+    // SOI
+    0xFF, 0xD8,
+
+    // DQT marker (required for jpeg_start_decompress to succeed)
+    0xFF, 0xDB,
+    // Length = 67 (2 + 1 + 64): precision/table byte + 64 quantization values
+    0x00, 0x43,
+    // Table 0, 8-bit precision
+    0x00,
+    // 64 quantization values (all 1s — minimal valid table)
+    0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
+    0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
+    0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
+    0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
+    0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
+    0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
+    0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
+    0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
+
+    // SOF0 marker
+    0xFF, 0xC0,
+    // Length = 11
+    0x00, 0x0B,
+    // Precision = 8
+    0x08,
+    // Height = 17000 (0x4268)
+    0x42, 0x68,
+    // Width = 17000 (0x4268)
+    0x42, 0x68,
+    // Number of components = 1 (grayscale)
+    0x01,
+    // Component 1: id=1, sampling=0x11, quant_table=0
+    0x01, 0x11, 0x00,
+
+    // DHT marker (minimal Huffman table for DC, required by libjpeg)
+    0xFF, 0xC4,
+    // Length = 31 (2 + 1 + 16 + 12 symbols)
+    0x00, 0x1F,
+    // DC table 0
+    0x00,
+    // Number of codes of each length 1-16 (12 codes total)
+    0x00, 0x01, 0x05, 0x01, 0x01, 0x01, 0x01, 0x01,
+    0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    // Symbol values
+    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+    0x08, 0x09, 0x0A, 0x0B,
+
+    // SOS marker
+    0xFF, 0xDA,
+    // Length = 8
+    0x00, 0x08,
+    // Number of components = 1
+    0x01,
+    // Component 1: DC table 0, AC table 0
+    0x01, 0x00,
+    // Spectral selection start, end, approximation
+    0x00, 0x3F, 0x00,
+
+    // Minimal scan data (a single zero byte) + EOI
+    0x00,
+    0xFF, 0xD9
+  };
+
+  ortc::Tensor<uint8_t> jpeg_tensor({static_cast<int64_t>(jpeg_oversize.size())}, jpeg_oversize.data());
+  ortc::Tensor<uint8_t> out_tensor{&CppAllocator::Instance()};
+  auto status = image_decoder.Compute(jpeg_tensor, out_tensor);
+
+  // Must be rejected — on libjpeg platforms this hits our dimension check;
+  // on macOS (CoreGraphics) and Windows (WIC) the platform decoder may reject the
+  // synthetic data before our check runs. Either way, the image must not be accepted.
+  std::cout << "[Expected rejection] JPEG 17000x17000: " << status.ToString() << std::endl;
+  ASSERT_FALSE(status.IsOk()) << "Oversized JPEG (17000x17000) should have been rejected but was accepted.";
+}