diff --git a/crates/psd/src/pixels.rs b/crates/psd/src/pixels.rs index 013d8217..d4948457 100644 --- a/crates/psd/src/pixels.rs +++ b/crates/psd/src/pixels.rs @@ -31,35 +31,58 @@ pub fn f32_to_bytes(samples: &[f32]) -> Vec { samples.iter().flat_map(|s| s.to_be_bytes()).collect() } +/// Allocates `len` bytes, reporting a machine without that much memory instead of aborting. +fn try_buffer(len: usize) -> Result> { + let mut out = Vec::new(); + out.try_reserve_exact(len).map_err(|_| PsdError::LimitExceeded("not enough memory for the converted plane"))?; + Ok(out) +} + /// Expands a 1-bit plane (MSB first, rows padded to bytes) to one byte per -/// pixel (0 or 1). -pub fn unpack_bits(bytes: &[u8], width: usize, height: usize) -> Vec { +/// pixel (0 or 1). Errors with [`PsdError::LimitExceeded`] when the plane +/// size overflows or cannot be allocated. +pub fn unpack_bits(bytes: &[u8], width: usize, height: usize) -> Result> { + unpack_bits_with(bytes, width, height, |bit| bit) +} + +/// Shared 1-bit expansion; `map` turns each bit (0 or 1) into the output byte. +fn unpack_bits_with(bytes: &[u8], width: usize, height: usize, map: impl Fn(u8) -> u8) -> Result> { + let len = width.checked_mul(height).ok_or(PsdError::LimitExceeded("plane size overflow"))?; let rb = width.div_ceil(8); - let mut out = Vec::with_capacity(width * height); + let mut out = try_buffer(len)?; for y in 0..height { for x in 0..width { - let b = bytes.get(y * rb + x / 8).copied().unwrap_or(0); - out.push((b >> (7 - (x % 8))) & 1); + let b = y.checked_mul(rb).and_then(|row| bytes.get(row + x / 8)).copied().unwrap_or(0); + out.push(map((b >> (7 - (x % 8))) & 1)); } } - out + Ok(out) } /// Converts one plane of samples at `depth` to 8-bit. pub fn plane_to_u8(bytes: &[u8], depth: u16, width: usize, height: usize) -> Result> { Ok(match depth { - 1 => unpack_bits(bytes, width, height).into_iter().map(|b| if b == 1 { 0 } else { 255 }).collect(), - 8 => bytes.to_vec(), - 16 => bytes.as_chunks::<2>().0.iter().map(|b| ((u32::from(u16::from_be_bytes(*b)) * 255 + 32767) / 65535) as u8).collect(), - 32 => bytes - .as_chunks::<4>() - .0 - .iter() - .map(|b| { + 1 => unpack_bits_with(bytes, width, height, |b| if b == 1 { 0 } else { 255 })?, + 8 => { + let mut out = try_buffer(bytes.len())?; + out.extend_from_slice(bytes); + out + } + 16 => { + let chunks = bytes.as_chunks::<2>().0; + let mut out = try_buffer(chunks.len())?; + out.extend(chunks.iter().map(|b| ((u32::from(u16::from_be_bytes(*b)) * 255 + 32767) / 65535) as u8)); + out + } + 32 => { + let chunks = bytes.as_chunks::<4>().0; + let mut out = try_buffer(chunks.len())?; + out.extend(chunks.iter().map(|b| { let v = f32::from_be_bytes(*b); if v.is_nan() { 0 } else { (v.clamp(0.0, 1.0) * 255.0).round() as u8 } - }) - .collect(), + })); + out + } d => return Err(PsdError::Unsupported(format!("depth {d}"))), }) } @@ -282,6 +305,10 @@ impl PsdFile { h.validate()?; let (w, hh) = (h.width as usize, h.height as usize); let n = w.checked_mul(hh).ok_or(PsdError::LimitExceeded("composite size overflow"))?; + // Same budget as `Layer::rgba8`: refuse before decoding or converting any plane. + if n.checked_mul(4).is_none_or(|len| len as u64 > crate::compression::MAX_DECODED_BYTES) { + return Err(PsdError::LimitExceeded("RGBA output exceeds MAX_DECODED_BYTES")); + } let cc = color_channel_count(h.color_mode)?; if usize::from(h.channels) < cc { return Err(PsdError::invalid("fewer channels than the color mode requires")); @@ -358,7 +385,7 @@ mod tests { #[test] fn unpack_bits_rows() { - assert_eq!(unpack_bits(&[0b1010_0000, 0b1000_0000], 3, 2), vec![1, 0, 1, 1, 0, 0]); + assert_eq!(unpack_bits(&[0b1010_0000, 0b1000_0000], 3, 2).unwrap(), vec![1, 0, 1, 1, 0, 0]); } #[test] @@ -369,6 +396,26 @@ mod tests { assert!(plane_to_u8(&[], 12, 0, 0).is_err()); } + /// A plane size that overflows or cannot be allocated is an error, not a panic or abort. + #[test] + fn plane_conversion_reports_absurd_sizes() { + assert!(matches!(unpack_bits(&[0], usize::MAX / 2, 4), Err(PsdError::LimitExceeded(_)))); + assert!(matches!(plane_to_u8(&[0], 1, usize::MAX / 2, 4), Err(PsdError::LimitExceeded(_)))); + assert!(matches!(unpack_bits(&[0], 1 << 40, 1 << 20), Err(PsdError::LimitExceeded(_)))); + } + + /// A bitmap composite whose RGBA output is over the decode budget is refused up front. + #[test] + fn composite_rgba8_checks_the_output_budget_first() { + let mut file = crate::testgen::small(Version::Psb, crate::Compression::Zip); + file.header.color_mode = ColorMode::Bitmap; + file.header.depth = 1; + file.header.channels = 1; + file.header.width = 60_000; + file.header.height = 60_000; + assert!(matches!(file.composite_rgba8(), Err(PsdError::LimitExceeded("RGBA output exceeds MAX_DECODED_BYTES")))); + } + #[test] fn rgb_interleave_checks_planes_and_preserves_alpha() { let color = [Cow::Borrowed(&[1, 2][..]), Cow::Borrowed(&[3, 4][..]), Cow::Borrowed(&[5, 6][..])];