From 33aa865825da292ad07ed6e6740f60900e93f2f3 Mon Sep 17 00:00:00 2001 From: Ilia Sokolov Date: Sat, 25 Jul 2026 21:16:44 +0200 Subject: [PATCH 1/2] Return errors for malformed JPEG input --- operators/vision/image_decoder.hpp | 227 +++++++++++++++++++---------- test/pp_api_test/test_imgcodec.cc | 64 +++++++- 2 files changed, 209 insertions(+), 82 deletions(-) diff --git a/operators/vision/image_decoder.hpp b/operators/vision/image_decoder.hpp index cf2beb25a..35e6bd1f4 100644 --- a/operators/vision/image_decoder.hpp +++ b/operators/vision/image_decoder.hpp @@ -3,7 +3,9 @@ #pragma once +#include #include +#include #include "png.h" #if _WIN32 @@ -11,6 +13,7 @@ #include #endif #include "jpeglib.h" +#include "jerror.h" #include "op_def_struct.h" #include "ext_status.h" @@ -23,6 +26,91 @@ static constexpr uint64_t kMaxPixelCount = 100'000'000; // 100 megapixels struct DecodeImage { OrtxStatus OnInit() { return {}; } + struct JpegErrorManager { + jpeg_error_mgr base; + jmp_buf jump_buffer; + char message[JMSG_LENGTH_MAX]{}; + + static void ErrorExit(j_common_ptr cinfo) { + auto* error = reinterpret_cast(cinfo->err); + (*cinfo->err->format_message)(cinfo, error->message); + longjmp(error->jump_buffer, 1); + } + }; + + class JMemorySourceManager : public jpeg_source_mgr { + public: + JMemorySourceManager(const uint8_t* encoded_image_data, const int64_t encoded_image_data_len) { + next_input_byte = reinterpret_cast(encoded_image_data); + bytes_in_buffer = static_cast(encoded_image_data_len); + init_source = &JMemorySourceManager::initSource; + fill_input_buffer = &JMemorySourceManager::fillInputBuffer; + skip_input_data = &JMemorySourceManager::skipInputData; + resync_to_restart = jpeg_resync_to_restart; + term_source = &JMemorySourceManager::termSource; + } + + static void initSource(j_decompress_ptr cinfo) { + // No initialization needed + } + + // This is an in-memory, non-suspending source. Asking for more bytes means + // the JPEG is truncated, so report a fatal libjpeg error immediately. + static boolean fillInputBuffer(j_decompress_ptr cinfo) { + auto* srcMgr = reinterpret_cast(cinfo->src); + srcMgr->extError = kOrtxErrorCorruptData; + ERREXIT(cinfo, JERR_INPUT_EOF); + return FALSE; + } + + static void skipInputData(j_decompress_ptr cinfo, long num_bytes) { + auto* srcMgr = reinterpret_cast(cinfo->src); + if (num_bytes > 0) { + size_t bytes_to_skip = static_cast(num_bytes); + if (bytes_to_skip > srcMgr->bytes_in_buffer) { + srcMgr->next_input_byte += srcMgr->bytes_in_buffer; + srcMgr->bytes_in_buffer = 0; + srcMgr->extError = kOrtxErrorCorruptData; + ERREXIT(cinfo, JERR_INPUT_EOF); + return; + } + srcMgr->next_input_byte += bytes_to_skip; + srcMgr->bytes_in_buffer -= bytes_to_skip; + } + } + + static void termSource(j_decompress_ptr cinfo) { + // No cleanup needed + } + + extError_t extError{kOrtxOK}; + }; + + // libjpeg mutates its state after setjmp. Keep that mutable state on the + // heap so automatic variables do not become indeterminate after longjmp. + struct JpegDecodeState { + JpegDecodeState(const uint8_t* data, int64_t size) : source(data, size) {} + + jpeg_decompress_struct cinfo{}; + JpegErrorManager error{}; + JMemorySourceManager source; + std::vector output_dimensions; + }; + + static void DestroyJpegState(JpegDecodeState* state) { + // jpeg_create_decompress initializes mem to null before doing work, so this + // also cleans up a partially created decompressor. + if (state->cinfo.mem != nullptr) { + jpeg_destroy_decompress(&state->cinfo); + } + delete state; + } + + static OrtxStatus JpegFailure(JpegDecodeState* state, const std::string& message) { + DestroyJpegState(state); + return {kOrtxErrorCorruptData, message}; + } + OrtxStatus DecodePNG(const uint8_t* encoded_image_data, const int64_t encoded_image_data_len, ortc::Tensor& output) const { // Decode the PNG image @@ -131,125 +219,102 @@ struct DecodeImage { if (png_sig_cmp(encoded_image_data, 0, 8) == 0) { return DecodePNG(encoded_image_data, encoded_image_data_len, output); } else { - // Initialize JPEG decompression object - jpeg_decompress_struct cinfo; - jpeg_error_mgr jerr; - cinfo.err = jpeg_std_error(&jerr); - jpeg_create_decompress(&cinfo); + auto* const state = + new JpegDecodeState(encoded_image_data, encoded_image_data_len); + state->cinfo.err = jpeg_std_error(&state->error.base); + state->error.base.error_exit = &JpegErrorManager::ErrorExit; + + if (setjmp(state->error.jump_buffer)) { + const char* diagnostic = + state->error.message[0] == '\0' + ? "unknown libjpeg error" + : state->error.message; + return JpegFailure( + state, + std::string("[ImageDecoder]: Failed to decode JPEG image: ") + diagnostic); + } - // Set up the custom memory source manager - JMemorySourceManager srcManager(encoded_image_data, encoded_image_data_len); - cinfo.src = &srcManager; + jpeg_create_decompress(&state->cinfo); + state->cinfo.src = &state->source; // Read the JPEG header to get image info - jpeg_read_header(&cinfo, TRUE); + if (jpeg_read_header(&state->cinfo, TRUE) != JPEG_HEADER_OK) { + return JpegFailure( + state, "[ImageDecoder]: Failed to decode JPEG image header."); + } // Security: explicitly reject CMYK/YCCK color spaces before decompression. // These have 4 channels and downstream code assumes 3 channels (CVE-class: CWE-122). - if (cinfo.jpeg_color_space == JCS_CMYK || cinfo.jpeg_color_space == JCS_YCCK) { - jpeg_destroy_decompress(&cinfo); + if (state->cinfo.jpeg_color_space == JCS_CMYK || + state->cinfo.jpeg_color_space == JCS_YCCK) { + DestroyJpegState(state); return {kOrtxErrorInvalidArgument, "[ImageDecoder]: Unsupported JPEG color space (CMYK/YCCK). Only RGB and grayscale are supported."}; } // Force RGB output to ensure consistent 3-channel output regardless of input // (e.g., grayscale JPEGs are expanded to RGB). - cinfo.out_color_space = JCS_RGB; + state->cinfo.out_color_space = JCS_RGB; // Start decompression - jpeg_start_decompress(&cinfo); + if (!jpeg_start_decompress(&state->cinfo)) { + return JpegFailure( + state, "[ImageDecoder]: Failed to start JPEG decompression."); + } // Dimension limit to prevent decompression bombs - if (cinfo.output_width > kMaxImageDimension || - cinfo.output_height > kMaxImageDimension || - static_cast(cinfo.output_width) * cinfo.output_height > kMaxPixelCount) { - jpeg_destroy_decompress(&cinfo); + if (state->cinfo.output_width > kMaxImageDimension || + state->cinfo.output_height > kMaxImageDimension || + static_cast(state->cinfo.output_width) * + state->cinfo.output_height > + kMaxPixelCount) { + DestroyJpegState(state); 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); + if (state->cinfo.output_components != 3) { + const int output_components = state->cinfo.output_components; + DestroyJpegState(state); return {kOrtxErrorInvalidArgument, "[ImageDecoder]: Unexpected JPEG output channels. Expected 3 (RGB), got " + - std::to_string(cinfo.output_components) + "."}; + std::to_string(output_components) + "."}; } // Allocate memory for the image - std::vector output_dimensions{cinfo.output_height, cinfo.output_width, cinfo.output_components}; - uint8_t* imageBuffer = output.Allocate(output_dimensions); + state->output_dimensions = { + state->cinfo.output_height, + state->cinfo.output_width, + state->cinfo.output_components}; + uint8_t* imageBuffer = output.Allocate(state->output_dimensions); // Read the image data - int row_stride = cinfo.output_width * cinfo.output_components; - while (cinfo.output_scanline < cinfo.output_height) { - uint8_t* row_ptr = imageBuffer + (cinfo.output_scanline * row_stride); - jpeg_read_scanlines(&cinfo, &row_ptr, 1); - if (srcManager.extError != kOrtxOK) { + int row_stride = + state->cinfo.output_width * state->cinfo.output_components; + while (state->cinfo.output_scanline < state->cinfo.output_height) { + uint8_t* row_ptr = + imageBuffer + (state->cinfo.output_scanline * row_stride); + if (jpeg_read_scanlines(&state->cinfo, &row_ptr, 1) != 1) { + state->source.extError = kOrtxErrorCorruptData; break; } } - if (srcManager.extError != kOrtxOK) { - jpeg_destroy_decompress(&cinfo); - return {kOrtxErrorInternal, "[ImageDecoder]: Failed to decode JPEG image."}; + if (state->source.extError != kOrtxOK) { + return JpegFailure( + state, "[ImageDecoder]: Failed to decode JPEG image."); } // Finish decompression - jpeg_finish_decompress(&cinfo); - jpeg_destroy_decompress(&cinfo); + if (!jpeg_finish_decompress(&state->cinfo)) { + return JpegFailure( + state, "[ImageDecoder]: Failed to finish JPEG decompression."); + } + DestroyJpegState(state); } return {}; } - - class JMemorySourceManager : public jpeg_source_mgr { - public: - // Constructor - JMemorySourceManager(const uint8_t* encoded_image_data, const int64_t encoded_image_data_len) { - // Initialize source fields - next_input_byte = reinterpret_cast(encoded_image_data); - bytes_in_buffer = static_cast(encoded_image_data_len); - init_source = &JMemorySourceManager::initSource; - fill_input_buffer = &JMemorySourceManager::fillInputBuffer; - skip_input_data = &JMemorySourceManager::skipInputData; - resync_to_restart = jpeg_resync_to_restart; - term_source = &JMemorySourceManager::termSource; - } - - // Initialize source (no-op) - static void initSource(j_decompress_ptr cinfo) { - // No initialization needed - } - - // Fill input buffer (not used here, always return FALSE) - static boolean fillInputBuffer(j_decompress_ptr cinfo) { - return FALSE; // Buffer is managed manually - } - - // Skip input data - static void skipInputData(j_decompress_ptr cinfo, long num_bytes) { - JMemorySourceManager* srcMgr = reinterpret_cast(cinfo->src); - if (num_bytes > 0) { - size_t bytes_to_skip = static_cast(num_bytes); - while (bytes_to_skip > srcMgr->bytes_in_buffer) { - bytes_to_skip -= srcMgr->bytes_in_buffer; - if (srcMgr->fillInputBuffer(cinfo)) { - // Error: buffer ran out - srcMgr->extError = kOrtxErrorCorruptData; - } - } - srcMgr->next_input_byte += bytes_to_skip; - srcMgr->bytes_in_buffer -= bytes_to_skip; - } - } - - // Terminate source (no-op) - static void termSource(j_decompress_ptr cinfo) { - // No cleanup needed - } - - extError_t extError{kOrtxOK}; // Error handler - }; }; } // namespace ort_extensions::internal diff --git a/test/pp_api_test/test_imgcodec.cc b/test/pp_api_test/test_imgcodec.cc index 05ebeecf2..51fc65f65 100644 --- a/test/pp_api_test/test_imgcodec.cc +++ b/test/pp_api_test/test_imgcodec.cc @@ -159,6 +159,68 @@ TEST(ImageDecoderTest, TestJpegEncoderDecoder) { ASSERT_NE(encodeOutputBuffer, nullptr); } +TEST(ImageDecoderTest, InvalidJpegReturnsErrorWithoutTerminatingProcess) { + ort_extensions::DecodeImage image_decoder; + image_decoder.Init(std::unordered_map>()); + + const std::vector> invalid_images = { + // Non-PNG bytes long enough to enter the JPEG path. + {0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09}, + // JPEG SOI followed by a truncated APP0 marker. + {0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x10, 0x4A, 0x46, 0x49, 0x46}, + // Marker length requests a skip beyond the supplied buffer. + {0xFF, 0xD8, 0xFF, 0xE1, 0x7F, 0xFF, 0x00, 0x00, 0xFF, 0xD9}, + }; + + for (auto encoded : invalid_images) { + ortc::Tensor input( + {static_cast(encoded.size())}, + encoded.data()); + ortc::Tensor output{&CppAllocator::Instance()}; + auto status = image_decoder.Compute(input, output); + + ASSERT_FALSE(status.IsOk()); +#if !OCOS_ENABLE_VENDOR_IMAGE_CODECS || (!defined(_WIN32) && !defined(__APPLE__)) + EXPECT_EQ(status.Code(), kOrtxErrorCorruptData); + EXPECT_NE(status.ToString().find("JPEG"), std::string::npos); +#endif + } +} + +TEST(ImageDecoderTest, ValidJpegStillDecodesAfterInvalidInput) { + ort_extensions::DecodeImage image_decoder; + image_decoder.Init(std::unordered_map>()); + + std::ifstream jpeg_file("data/processor/australia.jpg", std::ios::binary); + ASSERT_TRUE(jpeg_file.is_open()); + jpeg_file.seekg(0, std::ios::end); + std::vector valid(static_cast(jpeg_file.tellg())); + jpeg_file.seekg(0, std::ios::beg); + jpeg_file.read(reinterpret_cast(valid.data()), valid.size()); + + // Preserve the valid header and some entropy-coded scan data, then truncate. + // This exercises recovery after output allocation, not only header parsing. + std::vector truncated(valid.begin(), valid.begin() + valid.size() / 2); + for (int attempt = 0; attempt < 3; ++attempt) { + ortc::Tensor invalid_tensor( + {static_cast(truncated.size())}, truncated.data()); + ortc::Tensor invalid_output{&CppAllocator::Instance()}; + auto invalid_status = image_decoder.Compute(invalid_tensor, invalid_output); + ASSERT_FALSE(invalid_status.IsOk()); +#if !OCOS_ENABLE_VENDOR_IMAGE_CODECS || (!defined(_WIN32) && !defined(__APPLE__)) + EXPECT_EQ(invalid_status.Code(), kOrtxErrorCorruptData); + EXPECT_NE(invalid_status.ToString().find("JPEG"), std::string::npos); +#endif + } + + ortc::Tensor valid_tensor( + {static_cast(valid.size())}, valid.data()); + ortc::Tensor valid_output{&CppAllocator::Instance()}; + auto status = image_decoder.Compute(valid_tensor, valid_output); + ASSERT_TRUE(status.IsOk()) << status.ToString(); + EXPECT_EQ(valid_output.Shape(), std::vector({876, 1300, 3})); +} + #if OCOS_ENABLE_VENDOR_IMAGE_CODECS #if defined(_WIN32) || defined(__APPLE__) TEST(ImageDecoderTest, TestTiffDecoder) { @@ -336,4 +398,4 @@ TEST(ImageDecoderTest, TestJpegOversizeDimensionsRejected) { // 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."; -} \ No newline at end of file +} From afd1b3fc90d87cc3880f9a5e0208160e6cac4081 Mon Sep 17 00:00:00 2001 From: Ilia Sokolov Date: Thu, 20 Aug 2026 18:18:33 +0200 Subject: [PATCH 2/2] Address JPEG decoder review and CI failures --- operators/vision/image_decoder.hpp | 17 ++++++++++------- test/pp_api_test/test_imgcodec.cc | 10 ++++++---- 2 files changed, 16 insertions(+), 11 deletions(-) diff --git a/operators/vision/image_decoder.hpp b/operators/vision/image_decoder.hpp index 35e6bd1f4..4e976981d 100644 --- a/operators/vision/image_decoder.hpp +++ b/operators/vision/image_decoder.hpp @@ -5,6 +5,7 @@ #include #include +#include #include #include "png.h" @@ -26,13 +27,12 @@ static constexpr uint64_t kMaxPixelCount = 100'000'000; // 100 megapixels struct DecodeImage { OrtxStatus OnInit() { return {}; } - struct JpegErrorManager { - jpeg_error_mgr base; + struct JpegErrorManager : jpeg_error_mgr { jmp_buf jump_buffer; char message[JMSG_LENGTH_MAX]{}; static void ErrorExit(j_common_ptr cinfo) { - auto* error = reinterpret_cast(cinfo->err); + auto* error = static_cast(cinfo->err); (*cinfo->err->format_message)(cinfo, error->message); longjmp(error->jump_buffer, 1); } @@ -219,10 +219,13 @@ struct DecodeImage { if (png_sig_cmp(encoded_image_data, 0, 8) == 0) { return DecodePNG(encoded_image_data, encoded_image_data_len, output); } else { - auto* const state = - new JpegDecodeState(encoded_image_data, encoded_image_data_len); - state->cinfo.err = jpeg_std_error(&state->error.base); - state->error.base.error_exit = &JpegErrorManager::ErrorExit; + auto* const state = new (std::nothrow) JpegDecodeState(encoded_image_data, encoded_image_data_len); + if (state == nullptr) { + return {kOrtxErrorOutOfMemory, "[ImageDecoder]: Failed to allocate JPEG decoder state."}; + } + + state->cinfo.err = jpeg_std_error(&state->error); + state->error.error_exit = &JpegErrorManager::ErrorExit; if (setjmp(state->error.jump_buffer)) { const char* diagnostic = diff --git a/test/pp_api_test/test_imgcodec.cc b/test/pp_api_test/test_imgcodec.cc index 51fc65f65..92032c933 100644 --- a/test/pp_api_test/test_imgcodec.cc +++ b/test/pp_api_test/test_imgcodec.cc @@ -163,7 +163,7 @@ TEST(ImageDecoderTest, InvalidJpegReturnsErrorWithoutTerminatingProcess) { ort_extensions::DecodeImage image_decoder; image_decoder.Init(std::unordered_map>()); - const std::vector> invalid_images = { + std::vector> invalid_images = { // Non-PNG bytes long enough to enter the JPEG path. {0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09}, // JPEG SOI followed by a truncated APP0 marker. @@ -172,7 +172,7 @@ TEST(ImageDecoderTest, InvalidJpegReturnsErrorWithoutTerminatingProcess) { {0xFF, 0xD8, 0xFF, 0xE1, 0x7F, 0xFF, 0x00, 0x00, 0xFF, 0xD9}, }; - for (auto encoded : invalid_images) { + for (auto& encoded : invalid_images) { ortc::Tensor input( {static_cast(encoded.size())}, encoded.data()); @@ -187,6 +187,9 @@ TEST(ImageDecoderTest, InvalidJpegReturnsErrorWithoutTerminatingProcess) { } } +// Windows and macOS vendor codecs may accept truncated JPEGs. This test +// specifically verifies recovery from libjpeg's fatal error callback. +#if !OCOS_ENABLE_VENDOR_IMAGE_CODECS || (!defined(_WIN32) && !defined(__APPLE__)) TEST(ImageDecoderTest, ValidJpegStillDecodesAfterInvalidInput) { ort_extensions::DecodeImage image_decoder; image_decoder.Init(std::unordered_map>()); @@ -207,10 +210,8 @@ TEST(ImageDecoderTest, ValidJpegStillDecodesAfterInvalidInput) { ortc::Tensor invalid_output{&CppAllocator::Instance()}; auto invalid_status = image_decoder.Compute(invalid_tensor, invalid_output); ASSERT_FALSE(invalid_status.IsOk()); -#if !OCOS_ENABLE_VENDOR_IMAGE_CODECS || (!defined(_WIN32) && !defined(__APPLE__)) EXPECT_EQ(invalid_status.Code(), kOrtxErrorCorruptData); EXPECT_NE(invalid_status.ToString().find("JPEG"), std::string::npos); -#endif } ortc::Tensor valid_tensor( @@ -220,6 +221,7 @@ TEST(ImageDecoderTest, ValidJpegStillDecodesAfterInvalidInput) { ASSERT_TRUE(status.IsOk()) << status.ToString(); EXPECT_EQ(valid_output.Shape(), std::vector({876, 1300, 3})); } +#endif #if OCOS_ENABLE_VENDOR_IMAGE_CODECS #if defined(_WIN32) || defined(__APPLE__)