address review: FBX textures: key image cache by the source used, charge decoded size before decoding, propagate budget errors

crates/model/src/fbx.rs

@@ -859,50 +859,71 @@ fn parse_within(bytes: &[u8], resolve: &dyn Fn(&str) -> Option<Vec<u8>>, limit:
     let mut textures: Vec<Arc<Texture>> = vec![];
     let mut material_of: HashMap<i64, usize> = HashMap::new();
     let mut texture_of: HashMap<i64, Option<usize>> = HashMap::new();
-    // Decoded images by their source (embedded Video, then file name): Texture objects aliasing
-    // one image share its pixels instead of each keeping a copy.
-    let mut image_of: HashMap<(Vec<i64>, Vec<String>), Option<usize>> = HashMap::new();
-    let mut load_texture = |tid: i64, textures: &mut Vec<Arc<Texture>>| -> Option<usize> {
+    // Decoded images by the source actually used (embedded Video id, else the resolved file
+    // name): Texture objects aliasing one image share its pixels instead of each keeping a copy.
+    let mut image_of: HashMap<(Option<i64>, String), Option<usize>> = HashMap::new();
+    // Ok(None): no usable image (the material goes untextured). Err: the memory budget ran out.
+    let mut load_texture = |tid: i64, textures: &mut Vec<Arc<Texture>>| -> Result<Option<usize>> {
         if let Some(t) = texture_of.get(&tid) {
-            return *t;
+            return Ok(*t);
         }
-        let t = (|| {
-            let tex = by_id.get(&tid)?;
-            let videos: Vec<i64> = children_of(tid).filter(|l| kind(l.0) == Some("Video")).map(|l| l.0).collect();
-            let files: Vec<String> = ["RelativeFilename", "FileName"].iter().filter_map(|k| Some(tex.child(k)?.str_at(0)?.replace('\\', "/"))).collect();
-            let key = (videos, files);
-            if let Some(i) = image_of.get(&key) {
-                return *i;
-            }
-            let embedded = key.0.iter().find_map(|v| {
-                let content = by_id.get(v)?.child("Content")?;
-                match content.props.first()? {
-                    Prop::Bytes(b) if !b.is_empty() => Some(b.clone()),
-                    Prop::Str(_) => {
-                        let s: String = content.props.iter().filter_map(Prop::as_str).collect();
-                        base64(&s).filter(|b| !b.is_empty())
-                    }
-                    _ => None,
+        let Some(tex) = by_id.get(&tid) else {
+            texture_of.insert(tid, None);
+            return Ok(None);
+        };
+        let videos: Vec<i64> = children_of(tid).filter(|l| kind(l.0) == Some("Video")).map(|l| l.0).collect();
+        let files: Vec<String> = ["RelativeFilename", "FileName"].iter().filter_map(|k| Some(tex.child(k)?.str_at(0)?.replace('\\', "/"))).collect();
+        let embedded = videos.iter().find_map(|&v| {
+            let content = by_id.get(&v)?.child("Content")?;
+            let bytes = match content.props.first()? {
+                Prop::Bytes(b) if !b.is_empty() => Some(b.clone()),
+                Prop::Str(_) => {
+                    let s: String = content.props.iter().filter_map(Prop::as_str).collect();
+                    base64(&s).filter(|b| !b.is_empty())
                 }
-            });
-            let i = (|| {
-                let bytes = embedded.or_else(|| key.1.iter().find_map(|f| resolve(f).or_else(|| resolve(f.rsplit('/').next()?))))?;
-                let t = Texture::decode(&bytes).ok()?;
-                budget.take(t.data.len().checked_mul(16), "a texture").ok()?;
-                textures.push(Arc::new(t));
-                Some(textures.len() - 1)
-            })();
-            image_of.insert(key, i);
-            i
-        })();
+                _ => None,
+            }?;
+            Some(((Some(v), String::new()), bytes))
+        });
+        let source = embedded.or_else(|| {
+            files.iter().find_map(|f| {
+                let g = resolve(f).map(|b| (f.clone(), b)).or_else(|| {
+                    let base = f.rsplit('/').next()?;
+                    Some((base.to_string(), resolve(base)?))
+                })?;
+                Some(((None, g.0), g.1))
+            })
+        });
+        let t = match source {
+            None => None,
+            Some((key, bytes)) => match image_of.get(&key) {
+                Some(i) => *i,
+                None => {
+                    // Charge the decoded size before decoding, so a small file claiming huge
+                    // dimensions is refused before its pixels are allocated.
+                    let i = match Texture::dimensions(&bytes) {
+                        Ok((w, h)) => {
+                            budget.take((w as usize).checked_mul(h as usize).and_then(|n| n.checked_mul(16)), "a texture")?;
+                            Texture::decode(&bytes).ok().map(|t| {
+                                textures.push(Arc::new(t));
+                                textures.len() - 1
+                            })
+                        }
+                        Err(_) => None,
+                    };
+                    image_of.insert(key, i);
+                    i
+                }
+            },
+        };
         texture_of.insert(tid, t);
-        t
+        Ok(t)
     };
-    let mut material_index = |mid: i64, materials: &mut Vec<Material>, textures: &mut Vec<Arc<Texture>>| -> Option<usize> {
+    let mut material_index = |mid: i64, materials: &mut Vec<Material>, textures: &mut Vec<Arc<Texture>>| -> Result<Option<usize>> {
         if let Some(&m) = material_of.get(&mid) {
-            return Some(m);
+            return Ok(Some(m));
         }
-        let node = by_id.get(&mid)?;
+        let Some(node) = by_id.get(&mid) else { return Ok(None) };
         let mut m = material(short_name(node.str_at(1).unwrap_or("")), node);
         for (tid, prop) in children_of(mid).filter(|l| kind(l.0) == Some("Texture")) {
             let slot = match prop {
@@ -912,7 +933,7 @@ fn parse_within(bytes: &[u8], resolve: &dyn Fn(&str) -> Option<Vec<u8>>, limit:
                 _ => continue,
             };
             if slot.is_none() {
-                *slot = load_texture(tid, textures).map(|texture| TexRef { texture, uv_set: 0 });
+                *slot = load_texture(tid, textures)?.map(|texture| TexRef { texture, uv_set: 0 });
             }
         }
         if m.emissive_tex.is_some() && m.emissive == [0.0; 3] {
@@ -920,7 +941,7 @@ fn parse_within(bytes: &[u8], resolve: &dyn Fn(&str) -> Option<Vec<u8>>, limit:
         }
         materials.push(m);
         material_of.insert(mid, materials.len() - 1);
-        Some(materials.len() - 1)
+        Ok(Some(materials.len() - 1))
     };
 
     // One mesh per Model with Mesh geometry (material slots are per Model). Models instancing the
@@ -931,8 +952,10 @@ fn parse_within(bytes: &[u8], resolve: &dyn Fn(&str) -> Option<Vec<u8>>, limit:
         let Some(&ni) = node_of.get(&id) else { continue };
         let Some(gid) = children_of(id).find(|l| kind(l.0) == Some("Geometry")).map(|l| l.0) else { continue };
         let Some(geom) = by_id.get(&gid) else { continue };
-        let mats: Vec<usize> =
-            children_of(id).filter(|l| kind(l.0) == Some("Material")).filter_map(|l| material_index(l.0, &mut materials, &mut textures)).collect();
+        let mut mats: Vec<usize> = vec![];
+        for (mid, _) in children_of(id).filter(|l| kind(l.0) == Some("Material")) {
+            mats.extend(material_index(mid, &mut materials, &mut textures)?);
+        }
         let name = nodes.get(ni).map(|n| n.name.clone()).unwrap_or_default();
         let mesh_index = match mesh_of.get(&(gid, mats.clone())) {
             Some(&i) => i,

crates/model/src/fbx/tests.rs

@@ -455,3 +455,64 @@ Connections:  {"#;
     assert_eq!(parse_within(doc.as_bytes(), &resolve, values + mesh + 32).unwrap().textures.len(), 1);
     assert!(parse_within(doc.as_bytes(), &resolve, values + mesh + 31).unwrap_err().to_string().contains("memory limit"));
 }
+
+/// A PNG that is only a header claiming `w`×`h` pixels (an empty data chunk follows).
+fn png_header_only(w: u32, h: u32) -> Vec<u8> {
+    fn chunk(png: &mut Vec<u8>, ty: &[u8], data: &[u8]) {
+        let body = [ty, data].concat();
+        let mut crc = 0xFFFF_FFFFu32;
+        for &b in &body {
+            crc ^= u32::from(b);
+            for _ in 0..8 {
+                crc = if crc & 1 != 0 { (crc >> 1) ^ 0xEDB8_8320 } else { crc >> 1 };
+            }
+        }
+        png.extend((data.len() as u32).to_be_bytes());
+        png.extend(&body);
+        png.extend((!crc).to_be_bytes());
+    }
+    let mut png = vec![0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A];
+    chunk(&mut png, b"IHDR", &[w.to_be_bytes(), h.to_be_bytes(), [8, 6, 0, 0]].concat().iter().copied().chain([0]).collect::<Vec<u8>>());
+    chunk(&mut png, b"IDAT", &[]);
+    png
+}
+
+#[test]
+fn texture_budget_is_charged_before_decoding_and_its_error_propagates() {
+    let objects = r#"
+	Texture: 600, "Texture::Base", "" {
+		RelativeFilename: "maps\base.png"
+	}
+}
+Connections:  {"#;
+    let doc = ASCII.replace("}\nConnections:  {", objects).replace("C: \"OO\",500,300\n", "C: \"OO\",500,300\n\tC: \"OP\",600,400, \"DiffuseColor\"\n");
+    // 30000×30000 pixels would take ~14 GiB; the header alone must be refused against the budget.
+    let huge = png_header_only(30000, 30000);
+    let resolve = |_: &str| Some(huge.clone());
+    assert!(parse(doc.as_bytes(), &resolve).unwrap_err().to_string().contains("memory limit"));
+}
+
+#[test]
+fn texture_files_resolving_to_one_image_share_it() {
+    let png = {
+        let img = image::RgbaImage::from_pixel(2, 1, image::Rgba([0, 255, 0, 255]));
+        let mut b = std::io::Cursor::new(vec![]);
+        img.write_to(&mut b, image::ImageFormat::Png).unwrap();
+        b.into_inner()
+    };
+    let objects = r#"
+	Texture: 600, "Texture::Base", "" {
+		RelativeFilename: "maps\base.png"
+	}
+	Texture: 700, "Texture::Other spelling", "" {
+		RelativeFilename: "missing\base.png"
+		FileName: "maps\base.png"
+	}
+}
+Connections:  {"#;
+    let doc = ASCII
+        .replace("}\nConnections:  {", objects)
+        .replace("C: \"OO\",500,300\n", "C: \"OO\",500,300\n\tC: \"OP\",600,400, \"DiffuseColor\"\n\tC: \"OP\",700,500, \"DiffuseColor\"\n");
+    let resolve = |uri: &str| (uri == "maps/base.png").then(|| png.clone());
+    assert_eq!(parse(doc.as_bytes(), &resolve).unwrap().textures.len(), 1);
+}

crates/model/src/lib.rs

@@ -68,6 +68,11 @@ impl Texture {
     pub fn solid(c: [f32; 4]) -> Texture {
         Texture { width: 1, height: 1, data: vec![c], wrap: [Wrap::Repeat; 2] }
     }
+    /// The pixel size an image claims, read from its header without decoding it.
+    pub fn dimensions(bytes: &[u8]) -> Result<(u32, u32)> {
+        let reader = image::ImageReader::new(std::io::Cursor::new(bytes)).with_guessed_format().map_err(|e| perr(format!("image: {e}")))?;
+        reader.into_dimensions().map_err(|e| perr(format!("image: {e}")))
+    }
     /// Decode a PNG or JPEG.
     pub fn decode(bytes: &[u8]) -> Result<Texture> {
         let img = image::load_from_memory(bytes).map_err(|e| perr(format!("image: {e}")))?.to_rgba32f();