1 files changed, 37 insertions, 12 deletions
diff --git a/third_party/jpeg-xl/lib/jpegli/decode.cc b/third_party/jpeg-xl/lib/jpegli/decode.cc
index 9fdf68dd18..d967b787d3 100644
--- a/third_party/jpeg-xl/lib/jpegli/decode.cc
+++ b/third_party/jpeg-xl/lib/jpegli/decode.cc
@@ -54,6 +54,7 @@ void InitializeImage(j_decompress_ptr cinfo) {
   m->found_soi_ = false;
   m->found_dri_ = false;
   m->found_sof_ = false;
+  m->found_sos_ = false;
   m->found_eoi_ = false;
   m->icc_index_ = 0;
   m->icc_total_ = 0;
@@ -243,10 +244,14 @@ void PrepareForScan(j_decompress_ptr cinfo) {
   // Copy quantization tables into comp_info.
   for (int i = 0; i < cinfo->comps_in_scan; ++i) {
     jpeg_component_info* comp = cinfo->cur_comp_info[i];
+    int quant_tbl_idx = comp->quant_tbl_no;
+    JQUANT_TBL* quant_table = cinfo->quant_tbl_ptrs[quant_tbl_idx];
+    if (!quant_table) {
+      JPEGLI_ERROR("Quantization table with index %d not found", quant_tbl_idx);
+    }
     if (comp->quant_table == nullptr) {
       comp->quant_table = Allocate<JQUANT_TBL>(cinfo, 1, JPOOL_IMAGE);
-      memcpy(comp->quant_table, cinfo->quant_tbl_ptrs[comp->quant_tbl_no],
-             sizeof(JQUANT_TBL));
+      memcpy(comp->quant_table, quant_table, sizeof(JQUANT_TBL));
     }
   }
   if (cinfo->comps_in_scan == 1) {
@@ -723,16 +728,36 @@ void jpegli_calc_output_dimensions(j_decompress_ptr cinfo) {
       }
     }
   }
-  if (cinfo->out_color_space == JCS_GRAYSCALE) {
-    cinfo->out_color_components = 1;
-  } else if (cinfo->out_color_space == JCS_RGB ||
-             cinfo->out_color_space == JCS_YCbCr) {
-    cinfo->out_color_components = 3;
-  } else if (cinfo->out_color_space == JCS_CMYK ||
-             cinfo->out_color_space == JCS_YCCK) {
-    cinfo->out_color_components = 4;
-  } else {
-    cinfo->out_color_components = cinfo->num_components;
+  switch (cinfo->out_color_space) {
+    case JCS_GRAYSCALE:
+      cinfo->out_color_components = 1;
+      break;
+    case JCS_RGB:
+    case JCS_YCbCr:
+#ifdef JCS_EXTENSIONS
+    case JCS_EXT_RGB:
+    case JCS_EXT_BGR:
+#endif
+      cinfo->out_color_components = 3;
+      break;
+    case JCS_CMYK:
+    case JCS_YCCK:
+#ifdef JCS_EXTENSIONS
+    case JCS_EXT_RGBX:
+    case JCS_EXT_BGRX:
+    case JCS_EXT_XBGR:
+    case JCS_EXT_XRGB:
+#endif
+#ifdef JCS_ALPHA_EXTENSIONS
+    case JCS_EXT_RGBA:
+    case JCS_EXT_BGRA:
+    case JCS_EXT_ABGR:
+    case JCS_EXT_ARGB:
+#endif
+      cinfo->out_color_components = 4;
+      break;
+    default:
+      cinfo->out_color_components = cinfo->num_components;
   }
   cinfo->output_components =
       cinfo->quantize_colors ? 1 : cinfo->out_color_components;