fix(psd): retain pattern budget charges after parse failures (#8767)

Keep allocation charges monotonic when malformed pattern blocks fail, so later sections cannot reuse budget already consumed.

🤖 Generated with Codebuff

Co-authored-by: Codebuff <noreply@codebuff.com>

crates/psd/src/abr.rs

@@ -435,6 +435,7 @@ fn parse_v6_within(r: &mut Rd, out: &mut AbrFile, budget: &mut Budget) -> Result
         return Err(PsdError::Unsupported(format!("brush file subversion {}", out.subversion)));
     }
     let mut hierarchy = None;
+    let mut patterns = PatternBudget::default();
     while r.left() >= 12 {
         let sig = r.arr::<4>()?;
         if &sig != b"8BIM" {
@@ -458,16 +459,17 @@ fn parse_v6_within(r: &mut Rd, out: &mut AbrFile, budget: &mut Budget) -> Result
         match &key {
             b"samp" => read_samp(data, out.subversion, out, budget)?,
             b"patt" => {
-                // Patterns share the samples' budget, so neither can claim it twice.
+                // Patterns share the samples' budget and one padding budget across every `patt`
+                // section, so neither can be claimed twice and failed blocks still deduct their
+                // decoded charges.
                 let limit = u64::try_from(budget.0).unwrap_or(u64::MAX);
-                let mut patterns = PatternBudget::with_decoded_limit(limit);
+                patterns.cap_decoded_to(limit);
                 let left = patterns.decoded_left();
-                match parse_pattern_block_with(data, &mut patterns) {
-                    Ok(p) => {
-                        let spent = usize::try_from(left.saturating_sub(patterns.decoded_left())).unwrap_or(usize::MAX);
-                        budget.spend(spent)?;
-                        out.patterns.extend(p);
-                    }
+                let res = parse_pattern_block_with(data, &mut patterns);
+                let spent = usize::try_from(left.saturating_sub(patterns.decoded_left())).unwrap_or(usize::MAX);
+                budget.spend(spent)?;
+                match res {
+                    Ok(p) => out.patterns.extend(p),
                     Err(e) => out.warnings.push(format!("embedded patterns unreadable: {e}")),
                 }
             }
@@ -706,10 +708,39 @@ mod tests {
         assert_eq!((out.samples.len(), out.patterns.len(), left), (1, 2, 0));
         // One byte short: the patterns no longer fit after the sample (they used not to count).
         let (out, left) = parse_within(189);
-        assert_eq!((out.samples.len(), out.patterns.len(), left), (1, 0, 89));
+        assert_eq!((out.samples.len(), out.patterns.len(), left), (1, 0, 14));
         assert!(out.warnings.iter().any(|w| w.contains("embedded patterns unreadable")), "{:?}", out.warnings);
     }
 
+    #[test]
+    fn failed_pattern_section_does_not_refund_budget_to_later_sections() {
+        let p = PsdPattern {
+            mode: 3,
+            width: 5,
+            height: 3,
+            name: "p".into(),
+            id: "id".into(),
+            palette: None,
+            depth: 8,
+            channels: vec![vec![1; 15], vec![2; 15], vec![3; 15]],
+            alpha: None,
+        };
+        let valid = write_pattern_block(std::slice::from_ref(&p)).unwrap();
+        let mut malformed = valid.clone();
+        malformed.extend_from_slice(&100u32.to_be_bytes());
+
+        let mut body = 2u16.to_be_bytes().to_vec();
+        section(&mut body, b"patt", &malformed);
+        section(&mut body, b"patt", &valid);
+        let mut out = AbrFile::default();
+        let mut budget = Budget(60);
+        parse_v6_within(&mut Rd::new(&body), &mut out, &mut budget).unwrap();
+
+        assert!(out.patterns.is_empty(), "a failed block is discarded atomically");
+        assert_eq!(budget.0, 0, "the failed block's first 45 bytes and later partial 15-byte charge are retained");
+        assert_eq!(out.warnings.iter().filter(|w| w.contains("embedded patterns unreadable")).count(), 2);
+    }
+
     #[test]
     fn v1_v2_round_trip_raw_and_rle() {
         for version in [1u16, 2] {

crates/psd/src/patterns.rs

@@ -56,6 +56,12 @@ impl PatternBudget {
         PatternBudget { decoded_left: d.decoded_left.min(decoded), padded_left: d.padded_left.min(decoded) }
     }
 
+    /// Cap both decoded and padded bytes to the amount still available to a containing file.
+    pub fn cap_decoded_to(&mut self, decoded: u64) {
+        self.decoded_left = self.decoded_left.min(decoded);
+        self.padded_left = self.padded_left.min(decoded);
+    }
+
     /// Decoded bytes still available.
     pub fn decoded_left(&self) -> u64 {
         self.decoded_left
@@ -425,9 +431,8 @@ pub fn parse_pattern_block(data: &[u8]) -> Result<Vec<PsdPattern>> {
 }
 
 /// [`parse_pattern_block`] charging `budget`, so several blocks of one document share one limit.
-/// The budget is charged only for patterns that are returned.
+/// Charges are monotonic even when a later record in the block fails.
 pub fn parse_pattern_block_with(data: &[u8], budget: &mut PatternBudget) -> Result<Vec<PsdPattern>> {
-    let mut left = budget.clone();
     let mut r = Rd { b: data, p: 0 };
     let mut out = Vec::new();
     while r.p + 4 <= data.len() {
@@ -436,11 +441,10 @@ pub fn parse_pattern_block_with(data: &[u8], budget: &mut PatternBudget) -> Resu
             break;
         }
         let body = r.take(len)?;
-        out.push(read_pattern(&mut Rd { b: body, p: 0 }, &mut left)?);
+        out.push(read_pattern(&mut Rd { b: body, p: 0 }, budget)?);
         let pad = (4 - len % 4) % 4;
         r.p = (r.p + pad).min(data.len());
     }
-    *budget = left;
     Ok(out)
 }
 
@@ -729,6 +733,18 @@ mod tests {
         block
     }
 
+    #[test]
+    fn failed_pattern_block_keeps_partial_charge_for_the_next_block() {
+        let one = padded_gray_block(100);
+        let mut malformed = one.clone();
+        malformed.extend_from_slice(&100u32.to_be_bytes());
+        let mut budget = PatternBudget::with_decoded_limit(15_000);
+        assert!(parse_pattern_block_with(&malformed, &mut budget).is_err());
+        assert_eq!(budget.decoded_left(), 5_000);
+        assert!(parse_pattern_block_with(&one, &mut budget).is_err());
+        assert_eq!(budget.decoded_left(), 5_000);
+    }
+
     #[test]
     fn padding_budget_is_shared_by_every_pattern_of_a_parse() {
         // Each pattern pads 100x100 = 10000 bytes; a 15000-byte budget fits one, not two. Before
@@ -741,8 +757,9 @@ mod tests {
         let two = [one.clone(), one.clone()].concat();
         let mut budget = PatternBudget::with_decoded_limit(15_000);
         assert!(matches!(parse_pattern_block_with(&two, &mut budget), Err(PsdError::LimitExceeded(_))));
-        // A failed parse keeps nothing, so it charges nothing.
-        assert_eq!(budget.decoded_left(), 15_000);
+        // Charges are monotonic: the first pattern's 10000 bytes stay charged even when the second
+        // record fails.
+        assert_eq!(budget.decoded_left(), 5_000);
 
         // Two blocks of one document share a budget too.
         let mut budget = PatternBudget::with_decoded_limit(15_000);