From 269495cd3867e9feee1876ac641627ba23c7dabb Mon Sep 17 00:00:00 2001
From: "Dylan M. Taylor" <dylan@dylanmtaylor.com>
Date: Sun, 25 Oct 2026 15:45:31 -0400
Subject: [PATCH] security(heif): budget every reconstruction before it runs

The primary item and the selected thumbnail are scanned before Context::grid rebuilds their descriptors, every tile reference is counted (repeats included) before the size check dedupes them, and the Exif and XMP items are scanned before they are read.
---
 crates/heif/src/lib.rs | 48 +++++++++++++++++++++++++++---------------
 1 file changed, 31 insertions(+), 17 deletions(-)

diff --git a/crates/heif/src/lib.rs b/crates/heif/src/lib.rs
index 723ffe5d9..6491489fb 100644
--- a/crates/heif/src/lib.rs
+++ b/crates/heif/src/lib.rs
@@ -158,27 +158,28 @@ fn probe_unguarded(bytes: &[u8]) -> Result<Info, Error> {
     })?;
     let ctx = Context::open(bytes).map_err(err)?;
     let id = ctx.meta.primary;
+    // heic-rs 0.1.4 caps the size of what it assembles, not the work: an item built from other
+    // 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 (whose grid descriptor
+    // `Context::grid` rebuilds) first and then every tile reference, repeats included.
+    let mut visits = 0_usize;
+    item_visits(&ctx, id, 0, &mut visits)?;
     let mut coded = match ctx.grid(id).map_err(err)? {
         Some((_, tiles)) => tiles,
         None => vec![id],
     };
-    // A `dimg` listing may name the same tile more than once; reconstrutable (`iloc`
-    // construction method 2) items are moreover rebuilt for every extent that references
-    // them in heic-rs 0.1.4, so our probe would pay the reconstruction again per tile. Each
-    // unique item only once, then the total reconstruction (all of it assembled from `bytes`
-    // by heic-rs) is held below one copy of the file — above that a repeat or chain wants
-    // more bytes than exist (CodeRabbit #21/#50 on this sweep). Order doesn't matter here:
-    // this only checks that every tile's bytes lie in the file.
-    coded.sort_unstable();
-    coded.dedup();
-    // heic-rs 0.1.4 caps the size of what it assembles, not the work: an item built from other
-    // 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.
-    let mut visits = 0_usize;
-    for &item in &coded {
+    for &item in coded.iter().filter(|&&item| item != id) {
         item_visits(&ctx, item, 0, &mut visits)?;
     }
+    // A `dimg` listing may name the same tile more than once, and reconstructable items are
+    // moreover rebuilt for every extent that references them in heic-rs 0.1.4. For the size
+    // check each unique item counts once: the total reconstruction (all of it assembled from
+    // `bytes` by heic-rs) is held below one copy of the file, above which a repeat or chain
+    // wants more bytes than exist. Order doesn't matter here: this only checks that every
+    // tile's bytes lie in the file.
+    coded.sort_unstable();
+    coded.dedup();
     let mut reconstructed = 0_usize;
     for item in coded {
         reconstructed = reconstructed.saturating_add(ctx.item_data(item).map_err(err)?.len());
@@ -252,6 +253,15 @@ pub fn decode_thumbnail(bytes: &[u8], min_edge: u32, options: &Options) -> Resul
             }
         }
         let Some((id, _)) = best else { return Ok(None) };
+        // The same pre-scan as `probe`: the thumbnail's own descriptor and every tile reference
+        // are counted before heic-rs rebuilds any of them.
+        let mut visits = 0_usize;
+        item_visits(&ctx, id, 0, &mut visits)?;
+        if let Some((_, tiles)) = ctx.grid(id).map_err(err)? {
+            for &tile in &tiles {
+                item_visits(&ctx, tile, 0, &mut visits)?;
+            }
+        }
         let mut decoded = decode_item(&ctx, id, options, false)?;
         if decoded.icc.is_none() {
             decoded.icc = item_icc(&ctx, primary);
@@ -444,8 +454,12 @@ fn colour(p: &ItemProps<'_>, coded: &ItemProps<'_>) -> Colour {
 /// metadata item never fails a decode whose pixels came out fine.
 fn metadata(ctx: &Context<'_>) -> (Option<Vec<u8>>, Option<String>) {
     let id = ctx.meta.primary;
-    let exif = ctx.exif(id).ok().flatten().map(<[u8]>::to_vec);
-    let xmp = ctx.xmp_item(id).and_then(|x| ctx.item_data(x).ok()).and_then(|d| match d {
+    // Metadata is best effort, so an item that would cost too much to rebuild is skipped, not an error.
+    let mut visits = 0_usize;
+    let cheap = |item: Option<u32>, visits: &mut usize| item.is_none_or(|item| item_visits(ctx, item, 0, visits).is_ok());
+    let exif = if cheap(ctx.exif_item(id), &mut visits) { ctx.exif(id).ok().flatten().map(<[u8]>::to_vec) } else { None };
+    let xmp_item = ctx.xmp_item(id).filter(|&x| cheap(Some(x), &mut visits));
+    let xmp = xmp_item.and_then(|x| ctx.item_data(x).ok()).and_then(|d| match d {
         Cow::Borrowed(b) => String::from_utf8(b.to_vec()).ok(),
         Cow::Owned(b) => String::from_utf8(b).ok(),
     });