jpeglossless bug fixes - SOF3 P in [9, 15] and DHT before SOF3

extensions/nvjpeg/lossless_decoder.cpp

@@ -127,6 +127,39 @@ nvimgcodecStatus_t DecoderImpl::getMetadata(const nvimgcodecCodeStreamDesc_t* co
     return NVIMGCODEC_STATUS_IMPLEMENTATION_UNSUPPORTED;
 }
 
+// Returns true if a DHT marker (0xC4) appears before the first SOF3 marker (0xC3)
+// in the raw JPEG bitstream; nvjpegDecodeBatched silently zero-fills for such streams
+static bool has_dht_before_sof3(const unsigned char* data, size_t size)
+{
+    if (size < 4)
+        return false;
+    size_t i = 2; // skip SOI (0xFF 0xD8)
+    while (i + 1 < size) {
+        if (data[i] != 0xFF) {
+            ++i;
+            continue;
+        }
+        uint8_t code = data[i + 1];
+        // consume fill bytes (0xFF padding between marker bytes is legal per ITU-T T.81 B.1.1.2)
+        while (code == 0xFF && i + 2 < size) {
+            ++i;
+            code = data[i + 1];
+        }
+        if (code == 0xC3) return false; // SOF3 found first — stream order is OK
+        if (code == 0xC4) return true;  // DHT found before SOF3 — nvjpeg will zero-fill
+        if (code == 0xD9) break;        // EOI — malformed, no SOF3 found
+        if (code == 0xDA) break;        // SOS — malformed, no SOF3 found
+        // advance past this segment: 2 marker bytes + segment length (length field includes itself)
+        if (i + 3 >= size)
+            break;
+        uint16_t seg_len = (static_cast<uint16_t>(data[i + 2]) << 8) | data[i + 3];
+        if (seg_len < 2)
+            break; // malformed length
+        i += 2 + seg_len;
+    }
+    return false;
+}
+
 nvimgcodecProcessingStatus_t DecoderImpl::canDecode(const nvimgcodecImageDesc_t* image, const nvimgcodecCodeStreamDesc_t* code_stream,
     const nvimgcodecDecodeParams_t* params, int thread_idx)
 {
@@ -168,8 +201,18 @@ nvimgcodecProcessingStatus_t DecoderImpl::canDecode(const nvimgcodecImageDesc_t*
 
         if (image_info.plane_info[0].sample_type != NVIMGCODEC_SAMPLE_DATA_TYPE_UINT16)
             status |= NVIMGCODEC_PROCESSING_STATUS_SAMPLE_TYPE_UNSUPPORTED;
-
-        XM_CHECK_NULL(code_stream); 
+
+        // nvjpegDecodeBatched only correctly handles P=16 for UINT16 lossless streams
+        // For P in range [9,15], it silently zero-fills the output buffer
+        uint8_t encoded_precision = cs_image_info.plane_info[0].precision;
+        if (encoded_precision > 8 && encoded_precision < 16) {
+            NVIMGCODEC_LOG_INFO(framework_, plugin_id_,
+                "JPEG Lossless SOF3 precision=" << static_cast<int>(encoded_precision)
+                << " is not supported (only P=16 is handled correctly by nvjpeg lossless decoder)");
+            status |= NVIMGCODEC_PROCESSING_STATUS_SAMPLE_TYPE_UNSUPPORTED;
+        }
+
+        XM_CHECK_NULL(code_stream);
         nvimgcodecCodeStreamInfo_t codestream_info{NVIMGCODEC_STRUCTURE_TYPE_CODE_STREAM_INFO, sizeof(nvimgcodecCodeStreamInfo_t), nullptr};
         ret = code_stream->getCodeStreamInfo(code_stream->instance, &codestream_info);
         if (ret != NVIMGCODEC_STATUS_SUCCESS)
@@ -188,7 +231,6 @@ nvimgcodecProcessingStatus_t DecoderImpl::canDecode(const nvimgcodecImageDesc_t*
         if (status != NVIMGCODEC_PROCESSING_STATUS_SUCCESS)
             return status;
 
-
         auto* io_stream = code_stream->io_stream;
         XM_CHECK_NULL(io_stream);
 
@@ -203,6 +245,14 @@ nvimgcodecProcessingStatus_t DecoderImpl::canDecode(const nvimgcodecImageDesc_t*
         assert(encoded_stream_data != nullptr);
         assert(encoded_stream_data_size > 0);
 
+        // DHT before SOF3 check
+        if (has_dht_before_sof3(static_cast<const unsigned char*>(encoded_stream_data), encoded_stream_data_size)) {
+            NVIMGCODEC_LOG_INFO(framework_, plugin_id_,
+                "JPEG Lossless has DHT segment before SOF3 (non-standard ordering); "
+                "not supported by nvjpeg lossless decoder, falling back");
+            return NVIMGCODEC_PROCESSING_STATUS_CODEC_UNSUPPORTED;
+        }
+
         XM_CHECK_NVJPEG(nvjpegJpegStreamParse(
             handle_, static_cast<const unsigned char*>(encoded_stream_data), encoded_stream_data_size, 0, 0, nvjpeg_stream));
         int isSupported = -1;

test/extensions/nvjpeg_ext_lossless_decoder_test.cpp

@@ -101,7 +101,40 @@ class NvJpegExtLosslessDecoderTestSingleImage :
         NvJpegExtDecoderTestBase::TearDown();
     }
 
+};
+
+// 9-15 bit rejection tests
+class NvJpegExtLosslessDecoderRejectionTest :
+    public NvJpegExtDecoderTestBase,
+    public NvJpegTestBase,
+    public TestWithParam<std::tuple<const char*, nvimgcodecColorSpec_t, nvimgcodecSampleFormat_t, nvimgcodecChromaSubsampling_t, nvimgcodecProcessingStatus_t>>
+{
+  public:
+    using NvJpegTestBase::SetUpTestSuite;
+    using NvJpegTestBase::TearDownTestSuite;
+    virtual ~NvJpegExtLosslessDecoderRejectionTest() = default;
+
+  protected:
+    void SetUp() override
+    {
+        image_file_ = std::get<0>(GetParam());
+        color_spec_ = std::get<1>(GetParam());
+        sample_format_ = std::get<2>(GetParam());
+        chroma_subsampling_ = std::get<3>(GetParam());
+        expected_status_ = std::get<4>(GetParam());
+        NvJpegExtDecoderTestBase::SetUp();
+        NvJpegTestBase::SetUp();
+    }
+
     nvimgcodecColorSpec_t output_color_spec_ = NVIMGCODEC_COLORSPEC_UNCHANGED;
+    nvimgcodecProcessingStatus_t expected_status_ = NVIMGCODEC_PROCESSING_STATUS_UNKNOWN;
+
+    virtual void TearDown()
+    {
+        NvJpegTestBase::TearDown();
+        NvJpegExtDecoderTestBase::TearDown();
+    }
+
 };
 
 TEST_P(NvJpegExtLosslessDecoderTestSingleImage, LosslessJpegValidFormatAndParameters)
@@ -131,10 +164,63 @@ TEST_P(NvJpegExtLosslessDecoderTestSingleImage, LosslessJpegValidFormatAndParame
     ASSERT_EQ(NVIMGCODEC_PROCESSING_STATUS_SUCCESS, status);
 }
 
-static const char* css_lossless_filenames[] = {"/jpeg/lossless/cat-1245673_640_grayscale_16bit.jpg",
-    "/jpeg/lossless/cat-3449999_640_grayscale_12bit.jpg",
+TEST_P(NvJpegExtLosslessDecoderRejectionTest, LosslessJpegRejected)
+{
+#if defined(_WIN32) || defined(_WIN64)
+    if (CC_major < 7) {
+        GTEST_SKIP() << "On Windows, nvJPEG lossless requires sm_70 or higher to work.";
+    }
+#endif
+
+    LoadImageFromFilename(instance_, in_code_stream_, resources_dir + image_file_);
+    ASSERT_EQ(NVIMGCODEC_STATUS_SUCCESS, nvimgcodecCodeStreamGetImageInfo(in_code_stream_, &image_info_));
+    image_info_.plane_info[0].sample_type = NVIMGCODEC_SAMPLE_DATA_TYPE_UINT16;
+    image_info_.plane_info[0].precision = 16;
+    PrepareImageForFormat();
+
+    nvimgcodecImageInfo_t out_image_info(image_info_);
+    ASSERT_EQ(NVIMGCODEC_STATUS_SUCCESS, nvimgcodecImageCreate(instance_, &out_image_, &out_image_info));
+    streams_.push_back(in_code_stream_);
+    images_.push_back(out_image_);
+    ASSERT_EQ(NVIMGCODEC_STATUS_SUCCESS, nvimgcodecDecoderDecode(decoder_, streams_.data(), images_.data(), 1, &params_, &future_));
+    ASSERT_EQ(NVIMGCODEC_STATUS_SUCCESS, nvimgcodecFutureWaitForAll(future_));
+    cudaDeviceSynchronize();
+    nvimgcodecProcessingStatus_t status;
+    size_t status_size;
+    ASSERT_EQ(NVIMGCODEC_STATUS_SUCCESS, nvimgcodecFutureGetProcessingStatus(future_, &status, &status_size));
+    ASSERT_EQ(expected_status_, status)
+        << "file=" << image_file_
+        << " status=" << static_cast<int>(status)
+        << " expected=" << static_cast<int>(expected_status_);
+
+}
+
+// 8 and 16 bit
+static const char* css_lossless_filenames[] = {
+    "/jpeg/lossless/cat-1245673_640_grayscale_16bit.jpg",
     "/jpeg/lossless/cat-3449999_640_grayscale_16bit.jpg",
-    "/jpeg/lossless/cat-3449999_640_grayscale_8bit.jpg"};
+    "/jpeg/lossless/cat-3449999_640_grayscale_8bit.jpg"
+};
+
+// 9-15 bit
+static const char* css_lossless_9_15bit_filenames[] = {
+    "/jpeg/lossless/cat-3449999_640_grayscale_12bit.jpg"
+};
+
+// 16 bit (DICOM-derived)
+static const char* dicom_lossless_p16_filenames[] = {
+    "/jpeg/lossless/dicom/bad_sequence_19d910fdeb_frame_0000.jpg"
+};
+
+// DHT before SOF3 (DICOM-derived)
+static const char* dicom_lossless_dht_before_sof3_filenames[] = {
+    "/jpeg/lossless/dicom/CT_c79833361c_frame_0000.jpg"
+};
+
+// 9-15 bit (DICOM-derived)
+static const char* dicom_lossless_9_15bit_filenames[] = {
+    "/jpeg/lossless/dicom/MR_c032f52f64_frame_0000.jpg"
+};
 
 // clang-format off
 INSTANTIATE_TEST_SUITE_P(NVJPEG_LOSSLESS_DECODE_VARIOUS_CHROMA_WITH_VALID_SRGB_OUTPUT_FORMATS,
@@ -144,5 +230,36 @@ INSTANTIATE_TEST_SUITE_P(NVJPEG_LOSSLESS_DECODE_VARIOUS_CHROMA_WITH_VALID_SRGB_O
         Values(NVIMGCODEC_SAMPLEFORMAT_P_Y, NVIMGCODEC_SAMPLEFORMAT_I_UNCHANGED),
         Values(NVIMGCODEC_SAMPLING_NONE)));
 
+INSTANTIATE_TEST_SUITE_P(NVJPEG_LOSSLESS_DECODE_9_15BIT,
+    NvJpegExtLosslessDecoderRejectionTest,
+    Combine(::testing::ValuesIn(css_lossless_9_15bit_filenames),
+        Values(NVIMGCODEC_COLORSPEC_SRGB),
+        Values(NVIMGCODEC_SAMPLEFORMAT_P_Y, NVIMGCODEC_SAMPLEFORMAT_I_UNCHANGED),
+        Values(NVIMGCODEC_SAMPLING_NONE),
+        Values(NVIMGCODEC_PROCESSING_STATUS_SAMPLE_TYPE_UNSUPPORTED)));
+
+INSTANTIATE_TEST_SUITE_P(NVJPEG_LOSSLESS_DECODE_DICOM_P16,
+    NvJpegExtLosslessDecoderTestSingleImage,
+    Combine(::testing::ValuesIn(dicom_lossless_p16_filenames),
+        Values(NVIMGCODEC_COLORSPEC_SRGB),
+        Values(NVIMGCODEC_SAMPLEFORMAT_P_Y, NVIMGCODEC_SAMPLEFORMAT_I_UNCHANGED),
+        Values(NVIMGCODEC_SAMPLING_NONE)));
+
+INSTANTIATE_TEST_SUITE_P(NVJPEG_LOSSLESS_DECODE_DICOM_DHT_BEFORE_SOF3,
+    NvJpegExtLosslessDecoderRejectionTest,
+    Combine(::testing::ValuesIn(dicom_lossless_dht_before_sof3_filenames),
+        Values(NVIMGCODEC_COLORSPEC_SRGB),
+        Values(NVIMGCODEC_SAMPLEFORMAT_P_Y, NVIMGCODEC_SAMPLEFORMAT_I_UNCHANGED),
+        Values(NVIMGCODEC_SAMPLING_NONE),
+        Values(NVIMGCODEC_PROCESSING_STATUS_CODEC_UNSUPPORTED)));
+
+INSTANTIATE_TEST_SUITE_P(NVJPEG_LOSSLESS_DECODE_DICOM_9_15BIT,
+    NvJpegExtLosslessDecoderRejectionTest,
+    Combine(::testing::ValuesIn(dicom_lossless_9_15bit_filenames),
+        Values(NVIMGCODEC_COLORSPEC_SRGB),
+        Values(NVIMGCODEC_SAMPLEFORMAT_P_Y, NVIMGCODEC_SAMPLEFORMAT_I_UNCHANGED),
+        Values(NVIMGCODEC_SAMPLING_NONE),
+        Values(NVIMGCODEC_PROCESSING_STATUS_SAMPLE_TYPE_UNSUPPORTED)));
+
 // clang-format on
 }} // namespace nvimgcodec::test