diff --git a/crates/jpeg/src/lib.rs b/crates/jpeg/src/lib.rs index 88edd2533..44c7673be 100644 --- a/crates/jpeg/src/lib.rs +++ b/crates/jpeg/src/lib.rs @@ -37,7 +37,7 @@ //! which also reads the APP14 `Adobe` segment and passes its transform byte to [`decode`]. //! //! The decoder never panics on input: every read from the stream is bounds-checked, the work -//! is bounded by the frame size times the scan count (capped at `MAX_SCANS`), and +//! is bounded by 64 passes over the frame's blocks (`SCAN_PASSES`), and //! allocations are checked against the [`Limits`] first. //! Damaged entropy-coded data stops that scan; what was decoded is kept and the caller is told //! through [`Decoded::damaged`]. @@ -386,17 +386,19 @@ fn parse_frame(marker: u8, seg: &[u8], limits: &Limits) -> Result } let hmax = raw.iter().map(|c| c.1).max().unwrap_or(1); let vmax = raw.iter().map(|c| c.2).max().unwrap_or(1); - limits.check(width, height, nf as u64 * if precision > 8 { 2 } else { 1 })?; + let bytes_per_pixel = nf as u64 * if precision > 8 { 2 } else { 1 }; + limits.check(width, height, bytes_per_pixel)?; let mcux = width.div_ceil(8 * hmax); let mcuy = height.div_ceil(8 * vmax); - // Coefficients (4 bytes) and the reconstructed plane (4 bytes) per padded sample. - let mut internal: u64 = 0; + // `render` holds the coefficients (4 bytes) and the reconstructed plane (4 bytes) of every + // padded sample while it writes the output, so all three count against `max_alloc` together. + let mut total = (width as u64 * height as u64).saturating_mul(bytes_per_pixel); for &(_, h, v, _) in &raw { let samples = (mcux * h * 8) as u64 * (mcuy * v * 8) as u64; - internal = internal.saturating_add(samples.saturating_mul(8)); + total = total.saturating_add(samples.saturating_mul(8)); } - if internal > limits.max_alloc { - return Err(Error::LimitExceeded(format!("{internal} bytes of decoder buffers exceed max_alloc {}", limits.max_alloc))); + if total > limits.max_alloc { + return Err(Error::LimitExceeded(format!("{total} bytes of decoder buffers and output exceed max_alloc {}", limits.max_alloc))); } let mut comps = Vec::with_capacity(nf); for (id, h, v, tq) in raw { @@ -426,21 +428,22 @@ fn scan_end(b: &[u8], start: usize) -> usize { b.len() } -/// Past the next restart marker at or after `pos` (or at the next other marker, or the end). -fn skip_restart(d: &[u8], mut pos: usize) -> usize { +/// Past the next restart marker at or after `pos`, and `true`; or at the next other marker +/// (or the end) and `false`. +fn skip_restart(d: &[u8], mut pos: usize) -> (usize, bool) { while let Some(&v) = d.get(pos) { if v != 0xFF { pos += 1; continue; } match d.get(pos + 1) { - Some(0xD0..=0xD7) => return pos + 2, + Some(0xD0..=0xD7) => return (pos + 2, true), Some(0x00) => pos += 2, Some(0xFF) => pos += 1, - _ => return pos, + _ => return (pos, false), } } - pos + (pos, false) } /// Next data byte for the decoders: `FF 00` is a data `FF`; at a marker (or the end) zeros @@ -519,7 +522,7 @@ impl<'a> Bits<'a> { } fn restart(&mut self) { - self.pos = skip_restart(self.d, self.pos); + self.pos = skip_restart(self.d, self.pos).0; self.acc = 0; self.n = 0; self.pad = 0; @@ -581,16 +584,41 @@ struct Arith<'a> { c: u32, a: u32, ct: u32, + /// Zero bytes fed in since the start or the last restart because the data had ended. + pad: u32, + /// A restart was due but its marker was missing. + lost: bool, } -impl Arith<'_> { +impl<'a> Arith<'a> { + /// The most zero bytes a scan or restart interval may consume past the end of its data. + /// The decoder reads two bytes ahead (INITDEC), and encoders drop trailing zero bytes + /// (libjpeg's `finish_pass` does, as D.1.8 allows), so a valid stream consumes a few. + /// Past that the data was cut short, and the decoder would only decode the zeros. + const MAX_PAD: u32 = 8; + + fn new(d: &'a [u8]) -> Self { + let mut ar = Arith { d, pos: 0, c: 0, a: 0, ct: 0, pad: 0, lost: false }; + ar.init(); + ar + } + fn byte_in(&mut self) { - let v = next_byte(self.d, &mut self.pos).unwrap_or(0); + let v = next_byte(self.d, &mut self.pos).unwrap_or_else(|| { + self.pad = self.pad.saturating_add(1); + 0 + }); self.c = self.c.wrapping_add(u32::from(v) << 8); } + /// The scan's data ran out, or a restart marker is missing. + fn exhausted(&self) -> bool { + self.lost || self.pad > Self::MAX_PAD + } + /// INITDEC (D.2.7). fn init(&mut self) { + self.pad = 0; self.a = 0x10000; self.c = 0; self.byte_in(); @@ -650,7 +678,9 @@ impl Arith<'_> { } fn restart(&mut self) { - self.pos = skip_restart(self.d, self.pos); + let (pos, found) = skip_restart(self.d, self.pos); + self.pos = pos; + self.lost |= !found; self.init(); } } @@ -905,6 +935,9 @@ impl ScanDecoder for ArithScan<'_> { ArithMode::AcFirst => self.ac_first(ta, blk)?, ArithMode::AcRefine => self.ac_refine(ta, blk)?, } + if self.ar.exhausted() { + return Err(err("entropy-coded data ends before the scan does")); + } Ok(()) } } @@ -918,16 +951,25 @@ impl ArithScan<'_> { } } +/// The scan's MCU grid (A.2): MCUs per line, MCU lines, and the blocks the scan decodes. +fn scan_grid(frame: &Frame, scan: &Scan) -> Result<(usize, usize, usize), Error> { + if let [(ci, _, _)] = scan.comps[..] { + let c = frame.comps.get(ci).ok_or_else(|| err("bad scan component"))?; + let (mcux, mcuy) = (c.cw.div_ceil(8), c.ch.div_ceil(8)); + return Ok((mcux, mcuy, mcux.saturating_mul(mcuy))); + } + let mut per_mcu = 0usize; + for &(ci, _, _) in &scan.comps { + let c = frame.comps.get(ci).ok_or_else(|| err("bad scan component"))?; + per_mcu += c.h * c.v; + } + Ok((frame.mcux, frame.mcuy, frame.mcux.saturating_mul(frame.mcuy).saturating_mul(per_mcu))) +} + /// Walk the scan's MCUs in order (A.2), handing each block to `dec`; counts decoded MCUs. fn walk_scan(frame: &mut Frame, scan: &Scan, restart: usize, dec: &mut dyn ScanDecoder, done: &mut usize) -> Result<(), Error> { let single = scan.comps.len() == 1; - let (mcux, mcuy) = if single { - let &(ci, _, _) = scan.comps.first().ok_or_else(|| err("empty scan"))?; - let c = frame.comps.get(ci).ok_or_else(|| err("bad scan component"))?; - (c.cw.div_ceil(8), c.ch.div_ceil(8)) - } else { - (frame.mcux, frame.mcuy) - }; + let (mcux, mcuy, _) = scan_grid(frame, scan)?; let total = mcux.saturating_mul(mcuy); for m in 0..total { if restart > 0 && m > 0 && m % restart == 0 { @@ -990,7 +1032,7 @@ fn decode_scan(frame: &mut Frame, tables: &Tables, scan: &Scan, d: &[u8], done: (true, _, _) => ArithMode::AcRefine, }; let mut dec = ArithScan { - ar: Arith { d, pos: 0, c: 0, a: 0, ct: 0 }, + ar: Arith::new(d), mode, ss: scan.ss, se: scan.se, @@ -1004,7 +1046,6 @@ fn decode_scan(frame: &mut Frame, tables: &Tables, scan: &Scan, d: &[u8], done: pred: [0; 4], ctx: [0; 4], }; - dec.ar.init(); walk_scan(frame, scan, tables.restart, &mut dec, done) } else { let mut sel = Vec::with_capacity(scan.comps.len()); @@ -1018,11 +1059,12 @@ fn decode_scan(frame: &mut Frame, tables: &Tables, scan: &Scan, d: &[u8], done: } } -/// The most scans decoded from one stream; later ones are dropped and the image is reported -/// damaged. Every scan walks its whole MCU grid, and an arithmetic-coded scan does so even with -/// no data (the decoder is fed zeros, D.2.6), so without a cap a small file repeating empty -/// scan headers multiplies the work by the frame size. Real files have a few dozen at most. -const MAX_SCANS: usize = 500; +/// The blocks decoded from one stream are capped at this many times the frame's blocks; the +/// scan that would pass the cap and every later one are dropped, and the image is reported +/// damaged. Every scan walks its whole MCU grid, so without a cap a small file repeating scan +/// headers multiplies the work by the frame size. 64 passes allow a progressive file one scan +/// per coefficient for every component; real files make about 10. +const SCAN_PASSES: usize = 64; /// Parse the stream and decode every scan into coefficients. fn decode_coefficients(b: &[u8], limits: &Limits) -> Result<(Frame, Tables, bool), Error> { @@ -1034,6 +1076,7 @@ fn decode_coefficients(b: &[u8], limits: &Limits) -> Result<(Frame, Tables, bool let mut damaged = false; let mut mcus = 0usize; let mut scans = 0usize; + let mut blocks = 0usize; let mut i = 2; while let Some(&v) = b.get(i) { if v != 0xFF { @@ -1079,11 +1122,13 @@ fn decode_coefficients(b: &[u8], limits: &Limits) -> Result<(Frame, Tables, bool } 0xDA => { let f = frame.as_mut().ok_or_else(|| err("scan before the frame header"))?; - if scans >= MAX_SCANS { + let scan = parse_scan(seg, f)?; + let budget = f.comps.iter().map(|c| c.coefs.len()).sum::().saturating_mul(SCAN_PASSES); + blocks = blocks.saturating_add(scan_grid(f, &scan)?.2); + if blocks > budget { damaged = true; break; } - let scan = parse_scan(seg, f)?; let end = scan_end(b, i); let data = b.get(i..end).unwrap_or_default(); scans += 1; diff --git a/crates/jpeg/src/tests.rs b/crates/jpeg/src/tests.rs index 3929e6438..44c10a6bd 100644 --- a/crates/jpeg/src/tests.rs +++ b/crates/jpeg/src/tests.rs @@ -530,12 +530,12 @@ fn malformed_streams_return_errors_not_panics() { assert!(decode_default(&streams[0][..40]).is_err()); } -/// Scans past `MAX_SCANS` aren't decoded (each one walks the whole MCU grid, with or without -/// data): the image keeps what the earlier scans decoded and is reported damaged. +/// Scans past the block budget (`SCAN_PASSES` passes over a one-block frame) aren't decoded: +/// the image keeps what the earlier scans decoded and is reported damaged. #[test] -fn scans_past_the_limit_are_not_decoded() { +fn scans_past_the_block_budget_are_not_decoded() { let blocks = coefficients(1, 3, 500); - for (count, damaged) in [(MAX_SCANS, false), (MAX_SCANS + 1, true)] { + for (count, damaged) in [(SCAN_PASSES, false), (SCAN_PASSES + 1, true)] { let stream = arith_stream(8, 1, 1, &blocks, &vec![(0, 0, 0, 0); count], 0); let (frame, _, d) = decode_coefficients(&stream, &Limits::default()).unwrap(); assert_eq!(d, damaged, "{count} scans"); @@ -549,6 +549,13 @@ fn limits_are_checked_before_allocating() { segment(&mut s, 0xC9, &[8, 0xFF, 0xFF, 0xFF, 0xFF, 3, 1, 0x44, 0, 2, 0x11, 0, 3, 0x11, 0]); let small = Limits { max_alloc: 1 << 20, ..Limits::default() }; assert!(matches!(decode(&s, &small, None), Err(Error::LimitExceeded(_)))); + // 64×64 gray: 32 KiB of coefficients and plane plus 4 KiB of output, each under the limit + // alone but not together. + let stream = arith_stream(8, 8, 8, &coefficients(64, 1, 100), &[(0, 63, 0, 0)], 0); + for (max_alloc, ok) in [(36 << 10, true), ((36 << 10) - 1, false), (32 << 10, false)] { + let limits = Limits { max_alloc, ..Limits::default() }; + assert_eq!(decode(&stream, &limits, None).is_ok(), ok, "max_alloc {max_alloc}"); + } } #[test]