1 files changed, 61 insertions, 12 deletions

diff --git a/lib/extras/dec/jpg.cc b/lib/extras/dec/jpg.cc
index 6b92f4a..3c8a4bc 100644
--- a/lib/extras/dec/jpg.cc
+++ b/lib/extras/dec/jpg.cc
@@ -5,8 +5,10 @@
 
 #include "lib/extras/dec/jpg.h"
 
+#if JPEGXL_ENABLE_JPEG
 #include <jpeglib.h>
 #include <setjmp.h>
+#endif
 #include <stdint.h>
 
 #include <algorithm>
@@ -14,12 +16,14 @@
 #include <utility>
 #include <vector>
 
+#include "lib/extras/size_constraints.h"
 #include "lib/jxl/base/status.h"
 #include "lib/jxl/sanitizers.h"
 
 namespace jxl {
 namespace extras {
 
+#if JPEGXL_ENABLE_JPEG
 namespace {
 
 constexpr unsigned char kICCSignature[12] = {
@@ -160,12 +164,35 @@ void MyOutputMessage(j_common_ptr cinfo) {
 #endif
 }
 
+void UnmapColors(uint8_t* row, size_t xsize, int components,
+                 JSAMPARRAY colormap, size_t num_colors) {
+  JXL_CHECK(colormap != nullptr);
+  std::vector<uint8_t> tmp(xsize * components);
+  for (size_t x = 0; x < xsize; ++x) {
+    JXL_CHECK(row[x] < num_colors);
+    for (int c = 0; c < components; ++c) {
+      tmp[x * components + c] = colormap[c][row[x]];
+    }
+  }
+  memcpy(row, tmp.data(), tmp.size());
+}
+
 }  // namespace
+#endif
+
+bool CanDecodeJPG() {
+#if JPEGXL_ENABLE_JPEG
+  return true;
+#else
+  return false;
+#endif
+}
 
 Status DecodeImageJPG(const Span<const uint8_t> bytes,
-                      const ColorHints& color_hints,
-                      const SizeConstraints& constraints,
-                      PackedPixelFile* ppf) {
+                      const ColorHints& color_hints, PackedPixelFile* ppf,
+                      const SizeConstraints* constraints,
+                      const JPGDecompressParams* dparams) {
+#if JPEGXL_ENABLE_JPEG
   // Don't do anything for non-JPEG files (no need to report an error)
   if (!IsJPG(bytes)) return false;
 
@@ -176,10 +203,7 @@ Status DecodeImageJPG(const Span<const uint8_t> bytes,
   std::unique_ptr<JSAMPLE[]> row;
 
   const auto try_catch_block = [&]() -> bool {
-    jpeg_decompress_struct cinfo;
-    // cinfo is initialized by libjpeg, which we are not instrumenting with
-    // msan, therefore we need to initialize cinfo here.
-    msan::UnpoisonMemory(&cinfo, sizeof(cinfo));
+    jpeg_decompress_struct cinfo = {};
     // Setup error handling in jpeg library so we can deal with broken jpegs in
     // the fuzzer.
     jpeg_error_mgr jerr;
@@ -207,8 +231,7 @@ Status DecodeImageJPG(const Span<const uint8_t> bytes,
     if (read_header_result == JPEG_SUSPENDED) {
       return failure("truncated JPEG input");
     }
-    if (!VerifyDimensions(&constraints, cinfo.image_width,
-                          cinfo.image_height)) {
+    if (!VerifyDimensions(constraints, cinfo.image_width, cinfo.image_height)) {
       return failure("image too big");
     }
     // Might cause CPU-zip bomb.
@@ -252,12 +275,21 @@ Status DecodeImageJPG(const Span<const uint8_t> bytes,
     ppf->info.num_color_channels = nbcomp;
     ppf->info.orientation = JXL_ORIENT_IDENTITY;
 
+    if (dparams && dparams->num_colors > 0) {
+      cinfo.quantize_colors = TRUE;
+      cinfo.desired_number_of_colors = dparams->num_colors;
+      cinfo.two_pass_quantize = dparams->two_pass_quant;
+      cinfo.dither_mode = (J_DITHER_MODE)dparams->dither_mode;
+    }
+
     jpeg_start_decompress(&cinfo);
-    JXL_ASSERT(cinfo.output_components == nbcomp);
+    JXL_ASSERT(cinfo.out_color_components == nbcomp);
+    JxlDataType data_type =
+        ppf->info.bits_per_sample <= 8 ? JXL_TYPE_UINT8 : JXL_TYPE_UINT16;
 
     const JxlPixelFormat format{
         /*num_channels=*/static_cast<uint32_t>(nbcomp),
-        /*data_type=*/BITS_IN_JSAMPLE == 8 ? JXL_TYPE_UINT8 : JXL_TYPE_UINT16,
+        data_type,
         /*endianness=*/JXL_NATIVE_ENDIAN,
         /*align=*/0,
     };
@@ -265,9 +297,19 @@ Status DecodeImageJPG(const Span<const uint8_t> bytes,
     // Allocates the frame buffer.
     ppf->frames.emplace_back(cinfo.image_width, cinfo.image_height, format);
     const auto& frame = ppf->frames.back();
-    JXL_ASSERT(sizeof(JSAMPLE) * cinfo.output_components * cinfo.image_width <=
+    JXL_ASSERT(sizeof(JSAMPLE) * cinfo.out_color_components *
+                   cinfo.image_width <=
                frame.color.stride);
 
+    if (cinfo.quantize_colors) {
+      jxl::msan::UnpoisonMemory(cinfo.colormap, cinfo.out_color_components *
+                                                    sizeof(cinfo.colormap[0]));
+      for (int c = 0; c < cinfo.out_color_components; ++c) {
+        jxl::msan::UnpoisonMemory(
+            cinfo.colormap[c],
+            cinfo.actual_number_of_colors * sizeof(cinfo.colormap[c][0]));
+      }
+    }
     for (size_t y = 0; y < cinfo.image_height; ++y) {
       JSAMPROW rows[] = {reinterpret_cast<JSAMPLE*>(
           static_cast<uint8_t*>(frame.color.pixels()) +
@@ -275,6 +317,10 @@ Status DecodeImageJPG(const Span<const uint8_t> bytes,
       jpeg_read_scanlines(&cinfo, rows, 1);
       msan::UnpoisonMemory(rows[0], sizeof(JSAMPLE) * cinfo.output_components *
                                         cinfo.image_width);
+      if (dparams && dparams->num_colors > 0) {
+        UnmapColors(rows[0], cinfo.output_width, cinfo.out_color_components,
+                    cinfo.colormap, cinfo.actual_number_of_colors);
+      }
     }
 
     jpeg_finish_decompress(&cinfo);
@@ -283,6 +329,9 @@ Status DecodeImageJPG(const Span<const uint8_t> bytes,
   };
 
   return try_catch_block();
+#else
+  return false;
+#endif
 }
 
 }  // namespace extras