diff --git a/crates/jpeg/src/lib.rs b/crates/jpeg/src/lib.rs index f1d3be2bd..3885724c9 100644 --- a/crates/jpeg/src/lib.rs +++ b/crates/jpeg/src/lib.rs @@ -380,12 +380,27 @@ struct Bits<'a> { pos: usize, acc: u32, n: u32, + /// How many of the `n` buffered bits (always the last ones) are zero padding supplied at + /// a marker or the end of the data, not real entropy-coded bits. + pad: u32, + /// Set once a read consumed padding: the scan's data ended before the scan did. + exhausted: bool, } -impl Bits<'_> { +impl<'a> Bits<'a> { + fn new(d: &'a [u8]) -> Self { + Bits { d, pos: 0, acc: 0, n: 0, pad: 0, exhausted: false } + } + fn fill(&mut self) { while self.n <= 24 { - let v = next_byte(self.d, &mut self.pos).unwrap_or(0); + let v = match next_byte(self.d, &mut self.pos) { + Some(v) => v, + None => { + self.pad += 8; + 0 + } + }; self.acc |= u32::from(v) << (24 - self.n); self.n += 8; } @@ -403,9 +418,13 @@ impl Bits<'_> { if self.n < n { self.fill(); } + if n > self.n.saturating_sub(self.pad) { + self.exhausted = true; + } let v = self.acc >> (32 - n); self.acc <<= n; self.n -= n; + self.pad = self.pad.min(self.n); v } @@ -413,6 +432,7 @@ impl Bits<'_> { self.pos = skip_restart(self.d, self.pos); self.acc = 0; self.n = 0; + self.pad = 0; } } @@ -605,6 +625,9 @@ impl ScanDecoder for HuffmanScan<'_> { *slot = extend(self.bits.get(s), s); k += 1; } + if self.bits.exhausted { + return Err(err("entropy-coded data ends before the scan does")); + } Ok(()) } } @@ -900,7 +923,7 @@ fn decode_scan(frame: &mut Frame, tables: &Tables, scan: &Scan, d: &[u8], done: let ac = tables.ac.get(ta).and_then(Option::as_ref).ok_or_else(|| err("scan uses an undefined Huffman table"))?; sel.push((dc, ac)); } - let mut dec = HuffmanScan { bits: Bits { d, pos: 0, acc: 0, n: 0 }, tables: sel, pred: [0; 4] }; + let mut dec = HuffmanScan { bits: Bits::new(d), tables: sel, pred: [0; 4] }; walk_scan(frame, scan, tables.restart, &mut dec, done) } } diff --git a/crates/jpeg/src/tests.rs b/crates/jpeg/src/tests.rs index 54699dc39..3b5acadda 100644 --- a/crates/jpeg/src/tests.rs +++ b/crates/jpeg/src/tests.rs @@ -417,6 +417,44 @@ fn twelve_bit_huffman_block_decodes_to_sixteen_bit() { assert!(px.iter().all(|&v| v.abs_diff(want) <= 1), "{:?} vs {want}", &px[..4]); } +/// Where the entropy-coded data of the last scan starts. +fn last_scan_data(s: &[u8]) -> usize { + let sos = s.windows(2).rposition(|p| p == [0xFF, 0xDA]).unwrap(); + sos + 2 + usize::from(u16::from_be_bytes([s[sos + 2], s[sos + 3]])) +} + +/// A Huffman scan whose entropy data runs out is not a clean image, even with EOI present: +/// the zero padding supplied at the marker must not be decoded as real coefficients. +#[test] +fn huffman_scan_out_of_data_is_not_clean() { + let full = twelve_bit_huffman_block(); + let data = last_scan_data(&full); + for keep in 0..2 { + let mut s = full[..data + keep].to_vec(); + s.extend([0xFF, 0xD9]); + assert!(decode(&s, &Limits::default(), None).is_err(), "keep {keep}"); + if let Ok(img) = crate::decode(&s) { + assert!(!img.warnings.is_empty(), "keep {keep}: clean decode of missing data"); + } + } + // A multi-MCU scan cut in half, EOI kept, decodes as damaged. + let (w, h) = (64usize, 64usize); + let rgb: Vec = (0..w * h).flat_map(|i| [(i % w * 4) as u8, (i / w * 4) as u8, (i * 7) as u8]).collect(); + let mut out = Vec::new(); + jpeg_encoder::Encoder::new(&mut out, 90).encode(&rgb, w as u16, h as u16, jpeg_encoder::ColorType::Rgb).unwrap(); + let sof = out.windows(2).position(|p| p == [0xFF, 0xC0]).unwrap(); + out[sof + 1] = 0xC1; + let data = last_scan_data(&out); + let end = out.len() - 2; + let mut cut = out[..data + (end - data) / 2].to_vec(); + if cut.last() == Some(&0xFF) { + cut.pop(); + } + cut.extend([0xFF, 0xD9]); + assert!(!decode(&out, &Limits::default(), None).unwrap().damaged); + assert!(decode(&cut, &Limits::default(), None).unwrap().damaged); +} + /// A baseline file from our encoder, relabelled SOF1, decodes like zune decodes the original. #[test] fn extended_huffman_matches_zune_on_baseline_data() {