craft-heif: cap hvcC parameter sets and tiles, count alpha-grid tiles (port of lightcraft #3645) (#812)

* craft-heif: cap hvcC parameter sets and tiles, count alpha-grid tiles in the visit budget

* craft-heif: charge every tile reference of a thumbnail candidate to the visit budget

crates/heif/src/lib.rs

@@ -175,13 +175,22 @@ fn probe_unguarded(bytes: &[u8]) -> Result<Info, Error> {
     // items (`iloc` construction method 2) is rebuilt once per extent, so a chain of such items
     // multiplies at every level while each result stays small. Count those visits from the
     // `iloc` table before reconstructing anything: the primary item first (`heic_rs::probe` and
-    // `Context::grid` both rebuild its descriptor), then its alpha plane, then every tile
-    // reference, repeats included.
+    // `Context::grid` both rebuild its descriptor), then its alpha plane and, when that is a grid,
+    // its tiles, then every tile reference, repeats included.
     let mut visits = 0_usize;
     item_visits(&ctx, id, 0, &mut visits)?;
+    // `heic_rs::probe` reads an item without `ispe` from its parameter sets
+    check_item_parameter_sets(&ctx.props(id).map_err(err)?)?;
     let info = heic_rs::probe(bytes).map_err(open_error)?;
-    if let Some(alpha) = ctx.alpha_item(id).map_err(err)? {
+    let alpha = ctx.alpha_item(id).map_err(err)?;
+    if let Some(alpha) = alpha {
         item_visits(&ctx, alpha, 0, &mut visits)?;
+        // an alpha plane may be a grid too: its tiles are rebuilt when it is decoded
+        if let Some((_, tiles)) = ctx.grid(alpha).map_err(err)? {
+            for &tile in tiles.iter().filter(|&&tile| tile != alpha) {
+                item_visits(&ctx, tile, 0, &mut visits)?;
+            }
+        }
     }
     let mut coded = match ctx.grid(id).map_err(err)? {
         Some((_, tiles)) => tiles,
@@ -271,6 +280,28 @@ pub fn decode_thumbnail(bytes: &[u8], min_edge: u32, options: &Options) -> Resul
                 continue;
             }
             let Ok(p) = ctx.props(r.from) else { continue };
+            if check_item_parameter_sets(&p).is_err() {
+                continue;
+            }
+            // A grid whose tile would fail the same check must not win and then fail in
+            // `decode_item`, with no fallback to a larger valid thumbnail.
+            let grid = match ctx.grid(r.from) {
+                Ok(grid) => grid,
+                // an unreadable grid fails in `decode_item` too: not a candidate
+                Err(_) => continue,
+            };
+            if let Some((_, tiles)) = grid {
+                // every reference is rebuilt and decoded, repeats included, so each is charged;
+                // the properties are checked once per distinct tile
+                let mut checked = std::collections::HashSet::new();
+                let tiles_ok = tiles.iter().all(|&t| {
+                    item_visits(&ctx, t, 0, &mut candidate_visits).is_ok()
+                        && (!checked.insert(t) || ctx.props(t).is_ok_and(|tp| check_item_parameter_sets(&tp).is_ok()))
+                });
+                if !tiles_ok {
+                    continue;
+                }
+            }
             let Ok((w, h)) = ctx.coded_size(r.from, &p) else { continue };
             let edge = w.max(h);
             if edge >= min_edge && best.is_none_or(|(_, e)| edge < e) {
@@ -332,7 +363,9 @@ fn decode_item(ctx: &Context<'_>, id: u32, options: &Options, with_alpha: bool)
         return Err(Error::Limit(format!("the image carries {} transform properties, over the {MAX_TRANSFORMS}-property cap", p.transforms.len())));
     }
     let limit = options.max_pixels;
-    // The declared sizes are refused before anything is decoded or allocated.
+    // The declared sizes are refused before anything is decoded or allocated (an item without
+    // `ispe` takes its size from its parameter sets, checked first).
+    check_item_parameter_sets(&p)?;
     let (cw, ch) = ctx.coded_size(id, &p).map_err(err)?;
     heic_rs::image::check_pixels(cw, ch, limit).map_err(err)?;
     let (tw, th) = heic_rs::transform::transformed_size(cw, ch, &p.transforms).map_err(err)?;
@@ -396,6 +429,44 @@ const GRID_BUDGET: u64 = 2;
 /// most one of each, so anything approaching this is a reordered or hostile association list.
 const MAX_TRANSFORMS: usize = 12;
 
+/// The most parameter sets one `hvcC` may carry: one per id H.265 allows (16 VPS, 16 SPS, 64 PPS).
+/// heic-rs 0.1.4 replaces a repeated id by scanning every set kept so far, so more sets (repeats,
+/// or ids past the standard's range) only make that collection quadratic.
+const MAX_PARAMETER_SETS: usize = 16 + 16 + 64;
+/// Highest `pps_pic_parameter_set_id` H.265 allows (7.4.3.3.1).
+const MAX_PPS_ID: u32 = 63;
+/// Most tile columns and rows a picture may have: H.265's level limits (Table A.8, levels 6 to
+/// 6.2). heic-rs 0.1.4 builds its tile tables by walking the tile boundaries for every coding tree
+/// block, so its work grows with blocks × tiles; these keep that a small multiple of the blocks.
+const MAX_TILE_COLS: u32 = 20;
+const MAX_TILE_ROWS: u32 = 22;
+
+/// Refuses parameter sets that would make heic-rs's set collection or tile tables cost far more
+/// than the picture: too many sets, a PPS id out of range, or more tiles than any level allows.
+fn check_parameter_sets(sets: &[&[u8]]) -> Result<(), Error> {
+    if sets.len() > MAX_PARAMETER_SETS {
+        return Err(Error::Limit(format!("the hvcC record carries {} parameter sets, over the {MAX_PARAMETER_SETS}-set cap", sets.len())));
+    }
+    let layouts = vui::pps_layouts(sets).ok_or_else(|| Error::Malformed("unreadable picture parameter set".into()))?;
+    for p in layouts {
+        if p.id > MAX_PPS_ID {
+            return Err(Error::Malformed(format!("picture parameter set id {} is out of range", p.id)));
+        }
+        if p.tile_cols > MAX_TILE_COLS || p.tile_rows > MAX_TILE_ROWS {
+            return Err(Error::Limit(format!(
+                "a picture of {} × {} tiles, over the {MAX_TILE_COLS} × {MAX_TILE_ROWS} H.265 allows",
+                p.tile_cols, p.tile_rows
+            )));
+        }
+    }
+    Ok(())
+}
+
+/// [`check_parameter_sets`] of a coded item's `hvcC`, if it has one.
+fn check_item_parameter_sets(p: &ItemProps<'_>) -> Result<(), Error> {
+    p.hvcc.as_ref().map_or(Ok(()), |h| check_parameter_sets(&h.parameter_sets()))
+}
+
 /// The decoded pictures of an item: one, or a grid's tiles in `dimg` order.
 struct Coded<'a> {
     frames: Vec<Frame>,
@@ -420,15 +491,24 @@ impl<'a> Coded<'a> {
         // conformance window (or a slice picking another SPS) can make far larger than the
         // declared one, and a grid holds all its tiles at once. Both are refused before decoding.
         let mut total = 0u64;
+        // A grid may list one tile thousands of times: each distinct tile is read and checked
+        // once, so the work follows the file's items, not the grid's cell count.
+        let mut largest_of: std::collections::HashMap<u32, u64> = std::collections::HashMap::new();
         for t in &ids {
+            if let Some(largest) = largest_of.get(t) {
+                total = total.saturating_add(*largest);
+                continue;
+            }
             let tp = ctx.props(*t).map_err(err)?;
             let hvcc = tp.hvcc.as_ref().ok_or(heic_rs::Error::MissingBox("hvcC")).map_err(err)?;
+            check_parameter_sets(&hvcc.parameter_sets())?;
             let sizes = vui::coded_sizes(&hvcc.parameter_sets()).ok_or_else(|| Error::Malformed("unreadable sequence parameter set".into()))?;
             let mut largest = 0u64;
             for (w, h) in sizes {
                 heic_rs::image::check_pixels(w, h, limit).map_err(err)?;
                 largest = largest.max(u64::from(w) * u64::from(h));
             }
+            largest_of.insert(*t, largest);
             total = total.saturating_add(largest);
         }
         if total > limit.saturating_mul(GRID_BUDGET) {

crates/heif/src/testdata.rs

@@ -47,11 +47,29 @@ pub struct Spec<'a> {
     pub thumbnail: Option<Sampler<'a>>,
     /// Auxiliary images of these types the decoder must ignore (HDR gain map, depth…).
     pub ignored_aux: Vec<&'static str>,
+    /// When nonzero, the alpha image is instead a grid of `n` × `n` references to one alpha tile
+    /// (a hostile file: the same stored picture rebuilt for every reference).
+    pub alpha_tile_repeats: u8,
+    /// When nonzero, a thumbnail that is a grid of `n` × `n` references to one tile, written
+    /// after any `thumbnail` item (a hostile file, as with `alpha_tile_repeats`).
+    pub thumbnail_tile_repeats: u8,
 }
 
 impl<'a> Spec<'a> {
     pub fn new(width: u32, height: u32, sample: Sampler<'a>) -> Spec<'a> {
-        Spec { width, height, bit_depth: 8, sample, props: Vec::new(), exif: None, alpha: None, thumbnail: None, ignored_aux: Vec::new() }
+        Spec {
+            width,
+            height,
+            bit_depth: 8,
+            sample,
+            props: Vec::new(),
+            exif: None,
+            alpha: None,
+            thumbnail: None,
+            ignored_aux: Vec::new(),
+            alpha_tile_repeats: 0,
+            thumbnail_tile_repeats: 0,
+        }
     }
 }
 
@@ -524,7 +542,20 @@ pub fn build(spec: &Spec<'_>) -> Vec<u8> {
         }
         f.refs.push((b"auxl", id, vec![primary]));
     };
-    if let Some(alpha) = spec.alpha {
+    if let (Some(alpha), n @ 1..) = (spec.alpha, spec.alpha_tile_repeats) {
+        let tile = f.image(TILE, TILE, bd, true, &|x, y| [alpha(x, y), 0, 0]);
+        if let Some(it) = f.items.iter_mut().find(|i| i.id == tile) {
+            it.hidden = true;
+        }
+        let side = u32::from(n) * TILE;
+        let mut g = vec![0u8, 0, n - 1, n - 1];
+        g.extend_from_slice(&(side as u16).to_be_bytes());
+        g.extend_from_slice(&(side as u16).to_be_bytes());
+        let auxc = f.prop(full_box(b"auxC", 0, b"urn:mpeg:hevc:2015:auxid:1\0"));
+        let id = f.add(b"grid", g, vec![(auxc, true)]);
+        f.refs.push((b"dimg", id, vec![tile; usize::from(n) * usize::from(n)]));
+        f.refs.push((b"auxl", id, vec![primary]));
+    } else if let Some(alpha) = spec.alpha {
         aux(&mut f, "urn:mpeg:hevc:2015:auxid:1", &|x, y| [alpha(x, y), 0, 0]);
     }
     for urn in &spec.ignored_aux {
@@ -538,5 +569,18 @@ pub fn build(spec: &Spec<'_>) -> Vec<u8> {
         }
         f.refs.push((b"thmb", id, vec![primary]));
     }
+    if let n @ 1.. = spec.thumbnail_tile_repeats {
+        let tile = f.image(TILE, TILE, bd, false, spec.sample);
+        if let Some(it) = f.items.iter_mut().find(|i| i.id == tile) {
+            it.hidden = true;
+        }
+        let side = u32::from(n) * TILE;
+        let mut g = vec![0u8, 0, n - 1, n - 1];
+        g.extend_from_slice(&(side as u16).to_be_bytes());
+        g.extend_from_slice(&(side as u16).to_be_bytes());
+        let id = f.add(b"grid", g, Vec::new());
+        f.refs.push((b"dimg", id, vec![tile; usize::from(n) * usize::from(n)]));
+        f.refs.push((b"thmb", id, vec![primary]));
+    }
     f.write(primary)
 }

crates/heif/src/tests_synthetic.rs

@@ -139,6 +139,35 @@ fn alpha_comes_from_its_auxiliary_image() {
     assert_eq!(d.data[..3], [luma(0, 0) * 257; 3]);
 }
 
+#[test]
+fn an_alpha_grid_decodes_and_its_tile_references_share_the_visit_budget() {
+    let alpha = |x: u32, _: u32| (x * 4) as u16;
+    // a 2 × 2 alpha grid naming one tile four times: a valid file
+    let d = decoded(&Spec { alpha: Some(&alpha), alpha_tile_repeats: 2, ..Spec::new(64, 64, &grey) });
+    assert!(d.has_alpha);
+    assert_eq!((d.data[10 * 4 + 3], d.data[42 * 4 + 3]), (40 * 257, 40 * 257));
+    // 70 × 70 references to that tile: each would be rebuilt and decoded, so it is refused first
+    let bytes = build(&Spec { alpha: Some(&alpha), alpha_tile_repeats: 70, ..Spec::new(32, 32, &grey) });
+    for r in [probe(&bytes).map(|_| ()), decode(&bytes, &Options::default()).map(|_| ())] {
+        assert!(matches!(&r, Err(Error::Limit(m)) if m.contains("extent reads")), "{r:?}");
+    }
+}
+
+/// A thumbnail grid naming one tile thousands of times is not a candidate (each reference would be
+/// rebuilt and decoded). The candidates share one visit budget, so a valid thumbnail listed before
+/// it still wins; a small repeat count is fine.
+#[test]
+fn a_thumbnail_grid_charges_every_tile_reference() {
+    let ok = Spec { thumbnail_tile_repeats: 2, ..Spec::new(64, 64, &grey) };
+    let t = decode_thumbnail(&build(&ok), 1, &Options::default()).unwrap();
+    assert!(t.is_some(), "a 2 x 2 repeated-tile thumbnail is valid");
+    let hostile = Spec { thumbnail_tile_repeats: 70, ..Spec::new(64, 64, &grey) };
+    assert!(decode_thumbnail(&build(&hostile), 1, &Options::default()).unwrap().is_none());
+    let both = Spec { thumbnail_tile_repeats: 70, thumbnail: Some(&grey), ..Spec::new(64, 64, &grey) };
+    let t = decode_thumbnail(&build(&both), 1, &Options::default()).unwrap().expect("the plain thumbnail is the candidate");
+    assert_eq!((t.width, t.height), (32, 32));
+}
+
 #[test]
 fn limits_are_checked_before_decoding() {
     let bytes = build(&Spec::new(96, 64, &grey));

crates/heif/src/vui.rs

@@ -49,6 +49,55 @@ pub fn coded_sizes(parameter_sets: &[&[u8]]) -> Option<Vec<(u32, u32)>> {
         .collect()
 }
 
+/// NAL unit type of a picture parameter set.
+const PPS_NUT: u8 = 34;
+
+/// What [`pps_layouts`] reads of a picture parameter set.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct PpsLayout {
+    /// `pps_pic_parameter_set_id`.
+    pub id: u32,
+    /// `num_tile_columns_minus1 + 1` (1 without tiles).
+    pub tile_cols: u32,
+    /// `num_tile_rows_minus1 + 1` (1 without tiles).
+    pub tile_rows: u32,
+}
+
+/// The id and tile layout of every PPS among `parameter_sets`, read up to the tile counts
+/// (H.265 7.3.2.3.1). `None`: a PPS this reader can't follow.
+pub fn pps_layouts(parameter_sets: &[&[u8]]) -> Option<Vec<PpsLayout>> {
+    parameter_sets
+        .iter()
+        .filter(|n| n.first().is_some_and(|b| (b >> 1) & 0x3f == PPS_NUT))
+        .map(|pps| {
+            let rbsp = unescape(pps.get(2..)?);
+            pps_layout(&mut Bits { data: &rbsp, pos: 0 })
+        })
+        .collect()
+}
+
+fn pps_layout(r: &mut Bits) -> Option<PpsLayout> {
+    let id = r.ue()?; // pps_pic_parameter_set_id
+    r.ue()?; // pps_seq_parameter_set_id
+    // dependent slices, output flag, 3 extra slice header bits, sign hiding, cabac_init_present
+    r.skip(7)?;
+    r.ue()?; // num_ref_idx_l0_default_active_minus1
+    r.ue()?; // num_ref_idx_l1_default_active_minus1
+    r.se()?; // init_qp_minus26
+    r.skip(2)?; // constrained_intra_pred_flag, transform_skip_enabled_flag
+    if r.flag()? {
+        r.ue()?; // diff_cu_qp_delta_depth
+    }
+    r.se()?; // pps_cb_qp_offset
+    r.se()?; // pps_cr_qp_offset
+    // slice chroma qp offsets present, weighted pred, weighted bipred, transquant bypass
+    r.skip(4)?;
+    let tiles = r.flag()?; // tiles_enabled_flag
+    r.skip(1)?; // entropy_coding_sync_enabled_flag
+    let (tile_cols, tile_rows) = if tiles { (r.ue()?.checked_add(1)?, r.ue()?.checked_add(1)?) } else { (1, 1) };
+    Some(PpsLayout { id, tile_cols, tile_rows })
+}
+
 /// Reads an SPS up to and including the picture size, leaving `r` after it: width, height and
 /// `sps_max_sub_layers_minus1`.
 fn sps_size(r: &mut Bits) -> Option<(u32, u32, u32)> {
@@ -456,4 +505,63 @@ mod tests {
         }
         assert_eq!(signal(&[&[SPS_NUT << 1, 1, 0xff, 0xff, 0xff]]), None);
     }
+
+    /// A PPS NAL as encoders write it, with `tiles` (columns, rows) when tiles are enabled.
+    fn pps(id: u32, cu_qp_delta: bool, tiles: Option<(u32, u32)>) -> Vec<u8> {
+        let mut w = Writer::default();
+        w.ue(id);
+        w.ue(0); // sps id
+        w.u(0, 7);
+        w.ue(0);
+        w.ue(0);
+        w.ue(3); // init_qp_minus26 = -2
+        w.u(0, 2);
+        w.u(u32::from(cu_qp_delta), 1);
+        if cu_qp_delta {
+            w.ue(1);
+        }
+        w.ue(1); // cb offset +1
+        w.ue(2); // cr offset -1
+        w.u(0, 4);
+        w.u(u32::from(tiles.is_some()), 1);
+        w.u(0, 1);
+        if let Some((c, r)) = tiles {
+            w.ue(c - 1);
+            w.ue(r - 1);
+            w.u(1, 1); // uniform spacing
+            w.u(1, 1);
+        }
+        w.u(0, 3);
+        [vec![PPS_NUT << 1, 1], w.bytes()].concat()
+    }
+
+    #[test]
+    fn reads_pps_ids_and_tile_counts() {
+        let plain = pps(5, true, None);
+        let tiled = pps(63, false, Some((4, 3)));
+        let sps_nal = sps(|w| w.u(0, 1), |w| w.ue(0));
+        assert_eq!(
+            pps_layouts(&[&sps_nal, &plain, &tiled]),
+            Some(vec![PpsLayout { id: 5, tile_cols: 1, tile_rows: 1 }, PpsLayout { id: 63, tile_cols: 4, tile_rows: 3 }])
+        );
+        assert_eq!(pps_layouts(&[&[PPS_NUT << 1, 1]]), None);
+        assert!(crate::check_parameter_sets(&[&sps_nal, &plain, &tiled]).is_ok());
+        assert!(crate::check_parameter_sets(&[&pps(0, false, Some((20, 22)))]).is_ok());
+    }
+
+    #[test]
+    fn refuses_parameter_sets_that_multiply_decoder_work() {
+        // a PPS id past 63: heic-rs would keep every distinct one and rescan them all for each
+        let r = crate::check_parameter_sets(&[&pps(64, false, None)]);
+        assert!(matches!(r, Err(crate::Error::Malformed(_))), "{r:?}");
+        // thousands of PPS records (even with valid, repeated ids): quadratic collection
+        let one = pps(1, false, None);
+        let many: Vec<&[u8]> = vec![&one; 5000];
+        assert!(matches!(crate::check_parameter_sets(&many), Err(crate::Error::Limit(_))));
+        // more tiles than any H.265 level allows: tile tables cost blocks × tiles
+        for (c, r) in [(21, 1), (1, 23), (1024, 1024)] {
+            let r = crate::check_parameter_sets(&[&pps(0, false, Some((c, r)))]);
+            assert!(matches!(r, Err(crate::Error::Limit(_))), "{r:?}");
+        }
+    }
 }