1 // Malicious JPEG files can cause operations in the idct to overflow.
2 // One example is tests/crashtest/images/imagetestsuite/b0b8914cc5f7a6eff409f16d8cc236c5.jpg
3 // That's why wrapping operators are needed.
4 use crate::parser::Dimensions;
5 use std::{convert::TryFrom, num::Wrapping};
6
choose_idct_size(full_size: Dimensions, requested_size: Dimensions) -> usize7 pub(crate) fn choose_idct_size(full_size: Dimensions, requested_size: Dimensions) -> usize {
8 fn scaled(len: u16, scale: usize) -> u16 { ((len as u32 * scale as u32 - 1) / 8 + 1) as u16 }
9
10 for &scale in &[1, 2, 4] {
11 if scaled(full_size.width, scale) >= requested_size.width || scaled(full_size.height, scale) >= requested_size.height {
12 return scale;
13 }
14 }
15
16 8
17 }
18
19 #[test]
test_choose_idct_size()20 fn test_choose_idct_size() {
21 assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 200, height: 200}), 1);
22 assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 500, height: 500}), 1);
23 assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 684, height: 456}), 1);
24 assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 999, height: 456}), 1);
25 assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 684, height: 999}), 1);
26 assert_eq!(choose_idct_size(Dimensions{width: 500, height: 333}, Dimensions{width: 63, height: 42}), 1);
27
28 assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 685, height: 999}), 2);
29 assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 1000, height: 1000}), 2);
30 assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 1400, height: 1400}), 4);
31
32 assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 5472, height: 3648}), 8);
33 assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 16384, height: 16384}), 8);
34 assert_eq!(choose_idct_size(Dimensions{width: 1, height: 1}, Dimensions{width: 65535, height: 65535}), 8);
35 assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 16384, height: 16384}), 8);
36 }
37
dequantize_and_idct_block(scale: usize, coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8])38 pub(crate) fn dequantize_and_idct_block(scale: usize, coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8]) {
39 match scale {
40 8 => dequantize_and_idct_block_8x8(coefficients, quantization_table, output_linestride, output),
41 4 => dequantize_and_idct_block_4x4(coefficients, quantization_table, output_linestride, output),
42 2 => dequantize_and_idct_block_2x2(coefficients, quantization_table, output_linestride, output),
43 1 => dequantize_and_idct_block_1x1(coefficients, quantization_table, output_linestride, output),
44 _ => panic!("Unsupported IDCT scale {}/8", scale),
45 }
46 }
47
dequantize_and_idct_block_8x8( coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8] )48 pub fn dequantize_and_idct_block_8x8(
49 coefficients: &[i16],
50 quantization_table: &[u16; 64],
51 output_linestride: usize,
52 output: &mut [u8]
53 ) {
54 let output = output
55 .chunks_mut(output_linestride);
56 dequantize_and_idct_block_8x8_inner(coefficients, quantization_table, output)
57 }
58
59 // This is based on stb_image's 'stbi__idct_block'.
dequantize_and_idct_block_8x8_inner<'a, I>( coefficients: &[i16], quantization_table: &[u16; 64], output: I, ) where I: IntoIterator<Item = &'a mut [u8]>, I::IntoIter: ExactSizeIterator<Item = &'a mut [u8]>,60 fn dequantize_and_idct_block_8x8_inner<'a, I>(
61 coefficients: &[i16],
62 quantization_table: &[u16; 64],
63 output: I,
64 ) where
65 I: IntoIterator<Item = &'a mut [u8]>,
66 I::IntoIter: ExactSizeIterator<Item = &'a mut [u8]>,
67 {
68 let output = output.into_iter();
69 debug_assert!(
70 output.len() >= 8,
71 "Output iterator has the wrong length: {}",
72 output.len()
73 );
74
75 // optimizer hint to eliminate bounds checks within loops
76 assert!(coefficients.len() == 64);
77
78 let mut temp = [Wrapping(0); 64];
79
80 // columns
81 for i in 0..8 {
82 if coefficients[i + 8] == 0
83 && coefficients[i + 16] == 0
84 && coefficients[i + 24] == 0
85 && coefficients[i + 32] == 0
86 && coefficients[i + 40] == 0
87 && coefficients[i + 48] == 0
88 && coefficients[i + 56] == 0
89 {
90 let dcterm = dequantize(coefficients[i], quantization_table[i]) << 2;
91 temp[i] = dcterm;
92 temp[i + 8] = dcterm;
93 temp[i + 16] = dcterm;
94 temp[i + 24] = dcterm;
95 temp[i + 32] = dcterm;
96 temp[i + 40] = dcterm;
97 temp[i + 48] = dcterm;
98 temp[i + 56] = dcterm;
99 } else {
100 let s0 = dequantize(coefficients[i], quantization_table[i]);
101 let s1 = dequantize(coefficients[i + 8], quantization_table[i + 8]);
102 let s2 = dequantize(coefficients[i + 16], quantization_table[i + 16]);
103 let s3 = dequantize(coefficients[i + 24], quantization_table[i + 24]);
104 let s4 = dequantize(coefficients[i + 32], quantization_table[i + 32]);
105 let s5 = dequantize(coefficients[i + 40], quantization_table[i + 40]);
106 let s6 = dequantize(coefficients[i + 48], quantization_table[i + 48]);
107 let s7 = dequantize(coefficients[i + 56], quantization_table[i + 56]);
108
109 let Kernel {
110 xs: [x0, x1, x2, x3],
111 ts: [t0, t1, t2, t3],
112 } = kernel(
113 [s0, s1, s2, s3, s4, s5, s6, s7],
114 // constants scaled things up by 1<<12; let's bring them back
115 // down, but keep 2 extra bits of precision
116 512,
117 );
118
119 temp[i] = (x0 + t3) >> 10;
120 temp[i + 56] = (x0 - t3) >> 10;
121 temp[i + 8] = (x1 + t2) >> 10;
122 temp[i + 48] = (x1 - t2) >> 10;
123 temp[i + 16] = (x2 + t1) >> 10;
124 temp[i + 40] = (x2 - t1) >> 10;
125 temp[i + 24] = (x3 + t0) >> 10;
126 temp[i + 32] = (x3 - t0) >> 10;
127 }
128 }
129
130 for (chunk, output_chunk) in temp.chunks_exact(8).zip(output) {
131 let chunk = <&[_; 8]>::try_from(chunk).unwrap();
132
133 // constants scaled things up by 1<<12, plus we had 1<<2 from first
134 // loop, plus horizontal and vertical each scale by sqrt(8) so together
135 // we've got an extra 1<<3, so 1<<17 total we need to remove.
136 // so we want to round that, which means adding 0.5 * 1<<17,
137 // aka 65536. Also, we'll end up with -128 to 127 that we want
138 // to encode as 0..255 by adding 128, so we'll add that before the shift
139 const X_SCALE: i32 = 65536 + (128 << 17);
140
141 // TODO When the minimum rust version supports it
142 // let [s0, rest @ ..] = chunk;
143 let (s0, rest) = chunk.split_first().unwrap();
144 if *rest == [Wrapping(0); 7] {
145 let dcterm = stbi_clamp((stbi_fsh(*s0) + Wrapping(X_SCALE)) >> 17);
146 output_chunk[0] = dcterm;
147 output_chunk[1] = dcterm;
148 output_chunk[2] = dcterm;
149 output_chunk[3] = dcterm;
150 output_chunk[4] = dcterm;
151 output_chunk[5] = dcterm;
152 output_chunk[6] = dcterm;
153 output_chunk[7] = dcterm;
154 } else {
155 let Kernel {
156 xs: [x0, x1, x2, x3],
157 ts: [t0, t1, t2, t3],
158 } = kernel(*chunk, X_SCALE);
159
160 output_chunk[0] = stbi_clamp((x0 + t3) >> 17);
161 output_chunk[7] = stbi_clamp((x0 - t3) >> 17);
162 output_chunk[1] = stbi_clamp((x1 + t2) >> 17);
163 output_chunk[6] = stbi_clamp((x1 - t2) >> 17);
164 output_chunk[2] = stbi_clamp((x2 + t1) >> 17);
165 output_chunk[5] = stbi_clamp((x2 - t1) >> 17);
166 output_chunk[3] = stbi_clamp((x3 + t0) >> 17);
167 output_chunk[4] = stbi_clamp((x3 - t0) >> 17);
168 }
169 }
170 }
171
172 struct Kernel {
173 xs: [Wrapping<i32>; 4],
174 ts: [Wrapping<i32>; 4],
175 }
176
177 #[inline]
kernel_x([s0, s2, s4, s6]: [Wrapping<i32>; 4], x_scale: i32) -> [Wrapping<i32>; 4]178 fn kernel_x([s0, s2, s4, s6]: [Wrapping<i32>; 4], x_scale: i32) -> [Wrapping<i32>; 4] {
179 // Even `chunk` indicies
180 let (t2, t3);
181 {
182 let p2 = s2;
183 let p3 = s6;
184
185 let p1 = (p2 + p3) * stbi_f2f(0.5411961);
186 t2 = p1 + p3 * stbi_f2f(-1.847759065);
187 t3 = p1 + p2 * stbi_f2f(0.765366865);
188 }
189
190 let (t0, t1);
191 {
192 let p2 = s0;
193 let p3 = s4;
194
195 t0 = stbi_fsh(p2 + p3);
196 t1 = stbi_fsh(p2 - p3);
197 }
198
199 let x0 = t0 + t3;
200 let x3 = t0 - t3;
201 let x1 = t1 + t2;
202 let x2 = t1 - t2;
203
204 let x_scale = Wrapping(x_scale);
205
206 [x0 + x_scale, x1 + x_scale, x2 + x_scale, x3 + x_scale]
207 }
208
209 #[inline]
kernel_t([s1, s3, s5, s7]: [Wrapping<i32>; 4]) -> [Wrapping<i32>; 4]210 fn kernel_t([s1, s3, s5, s7]: [Wrapping<i32>; 4]) -> [Wrapping<i32>; 4] {
211 // Odd `chunk` indicies
212 let mut t0 = s7;
213 let mut t1 = s5;
214 let mut t2 = s3;
215 let mut t3 = s1;
216
217 let p3 = t0 + t2;
218 let p4 = t1 + t3;
219 let p1 = t0 + t3;
220 let p2 = t1 + t2;
221 let p5 = (p3 + p4) * stbi_f2f(1.175875602);
222
223 t0 *= stbi_f2f(0.298631336);
224 t1 *= stbi_f2f(2.053119869);
225 t2 *= stbi_f2f(3.072711026);
226 t3 *= stbi_f2f(1.501321110);
227
228 let p1 = p5 + p1 * stbi_f2f(-0.899976223);
229 let p2 = p5 + p2 * stbi_f2f(-2.562915447);
230 let p3 = p3 * stbi_f2f(-1.961570560);
231 let p4 = p4 * stbi_f2f(-0.390180644);
232
233 t3 += p1 + p4;
234 t2 += p2 + p3;
235 t1 += p2 + p4;
236 t0 += p1 + p3;
237
238 [t0, t1, t2, t3]
239 }
240
241 #[inline]
kernel([s0, s1, s2, s3, s4, s5, s6, s7]: [Wrapping<i32>; 8], x_scale: i32) -> Kernel242 fn kernel([s0, s1, s2, s3, s4, s5, s6, s7]: [Wrapping<i32>; 8], x_scale: i32) -> Kernel {
243 Kernel {
244 xs: kernel_x([s0, s2, s4, s6], x_scale),
245 ts: kernel_t([s1, s3, s5, s7]),
246 }
247 }
248
249 #[inline(always)]
dequantize(c: i16, q: u16) -> Wrapping<i32>250 fn dequantize(c: i16, q: u16) -> Wrapping<i32> {
251 Wrapping(i32::from(c) * i32::from(q))
252 }
253
254 // 4x4 and 2x2 IDCT based on Rakesh Dugad and Narendra Ahuja: "A Fast Scheme for Image Size Change in the Compressed Domain" (2001).
255 // http://sylvana.net/jpegcrop/jidctred/
dequantize_and_idct_block_4x4(coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8])256 fn dequantize_and_idct_block_4x4(coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8]) {
257 debug_assert_eq!(coefficients.len(), 64);
258 let mut temp = [Wrapping(0i32); 4 * 4];
259
260 const CONST_BITS: usize = 12;
261 const PASS1_BITS: usize = 2;
262 const FINAL_BITS: usize = CONST_BITS + PASS1_BITS + 3;
263
264 // columns
265 for i in 0..4 {
266 let s0 = Wrapping(coefficients[i + 8 * 0] as i32 * quantization_table[i + 8 * 0] as i32);
267 let s1 = Wrapping(coefficients[i + 8 * 1] as i32 * quantization_table[i + 8 * 1] as i32);
268 let s2 = Wrapping(coefficients[i + 8 * 2] as i32 * quantization_table[i + 8 * 2] as i32);
269 let s3 = Wrapping(coefficients[i + 8 * 3] as i32 * quantization_table[i + 8 * 3] as i32);
270
271 let x0 = (s0 + s2) << PASS1_BITS;
272 let x2 = (s0 - s2) << PASS1_BITS;
273
274 let p1 = (s1 + s3) * stbi_f2f(0.541196100);
275 let t0 = (p1 + s3 * stbi_f2f(-1.847759065) + Wrapping(512)) >> (CONST_BITS - PASS1_BITS);
276 let t2 = (p1 + s1 * stbi_f2f(0.765366865) + Wrapping(512)) >> (CONST_BITS - PASS1_BITS);
277
278 temp[i + 4 * 0] = x0 + t2;
279 temp[i + 4 * 3] = x0 - t2;
280 temp[i + 4 * 1] = x2 + t0;
281 temp[i + 4 * 2] = x2 - t0;
282 }
283
284 for i in 0 .. 4 {
285 let s0 = temp[i * 4 + 0];
286 let s1 = temp[i * 4 + 1];
287 let s2 = temp[i * 4 + 2];
288 let s3 = temp[i * 4 + 3];
289
290 let x0 = (s0 + s2) << CONST_BITS;
291 let x2 = (s0 - s2) << CONST_BITS;
292
293 let p1 = (s1 + s3) * stbi_f2f(0.541196100);
294 let t0 = p1 + s3 * stbi_f2f(-1.847759065);
295 let t2 = p1 + s1 * stbi_f2f(0.765366865);
296
297 // constants scaled things up by 1<<12, plus we had 1<<2 from first
298 // loop, plus horizontal and vertical each scale by sqrt(8) so together
299 // we've got an extra 1<<3, so 1<<17 total we need to remove.
300 // so we want to round that, which means adding 0.5 * 1<<17,
301 // aka 65536. Also, we'll end up with -128 to 127 that we want
302 // to encode as 0..255 by adding 128, so we'll add that before the shift
303 let x0 = x0 + Wrapping(1 << (FINAL_BITS - 1)) + Wrapping(128 << FINAL_BITS);
304 let x2 = x2 + Wrapping(1 << (FINAL_BITS - 1)) + Wrapping(128 << FINAL_BITS);
305
306 output[i * output_linestride + 0] = stbi_clamp((x0 + t2) >> FINAL_BITS);
307 output[i * output_linestride + 3] = stbi_clamp((x0 - t2) >> FINAL_BITS);
308 output[i * output_linestride + 1] = stbi_clamp((x2 + t0) >> FINAL_BITS);
309 output[i * output_linestride + 2] = stbi_clamp((x2 - t0) >> FINAL_BITS);
310 }
311 }
312
dequantize_and_idct_block_2x2(coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8])313 fn dequantize_and_idct_block_2x2(coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8]) {
314 debug_assert_eq!(coefficients.len(), 64);
315
316 const SCALE_BITS: usize = 3;
317
318 // Column 0
319 let s00 = Wrapping(coefficients[8 * 0] as i32 * quantization_table[8 * 0] as i32);
320 let s10 = Wrapping(coefficients[8 * 1] as i32 * quantization_table[8 * 1] as i32);
321
322 let x0 = s00 + s10;
323 let x2 = s00 - s10;
324
325 // Column 1
326 let s01 = Wrapping(coefficients[8 * 0 + 1] as i32 * quantization_table[8 * 0 + 1] as i32);
327 let s11 = Wrapping(coefficients[8 * 1 + 1] as i32 * quantization_table[8 * 1 + 1] as i32);
328
329 let x1 = s01 + s11;
330 let x3 = s01 - s11;
331
332 let x0 = x0 + Wrapping(1 << (SCALE_BITS - 1)) + Wrapping(128 << SCALE_BITS);
333 let x2 = x2 + Wrapping(1 << (SCALE_BITS - 1)) + Wrapping(128 << SCALE_BITS);
334
335 // Row 0
336 output[0] = stbi_clamp((x0 + x1) >> SCALE_BITS);
337 output[1] = stbi_clamp((x0 - x1) >> SCALE_BITS);
338
339 // Row 1
340 output[output_linestride + 0] = stbi_clamp((x2 + x3) >> SCALE_BITS);
341 output[output_linestride + 1] = stbi_clamp((x2 - x3) >> SCALE_BITS);
342 }
343
dequantize_and_idct_block_1x1(coefficients: &[i16], quantization_table: &[u16; 64], _output_linestride: usize, output: &mut [u8])344 fn dequantize_and_idct_block_1x1(coefficients: &[i16], quantization_table: &[u16; 64], _output_linestride: usize, output: &mut [u8]) {
345 debug_assert_eq!(coefficients.len(), 64);
346
347 let s0 = (Wrapping(coefficients[0] as i32 * quantization_table[0] as i32) + Wrapping(128 * 8)) / Wrapping(8);
348 output[0] = stbi_clamp(s0);
349 }
350
351 // take a -128..127 value and stbi__clamp it and convert to 0..255
stbi_clamp(x: Wrapping<i32>) -> u8352 fn stbi_clamp(x: Wrapping<i32>) -> u8
353 {
354 x.0.max(0).min(255) as u8
355 }
356
stbi_f2f(x: f32) -> Wrapping<i32>357 fn stbi_f2f(x: f32) -> Wrapping<i32> {
358 Wrapping((x * 4096.0 + 0.5) as i32)
359 }
360
stbi_fsh(x: Wrapping<i32>) -> Wrapping<i32>361 fn stbi_fsh(x: Wrapping<i32>) -> Wrapping<i32> {
362 x << 12
363 }
364
365 #[test]
test_dequantize_and_idct_block_8x8()366 fn test_dequantize_and_idct_block_8x8() {
367 let coefficients: [i16; 8 * 8] = [
368 -14, -39, 58, -2, 3, 3, 0, 1,
369 11, 27, 4, -3, 3, 0, 1, 0,
370 -6, -13, -9, -1, -2, -1, 0, 0,
371 -4, 0, -1, -2, 0, 0, 0, 0,
372 3, 0, 0, 0, 0, 0, 0, 0,
373 -3, -2, 0, 0, 0, 0, 0, 0,
374 0, 0, 0, 0, 0, 0, 0, 0,
375 0, 0, 0, 0, 0, 0, 0, 0];
376
377 let quantization_table: [u16; 8 * 8] = [
378 8, 6, 5, 8, 12, 20, 26, 31,
379 6, 6, 7, 10, 13, 29, 30, 28,
380 7, 7, 8, 12, 20, 29, 35, 28,
381 7, 9, 11, 15, 26, 44, 40, 31,
382 9, 11, 19, 28, 34, 55, 52, 39,
383 12, 18, 28, 32, 41, 52, 57, 46,
384 25, 32, 39, 44, 52, 61, 60, 51,
385 36, 46, 48, 49, 56, 50, 52, 50];
386 let output_linestride: usize = 8;
387 let mut output = [0u8; 8 * 8];
388 dequantize_and_idct_block_8x8(
389 &coefficients,
390 &quantization_table,
391 output_linestride,
392 &mut output);
393 let expected_output = [
394 118, 92, 110, 83, 77, 93, 144, 198,
395 172, 116, 114, 87, 78, 93, 146, 191,
396 194, 107, 91, 76, 71, 93, 160, 198,
397 196, 100, 80, 74, 67, 92, 174, 209,
398 182, 104, 88, 81, 68, 89, 178, 206,
399 105, 64, 59, 59, 63, 94, 183, 201,
400 35, 27, 28, 37, 72, 121, 203, 204,
401 37, 45, 41, 47, 98, 154, 223, 208];
402 assert_eq!(&output[..], &expected_output[..]);
403 }
404
405 #[test]
test_dequantize_and_idct_block_8x8_all_zero()406 fn test_dequantize_and_idct_block_8x8_all_zero() {
407 let mut output = [0u8; 8 * 8];
408 dequantize_and_idct_block_8x8(
409 &[0; 8*8],
410 &[666; 8*8],
411 8,
412 &mut output);
413 assert_eq!(&output[..], &[128; 8*8][..]);
414 }
415
416 #[test]
test_dequantize_and_idct_block_8x8_saturated()417 fn test_dequantize_and_idct_block_8x8_saturated() {
418 let mut output = [0u8; 8 * 8];
419 dequantize_and_idct_block_8x8(
420 &[std::i16::MAX; 8*8],
421 &[std::u16::MAX; 8*8],
422 8,
423 &mut output);
424 let expected = [
425 0, 0, 0, 255, 255, 0, 0, 255,
426 0, 0, 215, 0, 0, 255, 255, 0,
427 255, 255, 255, 255, 255, 0, 0, 255,
428 0, 0, 255, 0, 255, 0, 255, 255,
429 0, 0, 255, 255, 0, 255, 0, 0,
430 255, 255, 0, 255, 255, 255, 170, 0,
431 0, 255, 0, 0, 0, 0, 0, 255,
432 255, 255, 0, 255, 0, 255, 0, 0];
433 assert_eq!(&output[..], &expected[..]);
434 }
435