diff --git a/crates/engine/src/psd_import.rs b/crates/engine/src/psd_import.rs index b025a7a5..32b9038e 100644 --- a/crates/engine/src/psd_import.rs +++ b/crates/engine/src/psd_import.rs @@ -202,7 +202,8 @@ fn smart_object_layer(proj: &mut Project, cx: &mut Ctx, comp: &Comp, index: usiz let bb = so.placed_bounds(cw as f64, ch as f64); let area = ((bb[2] - bb[0]) * (bb[3] - bb[1])).max(1.0); let k = ((cw as f64 * ch as f64) / area).sqrt().clamp(1.0, 4.0); - let (w, h) = so.placed_size(cw as f64, ch as f64, k); + // At most 4× the document across and 2^24 pixels, however far the file moves a corner. + let (k, (w, h)) = so.bake_scale(cw as f64, ch as f64, k, (cx.psd.width, cx.psd.height)); (w, h, k, bb) } else { (cw, ch, 1.0, [0.0; 4]) diff --git a/crates/psd/src/lib.rs b/crates/psd/src/lib.rs index 1c0fd3d4..a7816599 100644 --- a/crates/psd/src/lib.rs +++ b/crates/psd/src/lib.rs @@ -56,6 +56,79 @@ fn decode_pixels(w: usize, h: usize) -> Result { } } +fn too_large(w: usize, h: usize) -> Error { + Error::Unsupported(format!("{w}×{h} pixels is too large to decode")) +} + +/// What one decode (the composite, or one layer) may hold in memory at once: 4 GiB, 1 GiB where +/// pointers are 32 bits wide (WebAssembly). [`MAX_DECODE_PIXELS`] bounds one image; this bounds +/// the sum of the buffers a decode keeps alive together. +#[cfg(target_pointer_width = "32")] +const DECODE_BUDGET: u64 = 1 << 30; +#[cfg(not(target_pointer_width = "32"))] +const DECODE_BUDGET: u64 = 4 << 30; + +/// The bytes the buffers of one decode hold right now, against a limit. A decode charges a buffer +/// before it allocates it (and releases it when it drops it), so a file that would need more +/// than the limit fails with an error instead of an out-of-memory abort. +struct Budget { + limit: u64, + live: u64, +} + +impl Budget { + fn new(limit: u64) -> Budget { + Budget { limit, live: 0 } + } + + fn charge(&mut self, bytes: u64) -> Result<()> { + self.check(bytes)?; + self.live += bytes; + Ok(()) + } + + /// Whether `bytes` more would still fit, without taking them. A decode checks its whole + /// expected peak first, so one that cannot fit fails before it has allocated anything. + fn check(&self, bytes: u64) -> Result<()> { + match self.live.checked_add(bytes) { + Some(total) if total <= self.limit => Ok(()), + _ => { + const MIB: u64 = 1 << 20; + Err(Error::Unsupported(format!( + "decoding needs {} MiB more with {} MiB in use, over the {} MiB memory budget", + bytes / MIB, + self.live / MIB, + self.limit / MIB + ))) + } + } + } + + fn release(&mut self, bytes: u64) { + self.live = self.live.saturating_sub(bytes); + } + + /// An empty vector with room for `n` items, charged. + fn vec(&mut self, n: usize) -> Result> { + self.charge((n as u64).saturating_mul(std::mem::size_of::() as u64))?; + let mut v = Vec::new(); + v.try_reserve_exact(n).map_err(|_| Error::Unsupported(format!("out of memory for {n} items")))?; + Ok(v) + } + + /// `n` copies of `value`, charged. + fn filled(&mut self, n: usize, value: T) -> Result> { + let mut v = self.vec::(n)?; + v.resize(n, value); + Ok(v) + } +} + +/// Inflate a zlib stream to at most `limit` bytes. +fn inflate(src: &[u8], limit: usize) -> Result> { + miniz_oxide::inflate::decompress_to_vec_zlib_with_limit(src, limit).map_err(|e| Error::Invalid(format!("zip: {e:?}"))) +} + /// Whether `bytes` start like a Photoshop document. pub fn is_psd(bytes: &[u8]) -> bool { bytes.len() >= 6 && &bytes[..4] == b"8BPS" && matches!(u16::from_be_bytes([bytes[4], bytes[5]]), 1 | 2) @@ -788,38 +861,59 @@ impl Psd { } } - /// Decode one channel of `w`×`h` samples at `offset` (compression field first). Values 0..1. - fn decode_channel(&self, offset: usize, len: usize, w: usize, h: usize) -> Result> { + /// The raw big-endian bytes of one channel of `w`×`h` samples at `offset` (compression field + /// first): `row_bytes(w) × h` of them with the compression and the prediction undone. The + /// buffer is charged to `budget`; the caller releases it once it is dropped. + fn channel_bytes(&self, offset: usize, len: usize, w: usize, h: usize, budget: &mut Budget) -> Result> { let end = offset.checked_add(len).filter(|&e| e <= self.data.len()).ok_or(Error::Truncated)?; - let mut r = Reader::at(&self.data[..end], offset); + let mut r = Reader::at(self.data.get(..end).ok_or(Error::Truncated)?, offset); let comp = r.u16()?; decode_pixels(w, h)?; let rowb = self.row_bytes(w); - let raw = match comp { - 0 => r.bytes(rowb * h)?.to_vec(), + let size = rowb.checked_mul(h).ok_or_else(|| too_large(w, h))?; + match comp { + 0 => { + let src = r.bytes(size)?; + let mut v = budget.vec::(size)?; + v.extend_from_slice(src); + Ok(v) + } 1 => { let mut counts = Vec::with_capacity(h.min(r.remaining() / 2)); for _ in 0..h { counts.push(if self.psb { r.u32()? as usize } else { r.u16()? as usize }); } - let mut out = Vec::with_capacity(rowb * h); + let mut out = budget.vec::(size)?; for c in counts { unpack_bits(r.bytes(c)?, rowb, &mut out)?; } - out + Ok(out) } 2 | 3 => { let rest = r.bytes(r.remaining())?; - let mut v = miniz_oxide::inflate::decompress_to_vec_zlib_with_limit(rest, rowb * h).map_err(|e| Error::Invalid(format!("zip: {e:?}")))?; - v.resize(rowb * h, 0); + budget.charge(size as u64)?; + let mut v = inflate(rest, size)?; + v.resize(size, 0); if comp == 3 { self.unpredict(&mut v, w, h); } - v + Ok(v) } - c => return Err(Error::Unsupported(format!("compression {c}"))), - }; - Ok(self.samples(&raw, w, h)) + c => Err(Error::Unsupported(format!("compression {c}"))), + } + } + + /// One channel as `w * h` floats (0..1), charged to `budget` (the caller releases it). + fn channel_floats(&self, offset: usize, len: usize, w: usize, h: usize, budget: &mut Budget) -> Result> { + let raw = self.channel_bytes(offset, len, w, h, budget)?; + let mut out = budget.filled(decode_pixels(w, h)?, 0.0f32)?; + self.for_each_sample(&raw, w, h, |i, v| { + if let Some(o) = out.get_mut(i) { + *o = v; + } + }); + budget.release(raw.len() as u64); + Ok(out) } fn row_bytes(&self, w: usize) -> usize { @@ -862,29 +956,30 @@ impl Psd { } } - fn samples(&self, raw: &[u8], w: usize, h: usize) -> Vec { + /// Hand every sample of `raw` (`w`×`h`, row-major, in order) to `f` as a float. Samples a short + /// `raw` does not have are not reported. + fn for_each_sample(&self, raw: &[u8], w: usize, h: usize, mut f: impl FnMut(usize, f32)) { + let n = w.saturating_mul(h); match self.depth { 1 => { let rowb = w.div_ceil(8); - let mut out = Vec::with_capacity(w * h); for y in 0..h { for x in 0..w { let b = raw.get(y * rowb + x / 8).copied().unwrap_or(0); // 1 = black. - out.push(if b & (0x80 >> (x % 8)) != 0 { 0.0 } else { 1.0 }); + f(y * w + x, if b & (0x80 >> (x % 8)) != 0 { 0.0 } else { 1.0 }); } } - out } - 16 => raw.as_chunks::<2>().0.iter().take(w * h).map(|c| u16::from_be_bytes([c[0], c[1]]) as f32 / 65535.0).collect(), - 32 => raw.as_chunks::<4>().0.iter().take(w * h).map(|c| f32::from_be_bytes([c[0], c[1], c[2], c[3]])).collect(), - _ => raw.iter().take(w * h).map(|&b| b as f32 / 255.0).collect(), + 16 => raw.as_chunks::<2>().0.iter().take(n).enumerate().for_each(|(i, c)| f(i, u16::from_be_bytes([c[0], c[1]]) as f32 / 65535.0)), + 32 => raw.as_chunks::<4>().0.iter().take(n).enumerate().for_each(|(i, c)| f(i, f32::from_be_bytes([c[0], c[1], c[2], c[3]]))), + _ => raw.iter().take(n).enumerate().for_each(|(i, &b)| f(i, b as f32 / 255.0)), } } - /// Convert colour channels (in the document's mode) to RGB. - fn to_rgb(&self, ch: &[Option>], i: usize) -> [f32; 3] { - let g = |k: usize| ch.get(k).and_then(|c| c.as_ref()).map_or(0.0, |c| c[i]); + /// Convert colour values (the document mode's channels, missing ones 0) to RGB. + fn rgb_from(&self, v: [f32; 4]) -> [f32; 3] { + let g = |k: usize| v.get(k).copied().unwrap_or(0.0); let rgb = match self.mode { ColorMode::Rgb => [g(0), g(1), g(2)], ColorMode::Cmyk => { @@ -902,55 +997,101 @@ impl Psd { if self.depth == 32 && self.mode != ColorMode::Lab { rgb.map(srgb_encode) } else { rgb } } + /// Write merged-image channel `c` (raw bytes) into the result: the first three colour channels + /// and the alpha channel into their slot of a pixel, the fourth CMYK channel into `black`. + fn store_merged_channel(&self, c: usize, raw: &[u8], data: &mut [[f32; 4]], black: &mut Option>) { + let (w, h, cc) = (self.width as usize, self.height as usize, self.mode.color_channels()); + if c < cc.min(3) { + self.for_each_sample(raw, w, h, |i, v| { + if let Some(p) = data.get_mut(i) { + p[c] = v; + } + }); + } else if c == 3 && c < cc { + if let Some(k) = black { + self.for_each_sample(raw, w, h, |i, v| { + if let Some(o) = k.get_mut(i) { + *o = v; + } + }); + } + } else if c == cc && self.merged_alpha { + self.for_each_sample(raw, w, h, |i, v| { + if let Some(p) = data.get_mut(i) { + p[3] = v; + } + }); + } + } + /// The merged (flattened) image. pub fn composite(&self) -> Result { + self.composite_with_budget(DECODE_BUDGET) + } + + /// [`Self::composite`] with at most `limit` bytes live at once. The channels are written + /// straight into the result, so the peak is the result (16 bytes a pixel), the CMYK black + /// channel, and the one buffer a channel (or, for ZIP, all channels) decodes through. + pub(crate) fn composite_with_budget(&self, limit: u64) -> Result { + let mut budget = Budget::new(limit); let (w, h) = (self.width as usize, self.height as usize); - decode_pixels(w, h)?; + let n = decode_pixels(w, h)?; let mut r = Reader::at(&self.data, self.merged_offset); let comp = r.u16()?; let nch = self.channels as usize; let cc = self.mode.color_channels(); let want = (cc + usize::from(self.merged_alpha)).min(nch); let rowb = self.row_bytes(w); - let mut planes: Vec>> = vec![None; nch]; + let size = rowb.checked_mul(h).ok_or_else(|| too_large(w, h))?; + // The peak: the result, the CMYK black channel and what the channels decode through. + let black_bytes = if self.mode == ColorMode::Cmyk && want > 3 { 4 * n as u64 } else { 0 }; + let transient = if matches!(comp, 2 | 3) { (size as u64).saturating_mul(want as u64) } else { size as u64 }; + budget.check((16 * n as u64).saturating_add(black_bytes).saturating_add(transient))?; + // Opaque black until a channel says otherwise. + let mut data = budget.filled(n, [0.0f32, 0.0, 0.0, 1.0])?; + let mut black = if self.mode == ColorMode::Cmyk && want > 3 { Some(budget.filled(n, 0.0f32)?) } else { None }; match comp { 0 => { - for p in planes.iter_mut().take(want) { - *p = Some(self.samples(r.bytes(rowb * h)?, w, h)); + for c in 0..want { + self.store_merged_channel(c, r.bytes(size)?, &mut data, &mut black); } } 1 => { - let mut counts = Vec::with_capacity((nch * h).min(r.remaining() / 2)); - for _ in 0..nch * h { + let total = nch.saturating_mul(h); + let mut counts = Vec::with_capacity(total.min(r.remaining() / 2)); + for _ in 0..total { counts.push(if self.psb { r.u32()? as usize } else { r.u16()? as usize }); } - for (c, p) in planes.iter_mut().enumerate().take(want) { - let mut out = Vec::with_capacity(rowb * h); + for c in 0..want { + let mut out = budget.vec::(size)?; for row in 0..h { - unpack_bits(r.bytes(counts[c * h + row])?, rowb, &mut out)?; + let count = *counts.get(c * h + row).ok_or(Error::Truncated)?; + unpack_bits(r.bytes(count)?, rowb, &mut out)?; } - *p = Some(self.samples(&out, w, h)); + self.store_merged_channel(c, &out, &mut data, &mut black); + budget.release(size as u64); } } 2 | 3 => { let rest = r.bytes(r.remaining())?; - let v = miniz_oxide::inflate::decompress_to_vec_zlib_with_limit(rest, want * rowb * h).map_err(|e| Error::Invalid(format!("zip: {e:?}")))?; - for (c, p) in planes.iter_mut().enumerate().take(want) { - let mut plane = v.get(c * rowb * h..(c + 1) * rowb * h).ok_or(Error::Truncated)?.to_vec(); + // One stream holds every channel: what it may decode to is what they need. + let total = size.checked_mul(want).ok_or_else(|| too_large(w, h))?; + budget.charge(total as u64)?; + let mut v = inflate(rest, total)?; + for c in 0..want { + let plane = v.get_mut(c * size..(c + 1) * size).ok_or(Error::Truncated)?; if comp == 3 { - self.unpredict(&mut plane, w, h); + self.unpredict(plane, w, h); } - *p = Some(self.samples(&plane, w, h)); + self.store_merged_channel(c, plane, &mut data, &mut black); } } c => return Err(Error::Unsupported(format!("compression {c}"))), } - let alpha = if self.merged_alpha && nch > cc { planes[cc].take() } else { None }; - let mut data = Vec::with_capacity(w * h); - for i in 0..w * h { - let [r, g, b] = self.to_rgb(&planes, i); - let a = alpha.as_ref().map_or(1.0, |a| a[i]); - data.push([r, g, b, a]); + for (i, p) in data.iter_mut().enumerate() { + let k = black.as_ref().and_then(|k| k.get(i)).copied().unwrap_or(0.0); + let [r, g, b] = self.rgb_from([p[0], p[1], p[2], k]); + *p = [r, g, b, p[3]]; } Ok(Pixels { width: w as u32, height: h as u32, data }) } @@ -958,6 +1099,14 @@ impl Psd { /// One layer's pixels with its layer mask (and vector-less fill colour) applied to alpha. /// `canvas`: document-sized with the layer in place; otherwise the layer's own bounds. pub fn layer_pixels(&self, index: usize, canvas: bool) -> Result { + self.layer_pixels_with_budget(index, canvas, DECODE_BUDGET) + } + + /// [`Self::layer_pixels`] with at most `limit` bytes live at once: the result (16 bytes a + /// pixel) with the channels written straight into it, the CMYK black channel, the one channel + /// being decoded, and the mask. + pub(crate) fn layer_pixels_with_budget(&self, index: usize, canvas: bool, limit: u64) -> Result { + let mut budget = Budget::new(limit); let l = self.layers.get(index).ok_or_else(|| Error::Invalid(format!("no layer {index}")))?; let doc = Rect { top: 0, left: 0, bottom: self.height as i32, right: self.width as i32 }; // Fill layers often have no pixel data: they cover the canvas (or the mask). @@ -971,22 +1120,26 @@ impl Psd { fill_rect }; let (ow, oh) = (out_rect.width() as usize, out_rect.height() as usize); - let mut out = vec![[0.0f32; 4]; decode_pixels(ow, oh)?]; + let n = decode_pixels(ow, oh)?; let src = fill_rect; let (sw, sh) = (src.width() as usize, src.height() as usize); - if sw > 0 && sh > 0 { - let mut planes: Vec>> = vec![None; self.mode.color_channels()]; - let mut alpha = None; - if l.has_pixels() { - for c in &l.channels { - match c.id { - -1 => alpha = Some(self.decode_channel(c.offset, c.len, sw, sh)?), - id if id >= 0 && (id as usize) < planes.len() => planes[id as usize] = Some(self.decode_channel(c.offset, c.len, sw, sh)?), - _ => {} - } - } + // The peak: the result, then either the CMYK black channel with the channel being decoded, + // or the mask as floats with its bytes. Checked before anything is allocated. + let plane = |w: usize, h: usize| (self.row_bytes(w) as u64).saturating_mul(h as u64); + let mut pixels_peak = 0u64; + if l.has_pixels() && sw > 0 && sh > 0 { + pixels_peak = plane(sw, sh).saturating_add(if self.mode == ColorMode::Cmyk { 4 * n as u64 } else { 0 }); + } + let mut mask_peak = 0u64; + if let Some(m) = l.mask.filter(|m| !m.disabled) { + let (mw, mh) = (m.rect.width() as usize, m.rect.height() as usize); + if mw > 0 && mh > 0 && l.channels.iter().any(|c| c.id == -2) { + mask_peak = plane(mw, mh).saturating_add((mw as u64).saturating_mul(mh as u64).saturating_mul(4)); } - let fill = l.fill.map(|c| c.map(|v| v as f32)); + } + budget.check((16 * n as u64).saturating_add(pixels_peak.max(mask_peak)))?; + let mut out = budget.filled(n, [0.0f32; 4])?; + if sw > 0 && sh > 0 { // The source rows / columns that land inside the output (64-bit: file bounds can span // the whole i32 range). let span = |s0: i32, n: usize, o0: i32, on: usize| { @@ -997,16 +1150,77 @@ impl Psd { }; let (ys, oy) = span(src.top, sh, out_rect.top, oh); let (xs, ox) = span(src.left, sw, out_rect.left, ow); - for y in ys { - let dy = (y as i64 - oy) as usize; - for x in xs.clone() { - let dx = (x as i64 - ox) as usize; - let i = y * sw + x; - let (rgb, a) = match fill { - Some(f) if !l.has_pixels() => (f, 1.0), - _ => (self.to_rgb(&planes, i), alpha.as_ref().map_or(1.0, |a| a[i])), - }; - out[dy * ow + dx] = [rgb[0], rgb[1], rgb[2], a]; + let dest = |y: usize, x: usize| (y as i64 - oy) as usize * ow + (x as i64 - ox) as usize; + let fill = l.fill.map(|c| c.map(|v| v as f32)); + match fill { + Some(f) if !l.has_pixels() => { + for y in ys { + for x in xs.clone() { + if let Some(p) = out.get_mut(dest(y, x)) { + *p = [f[0], f[1], f[2], 1.0]; + } + } + } + } + _ => { + // Opaque unless there is an alpha channel. + for y in ys.clone() { + for x in xs.clone() { + if let Some(p) = out.get_mut(dest(y, x)) { + p[3] = 1.0; + } + } + } + let cc = self.mode.color_channels(); + let mut black: Option> = None; + if l.has_pixels() { + for c in &l.channels { + let slot = match c.id { + -1 => 3, + id if id >= 0 && (id as usize) < cc => id as usize, + _ => continue, + }; + let raw = self.channel_bytes(c.offset, c.len, sw, sh, &mut budget)?; + // The fourth CMYK channel is the black plate, not a slot of the result. + let to_black = c.id == 3; + if to_black && black.is_none() { + black = Some(budget.filled(out.len(), 0.0f32)?); + } + // Samples arrive in row-major order of the source. + let (mut sy, mut sx) = (0usize, 0usize); + self.for_each_sample(&raw, sw, sh, |_, v| { + if ys.contains(&sy) && xs.contains(&sx) { + let d = dest(sy, sx); + if to_black { + if let Some(o) = black.as_mut().and_then(|k| k.get_mut(d)) { + *o = v; + } + } else if let Some(p) = out.get_mut(d) { + p[slot] = v; + } + } + sx += 1; + if sx == sw { + sx = 0; + sy += 1; + } + }); + budget.release(raw.len() as u64); + } + } + for y in ys { + for x in xs.clone() { + let d = dest(y, x); + let k = black.as_ref().and_then(|k| k.get(d)).copied().unwrap_or(0.0); + if let Some(p) = out.get_mut(d) { + let [r, g, b] = self.rgb_from([p[0], p[1], p[2], k]); + *p = [r, g, b, p[3]]; + } + } + } + if let Some(k) = black { + budget.release(k.len() as u64 * 4); + } } } } @@ -1015,7 +1229,7 @@ impl Psd { let mc = l.channels.iter().find(|c| c.id == -2); let (mw, mh) = (m.rect.width() as usize, m.rect.height() as usize); let mask = match mc { - Some(c) if mw > 0 && mh > 0 => Some(self.decode_channel(c.offset, c.len, mw, mh)?), + Some(c) if mw > 0 && mh > 0 => Some(self.channel_floats(c.offset, c.len, mw, mh, &mut budget)?), _ => None, }; let def = m.default_color as f32 / 255.0; @@ -1025,13 +1239,15 @@ impl Psd { let px = out_rect.left + x as i32; let inside = py >= m.rect.top && py < m.rect.bottom && px >= m.rect.left && px < m.rect.right; let mut v = match (&mask, inside) { - (Some(mk), true) => mk[(py - m.rect.top) as usize * mw + (px - m.rect.left) as usize], + (Some(mk), true) => mk.get((py - m.rect.top) as usize * mw + (px - m.rect.left) as usize).copied().unwrap_or(def), _ => def, }; if m.inverted { v = 1.0 - v; } - out[y * ow + x][3] *= v; + if let Some(p) = out.get_mut(y * ow + x) { + p[3] *= v; + } } } } @@ -1147,5 +1363,7 @@ fn lab_to_rgb(l: f32, a: f32, b: f32) -> [f32; 3] { [r, g, bb].map(|v| srgb_encode(v.clamp(0.0, 1.0))) } +#[cfg(test)] +mod memory_tests; #[cfg(test)] mod tests; diff --git a/crates/psd/src/memory_tests.rs b/crates/psd/src/memory_tests.rs new file mode 100644 index 00000000..8fbe8307 --- /dev/null +++ b/crates/psd/src/memory_tests.rs @@ -0,0 +1,398 @@ +//! Tests for the decode memory budget, with the previous decoder (one `Vec` per channel and a +//! result vector, no budget) kept as the reference the new one must reproduce bit for bit. + +use super::write::*; +use super::*; + +/// The decoder as it was before the budget: copied unchanged except for the names. +impl Psd { + /// Decode one channel of `w`×`h` samples at `offset` (compression field first). Values 0..1. + fn ref_decode_channel(&self, offset: usize, len: usize, w: usize, h: usize) -> Result> { + let end = offset.checked_add(len).filter(|&e| e <= self.data.len()).ok_or(Error::Truncated)?; + let mut r = Reader::at(&self.data[..end], offset); + let comp = r.u16()?; + decode_pixels(w, h)?; + let rowb = self.row_bytes(w); + let raw = match comp { + 0 => r.bytes(rowb * h)?.to_vec(), + 1 => { + let mut counts = Vec::with_capacity(h.min(r.remaining() / 2)); + for _ in 0..h { + counts.push(if self.psb { r.u32()? as usize } else { r.u16()? as usize }); + } + let mut out = Vec::with_capacity(rowb * h); + for c in counts { + unpack_bits(r.bytes(c)?, rowb, &mut out)?; + } + out + } + 2 | 3 => { + let rest = r.bytes(r.remaining())?; + let mut v = miniz_oxide::inflate::decompress_to_vec_zlib_with_limit(rest, rowb * h).map_err(|e| Error::Invalid(format!("zip: {e:?}")))?; + v.resize(rowb * h, 0); + if comp == 3 { + self.unpredict(&mut v, w, h); + } + v + } + c => return Err(Error::Unsupported(format!("compression {c}"))), + }; + Ok(self.ref_samples(&raw, w, h)) + } + + fn ref_samples(&self, raw: &[u8], w: usize, h: usize) -> Vec { + match self.depth { + 1 => { + let rowb = w.div_ceil(8); + let mut out = Vec::with_capacity(w * h); + for y in 0..h { + for x in 0..w { + let b = raw.get(y * rowb + x / 8).copied().unwrap_or(0); + // 1 = black. + out.push(if b & (0x80 >> (x % 8)) != 0 { 0.0 } else { 1.0 }); + } + } + out + } + 16 => raw.as_chunks::<2>().0.iter().take(w * h).map(|c| u16::from_be_bytes([c[0], c[1]]) as f32 / 65535.0).collect(), + 32 => raw.as_chunks::<4>().0.iter().take(w * h).map(|c| f32::from_be_bytes([c[0], c[1], c[2], c[3]])).collect(), + _ => raw.iter().take(w * h).map(|&b| b as f32 / 255.0).collect(), + } + } + + /// Convert colour channels (in the document's mode) to RGB. + fn ref_to_rgb(&self, ch: &[Option>], i: usize) -> [f32; 3] { + let g = |k: usize| ch.get(k).and_then(|c| c.as_ref()).map_or(0.0, |c| c[i]); + let rgb = match self.mode { + ColorMode::Rgb => [g(0), g(1), g(2)], + ColorMode::Cmyk => { + // Stored inverted (1 = no ink). + let k = g(3); + [g(0) * k, g(1) * k, g(2) * k] + } + ColorMode::Indexed if self.palette.len() >= 768 => { + let idx = (g(0) * 255.0).round() as usize & 255; + [self.palette[idx] as f32 / 255.0, self.palette[256 + idx] as f32 / 255.0, self.palette[512 + idx] as f32 / 255.0] + } + ColorMode::Lab => lab_to_rgb(g(0) * 100.0, g(1) * 255.0 - 128.0, g(2) * 255.0 - 128.0), + _ => [g(0); 3], + }; + if self.depth == 32 && self.mode != ColorMode::Lab { rgb.map(srgb_encode) } else { rgb } + } + + /// The merged (flattened) image. + fn ref_composite(&self) -> Result { + let (w, h) = (self.width as usize, self.height as usize); + decode_pixels(w, h)?; + let mut r = Reader::at(&self.data, self.merged_offset); + let comp = r.u16()?; + let nch = self.channels as usize; + let cc = self.mode.color_channels(); + let want = (cc + usize::from(self.merged_alpha)).min(nch); + let rowb = self.row_bytes(w); + let mut planes: Vec>> = vec![None; nch]; + match comp { + 0 => { + for p in planes.iter_mut().take(want) { + *p = Some(self.ref_samples(r.bytes(rowb * h)?, w, h)); + } + } + 1 => { + let mut counts = Vec::with_capacity((nch * h).min(r.remaining() / 2)); + for _ in 0..nch * h { + counts.push(if self.psb { r.u32()? as usize } else { r.u16()? as usize }); + } + for (c, p) in planes.iter_mut().enumerate().take(want) { + let mut out = Vec::with_capacity(rowb * h); + for row in 0..h { + unpack_bits(r.bytes(counts[c * h + row])?, rowb, &mut out)?; + } + *p = Some(self.ref_samples(&out, w, h)); + } + } + 2 | 3 => { + let rest = r.bytes(r.remaining())?; + let v = miniz_oxide::inflate::decompress_to_vec_zlib_with_limit(rest, want * rowb * h).map_err(|e| Error::Invalid(format!("zip: {e:?}")))?; + for (c, p) in planes.iter_mut().enumerate().take(want) { + let mut plane = v.get(c * rowb * h..(c + 1) * rowb * h).ok_or(Error::Truncated)?.to_vec(); + if comp == 3 { + self.unpredict(&mut plane, w, h); + } + *p = Some(self.ref_samples(&plane, w, h)); + } + } + c => return Err(Error::Unsupported(format!("compression {c}"))), + } + let alpha = if self.merged_alpha && nch > cc { planes[cc].take() } else { None }; + let mut data = Vec::with_capacity(w * h); + for i in 0..w * h { + let [r, g, b] = self.ref_to_rgb(&planes, i); + let a = alpha.as_ref().map_or(1.0, |a| a[i]); + data.push([r, g, b, a]); + } + Ok(Pixels { width: w as u32, height: h as u32, data }) + } + + /// One layer's pixels with its layer mask (and vector-less fill colour) applied to alpha. + /// `canvas`: document-sized with the layer in place; otherwise the layer's own bounds. + fn ref_layer_pixels(&self, index: usize, canvas: bool) -> Result { + let l = self.layers.get(index).ok_or_else(|| Error::Invalid(format!("no layer {index}")))?; + let doc = Rect { top: 0, left: 0, bottom: self.height as i32, right: self.width as i32 }; + // Fill layers often have no pixel data: they cover the canvas (or the mask). + let fill_rect = + if l.fill.is_some() && !l.has_pixels() { l.mask.filter(|m| !m.rect.is_empty() && m.default_color == 0).map_or(doc, |m| m.rect) } else { l.rect }; + let out_rect = if canvas { + doc + } else if fill_rect.is_empty() { + Rect { top: 0, left: 0, bottom: 1, right: 1 } + } else { + fill_rect + }; + let (ow, oh) = (out_rect.width() as usize, out_rect.height() as usize); + let mut out = vec![[0.0f32; 4]; decode_pixels(ow, oh)?]; + let src = fill_rect; + let (sw, sh) = (src.width() as usize, src.height() as usize); + if sw > 0 && sh > 0 { + let mut planes: Vec>> = vec![None; self.mode.color_channels()]; + let mut alpha = None; + if l.has_pixels() { + for c in &l.channels { + match c.id { + -1 => alpha = Some(self.ref_decode_channel(c.offset, c.len, sw, sh)?), + id if id >= 0 && (id as usize) < planes.len() => planes[id as usize] = Some(self.ref_decode_channel(c.offset, c.len, sw, sh)?), + _ => {} + } + } + } + let fill = l.fill.map(|c| c.map(|v| v as f32)); + // The source rows / columns that land inside the output (64-bit: file bounds can span + // the whole i32 range). + let span = |s0: i32, n: usize, o0: i32, on: usize| { + let off = i64::from(o0) - i64::from(s0); + let lo = off.clamp(0, n as i64); + let hi = (off + on as i64).clamp(lo, n as i64); + (lo as usize..hi as usize, off) + }; + let (ys, oy) = span(src.top, sh, out_rect.top, oh); + let (xs, ox) = span(src.left, sw, out_rect.left, ow); + for y in ys { + let dy = (y as i64 - oy) as usize; + for x in xs.clone() { + let dx = (x as i64 - ox) as usize; + let i = y * sw + x; + let (rgb, a) = match fill { + Some(f) if !l.has_pixels() => (f, 1.0), + _ => (self.ref_to_rgb(&planes, i), alpha.as_ref().map_or(1.0, |a| a[i])), + }; + out[dy * ow + dx] = [rgb[0], rgb[1], rgb[2], a]; + } + } + } + // Layer mask: multiplies alpha (also outside the mask rect, by its default colour). + if let Some(m) = l.mask.filter(|m| !m.disabled) { + let mc = l.channels.iter().find(|c| c.id == -2); + let (mw, mh) = (m.rect.width() as usize, m.rect.height() as usize); + let mask = match mc { + Some(c) if mw > 0 && mh > 0 => Some(self.ref_decode_channel(c.offset, c.len, mw, mh)?), + _ => None, + }; + let def = m.default_color as f32 / 255.0; + for y in 0..oh { + let py = out_rect.top + y as i32; + for x in 0..ow { + let px = out_rect.left + x as i32; + let inside = py >= m.rect.top && py < m.rect.bottom && px >= m.rect.left && px < m.rect.right; + let mut v = match (&mask, inside) { + (Some(mk), true) => mk[(py - m.rect.top) as usize * mw + (px - m.rect.left) as usize], + _ => def, + }; + if m.inverted { + v = 1.0 - v; + } + out[y * ow + x][3] *= v; + } + } + } + Ok(Pixels { width: ow as u32, height: oh as u32, data: out }) + } +} + +// --------------------------------------------------------------------------------------------- +// Tests + +fn grad(w: u32, h: u32) -> Vec<[f32; 4]> { + (0..w * h) + .map(|i| [(i % w) as f32 / (w - 1).max(1) as f32, (i / w) as f32 / (h - 1).max(1) as f32, 0.25 + (i % 7) as f32 / 16.0, 0.3 + (i % 5) as f32 / 8.0]) + .collect() +} + +/// Layers covering: full canvas pixels with alpha, a partly off-canvas layer, a masked layer with +/// an inverted mask and a default colour, a solid layer. +fn doc(depth: u16, mode: ColorMode, rle: bool) -> WDoc { + let mut d = WDoc::new(40, 30); + d.depth = depth; + d.mode = mode; + d.rle = rle; + let bg = WLayer::pixels("bg", Rect::new(0, 0, 40, 30), grad(40, 30)); + let off = WLayer::pixels("off", Rect::new(-5, -4, 18, 13), grad(18, 13)); + let solid = WLayer::solid("solid", Rect::new(30, 20, 15, 15), [0.2, 0.6, 0.9, 1.0]); + let mut masked = WLayer::pixels("masked", Rect::new(10, 8, 12, 12), grad(12, 12)); + masked.mask = + Some(WMask { rect: Rect::new(12, 10, 12, 12), data: (0..144).map(|i| (i % 12) as f32 / 11.0).collect(), default_color: 255, disabled: false }); + d.layers = vec![bg, off, solid, masked]; + d +} + +fn assert_same_as_reference(p: &Psd, what: &str) { + assert_eq!(p.composite().unwrap(), p.ref_composite().unwrap(), "composite {what}"); + for i in 0..p.layers.len() { + for canvas in [false, true] { + assert_eq!(p.layer_pixels(i, canvas).unwrap(), p.ref_layer_pixels(i, canvas).unwrap(), "layer {i} canvas={canvas} {what}"); + } + } +} + +#[test] +fn pixels_match_the_one_buffer_per_channel_reference() { + for depth in [8u16, 16] { + for mode in [ColorMode::Rgb, ColorMode::Grayscale, ColorMode::Cmyk] { + for rle in [true, false] { + let p = Psd::parse(write(&doc(depth, mode, rle))).unwrap(); + assert_same_as_reference(&p, &format!("{depth}-bit {mode:?} rle={rle}")); + } + } + } +} + +/// A copy of `bytes` with a different size, depth and merged-image section (`comp`, then `body`). +fn with_merged(bytes: &[u8], size: Option<(u32, u32)>, depth: Option, comp: u16, body: &[u8]) -> Vec { + let offset = Psd::parse(bytes.to_vec()).unwrap().merged_offset; + let mut v = bytes[..offset].to_vec(); + if let Some((w, h)) = size { + v[14..18].copy_from_slice(&h.to_be_bytes()); + v[18..22].copy_from_slice(&w.to_be_bytes()); + } + if let Some(d) = depth { + v[22..24].copy_from_slice(&d.to_be_bytes()); + } + v.extend_from_slice(&comp.to_be_bytes()); + v.extend_from_slice(body); + v +} + +/// Channel planes for a 6×5 image: values that exercise the whole range. +fn planes(nch: usize, bytes_per_sample: usize, n: usize) -> Vec { + let mut v = vec![]; + for c in 0..nch { + for i in 0..n { + let t = ((i * 37 + c * 91) % 251) as f32 / 250.0; + match bytes_per_sample { + 1 => v.push((t * 255.0) as u8), + 2 => v.extend_from_slice(&((t * 65535.0) as u16).to_be_bytes()), + _ => v.extend_from_slice(&(t * 1.5 - 0.1).to_be_bytes()), + } + } + } + v +} + +#[test] +fn merged_image_matches_the_reference_for_32_bit_and_zip() { + for mode in [ColorMode::Rgb, ColorMode::Grayscale, ColorMode::Cmyk] { + let mut small = WDoc::new(6, 5); + small.mode = mode; + small.rle = false; + let base = write(&small); + let info = Psd::parse(base.clone()).unwrap(); + // The merged image only reads the colour channels and, when flagged, the alpha channel. + let nch = (mode.color_channels() + usize::from(info.merged_alpha)).min(info.channels as usize); + let n = 30; + // 32-bit float, raw. + let p = Psd::parse(with_merged(&base, None, Some(32), 0, &planes(nch, 4, n))).unwrap(); + assert_eq!(p.depth, 32); + assert_eq!(p.composite().unwrap(), p.ref_composite().unwrap(), "32-bit {mode:?}"); + // ZIP without and with prediction, 8 and 16 bit; 32-bit ZIP without prediction. + for (depth, bps) in [(8u16, 1usize), (16, 2), (32, 4)] { + let raw = planes(nch, bps, n); + let zip = miniz_oxide::deflate::compress_to_vec_zlib(&raw, 6); + let p = Psd::parse(with_merged(&base, None, Some(depth), 2, &zip)).unwrap(); + assert_eq!(p.composite().unwrap(), p.ref_composite().unwrap(), "zip {depth}-bit {mode:?}"); + if depth == 32 { + continue; + } + // Prediction: each row stores differences to its left neighbour. + let mut pred = raw.clone(); + for plane in pred.chunks_exact_mut(n * bps) { + for row in plane.chunks_exact_mut(6 * bps) { + if bps == 1 { + for i in (1..6).rev() { + row[i] = row[i].wrapping_sub(row[i - 1]); + } + } else { + for i in (1..6).rev() { + let a = u16::from_be_bytes([row[i * 2], row[i * 2 + 1]]); + let b = u16::from_be_bytes([row[i * 2 - 2], row[i * 2 - 1]]); + row[i * 2..i * 2 + 2].copy_from_slice(&a.wrapping_sub(b).to_be_bytes()); + } + } + } + } + let zip = miniz_oxide::deflate::compress_to_vec_zlib(&pred, 6); + let p = Psd::parse(with_merged(&base, None, Some(depth), 3, &zip)).unwrap(); + assert_eq!(p.composite().unwrap(), p.ref_composite().unwrap(), "zip+prediction {depth}-bit {mode:?}"); + } + } +} + +#[test] +fn budget_counts_the_sum_of_the_buffers_of_a_decode() { + let mut b = Budget::new(100); + b.charge(60).unwrap(); + assert!(b.charge(41).is_err(), "60 + 41 is over 100 although each alone fits"); + b.check(40).unwrap(); + assert!(b.check(41).is_err()); + b.release(30); + b.charge(70).unwrap(); + assert!(b.vec::(1).unwrap_err().to_string().contains("memory budget")); + assert!(Budget::new(u64::MAX).vec::(usize::MAX / 2).is_err(), "a size that cannot be reserved is an error, not an abort"); + + // The same for a real decode: the result (16 bytes a pixel) and a channel's bytes together. + let p = Psd::parse(write(&doc(8, ColorMode::Rgb, true))).unwrap(); + let n = 40 * 30; + let peak = 16 * n + n; + assert!(p.composite_with_budget(peak).is_ok()); + for limit in [peak - 1, 16 * n, 16 * n - 1, 1000, 0] { + let e = p.composite_with_budget(limit).unwrap_err(); + assert!(e.to_string().contains("memory budget"), "{limit}: {e}"); + } + // A layer: the result plus its channel; the mask comes after the channels are gone. + let bg = p.layers.iter().position(|l| l.name == "bg").unwrap(); + assert!(p.layer_pixels_with_budget(bg, true, 16 * n + n).is_ok()); + assert!(p.layer_pixels_with_budget(bg, true, 16 * n).unwrap_err().to_string().contains("memory budget")); + assert_eq!(p.layer_pixels_with_budget(bg, true, 16 * n + n).unwrap(), p.layer_pixels(bg, true).unwrap()); +} + +#[test] +fn a_huge_header_with_a_tiny_zip_stream_is_refused_before_it_allocates() { + // 16384 × 16384, RGB + alpha: a few hundred bytes of zlib that inflate to gigabytes of zeros. + let mut small = WDoc::new(2, 2); + small.rle = false; + let base = write(&small); + let zeros = vec![0u8; 1 << 20]; + let stream = miniz_oxide::deflate::compress_to_vec_zlib(&zeros, 9); + let bytes = with_merged(&base, Some((16384, 16384)), None, 2, &stream); + assert!(bytes.len() < 8192); + let p = Psd::parse(bytes).unwrap(); + assert_eq!((p.width, p.height), (16384, 16384)); + let e = p.composite().unwrap_err(); + assert!(e.to_string().contains("memory budget"), "{e}"); + // The same file with the compression types that read straight from the file. + for comp in [0u16, 1, 3] { + let p = Psd::parse(with_merged(&base, Some((16384, 16384)), None, comp, &stream)).unwrap(); + let e = p.composite().unwrap_err(); + assert!(e.to_string().contains("memory budget"), "comp {comp}: {e}"); + } + // And a smaller image under a small budget: the stream is not inflated. + let p = Psd::parse(with_merged(&base, Some((4096, 4096)), None, 2, &stream)).unwrap(); + assert!(p.composite_with_budget(1 << 20).unwrap_err().to_string().contains("memory budget")); +} diff --git a/crates/psd/src/placed.rs b/crates/psd/src/placed.rs index 156a8e27..25490e4d 100644 --- a/crates/psd/src/placed.rs +++ b/crates/psd/src/placed.rs @@ -12,6 +12,10 @@ use crate::descriptor::{DValue, Descriptor}; use crate::{Pixels, SmartObject}; +/// The most pixels a baked smart object may have: 2^24 (256 MiB as float RGBA). The size comes +/// from the file's transform, which is untrusted. +pub const MAX_BAKE_PIXELS: u64 = 1 << 24; + /// A placed layer's warp. #[derive(Clone, Debug, Default, PartialEq)] pub struct Warp { @@ -252,9 +256,27 @@ impl SmartObject { (((bb[2] - bb[0]) * k).ceil().clamp(1.0, 16384.0) as u32, ((bb[3] - bb[1]) * k).ceil().clamp(1.0, 16384.0) as u32) } + /// The scale and size to bake a `w`×`h` content at, wanting `k` pixels per document pixel but + /// never more than [`MAX_BAKE_PIXELS`] pixels, nor more than 4× the `doc` size across either + /// way: a quad corner far outside the document lowers `k` (possibly below 1) instead of asking + /// for a gigantic image. Returns the scale actually used and the size in pixels. + pub fn bake_scale(&self, w: f64, h: f64, k: f64, doc: (u32, u32)) -> (f64, (u32, u32)) { + let bb = self.placed_bounds(w, h); + let (bw, bh) = ((bb[2] - bb[0]).max(1.0), (bb[3] - bb[1]).max(1.0)); + if !(bw.is_finite() && bh.is_finite() && k.is_finite() && k > 0.0) { + return (1.0, (1, 1)); + } + let (max_w, max_h) = ((f64::from(doc.0) * 4.0).max(1.0), (f64::from(doc.1) * 4.0).max(1.0)); + let capped = k.min(max_w / bw).min(max_h / bh).min((MAX_BAKE_PIXELS as f64 / (bw * bh)).sqrt()); + // Rounding up could step over a cap, so a lowered scale rounds down. + let side = |len: f64| if capped < k { (len * capped).floor() } else { (len * capped).ceil() }.clamp(1.0, 16384.0) as u32; + (capped, (side(bw), side(bh))) + } + /// Bake the content (straight RGBA) as placed: warped and mapped by the quad, rendered at /// `k` pixels per document pixel (into `size` pixels when given; see [`Self::placed_size`]). - /// Returns the pixels and their document-space origin. + /// Returns the pixels and their document-space origin. `None` when the content is empty or + /// the baked image would exceed [`MAX_BAKE_PIXELS`]. pub fn render_placed(&self, content: &Pixels, k: f64, size: Option<(u32, u32)>) -> Option<(Pixels, [f64; 2])> { let (cw, ch) = (content.width as f64, content.height as f64); if content.width == 0 || content.height == 0 { @@ -266,6 +288,7 @@ impl SmartObject { None => self.placed_size(cw, ch, k), }; let (ow, oh) = (ow as usize, oh as usize); + let count = ow.checked_mul(oh).filter(|&c| c as u64 <= MAX_BAKE_PIXELS)?; // A forward grid, filled triangle by triangle with inverse barycentric sampling. let n = 48usize; let mut dst = Vec::with_capacity((n + 1) * (n + 1)); @@ -278,7 +301,9 @@ impl SmartObject { src.push(s); } } - let mut out = vec![[0.0f32; 4]; ow * oh]; + let mut out: Vec<[f32; 4]> = Vec::new(); + out.try_reserve_exact(count).ok()?; + out.resize(count, [0.0; 4]); let sample = |x: f64, y: f64| -> [f32; 4] { let (fx, fy) = ((x - 0.5).clamp(0.0, cw - 1.0), (y - 0.5).clamp(0.0, ch - 1.0)); let (x0, y0) = (fx.floor() as usize, fy.floor() as usize); @@ -407,4 +432,34 @@ mod tests { assert!(at(35, 10)[1] > 0.9, "{:?}", at(35, 10)); assert_eq!(at(38, 1)[3], 0.0, "outside the trapezoid"); } + #[test] + fn an_absurd_quad_corner_bakes_within_four_times_the_document_and_2_pow_24_pixels() { + let doc = (1920u32, 1080u32); + for far in [1e4, 1e6, 1e12] { + let s = SmartObject { uuid: "x".into(), quad: [[0.0, 0.0], [far, 5.0], [far, far], [0.0, 20.0]], size: None, warp: None }; + assert!(s.needs_bake()); + // What the importer asks for: the unbounded size is the 16384² clamp, 4 GiB as float RGBA. + assert_eq!(s.placed_size(100.0, 100.0, 4.0), (16384, 16384)); + let (k, (w, h)) = s.bake_scale(100.0, 100.0, 4.0, doc); + assert!(k > 0.0 && k <= 4.0, "{k}"); + assert!(w <= doc.0 * 4 && h <= doc.1 * 4, "{w}x{h}"); + assert!(u64::from(w) * u64::from(h) <= MAX_BAKE_PIXELS, "{w}x{h}"); + } + // A quad inside the document keeps the scale and the size it always had. + let s = SmartObject { uuid: "x".into(), quad: [[0.0, 0.0], [400.0, 50.0], [400.0, 150.0], [0.0, 200.0]], size: None, warp: None }; + let (k, size) = s.bake_scale(100.0, 100.0, 2.0, doc); + assert_eq!((k, size), (2.0, s.placed_size(100.0, 100.0, 2.0))); + // Not-a-number corners stay small too. + let s = SmartObject { uuid: "x".into(), quad: [[f64::NAN; 2]; 4], size: None, warp: None }; + let (_, (w, h)) = s.bake_scale(100.0, 100.0, 4.0, doc); + assert!(u64::from(w) * u64::from(h) <= MAX_BAKE_PIXELS); + } + + #[test] + fn render_placed_refuses_an_output_over_the_pixel_cap() { + let content = Pixels { width: 2, height: 1, data: vec![[1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 1.0]] }; + let s = SmartObject { uuid: "x".into(), quad: [[0.0, 0.0], [40.0, 5.0], [40.0, 15.0], [0.0, 20.0]], size: None, warp: None }; + assert!(s.render_placed(&content, 1.0, Some((16384, 16384))).is_none()); + assert!(s.render_placed(&content, 1.0, Some((u32::MAX, u32::MAX))).is_none()); + } }