Bound PSD layer and mask import memory before decoding (#10503)

* Bound PSD import pixel storage before decoding

* Include ZIP inflater workspace in PSD preflight

* PSD import: budget unknown-mode layers, stop after a fatal layer error, keep best-effort for non-allocation limits (#8593 review)

crates/psd/src/compression.rs

@@ -210,7 +210,7 @@ pub fn decode_planes(compression: Compression, data: &[u8], layout: &PlaneLayout
                 return Err(PsdError::UnexpectedEof { offset: data.len(), needed: total - data.len() });
             }
             let mut out = Vec::new();
-            out.try_reserve_exact(total).map_err(|_| PsdError::LimitExceeded("not enough memory for the decoded channel data"))?;
+            out.try_reserve_exact(total).map_err(|_| PsdError::OutOfMemory("not enough memory for the decoded channel data"))?;
             out.extend_from_slice(&data[..total]);
             Ok(out)
         }
@@ -311,7 +311,7 @@ thread_local! {
 // Keep only the requested pixels, but consume the trailer to verify Adler-32.
 // A bounded surplus tolerates padding without unbounded decompression work.
 fn inflate_checked(data: &[u8], expected: u64, keep: bool) -> Result<Vec<u8>> {
-    const OOM: PsdError = PsdError::LimitExceeded("not enough memory for the decoded channel data");
+    const OOM: PsdError = PsdError::OutOfMemory("not enough memory for the decoded channel data");
     #[cfg(test)]
     INFLATES.with(|n| n.set(n.get() + 1));
     let mut decoder = flate2::Decompress::new(true);
@@ -364,7 +364,7 @@ fn decode_rle(data: &[u8], layout: &PlaneLayout, total: usize) -> Result<Vec<u8>
     }
     // Up to MAX_DECODED_BYTES: report a machine without that much memory instead of aborting.
     let mut out = Vec::new();
-    out.try_reserve_exact(total).map_err(|_| PsdError::LimitExceeded("not enough memory for the decoded channel data"))?;
+    out.try_reserve_exact(total).map_err(|_| PsdError::OutOfMemory("not enough memory for the decoded channel data"))?;
     for c in counts {
         let row = r.bytes(c)?;
         packbits::decode_into(row, rb, &mut out)?;
@@ -450,7 +450,7 @@ pub fn zip_compress(data: &[u8]) -> Vec<u8> {
 /// bytes, so a stream declaring a huge size reports an error on a machine
 /// without that much memory instead of aborting.
 pub fn zip_decompress(data: &[u8], expected: usize) -> Result<Vec<u8>> {
-    const OOM: PsdError = PsdError::LimitExceeded("not enough memory for the decoded channel data");
+    const OOM: PsdError = PsdError::OutOfMemory("not enough memory for the decoded channel data");
     const CHUNK: usize = 1 << 20;
     let mut dec = flate2::read::ZlibDecoder::new(data).take(u64::try_from(expected).unwrap_or(u64::MAX));
     let mut out = Vec::new();

crates/psd/src/error.rs

@@ -31,6 +31,10 @@ pub enum PsdError {
     /// A dimension, count or allocation exceeded the configured safety limits.
     #[error("limit exceeded: {0}")]
     LimitExceeded(&'static str),
+    /// A buffer could not be allocated: the system has no memory for it. Unlike
+    /// [`PsdError::LimitExceeded`] the data may be valid, so the caller can retry smaller.
+    #[error("out of memory: {0}")]
+    OutOfMemory(&'static str),
     /// Structurally invalid data.
     #[error("invalid data: {0}")]
     Invalid(String),

crates/psd/src/file.rs

@@ -132,6 +132,17 @@ impl PsdFile {
         data: &[u8],
         want: impl FnOnce(&PsdFile) -> bool,
         recover: impl FnOnce(&PsdFile) -> bool,
+    ) -> Result<(Self, Option<Result<Vec<u8>>>)> {
+        Self::from_bytes_and_merged_preflight(data, |file| Ok(want(file)), recover)
+    }
+
+    /// Parses with a fallible pixel preflight after layer records are available and before
+    /// the merged image is read, validated or decoded. A preflight error is always fatal;
+    /// `recover` only handles damage encountered after an accepted preflight.
+    pub fn from_bytes_and_merged_preflight(
+        data: &[u8],
+        want: impl FnOnce(&PsdFile) -> Result<bool>,
+        recover: impl FnOnce(&PsdFile) -> bool,
     ) -> Result<(Self, Option<Result<Vec<u8>>>)> {
         let mut r = Reader::new(data);
         let header = Header::read(&mut r)?;
@@ -202,7 +213,7 @@ impl PsdFile {
             image_data,
             write_as_psb: false,
         };
-        let want = want(&file);
+        let want = want(&file)?;
         let merged = file.read_merged_section(&mut r, want, recover)?;
         Ok((file, merged))
     }
@@ -227,13 +238,13 @@ impl PsdFile {
         let merged = match merged {
             Ok(merged) => merged,
             Err(error) => {
-                if matches!(error, PsdError::LimitExceeded(_)) || !recover(file) {
+                if matches!(error, PsdError::LimitExceeded(_) | PsdError::OutOfMemory(_)) || !recover(file) {
                     return Err(error);
                 }
                 let compression = Compression::from_u16(encoded.u16()?);
                 let source = encoded.peek_rest();
                 let mut data = Vec::new();
-                data.try_reserve_exact(source.len()).map_err(|_| PsdError::LimitExceeded("not enough memory to retain the damaged merged image"))?;
+                data.try_reserve_exact(source.len()).map_err(|_| PsdError::OutOfMemory("not enough memory to retain the damaged merged image"))?;
                 data.extend_from_slice(source);
                 file.image_data = ImageData::new(compression, data);
                 Some(Err(error))
@@ -259,10 +270,18 @@ impl PsdFile {
         r.u16()?;
         let data = r.peek_rest();
         let decoded = decode_planes(compression, data, &layout);
+        // An allocation failure is the open's, before the encoded image is copied under the same
+        // memory pressure.
+        if matches!(decoded, Err(PsdError::OutOfMemory(_))) {
+            return decoded.map(Ok);
+        }
         if decoded.is_err() {
             validate_planes(compression, data, &layout)?;
         }
-        self.image_data = ImageData::new(compression, data.to_vec());
+        let mut copy = Vec::new();
+        copy.try_reserve_exact(data.len()).map_err(|_| PsdError::OutOfMemory("not enough memory to retain the merged image"))?;
+        copy.extend_from_slice(data);
+        self.image_data = ImageData::new(compression, copy);
         r.skip(data.len())?;
         Ok(decoded)
     }
@@ -281,7 +300,7 @@ impl PsdFile {
             .saturating_add(usize::try_from(self.image_data.encoded_len()).unwrap_or(usize::MAX))
             .saturating_add(usize::try_from(widened).unwrap_or(usize::MAX));
         let mut out = Vec::new();
-        out.try_reserve(hint).map_err(|_| PsdError::LimitExceeded("not enough memory to serialize the file"))?;
+        out.try_reserve(hint).map_err(|_| PsdError::OutOfMemory("not enough memory to serialize the file"))?;
         if widen {
             Header { version: v, ..self.header.clone() }.write(&mut out);
         } else {
@@ -848,4 +867,22 @@ mod write_tests {
             assert!(!other.merged_has_alpha(), "{}", String::from_utf8_lossy(&key));
         }
     }
+    #[test]
+    fn preflight_refusal_precedes_merged_validation_and_cannot_be_recovered() {
+        let file = testgen::small(Version::Psd, Compression::Zip);
+        let mut bytes = file.to_bytes().unwrap();
+        bytes.truncate(bytes.len() - file.image_data.data.len() - 2);
+        let recover_called = std::cell::Cell::new(false);
+        let preflight = |_: &PsdFile| Err(PsdError::LimitExceeded("test import budget"));
+        let recover = |_: &PsdFile| {
+            recover_called.set(true);
+            true
+        };
+        assert!(matches!(PsdFile::from_bytes_and_merged_preflight(&bytes, preflight, recover), Err(PsdError::LimitExceeded("test import budget"))));
+        assert!(matches!(
+            PsdFile::from_source_and_merged_preflight(crate::Source::new(bytes), preflight, recover),
+            Err(PsdError::LimitExceeded("test import budget"))
+        ));
+        assert!(!recover_called.get());
+    }
 }

crates/psd/src/lazy.rs

@@ -131,6 +131,17 @@ impl PsdFile {
         source: Source,
         want: impl FnOnce(&PsdFile) -> bool,
         recover: impl FnOnce(&PsdFile) -> bool,
+    ) -> Result<(Self, Option<Result<Vec<u8>>>)> {
+        Self::from_source_and_merged_preflight(source, |file| Ok(want(file)), recover)
+    }
+
+    /// Parses with a fallible pixel preflight after layer records are available and before
+    /// the merged image is read, validated or decoded. A preflight error is always fatal;
+    /// `recover` only handles damage encountered after an accepted preflight.
+    pub fn from_source_and_merged_preflight(
+        source: Source,
+        want: impl FnOnce(&PsdFile) -> Result<bool>,
+        recover: impl FnOnce(&PsdFile) -> bool,
     ) -> Result<(Self, Option<Result<Vec<u8>>>)> {
         let mut c = Cursor::new(&source);
         let header = Header::read(&mut Reader::new(&c.bytes(26)?))?;
@@ -217,7 +228,7 @@ impl PsdFile {
             image_data: ImageData::new(Compression::Raw, Vec::new()),
             write_as_psb: false,
         };
-        if !want(&file) {
+        if !want(&file)? {
             let compression = Compression::from_u16(c.u16()?);
             file.image_data = ImageData { compression, data: Vec::new(), stored: Some(c.stored_rest()) };
             return Ok((file, None));

crates/psd/src/patterns.rs

@@ -377,7 +377,7 @@ fn place_channel_plane(mut plane: Vec<u8>, source: &PlaneLayout, x: usize, y: us
 
 fn zeroed_pattern_plane(len: usize) -> Result<Vec<u8>> {
     let mut out = Vec::new();
-    out.try_reserve_exact(len).map_err(|_| PsdError::LimitExceeded("not enough memory for the padded pattern channel"))?;
+    out.try_reserve_exact(len).map_err(|_| PsdError::OutOfMemory("not enough memory for the padded pattern channel"))?;
     out.resize(len, 0);
     Ok(out)
 }

crates/psd/src/patterns/tests.rs

@@ -267,7 +267,7 @@ fn large_channel_padding_within_the_existing_budget_is_allowed() {
 #[test]
 fn padded_plane_allocation_failure_is_an_error() {
     // Capacity overflow is deterministic and never asks the OS for huge memory.
-    assert!(matches!(zeroed_pattern_plane(usize::MAX), Err(PsdError::LimitExceeded("not enough memory for the padded pattern channel"))));
+    assert!(matches!(zeroed_pattern_plane(usize::MAX), Err(PsdError::OutOfMemory("not enough memory for the padded pattern channel"))));
 }
 
 #[test]

crates/psd/src/pixels.rs

@@ -34,7 +34,7 @@ pub fn f32_to_bytes(samples: &[f32]) -> Vec<u8> {
 /// 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"))?;
+    out.try_reserve_exact(len).map_err(|_| PsdError::OutOfMemory("not enough memory for the converted plane"))?;
     Ok(out)
 }
 
@@ -128,7 +128,7 @@ fn interleave(mode: ColorMode, color: &[Cow<'_, [u8]>], alpha: Option<&[u8]>, pa
         .filter(|&len| len as u64 <= crate::compression::MAX_DECODED_BYTES)
         .ok_or(PsdError::LimitExceeded("RGBA output exceeds MAX_DECODED_BYTES"))?;
     let mut out = Vec::new();
-    out.try_reserve_exact(len).map_err(|_| PsdError::LimitExceeded("not enough memory for RGBA output"))?;
+    out.try_reserve_exact(len).map_err(|_| PsdError::OutOfMemory("not enough memory for RGBA output"))?;
     out.resize(len, 255);
     if mode == ColorMode::Rgb {
         let [r, g, b] = color else { return Err(PsdError::invalid("RGB requires three channels")) };
@@ -439,7 +439,7 @@ mod tests {
     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(_))));
+        assert!(matches!(unpack_bits(&[0], 1 << 40, 1 << 20), Err(PsdError::LimitExceeded(_) | PsdError::OutOfMemory(_))));
     }
 
     /// A bitmap composite whose RGBA output is over the decode budget is refused up front.

crates/psd/src/rows.rs

@@ -98,7 +98,7 @@ impl ChannelData {
                 }
                 let table = bytes.get(0, table_len)?;
                 let mut offsets = Vec::new();
-                offsets.try_reserve_exact(height + 1).map_err(|_| PsdError::LimitExceeded("not enough memory for the RLE row table"))?;
+                offsets.try_reserve_exact(height + 1).map_err(|_| PsdError::OutOfMemory("not enough memory for the RLE row table"))?;
                 let mut at = table_len;
                 offsets.push(at);
                 for c in table.chunks_exact(size as usize) {
@@ -150,7 +150,7 @@ impl RowDecoder<'_> {
             return Err(PsdError::LimitExceeded("decoded channel data exceeds MAX_DECODED_BYTES"));
         }
         let mut out = Vec::new();
-        out.try_reserve_exact(n).map_err(|_| PsdError::LimitExceeded("not enough memory for the decoded channel data"))?;
+        out.try_reserve_exact(n).map_err(|_| PsdError::OutOfMemory("not enough memory for the decoded channel data"))?;
         match &self.kind {
             Kind::Empty => {}
             Kind::Raw(bytes) => out.extend_from_slice(&bytes.get((r0 * self.row_bytes) as u64, n as u64)?),

crates/psd/src/source.rs

@@ -173,7 +173,7 @@ impl Source {
             return Err(out_of_range(offset, n, self.0.len()));
         }
         let mut out = Vec::new();
-        out.try_reserve_exact(n).map_err(|_| PsdError::LimitExceeded("not enough memory to read the data"))?;
+        out.try_reserve_exact(n).map_err(|_| PsdError::OutOfMemory("not enough memory to read the data"))?;
         out.resize(n, 0);
         self.0.read_at(offset, &mut out)?;
         Ok(out)

crates/psd/src/tiff.rs

@@ -158,7 +158,7 @@ impl ImageSourceData {
             // Pad to 4 without cloning the layers (and all their pixel data) just to set `padding`.
             let pad = info.pad_len(version, 4)?;
             let mut data = Vec::new();
-            data.try_reserve(info.channel_data_len().saturating_add(pad)).map_err(|_| PsdError::LimitExceeded("not enough memory for the layer data"))?;
+            data.try_reserve(info.channel_data_len().saturating_add(pad)).map_err(|_| PsdError::OutOfMemory("not enough memory for the layer data"))?;
             info.write_unpadded(&mut data, version)?;
             data.resize(data.len() + pad, 0);
             padded_block(self.layer_info_key, data).write(&mut body, version, 4)?;