diff --git a/crates/format/src/image_budget_tests.rs b/crates/format/src/image_budget_tests.rs new file mode 100644 index 00000000..58bb4da1 --- /dev/null +++ b/crates/format/src/image_budget_tests.rs @@ -0,0 +1,93 @@ +//! The images of a native file within a budget (#189): however many members each within the +//! per-member limit a file holds, the images read from it together stay within a byte and a count +//! budget. The document's own images are read first, and the ones past the budget are reported. + +use std::io::Write; + +use figcraft_doc::Document; + +use super::*; + +fn png(k: u8) -> Vec { + let img = image::RgbaImage::from_pixel(2, 2, image::Rgba([k, 7, 0, 255])); + let mut out = Vec::new(); + img.write_to(&mut std::io::Cursor::new(&mut out), image::ImageFormat::Png).unwrap(); + out +} + +/// A saved document with `n` images, and their hashes and total size. +fn file_with_images(n: u8) -> (Vec, Vec, u64) { + let mut d = Document::new("File"); + let mut hashes = Vec::new(); + let mut total = 0; + for k in 0..n { + let bytes = png(k); + total += bytes.len() as u64; + hashes.push(add_image(&mut d, bytes).unwrap().0); + } + (save_bytes(&d, None).unwrap(), hashes, total) +} + +/// `bytes` with one more member, `name` holding `data`. +fn with_member(bytes: &[u8], name: &str, data: &[u8]) -> Vec { + let mut z = zip::ZipArchive::new(std::io::Cursor::new(bytes)).unwrap(); + let names: Vec = z.file_names().map(str::to_string).collect(); + let mut buf = std::io::Cursor::new(Vec::new()); + let mut w = zip::ZipWriter::new(&mut buf); + let stored = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored); + for n in names { + let member = read_member(&mut z, &n, u64::MAX).unwrap(); + w.start_file(n, stored).unwrap(); + w.write_all(&member).unwrap(); + } + w.start_file(name, stored).unwrap(); + w.write_all(data).unwrap(); + w.finish().unwrap(); + buf.into_inner() +} + +#[test] +fn images_exactly_within_the_budget_all_load() { + let (bytes, hashes, total) = file_with_images(3); + let loaded = load_within(&bytes, ImageBudget { images: 3, bytes: total }).unwrap(); + assert!(loaded.issues.is_empty(), "{:?}", loaded.issues); + assert!(hashes.iter().all(|h| loaded.doc.images.contains_key(h))); +} + +#[test] +fn images_past_the_byte_budget_are_left_out_and_reported() { + let (bytes, hashes, total) = file_with_images(3); + let loaded = load_within(&bytes, ImageBudget { images: 100, bytes: total - 1 }).unwrap(); + let held: u64 = loaded.doc.images.values().map(|i| i.bytes.len() as u64).sum(); + assert!(held < total, "{held} bytes held"); + assert_eq!(loaded.doc.images.len(), 2); + let lost: Vec<&String> = hashes.iter().filter(|h| !loaded.doc.images.contains_key(*h)).collect(); + assert_eq!(lost.len(), 1); + assert!(loaded.issues.iter().any(|i| i.contains(lost[0].as_str()) && i.contains("together")), "{:?}", loaded.issues); +} + +#[test] +fn images_past_the_count_budget_are_left_out_and_reported() { + let (bytes, _, _) = file_with_images(4); + let loaded = load_within(&bytes, ImageBudget { images: 3, bytes: u64::MAX }).unwrap(); + assert_eq!(loaded.doc.images.len(), 3); + assert!(loaded.issues.iter().any(|i| i.contains("more than 3 images")), "{:?}", loaded.issues); +} + +#[test] +fn the_documents_own_images_are_read_before_others() { + let (bytes, hashes, total) = file_with_images(2); + // An image the document doesn't list, named to come first in the archive's order. + let extra = with_member(&bytes, "images/0000", &png(99)); + let loaded = load_within(&extra, ImageBudget { images: 2, bytes: total }).unwrap(); + assert!(hashes.iter().all(|h| loaded.doc.images.contains_key(h)), "{:?}", loaded.issues); + assert!(!loaded.issues.iter().any(|i| hashes.iter().any(|h| i.contains(h.as_str()))), "{:?}", loaded.issues); +} + +#[test] +fn a_file_within_the_shipped_budget_loads_as_before() { + let (bytes, hashes, _) = file_with_images(3); + let loaded = load_bytes_with_versions(&bytes).unwrap(); + assert!(loaded.issues.is_empty(), "{:?}", loaded.issues); + assert!(hashes.iter().all(|h| loaded.doc.images.contains_key(h))); +} diff --git a/crates/format/src/lib.rs b/crates/format/src/lib.rs index 56798d14..f88416e2 100644 --- a/crates/format/src/lib.rs +++ b/crates/format/src/lib.rs @@ -47,6 +47,8 @@ pub const MAX_VERSION_DESCRIPTION: usize = 2000; /// Largest version index and total size of the version snapshots read from one file. const MAX_VERSION_INDEX: u64 = 16 * 1024 * 1024; const MAX_VERSION_BYTES: u64 = 1024 * 1024 * 1024; +/// Most images, and image bytes together, read from one file (members past it are left out). +const IMAGE_BUDGET: ImageBudget = ImageBudget { images: 20_000, bytes: 4 * 1024 * 1024 * 1024 }; const BRANCH_INDEX: &str = "branches/index.json"; /// Most branches a file keeps: making more is refused, extra entries in a file are skipped. pub const MAX_BRANCHES: usize = 100; @@ -453,8 +455,20 @@ pub(crate) fn read_member(z: &mut zip::ZipArchive, name: &str Ok(data) } +/// How many images, and how many bytes of them together, one file may give. +#[derive(Clone, Copy, Debug)] +pub(crate) struct ImageBudget { + pub images: usize, + pub bytes: u64, +} + /// Read a `.figcraft` archive (or bare JSON) with its version history. pub fn load_bytes_with_versions(bytes: &[u8]) -> Result { + load_within(bytes, IMAGE_BUDGET) +} + +/// [`load_bytes_with_versions`] reading at most `budget` of images. +pub(crate) fn load_within(bytes: &[u8], budget: ImageBudget) -> Result { let trimmed = bytes.iter().position(|b| !b.is_ascii_whitespace()).and_then(|i| bytes.get(i)); if trimmed == Some(&b'{') { let text = std::str::from_utf8(bytes).map_err(|e| FormatError::Corrupt(e.to_string()))?; @@ -464,18 +478,7 @@ pub fn load_bytes_with_versions(bytes: &[u8]) -> Result { let mut z = zip::ZipArchive::new(std::io::Cursor::new(bytes)).map_err(|e| FormatError::Corrupt(e.to_string()))?; let text = String::from_utf8(read_member(&mut z, "document.json", MAX_MEMBER)?).map_err(|e| FormatError::Corrupt(e.to_string()))?; let (mut doc, mut issues, own_images) = parse_json(&text)?; - let names: Vec = z.file_names().filter(|n| n.starts_with("images/")).map(str::to_string).collect(); - let mut images = BTreeMap::new(); - let mut unread = Vec::new(); - for name in names { - let hash = name.trim_start_matches("images/").to_string(); - match read_member(&mut z, &name, MAX_MEMBER) { - Ok(data) => { - images.insert(hash, image_data(data)); - } - Err(e) => unread.push((hash, e)), - } - } + let (images, unread) = read_images(&mut z, own_images.as_ref(), budget); report_lost_images(&mut issues, &unread, own_images.as_ref(), &images); match library::read(&mut z, &images, u64::MAX) { Ok(lib) => doc.published = lib.map(Arc::new), @@ -497,6 +500,41 @@ pub fn load_bytes_with_versions(bytes: &[u8]) -> Result { Ok(Loaded { doc, issues, versions, branches }) } +/// The archive's `images/` members by hash, within `budget`, and the ones that couldn't be read +/// with why. The document's own images (`own`) are read first, so a file past the budget loses +/// the images only its history uses before its own. +fn read_images( + z: &mut zip::ZipArchive, + own: Option<&BTreeSet>, + budget: ImageBudget, +) -> (BTreeMap, Vec<(String, FormatError)>) { + let mut hashes: Vec = z.file_names().filter_map(|n| n.strip_prefix("images/")).map(str::to_string).collect(); + hashes.sort_by_key(|h| !own.is_some_and(|o| o.contains(h))); + let mut images = BTreeMap::new(); + let mut unread = Vec::new(); + let mut left = budget.bytes; + for (i, hash) in hashes.into_iter().enumerate() { + let name = format!("images/{hash}"); + if i >= budget.images { + unread.push((hash, FormatError::Corrupt(format!("the file has more than {} images", budget.images)))); + continue; + } + if z.by_name(&name).is_ok_and(|f| f.size() > left) { + let mb = budget.bytes / (1024 * 1024); + unread.push((hash, FormatError::Corrupt(format!("the file's images together pass {mb} MB")))); + continue; + } + match read_member(z, &name, left.min(MAX_MEMBER)) { + Ok(data) => { + left = left.saturating_sub(data.len() as u64); + images.insert(hash, image_data(data)); + } + Err(e) => unread.push((hash, e)), + } + } + (images, unread) +} + /// Most lost images reported one by one (the rest as a count). const MAX_IMAGE_ISSUES: usize = 20; @@ -712,6 +750,8 @@ pub fn load_file(path: &std::path::Path) -> Result<(Document, Vec), Form #[cfg(test)] mod branch_tests; #[cfg(test)] +mod image_budget_tests; +#[cfg(test)] mod image_tests; #[cfg(test)] mod large_tests; diff --git a/crates/render/src/image_budget_tests.rs b/crates/render/src/image_budget_tests.rs new file mode 100644 index 00000000..6228efc5 --- /dev/null +++ b/crates/render/src/image_budget_tests.rs @@ -0,0 +1,125 @@ +//! Decoded images within a memory budget (#203): an image too large is refused from its header, +//! before its pixels are decoded, and the decoded-image caches hold a bounded number of bytes +//! however many images, filters and colour profiles a document uses. + +use std::sync::Arc; + +use figcraft_color::ColorProfile; +use tiny_skia::Pixmap; + +use super::*; + +fn png(w: u32, h: u32) -> Vec { + let img = image::RgbaImage::from_pixel(w, h, image::Rgba([200, 40, 40, 255])); + let mut out = Vec::new(); + img.write_to(&mut std::io::Cursor::new(&mut out), image::ImageFormat::Png).unwrap(); + out +} + +/// A PNG whose header claims `w` × `h` but which holds no pixel data at all. +fn png_header_only(w: u32, h: u32) -> Vec { + let mut ihdr = Vec::new(); + ihdr.extend_from_slice(&w.to_be_bytes()); + ihdr.extend_from_slice(&h.to_be_bytes()); + ihdr.extend_from_slice(&[8, 6, 0, 0, 0]); + let mut out = b"\x89PNG\r\n\x1a\n".to_vec(); + out.extend_from_slice(&13u32.to_be_bytes()); + let chunk: Vec = b"IHDR".iter().copied().chain(ihdr).collect(); + out.extend_from_slice(&chunk); + out.extend_from_slice(&crc32(&chunk).to_be_bytes()); + out +} + +fn crc32(data: &[u8]) -> u32 { + let mut crc = !0u32; + for b in data { + crc ^= u32::from(*b); + for _ in 0..8 { + crc = if crc & 1 == 1 { (crc >> 1) ^ 0xEDB8_8320 } else { crc >> 1 }; + } + } + !crc +} + +#[test] +fn an_image_within_the_limits_decodes() { + let (p, _) = images::decode_within(&png(30, 20), ColorProfile::Srgb, 30, 30 * 20 * 4).unwrap(); + assert_eq!((p.width(), p.height()), (30, 20)); +} + +#[test] +fn an_image_wider_or_taller_than_the_limit_is_refused() { + assert!(images::decode_within(&png(30, 20), ColorProfile::Srgb, 29, u64::MAX).is_none()); + assert!(images::decode_within(&png(20, 30), ColorProfile::Srgb, 29, u64::MAX).is_none()); +} + +#[test] +fn an_image_whose_pixels_pass_the_byte_budget_is_refused() { + assert!(images::decode_within(&png(30, 20), ColorProfile::Srgb, 1000, 30 * 20 * 4 - 1).is_none()); +} + +#[test] +fn a_huge_image_is_refused_from_its_header_alone() { + // 16 000 × 16 000 is within the side limit; its pixels (1 GB) aren't within a 1 MB budget, and + // the file has none to decode: only the header is read. + assert!(images::decode_within(&png_header_only(16_000, 16_000), ColorProfile::Srgb, 16_384, 1 << 20).is_none()); + // Past the side limit, the shipped decode refuses it the same way. + assert!(images::decode(&png_header_only(20_000, 10)).is_none()); +} + +fn key(name: &str) -> images::Key { + (name.to_string(), [0; 7], ColorProfile::Srgb) +} + +fn pixmap(side: u32) -> Arc { + Arc::new(Pixmap::new(side, side).unwrap()) +} + +#[test] +fn the_image_cache_keeps_within_its_byte_budget() { + // 10 × 10 pixmaps are 400 bytes; 1000 bytes hold two. + let mut c = images::PixmapCache::default(); + for name in ["a", "b", "c", "d", "e"] { + c.insert(key(name), pixmap(10), 1000); + assert!(c.bytes() <= 1000, "{} bytes held", c.bytes()); + assert!(c.get(&key(name)).is_some(), "{name} was just added"); + } +} + +#[test] +fn a_full_image_cache_drops_the_least_recently_used_image_only() { + // 1000 bytes hold two 400-byte pixmaps: the third drops the one used longest ago. + let mut c = images::PixmapCache::default(); + c.insert(key("a"), pixmap(10), 1000); + c.insert(key("b"), pixmap(10), 1000); + assert!(c.get(&key("a")).is_some()); + c.insert(key("c"), pixmap(10), 1000); + assert!(c.get(&key("a")).is_some(), "used since b: kept"); + assert!(c.get(&key("b")).is_none(), "used longest ago: dropped"); + assert!(c.get(&key("c")).is_some()); + assert_eq!(c.bytes(), 800); +} + +#[test] +fn an_image_larger_than_the_budget_is_held_alone() { + let mut c = images::PixmapCache::default(); + c.insert(key("small"), pixmap(10), 1000); + c.insert(key("big"), pixmap(20), 1000); + assert!(c.get(&key("big")).is_some()); + assert!(c.get(&key("small")).is_none()); + assert_eq!(c.bytes(), 20 * 20 * 4); + // The next image takes its place. + c.insert(key("next"), pixmap(10), 1000); + assert_eq!(c.bytes(), 400); +} + +#[test] +fn filter_and_profile_variants_count_against_one_budget() { + let mut c = images::PixmapCache::default(); + for (i, profile) in [ColorProfile::Srgb, ColorProfile::DisplayP3].into_iter().enumerate() { + for f in 0..4 { + c.insert(("img".to_string(), [f, 0, 0, 0, 0, 0, i as i32], profile), pixmap(10), 1000); + assert!(c.bytes() <= 1000, "{} bytes held", c.bytes()); + } + } +} diff --git a/crates/render/src/images.rs b/crates/render/src/images.rs index abdb9e1f..f4f9c5d9 100644 --- a/crates/render/src/images.rs +++ b/crates/render/src/images.rs @@ -13,11 +13,69 @@ use figcraft_color::{ColorProfile, ImageFilters, PaintKind}; use figcraft_doc::{Document, ImageData, NodeId}; use tiny_skia::Pixmap; -type Key = (String, [i32; 7], ColorProfile); +pub(crate) type Key = (String, [i32; 7], ColorProfile); -fn cache() -> &'static Mutex>> { - static C: OnceLock>>> = OnceLock::new(); - C.get_or_init(|| Mutex::new(HashMap::new())) +/// The largest image side decoded. +const MAX_SIDE: u32 = 16_384; +/// The most bytes one image's decoding may allocate: the image crate's own default limit (a +/// larger image is refused from its header). +const MAX_DECODED_BYTES: u64 = 512 << 20; +/// The most bytes of pixels each decoded-image cache holds (one image larger than this is held +/// alone). +const MAX_CACHED_BYTES: u64 = 512 << 20; +/// The most entries each decoded-image cache holds. +const MAX_CACHED_IMAGES: usize = 256; + +/// Decoded pixmaps by image, filters and profile, holding at most a byte budget of pixels: adding +/// one past the budget (or past [`MAX_CACHED_IMAGES`]) drops the least recently used first, so +/// the images on screen stay while they fit. +#[derive(Default)] +pub(crate) struct PixmapCache { + /// Each pixmap and when it was last used. + map: HashMap, u64)>, + bytes: u64, + tick: u64, +} + +impl PixmapCache { + pub(crate) fn get(&mut self, key: &Key) -> Option> { + self.tick += 1; + let tick = self.tick; + self.map.get_mut(key).map(|(p, used)| { + *used = tick; + p.clone() + }) + } + + /// Add `p` under `key`, keeping the cache within `max_bytes` (or holding `p` alone). + pub(crate) fn insert(&mut self, key: Key, p: Arc, max_bytes: u64) { + let size = p.data().len() as u64; + self.remove(&key); + while self.bytes.saturating_add(size) > max_bytes || self.map.len() >= MAX_CACHED_IMAGES { + let Some(oldest) = self.map.iter().min_by_key(|(_, (_, used))| *used).map(|(k, _)| k.clone()) else { break }; + self.remove(&oldest); + } + self.tick += 1; + self.map.insert(key, (p, self.tick)); + self.bytes = self.bytes.saturating_add(size); + } + + fn remove(&mut self, key: &Key) { + if let Some((p, _)) = self.map.remove(key) { + self.bytes = self.bytes.saturating_sub(p.data().len() as u64); + } + } + + /// The bytes of pixels held. + #[cfg(test)] + pub(crate) fn bytes(&self) -> u64 { + self.bytes + } +} + +fn cache() -> &'static Mutex { + static C: OnceLock> = OnceLock::new(); + C.get_or_init(|| Mutex::new(PixmapCache::default())) } fn fkey(f: &ImageFilters) -> [i32; 7] { @@ -34,7 +92,23 @@ pub fn decode(bytes: &[u8]) -> Option { /// image's own profile into `target` (`false`), or as sRGB when it has none (`true`: the sRGB → /// `target` conversion is still to do). fn decode_in(bytes: &[u8], target: ColorProfile) -> Option<(Pixmap, bool)> { - let mut dec = image::ImageReader::new(std::io::Cursor::new(bytes)).with_guessed_format().ok()?.into_decoder().ok()?; + decode_within(bytes, target, MAX_SIDE, MAX_DECODED_BYTES) +} + +/// [`decode_in`] for images at most `max_side` wide and tall whose decoded pixels take at most +/// `max_bytes`; larger ones are refused from their header, before any pixels are decoded. +pub(crate) fn decode_within(bytes: &[u8], target: ColorProfile, max_side: u32, max_bytes: u64) -> Option<(Pixmap, bool)> { + use image::ImageDecoder; + let mut reader = image::ImageReader::new(std::io::Cursor::new(bytes)).with_guessed_format().ok()?; + let mut limits = image::Limits::default(); + limits.max_image_width = Some(max_side); + limits.max_image_height = Some(max_side); + limits.max_alloc = Some(max_bytes); + reader.limits(limits); + let mut dec = reader.into_decoder().ok()?; + if dec.total_bytes() > max_bytes { + return None; + } let source = crate::image_profile::Source::read(bytes, &mut dec); let img = image::DynamicImage::from_decoder(dec).ok()?; match source.convert(&img, target) { @@ -66,7 +140,7 @@ pub fn profile_name(hash: &str, data: Option<&ImageData>) -> Option { /// Straight RGBA → a premultiplied pixmap (max 16k per side). fn premultiplied(rgba: &image::RgbaImage) -> Option { let (w, h) = rgba.dimensions(); - if w == 0 || h == 0 || w > 16_384 || h > 16_384 { + if w == 0 || h == 0 || w > MAX_SIDE || h > MAX_SIDE { return None; } let mut p = Pixmap::new(w, h)?; @@ -87,17 +161,13 @@ fn premultiplied(rgba: &image::RgbaImage) -> Option { pub fn get(hash: &str, data: Option<&ImageData>, filters: &ImageFilters, profile: ColorProfile) -> Option> { let key = (hash.to_string(), fkey(filters), profile); if let Some(p) = cache().lock().unwrap_or_else(std::sync::PoisonError::into_inner).get(&key) { - return Some(p.clone()); + return Some(p); } let (mut p, srgb) = decode_in(&data?.bytes, profile)?; // Images with their own profile are in `profile` already (sRGB → sRGB is the identity). finish(&mut p, filters, if srgb { profile } else { ColorProfile::Srgb }); let p = Arc::new(p); - let mut c = cache().lock().unwrap_or_else(std::sync::PoisonError::into_inner); - if c.len() > 256 { - c.clear(); - } - c.insert(key, p.clone()); + cache().lock().unwrap_or_else(std::sync::PoisonError::into_inner).insert(key, p.clone(), MAX_CACHED_BYTES); Some(p) } @@ -166,7 +236,7 @@ pub fn decode_frames(bytes: &[u8]) -> Option> { let dec = image::codecs::gif::GifDecoder::new(std::io::Cursor::new(bytes)).ok()?; let (w, h) = dec.dimensions(); let frame_bytes = u64::from(w) * u64::from(h) * 4; - if frame_bytes == 0 || w > 16_384 || h > 16_384 { + if frame_bytes == 0 || w > MAX_SIDE || h > MAX_SIDE { return None; } let most = (MAX_ANIMATION_BYTES / frame_bytes).min(MAX_FRAMES as u64) as usize; @@ -304,22 +374,18 @@ pub fn pad(pm: &Pixmap) -> Option { Some(out) } -fn padded_cache() -> &'static Mutex>> { - static C: OnceLock>>> = OnceLock::new(); - C.get_or_init(|| Mutex::new(HashMap::new())) +fn padded_cache() -> &'static Mutex { + static C: OnceLock> = OnceLock::new(); + C.get_or_init(|| Mutex::new(PixmapCache::default())) } /// `pm` with a transparent 1-pixel border (cached). pub fn padded(hash: &str, filters: &ImageFilters, profile: ColorProfile, pm: &Pixmap) -> Option> { let key = (hash.to_string(), fkey(filters), profile); if let Some(p) = padded_cache().lock().unwrap_or_else(std::sync::PoisonError::into_inner).get(&key) { - return Some(p.clone()); + return Some(p); } let out = Arc::new(pad(pm)?); - let mut c = padded_cache().lock().unwrap_or_else(std::sync::PoisonError::into_inner); - if c.len() > 256 { - c.clear(); - } - c.insert(key, out.clone()); + padded_cache().lock().unwrap_or_else(std::sync::PoisonError::into_inner).insert(key, out.clone(), MAX_CACHED_BYTES); Some(out) } diff --git a/crates/render/src/lib.rs b/crates/render/src/lib.rs index c2bf9dda..904e6dce 100644 --- a/crates/render/src/lib.rs +++ b/crates/render/src/lib.rs @@ -442,6 +442,8 @@ mod grain_tests; #[cfg(test)] mod group_effects_tests; #[cfg(test)] +mod image_budget_tests; +#[cfg(test)] mod image_profile_tests; #[cfg(test)] mod lod_tests;