jpeg: bound decoder memory and work

Check coefficients, planes and output against max_alloc together. Stop an
arithmetic scan once it has consumed more than 8 zero bytes past its data or
misses a restart marker, and cap the decoded blocks at 64 passes over the
frame instead of capping the scan count.

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<Frame, Error>
     }
     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::<usize>().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;

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]