1 // Copyright (c) the JPEG XL Project Authors. All rights reserved.
2 //
3 // Use of this source code is governed by a BSD-style
4 // license that can be found in the LICENSE file.
5 
6 #include "lib/jxl/jpeg/jpeg_data.h"
7 
8 #include "lib/jxl/base/status.h"
9 
10 namespace jxl {
11 namespace jpeg {
12 
13 namespace {
14 enum JPEGComponentType : uint32_t {
15   kGray = 0,
16   kYCbCr = 1,
17   kRGB = 2,
18   kCustom = 3,
19 };
20 
21 struct JPEGInfo {
22   size_t num_app_markers = 0;
23   size_t num_com_markers = 0;
24   size_t num_scans = 0;
25   size_t num_intermarker = 0;
26   bool has_dri = false;
27 };
28 
VisitMarker(uint8_t * marker,Visitor * visitor,JPEGInfo * info)29 Status VisitMarker(uint8_t* marker, Visitor* visitor, JPEGInfo* info) {
30   uint32_t marker32 = *marker - 0xc0;
31   JXL_RETURN_IF_ERROR(visitor->Bits(6, 0x00, &marker32));
32   *marker = marker32 + 0xc0;
33   if ((*marker & 0xf0) == 0xe0) {
34     info->num_app_markers++;
35   }
36   if (*marker == 0xfe) {
37     info->num_com_markers++;
38   }
39   if (*marker == 0xda) {
40     info->num_scans++;
41   }
42   // We use a fake 0xff marker to signal intermarker data.
43   if (*marker == 0xff) {
44     info->num_intermarker++;
45   }
46   if (*marker == 0xdd) {
47     info->has_dri = true;
48   }
49   return true;
50 }
51 
52 }  // namespace
53 
VisitFields(Visitor * visitor)54 Status JPEGData::VisitFields(Visitor* visitor) {
55   bool is_gray = components.size() == 1;
56   JXL_RETURN_IF_ERROR(visitor->Bool(false, &is_gray));
57   if (visitor->IsReading()) {
58     components.resize(is_gray ? 1 : 3);
59   }
60   JPEGInfo info;
61   if (visitor->IsReading()) {
62     uint8_t marker = 0xc0;
63     do {
64       JXL_RETURN_IF_ERROR(VisitMarker(&marker, visitor, &info));
65       marker_order.push_back(marker);
66       if (marker_order.size() > 16384) {
67         return JXL_FAILURE("Too many markers: %zu\n", marker_order.size());
68       }
69     } while (marker != 0xd9);
70   } else {
71     if (marker_order.size() > 16384) {
72       return JXL_FAILURE("Too many markers: %zu\n", marker_order.size());
73     }
74     for (size_t i = 0; i < marker_order.size(); i++) {
75       JXL_RETURN_IF_ERROR(VisitMarker(&marker_order[i], visitor, &info));
76     }
77     if (!marker_order.empty()) {
78       // Last marker should always be EOI marker.
79       JXL_CHECK(marker_order.back() == 0xd9);
80     }
81   }
82 
83   // Size of the APP and COM markers.
84   if (visitor->IsReading()) {
85     app_data.resize(info.num_app_markers);
86     app_marker_type.resize(info.num_app_markers);
87     com_data.resize(info.num_com_markers);
88     scan_info.resize(info.num_scans);
89   }
90   JXL_ASSERT(app_data.size() == info.num_app_markers);
91   JXL_ASSERT(app_marker_type.size() == info.num_app_markers);
92   JXL_ASSERT(com_data.size() == info.num_com_markers);
93   JXL_ASSERT(scan_info.size() == info.num_scans);
94   for (size_t i = 0; i < app_data.size(); i++) {
95     auto& app = app_data[i];
96     // Encodes up to 8 different values.
97     JXL_RETURN_IF_ERROR(
98         visitor->U32(Val(0), Val(1), BitsOffset(1, 2), BitsOffset(2, 4), 0,
99                      reinterpret_cast<uint32_t*>(&app_marker_type[i])));
100     if (app_marker_type[i] != AppMarkerType::kUnknown &&
101         app_marker_type[i] != AppMarkerType::kICC &&
102         app_marker_type[i] != AppMarkerType::kExif &&
103         app_marker_type[i] != AppMarkerType::kXMP) {
104       return JXL_FAILURE("Unknown app marker type %u",
105                          static_cast<uint32_t>(app_marker_type[i]));
106     }
107     uint32_t len = app.size() - 1;
108     JXL_RETURN_IF_ERROR(visitor->Bits(16, 0, &len));
109     if (visitor->IsReading()) app.resize(len + 1);
110     if (app.size() < 3) {
111       return JXL_FAILURE("Invalid marker size: %zu\n", app.size());
112     }
113   }
114   for (auto& com : com_data) {
115     uint32_t len = com.size() - 1;
116     JXL_RETURN_IF_ERROR(visitor->Bits(16, 0, &len));
117     if (visitor->IsReading()) com.resize(len + 1);
118     if (com.size() < 3) {
119       return JXL_FAILURE("Invalid marker size: %zu\n", com.size());
120     }
121   }
122 
123   uint32_t num_quant_tables = quant.size();
124   JXL_RETURN_IF_ERROR(
125       visitor->U32(Val(1), Val(2), Val(3), Val(4), 2, &num_quant_tables));
126   if (num_quant_tables == 4) {
127     return JXL_FAILURE("Invalid number of quant tables");
128   }
129   if (visitor->IsReading()) {
130     quant.resize(num_quant_tables);
131   }
132   for (size_t i = 0; i < num_quant_tables; i++) {
133     if (quant[i].precision > 1) {
134       return JXL_FAILURE(
135           "Quant tables with more than 16 bits are not supported");
136     }
137     JXL_RETURN_IF_ERROR(visitor->Bits(1, 0, &quant[i].precision));
138     JXL_RETURN_IF_ERROR(visitor->Bits(2, i, &quant[i].index));
139     JXL_RETURN_IF_ERROR(visitor->Bool(true, &quant[i].is_last));
140   }
141 
142   JPEGComponentType component_type =
143       components.size() == 1 && components[0].id == 1
144           ? JPEGComponentType::kGray
145           : components.size() == 3 && components[0].id == 1 &&
146                     components[1].id == 2 && components[2].id == 3
147                 ? JPEGComponentType::kYCbCr
148                 : components.size() == 3 && components[0].id == 'R' &&
149                           components[1].id == 'G' && components[2].id == 'B'
150                       ? JPEGComponentType::kRGB
151                       : JPEGComponentType::kCustom;
152   JXL_RETURN_IF_ERROR(
153       visitor->Bits(2, JPEGComponentType::kYCbCr,
154                     reinterpret_cast<uint32_t*>(&component_type)));
155   uint32_t num_components;
156   if (component_type == JPEGComponentType::kGray) {
157     num_components = 1;
158   } else if (component_type != JPEGComponentType::kCustom) {
159     num_components = 3;
160   } else {
161     num_components = components.size();
162     JXL_RETURN_IF_ERROR(
163         visitor->U32(Val(1), Val(2), Val(3), Val(4), 3, &num_components));
164     if (num_components != 1 && num_components != 3) {
165       return JXL_FAILURE("Invalid number of components: %u", num_components);
166     }
167   }
168   if (visitor->IsReading()) {
169     components.resize(num_components);
170   }
171   if (component_type == JPEGComponentType::kCustom) {
172     for (size_t i = 0; i < components.size(); i++) {
173       JXL_RETURN_IF_ERROR(visitor->Bits(8, 0, &components[i].id));
174     }
175   } else if (component_type == JPEGComponentType::kGray) {
176     components[0].id = 1;
177   } else if (component_type == JPEGComponentType::kRGB) {
178     components[0].id = 'R';
179     components[1].id = 'G';
180     components[2].id = 'B';
181   } else {
182     components[0].id = 1;
183     components[1].id = 2;
184     components[2].id = 3;
185   }
186   size_t used_tables = 0;
187   for (size_t i = 0; i < components.size(); i++) {
188     JXL_RETURN_IF_ERROR(visitor->Bits(2, 0, &components[i].quant_idx));
189     if (components[i].quant_idx >= quant.size()) {
190       return JXL_FAILURE("Invalid quant table for component %zu: %u\n", i,
191                          components[i].quant_idx);
192     }
193     used_tables |= 1U << components[i].quant_idx;
194   }
195   if (used_tables + 1 != 1U << quant.size()) {
196     return JXL_FAILURE(
197         "Not all quant tables are used (%zu tables, %zx used table mask)",
198         quant.size(), used_tables);
199   }
200 
201   uint32_t num_huff = huffman_code.size();
202   JXL_RETURN_IF_ERROR(visitor->U32(Val(4), BitsOffset(3, 2), BitsOffset(4, 10),
203                                    BitsOffset(6, 26), 4, &num_huff));
204   if (visitor->IsReading()) {
205     huffman_code.resize(num_huff);
206   }
207   for (JPEGHuffmanCode& hc : huffman_code) {
208     bool is_ac = hc.slot_id >> 4;
209     uint32_t id = hc.slot_id & 0xF;
210     JXL_RETURN_IF_ERROR(visitor->Bool(false, &is_ac));
211     JXL_RETURN_IF_ERROR(visitor->Bits(2, 0, &id));
212     hc.slot_id = (static_cast<uint32_t>(is_ac) << 4) | id;
213     JXL_RETURN_IF_ERROR(visitor->Bool(true, &hc.is_last));
214     size_t num_symbols = 0;
215     for (size_t i = 0; i <= 16; i++) {
216       JXL_RETURN_IF_ERROR(visitor->U32(Val(0), Val(1), BitsOffset(3, 2),
217                                        Bits(8), 0, &hc.counts[i]));
218       num_symbols += hc.counts[i];
219     }
220     if (num_symbols < 1) {
221       // Actually, at least 2 symbols are required, since one of them is EOI.
222       return JXL_FAILURE("Empty Huffman table");
223     }
224     if (num_symbols > hc.values.size()) {
225       return JXL_FAILURE("Huffman code too large (%zu)", num_symbols);
226     }
227     // Presence flags for 4 * 64 + 1 values.
228     uint64_t value_slots[5] = {};
229     for (size_t i = 0; i < num_symbols; i++) {
230       // Goes up to 256, included. Might have the same symbol appear twice...
231       JXL_RETURN_IF_ERROR(visitor->U32(Bits(2), BitsOffset(2, 4),
232                                        BitsOffset(4, 8), BitsOffset(8, 1), 0,
233                                        &hc.values[i]));
234       value_slots[hc.values[i] >> 6] |= (uint64_t)1 << (hc.values[i] & 0x3F);
235     }
236     if (hc.values[num_symbols - 1] != kJpegHuffmanAlphabetSize) {
237       return JXL_FAILURE("Missing EOI symbol");
238     }
239     // Last element, denoting EOI, have to be 1 after the loop.
240     JXL_ASSERT(value_slots[4] == 1);
241     size_t num_values = 1;
242     for (size_t i = 0; i < 4; ++i) num_values += hwy::PopCount(value_slots[i]);
243     if (num_values != num_symbols) {
244       return JXL_FAILURE("Duplicate Huffman symbols");
245     }
246     if (!is_ac) {
247       bool only_dc = ((value_slots[0] >> kJpegDCAlphabetSize) | value_slots[1] |
248                       value_slots[2] | value_slots[3]) == 0;
249       if (!only_dc) return JXL_FAILURE("Huffman symbols out of DC range");
250     }
251   }
252 
253   for (auto& scan : scan_info) {
254     JXL_RETURN_IF_ERROR(
255         visitor->U32(Val(1), Val(2), Val(3), Val(4), 1, &scan.num_components));
256     if (scan.num_components >= 4) {
257       return JXL_FAILURE("Invalid number of components in SOS marker");
258     }
259     JXL_RETURN_IF_ERROR(visitor->Bits(6, 0, &scan.Ss));
260     JXL_RETURN_IF_ERROR(visitor->Bits(6, 63, &scan.Se));
261     JXL_RETURN_IF_ERROR(visitor->Bits(4, 0, &scan.Al));
262     JXL_RETURN_IF_ERROR(visitor->Bits(4, 0, &scan.Ah));
263     for (size_t i = 0; i < scan.num_components; i++) {
264       JXL_RETURN_IF_ERROR(visitor->Bits(2, 0, &scan.components[i].comp_idx));
265       if (scan.components[i].comp_idx >= components.size()) {
266         return JXL_FAILURE("Invalid component idx in SOS marker");
267       }
268       JXL_RETURN_IF_ERROR(visitor->Bits(2, 0, &scan.components[i].ac_tbl_idx));
269       JXL_RETURN_IF_ERROR(visitor->Bits(2, 0, &scan.components[i].dc_tbl_idx));
270     }
271     // TODO(veluca): actually set and use this value.
272     JXL_RETURN_IF_ERROR(visitor->U32(Val(0), Val(1), Val(2), BitsOffset(3, 3),
273                                      kMaxNumPasses - 1,
274                                      &scan.last_needed_pass));
275   }
276 
277   // From here on, this is data that is not strictly necessary to get a valid
278   // JPEG, but necessary for bit-exact JPEG reconstruction.
279   if (info.has_dri) {
280     JXL_RETURN_IF_ERROR(visitor->Bits(16, 0, &restart_interval));
281   }
282 
283   uint64_t padding_spot_limit = scan_info.size();
284 
285   for (auto& scan : scan_info) {
286     uint32_t num_reset_points = scan.reset_points.size();
287     JXL_RETURN_IF_ERROR(visitor->U32(Val(0), BitsOffset(2, 1), BitsOffset(4, 4),
288                                      BitsOffset(16, 20), 0, &num_reset_points));
289     if (visitor->IsReading()) {
290       scan.reset_points.resize(num_reset_points);
291     }
292     int last_block_idx = -1;
293     for (auto& block_idx : scan.reset_points) {
294       block_idx -= last_block_idx + 1;
295       JXL_RETURN_IF_ERROR(visitor->U32(Val(0), BitsOffset(3, 1),
296                                        BitsOffset(5, 9), BitsOffset(28, 41), 0,
297                                        &block_idx));
298       block_idx += last_block_idx + 1;
299       if (static_cast<int>(block_idx) < last_block_idx + 1) {
300         return JXL_FAILURE("Invalid block ID: %u, last block was %d", block_idx,
301                            last_block_idx);
302       }
303       // TODO(eustas): better upper boundary could be given at this point; also
304       //               it could be applied during reset_points reading.
305       if (block_idx > (1u << 30)) {
306         // At most 8K x 8K x num_channels blocks are expected. That is,
307         // typically, 1.5 * 2^27. 2^30 should be sufficient for any sane
308         // image.
309         return JXL_FAILURE("Invalid block ID: %u", block_idx);
310       }
311       last_block_idx = block_idx;
312     }
313 
314     uint32_t num_extra_zero_runs = scan.extra_zero_runs.size();
315     JXL_RETURN_IF_ERROR(visitor->U32(Val(0), BitsOffset(2, 1), BitsOffset(4, 4),
316                                      BitsOffset(16, 20), 0,
317                                      &num_extra_zero_runs));
318     if (visitor->IsReading()) {
319       scan.extra_zero_runs.resize(num_extra_zero_runs);
320     }
321     last_block_idx = -1;
322     for (size_t i = 0; i < scan.extra_zero_runs.size(); ++i) {
323       uint32_t& block_idx = scan.extra_zero_runs[i].block_idx;
324       JXL_RETURN_IF_ERROR(visitor->U32(
325           Val(1), BitsOffset(2, 2), BitsOffset(4, 5), BitsOffset(8, 20), 1,
326           &scan.extra_zero_runs[i].num_extra_zero_runs));
327       block_idx -= last_block_idx + 1;
328       JXL_RETURN_IF_ERROR(visitor->U32(Val(0), BitsOffset(3, 1),
329                                        BitsOffset(5, 9), BitsOffset(28, 41), 0,
330                                        &block_idx));
331       block_idx += last_block_idx + 1;
332       if (static_cast<int>(block_idx) < last_block_idx + 1) {
333         return JXL_FAILURE("Invalid block ID: %u, last block was %d", block_idx,
334                            last_block_idx);
335       }
336       if (block_idx > (1u << 30)) {
337         // At most 8K x 8K x num_channels blocks are expected. That is,
338         // typically, 1.5 * 2^27. 2^30 should be sufficient for any sane
339         // image.
340         return JXL_FAILURE("Invalid block ID: %u", block_idx);
341       }
342       last_block_idx = block_idx;
343     }
344 
345     if (restart_interval > 0) {
346       int MCUs_per_row = 0;
347       int MCU_rows = 0;
348       CalculateMcuSize(scan, &MCUs_per_row, &MCU_rows);
349       padding_spot_limit += DivCeil(MCU_rows * MCUs_per_row, restart_interval);
350     }
351   }
352   std::vector<uint32_t> inter_marker_data_sizes;
353   inter_marker_data_sizes.reserve(info.num_intermarker);
354   for (size_t i = 0; i < info.num_intermarker; ++i) {
355     uint32_t len = visitor->IsReading() ? 0 : inter_marker_data[i].size();
356     JXL_RETURN_IF_ERROR(visitor->Bits(16, 0, &len));
357     if (visitor->IsReading()) inter_marker_data_sizes.emplace_back(len);
358   }
359   uint32_t tail_data_len = tail_data.size();
360   JXL_RETURN_IF_ERROR(visitor->U32(Val(0), BitsOffset(8, 1),
361                                    BitsOffset(16, 257), BitsOffset(22, 65793),
362                                    0, &tail_data_len));
363 
364   JXL_RETURN_IF_ERROR(visitor->Bool(false, &has_zero_padding_bit));
365   if (has_zero_padding_bit) {
366     uint32_t nbit = padding_bits.size();
367     JXL_RETURN_IF_ERROR(visitor->Bits(24, 0, &nbit));
368     if (nbit > 7 * padding_spot_limit) {
369       return JXL_FAILURE("Number of padding bits does not correspond to image");
370     }
371     // TODO(eustas): check that that much bits of input are available.
372     if (visitor->IsReading()) {
373       padding_bits.resize(nbit);
374     }
375     // TODO(eustas): read in (8-64?) bit groups to reduce overhead.
376     for (uint8_t& bit : padding_bits) {
377       bool bbit = bit;
378       JXL_RETURN_IF_ERROR(visitor->Bool(false, &bbit));
379       bit = bbit;
380     }
381   }
382 
383   // Apply postponed actions.
384   if (visitor->IsReading()) {
385     tail_data.resize(tail_data_len);
386     JXL_ASSERT(inter_marker_data_sizes.size() == info.num_intermarker);
387     inter_marker_data.reserve(info.num_intermarker);
388     for (size_t i = 0; i < info.num_intermarker; ++i) {
389       inter_marker_data.emplace_back(inter_marker_data_sizes[i]);
390     }
391   }
392 
393   return true;
394 }
395 
CalculateMcuSize(const JPEGScanInfo & scan,int * MCUs_per_row,int * MCU_rows) const396 void JPEGData::CalculateMcuSize(const JPEGScanInfo& scan, int* MCUs_per_row,
397                                 int* MCU_rows) const {
398   const bool is_interleaved = (scan.num_components > 1);
399   const JPEGComponent& base_component = components[scan.components[0].comp_idx];
400   // h_group / v_group act as numerators for converting number of blocks to
401   // number of MCU. In interleaved mode it is 1, so MCU is represented with
402   // max_*_samp_factor blocks. In non-interleaved mode we choose numerator to
403   // be the samping factor, consequently MCU is always represented with single
404   // block.
405   const int h_group = is_interleaved ? 1 : base_component.h_samp_factor;
406   const int v_group = is_interleaved ? 1 : base_component.v_samp_factor;
407   int max_h_samp_factor = 1;
408   int max_v_samp_factor = 1;
409   for (const auto& c : components) {
410     max_h_samp_factor = std::max(c.h_samp_factor, max_h_samp_factor);
411     max_v_samp_factor = std::max(c.v_samp_factor, max_v_samp_factor);
412   }
413   *MCUs_per_row = DivCeil(width * h_group, 8 * max_h_samp_factor);
414   *MCU_rows = DivCeil(height * v_group, 8 * max_v_samp_factor);
415 }
416 
SetJPEGDataFromICC(const PaddedBytes & icc,jpeg::JPEGData * jpeg_data)417 Status SetJPEGDataFromICC(const PaddedBytes& icc, jpeg::JPEGData* jpeg_data) {
418   size_t icc_pos = 0;
419   for (size_t i = 0; i < jpeg_data->app_data.size(); i++) {
420     if (jpeg_data->app_marker_type[i] != jpeg::AppMarkerType::kICC) {
421       continue;
422     }
423     size_t len = jpeg_data->app_data[i].size() - 17;
424     if (icc_pos + len > icc.size()) {
425       return JXL_FAILURE(
426           "ICC length is less than APP markers: requested %zu more bytes, "
427           "%zu available",
428           len, icc.size() - icc_pos);
429     }
430     memcpy(&jpeg_data->app_data[i][17], icc.data() + icc_pos, len);
431     icc_pos += len;
432   }
433   if (icc_pos != icc.size() && icc_pos != 0) {
434     return JXL_FAILURE("ICC length is more than APP markers");
435   }
436   return true;
437 }
438 
439 }  // namespace jpeg
440 }  // namespace jxl
441