heap-buffer-overflow in jpegli's TokenizeACRefinementScan

[roticv]

Describe the bug Under certain conditions, num_refinement_bits is computed wrongly causing a heap-buffer-overflow in jpegli's TokenizeACRefinementScan in https://github.com/google/jpegli/blob/main/lib/jpegli/entropy_coding.cc#L255

To Reproduce There is a need to use a specific scan_info to trigger it. I was able to reproduce it with the following patch:

--- a/jpeg-xl/lib/extras/enc/jpegli.cc
+++ b/jpeg-xl/lib/extras/enc/jpegli.cc
@@ -41,6 +41,19 @@

 namespace {

+constexpr auto kScanConfig = to_array<jpeg_scan_info>({
+    {3, {0, 1, 2}, 0, 0, 0, 1},  //
+    {1, {0}, 1, 5, 0, 2},        //
+    {1, {2}, 1, 63, 0, 1},       //
+    {1, {1}, 1, 63, 0, 1},       //
+    {1, {0}, 6, 63, 0, 2},       //
+    {1, {0}, 1, 63, 2, 1},       //
+    {3, {0, 1, 2}, 0, 0, 1, 0},  //
+    {1, {2}, 1, 63, 1, 0},       //
+    {1, {1}, 1, 63, 1, 0},       //
+    {1, {0}, 1, 63, 1, 0}        //
+});
+
 void MyErrorExit(j_common_ptr cinfo) {
   jmp_buf* env = static_cast<jmp_buf*>(cinfo->client_data);
   (*cinfo->err->output_message)(cinfo);
@@ -448,7 +461,9 @@
     } else {
       jpegli_set_distance(&cinfo, jpeg_settings.distance, TRUE);
     }
-    jpegli_set_progressive_level(&cinfo, jpeg_settings.progressive_level);
+    //jpegli_set_progressive_level(&cinfo, jpeg_settings.progressive_level);
+    cinfo.scan_info = kScanConfig.data();
+    cinfo.num_scans = kScanConfig.size();
     cinfo.optimize_coding = TO_JXL_BOOL(jpeg_settings.optimize_coding);
     if (!jpeg_settings.app_data.empty()) {
       // Make sure jpegli_start_compress() does not write any APP markers.

Steps to reproduce the behavior:

1. Apply the above patch
2. Run build cjpegli with ASAN and run with some sample images (Will provide example image if needed)

Expected behavior num_refinement_bits is computed correctly and there is no heap-buffer-overflow

Environment

• Centos 5
• Compiler: LLVM 15
• CPU type: x86_64
• jpeg-xl version 0.10.2

[roticv]

@haixia-meta proposed a possible fix:

@@ -441,18 +441,19 @@
   std::vector<int> processed(cinfo->num_scans);
   size_t max_refinement_tokens = 0;
   size_t num_refinement_bits = 0;
-  int num_refinement_scans[DCTSIZE2] = {};
+  int num_refinement_scans[kMaxComponents][DCTSIZE2] = {};
   int max_num_refinement_scans = 0;
   for (int i = 0; i < cinfo->num_scans; ++i) {
     const jpeg_scan_info* si = &cinfo->scan_info[i];
     ScanTokenInfo* sti = &m->scan_token_info[i];
     if (si->Ss > 0 && si->Ah == 0 && si->Al > 0) {
       int offset = m->ac_ctx_offset[i];
+      int comp_idx = si->component_index[0];
       TokenizeScan(cinfo, i, offset, sti);
       processed[i] = 1;
       max_refinement_tokens += sti->num_future_nonzeros;
       for (int k = si->Ss; k <= si->Se; ++k) {
-        num_refinement_scans[k] = si->Al;
+        num_refinement_scans[comp_idx][k] = si->Al;
       }
       max_num_refinement_scans = std::max(max_num_refinement_scans, si->Al);
       num_refinement_bits += sti->num_nonzeros;
@@ -475,15 +476,17 @@
     size_t new_refinement_bits = 0;
     for (int i = 0; i < cinfo->num_scans; ++i) {
       const jpeg_scan_info* si = &cinfo->scan_info[i];
+      int comp_idx = si->component_index[0];
       ScanTokenInfo* sti = &m->scan_token_info[i];
       if (si->Ss > 0 && si->Ah > 0 &&
-          si->Ah == num_refinement_scans[si->Ss] - j) {
+          si->Ah == num_refinement_scans[comp_idx][si->Ss] - j) {
         int offset = m->ac_ctx_offset[i];
         TokenizeScan(cinfo, i, offset, sti);
         processed[i] = 1;
         new_refinement_bits += sti->num_nonzeros;
       }
     }
+    printf("%llu ?= %llu\n", m->next_refinement_bit, refinement_bits + num_refinement_bits);
     JXL_DASSERT(m->next_refinement_bit ==
                 refinement_bits + num_refinement_bits);
     num_refinement_bits += new_refinement_bits;

and he explains the problem as

[the code] only has a single array int num_refinement_scans[DCTSIZE2] and is supposed to find out for each DCT coefficient how many refinement scans there are. the problem is that there are 3 color components and this can differ between component index. for example scan #2 and #3 has Ah = 0 and Al = 1 for chroma components, but the least significant bit is refined at scans #7 and #8.