From afa0a86341d76c13aede6a87df57d37665e6b4dc Mon Sep 17 00:00:00 2001 From: FredApps <153173854+FredApps@users.noreply.github.com> Date: Tue, 3 Nov 2026 13:06:09 +0100 Subject: [PATCH] psd: widen staged RLE data to PSB in short reads, never reading it back whole --- crates/psd/src/compression.rs | 15 +++++++--- crates/psd/src/file.rs | 52 +++++++++++++++++++++++++++++++---- crates/psd/src/io.rs | 39 ++++++++++++++++++++++---- crates/psd/src/layer.rs | 8 +++--- 4 files changed, 95 insertions(+), 19 deletions(-) diff --git a/crates/psd/src/compression.rs b/crates/psd/src/compression.rs index 5a49284c44..73c71bab1d 100644 --- a/crates/psd/src/compression.rs +++ b/crates/psd/src/compression.rs @@ -228,11 +228,18 @@ pub fn decode_planes(compression: Compression, data: &[u8], layout: &PlaneLayout /// Bytes a PSD RLE stream (two-byte row counts) of `rows` rows grows by when written as PSB /// (four-byte row counts); an error when `data` is too short to hold the count table. pub(crate) fn rle_widening(data: &[u8], rows: usize) -> Result { - let table = rows.checked_mul(2).ok_or(PsdError::LimitExceeded("row count overflow"))?; - if data.len() < table { - return Err(PsdError::UnexpectedEof { offset: data.len(), needed: table - data.len() }); + rle_widening_len(data.len() as u64, rows) +} + +/// [`rle_widening`] from the stream's length alone, so data still in its source (or staged on +/// the scratch disk) is not read to measure it. +pub(crate) fn rle_widening_len(len: u64, rows: usize) -> Result { + let table = rows.checked_mul(2).ok_or(PsdError::LimitExceeded("row count overflow"))? as u64; + if len < table { + let offset = usize::try_from(len).unwrap_or(usize::MAX); + return Err(PsdError::UnexpectedEof { offset, needed: usize::try_from(table - len).unwrap_or(usize::MAX) }); } - Ok(table as u64) + Ok(table) } /// Encodes planar data (excluding the 2-byte compression field). diff --git a/crates/psd/src/file.rs b/crates/psd/src/file.rs index babcef4f58..910d19563a 100644 --- a/crates/psd/src/file.rs +++ b/crates/psd/src/file.rs @@ -2,7 +2,7 @@ use std::io::{Seek, Write}; -use crate::compression::{Compression, MAX_DECODED_BYTES, decode_planes, rle_widening, validate_planes}; +use crate::compression::{Compression, MAX_DECODED_BYTES, decode_planes, rle_widening_len, validate_planes}; use crate::error::{PsdError, Result}; use crate::header::{ColorMode, Header, Version}; use crate::image_data::ImageData; @@ -384,7 +384,7 @@ impl PsdFile { })?; match self.merged_widened_rows(widen)? { 0 => out.put_encoded(&self.image_data.data, self.image_data.stored.as_ref()), - rows => out.put_widened_rle(&self.image_data.encoded()?, rows), + rows => out.put_widened_encoded(&self.image_data.data, self.image_data.stored.as_ref(), rows), } } @@ -414,13 +414,13 @@ impl PsdFile { for c in &layer.channels { match layer.widened_rows(c, true)? { 0 => {} - rows => n = n.saturating_add(rle_widening(&c.encoded()?, rows)?), + rows => n = n.saturating_add(rle_widening_len(c.encoded_len(), rows)?), } } } match self.merged_widened_rows(true)? { 0 => Ok(n), - rows => Ok(n.saturating_add(rle_widening(&self.image_data.encoded()?, rows)?)), + rows => Ok(n.saturating_add(rle_widening_len(self.image_data.encoded_len(), rows)?)), } } @@ -463,7 +463,7 @@ impl PsdFile { } let merged = match self.merged_widened_rows(widen)? { 0 => 0, - rows => rle_widening(&self.image_data.encoded()?, rows)?, + rows => rle_widening_len(self.image_data.encoded_len(), rows)?, }; Ok(n.saturating_add(2).saturating_add(self.image_data.encoded_len()).saturating_add(merged)) } @@ -747,6 +747,48 @@ mod write_tests { assert!(*source.1.lock().unwrap() <= 7, "read {} bytes at once", source.1.lock().unwrap()); } + /// #5333: RLE data still in its source widens to PSB in short reads too, the count table and + /// then the packed rows, one consecutive pass, never the whole merged image or channel at once. + #[test] + fn source_backed_rle_widens_in_short_reads() { + struct Reads(Vec, std::sync::Mutex<(usize, Vec<(u64, usize)>)>); + impl crate::source::ReadAt for Reads { + fn len(&self) -> u64 { + self.0.len() as u64 + } + fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<()> { + let mut seen = self.1.lock().unwrap(); + seen.0 = seen.0.max(buf.len()); + seen.1.push((offset, buf.len())); + self.0.as_slice().read_at(offset, buf) + } + } + let file = testgen::small(Version::Psd, Compression::Rle); + assert_eq!(file.image_data.compression, Compression::Rle); + let bytes = file.to_bytes().unwrap(); + let mut promoted = file.clone(); + promoted.promote_to_psb().unwrap(); + let want = promoted.to_bytes().unwrap(); + let source = std::sync::Arc::new(Reads(bytes, std::sync::Mutex::new((0, Vec::new())))); + let mut lazy = PsdFile::from_source(crate::Source::from_arc(source.clone())).unwrap(); + assert!(lazy.image_data.stored.is_some(), "the merged image stays in the source"); + lazy.promote_to_psb().unwrap(); + *source.1.lock().unwrap() = (0, Vec::new()); + let mut w = std::io::Cursor::new(Vec::new()); + lazy.write_to(&mut w).unwrap(); + assert!(w.into_inner() == want); + let seen = source.1.lock().unwrap(); + assert!(seen.0 <= 7, "read {} bytes at once", seen.0); + let merged = lazy.image_data.stored.as_ref().unwrap(); + let reads: Vec<_> = seen.1.iter().filter(|(at, _)| (merged.offset..merged.offset + merged.len).contains(at)).collect(); + let mut next = merged.offset; + for &&(at, n) in &reads { + assert_eq!(at, next, "the merged image is read in one consecutive pass"); + next += n as u64; + } + assert_eq!(next, merged.offset + merged.len); + } + /// #6607: pixel data goes to the stream as it is, never through a staging copy, so a save /// holds the model and a short buffer, not a second copy of the file. #[test] diff --git a/crates/psd/src/io.rs b/crates/psd/src/io.rs index bf4468711a..8537b74649 100644 --- a/crates/psd/src/io.rs +++ b/crates/psd/src/io.rs @@ -275,16 +275,43 @@ impl<'w, W: Write + Seek + ?Sized> Sink<'w, W> { let table_len = usize::try_from(crate::compression::rle_widening(data, rows)?).map_err(|_| PsdError::LimitExceeded("row count overflow"))?; let (table, packed) = data.split_at_checked(table_len).ok_or_else(|| PsdError::invalid("RLE row-count table out of range"))?; for part in table.chunks(STAGE / 2) { - self.small(|o| { - for count in part.as_chunks::<2>().0 { - o.extend_from_slice(&[0, 0, count[0], count[1]]); - } - Ok(()) - })?; + self.widen_counts(part)?; } self.put(packed) } + /// [`Self::put_widened_rle`] for data held in memory (`data`) or, when `stored`, still in its + /// source: the count table and then the packed rows are read in [`COPY_CHUNK`]-sized parts + /// (#5333), so a staged merged image or channel is never read back whole to widen it. + pub(crate) fn put_widened_encoded(&mut self, data: &[u8], stored: Option<&crate::source::Stored>, rows: usize) -> std::io::Result<()> { + let Some(s) = stored else { return self.put_widened_rle(data, rows) }; + let table_len = crate::compression::rle_widening_len(s.len, rows)?; + // Whole counts per part; the parts follow each other, so a staged range is read in one pass. + let table_part = (COPY_CHUNK / 2).max(1) * 2; + let mut at = 0u64; + while at < s.len { + let n = if at < table_len { (table_len - at).min(table_part) } else { (s.len - at).min(COPY_CHUNK) }; + let part = s.read_part(at, n)?; + if at < table_len { + self.widen_counts(&part)? + } else { + self.put(&part)? + } + at += n; + } + Ok(()) + } + + /// Two-byte RLE row counts as four-byte ones (`00 00` and the count), staged. + fn widen_counts(&mut self, counts: &[u8]) -> std::io::Result<()> { + self.small(|o| { + for count in counts.as_chunks::<2>().0 { + o.extend_from_slice(&[0, 0, count[0], count[1]]); + } + Ok(()) + }) + } + /// Appends encoded bytes held in memory (`data`) or, when `stored`, still in their source: /// those are copied through in [`COPY_CHUNK`] pieces, so a channel or merged image kept on /// disk (an imported file's, or one an exporter spilled to scratch) is never read whole. diff --git a/crates/psd/src/layer.rs b/crates/psd/src/layer.rs index 3b492e6a04..90f6550863 100644 --- a/crates/psd/src/layer.rs +++ b/crates/psd/src/layer.rs @@ -3,7 +3,7 @@ use std::borrow::Cow; use crate::blend::BlendMode; -use crate::compression::{Compression, PlaneLayout, decode_planes, encode_planes, rle_widening}; +use crate::compression::{Compression, PlaneLayout, decode_planes, encode_planes, rle_widening_len}; use crate::error::{PsdError, Result}; use crate::header::{MAX_DIMENSION, Version}; use crate::io::{Reader, Sink, WriteExt, append_with, decode_legacy_name, read_pascal, write_pascal}; @@ -692,7 +692,7 @@ impl LayerRecord { for c in &self.channels { out.put_i16(c.id); let rows = self.widened_rows(c, widen)?; - let grown = if rows == 0 { 0 } else { rle_widening(&c.encoded()?, rows)? }; + let grown = if rows == 0 { 0 } else { rle_widening_len(c.encoded_len(), rows)? }; out.put_len(c.stored_len().saturating_add(grown), version.is_psb())?; } out.put(b"8BIM"); @@ -783,7 +783,7 @@ impl LayerInfo { for l in &self.layers { for c in l.channels.iter().filter(|c| c.compression.is_some()) { let rows = l.widened_rows(c, widen)?; - let grown = if rows == 0 { 0 } else { rle_widening(&c.encoded()?, rows)? }; + let grown = if rows == 0 { 0 } else { rle_widening_len(c.encoded_len(), rows)? }; n = n.saturating_add(2).saturating_add(c.encoded_len()).saturating_add(grown); } } @@ -860,7 +860,7 @@ impl LayerInfo { })?; match l.widened_rows(c, widen)? { 0 => out.put_encoded(&c.data, c.stored.as_ref())?, - rows => out.put_widened_rle(&c.encoded()?, rows)?, + rows => out.put_widened_encoded(&c.data, c.stored.as_ref(), rows)?, } } }