PSD: composite_rgba8 checks the output budget and converts planes fallibly (#3639) (#3651)

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

crates/psd/src/pixels.rs

@@ -31,35 +31,58 @@ pub fn f32_to_bytes(samples: &[f32]) -> Vec<u8> {
     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<Vec<u8>> {
+    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<u8> {
+/// 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<Vec<u8>> {
+    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<Vec<u8>> {
+    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<Vec<u8>> {
     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][..])];