diff --git a/crates/psd/src/limits.rs b/crates/psd/src/limits.rs index c1c23b390..6db366c00 100644 --- a/crates/psd/src/limits.rs +++ b/crates/psd/src/limits.rs @@ -107,7 +107,7 @@ impl Budget { /// Inflate a zlib stream to at most `limit` bytes. pub(crate) fn inflate(src: &[u8], limit: usize) -> Result> { - miniz_oxide::inflate::decompress_to_vec_zlib_with_limit(src, limit).map_err(|e| Error::Invalid(format!("zip: {e:?}"))) + miniz_oxide::inflate::decompress_to_vec_zlib_with_limit(src, limit).map_err(|e| Error::Invalid(format!("zip: {:?}", e.status))) } /// Give back what a channel buffer from `Psd::channel_bytes` / `Psd::decode_plane` was charged: diff --git a/crates/psd/src/memory_tests.rs b/crates/psd/src/memory_tests.rs index 108046b66..96ac5c70c 100644 --- a/crates/psd/src/memory_tests.rs +++ b/crates/psd/src/memory_tests.rs @@ -28,10 +28,11 @@ impl Psd { } 2 | 3 => { let rest = r.bytes(r.remaining())?; - let mut v = miniz_oxide::inflate::decompress_to_vec_zlib_with_limit(rest, rowb * h).map_err(|e| Error::Invalid(format!("zip: {e:?}")))?; + let mut v = + miniz_oxide::inflate::decompress_to_vec_zlib_with_limit(rest, rowb * h).map_err(|e| Error::Invalid(format!("zip: {:?}", e.status)))?; v.resize(rowb * h, 0); if comp == 3 { - self.unpredict(&mut v, w, h); + self.unpredict(&mut v, w, h, &mut Budget::new(u64::MAX))?; } v } @@ -128,11 +129,12 @@ impl Psd { } 2 | 3 => { let rest = r.bytes(r.remaining())?; - let v = miniz_oxide::inflate::decompress_to_vec_zlib_with_limit(rest, want * rowb * h).map_err(|e| Error::Invalid(format!("zip: {e:?}")))?; + let v = miniz_oxide::inflate::decompress_to_vec_zlib_with_limit(rest, want * rowb * h) + .map_err(|e| Error::Invalid(format!("zip: {:?}", e.status)))?; for (c, p) in planes.iter_mut().enumerate().take(want) { let mut plane = v.get(c * rowb * h..(c + 1) * rowb * h).ok_or(Error::Truncated)?.to_vec(); if comp == 3 { - self.unpredict(&mut plane, w, h); + self.unpredict(&mut plane, w, h, &mut Budget::new(u64::MAX))?; } *p = Some(self.ref_samples(&plane, w, h)); } diff --git a/crates/psd/src/pixels.rs b/crates/psd/src/pixels.rs index 21a5c8845..290a7b6f4 100644 --- a/crates/psd/src/pixels.rs +++ b/crates/psd/src/pixels.rs @@ -37,7 +37,7 @@ impl Psd { PlaneStream::Rle(counts) => Cow::Owned(Self::unpack_plane(&counts, 0, r, rowb, h, budget)?), PlaneStream::Zip(mut v) => { v.resize(size, 0); - self.zip_plane(comp, &mut v, 0, w, h)?; + self.zip_plane(comp, &mut v, 0, w, h, budget)?; Cow::Owned(v) } }) @@ -90,7 +90,7 @@ impl Psd { Ok(match planes { PlaneStream::Raw => Cow::Borrowed(r.bytes(size)?), PlaneStream::Rle(counts) => Cow::Owned(Self::unpack_plane(counts, c, r, rowb, h, budget)?), - PlaneStream::Zip(v) => Cow::Borrowed(self.zip_plane(comp, v, c, w, h)?), + PlaneStream::Zip(v) => Cow::Borrowed(self.zip_plane(comp, v, c, w, h, budget)?), }) } @@ -106,11 +106,11 @@ impl Psd { } /// Plane `c` of the inflated ZIP data of consecutive planes, un-predicted in place (`comp` 3). - pub(crate) fn zip_plane<'v>(&self, comp: u16, v: &'v mut [u8], c: usize, w: usize, h: usize) -> Result<&'v [u8]> { + pub(crate) fn zip_plane<'v>(&self, comp: u16, v: &'v mut [u8], c: usize, w: usize, h: usize, budget: &mut Budget) -> Result<&'v [u8]> { let rowb = self.row_bytes(w); let plane = v.get_mut(c * rowb * h..(c + 1) * rowb * h).ok_or(Error::Truncated)?; if comp == 3 { - self.unpredict(plane, w, h); + self.unpredict(plane, w, h, budget)?; } Ok(plane) } @@ -132,7 +132,8 @@ impl Psd { if self.depth == 1 { w.div_ceil(8) } else { w * self.bytes_per_sample() } } - pub(crate) fn unpredict(&self, v: &mut [u8], w: usize, h: usize) { + /// Undo ZIP prediction in place; the 32-bit scratch row is charged to `budget`. + pub(crate) fn unpredict(&self, v: &mut [u8], w: usize, h: usize, budget: &mut Budget) -> Result<()> { match self.depth { 16 => { for row in v.chunks_exact_mut(w * 2).take(h) { @@ -145,8 +146,9 @@ impl Psd { } } 32 => { - // One scratch row for the whole plane. - let mut planar = vec![0u8; w * 4]; + // One scratch row for the whole plane, charged: a wide, short layer passes the + // decode's peak check and would otherwise allocate it on top. + let mut planar = budget.filled(w.checked_mul(4).ok_or_else(|| too_large(w, h))?, 0u8)?; #[cfg(test)] counters::bump(&counters::SCRATCH); for row in v.chunks_exact_mut(w * 4).take(h) { @@ -161,6 +163,7 @@ impl Psd { } } } + budget.release(planar.len() as u64); } _ => { for row in v.chunks_exact_mut(w).take(h) { @@ -170,6 +173,7 @@ impl Psd { } } } + Ok(()) } /// Hand every sample of `raw` (`w`×`h`, row-major, in order) to `f` as a float. Samples a short @@ -483,15 +487,17 @@ pub(crate) fn unpack_bits(src: &[u8], row: usize, out: &mut Vec) -> Result<( while i < src.len() && out.len() - start < row { let n = src[i] as i8; i += 1; + // A packet is written only up to the row's end, so the reserved plane never grows. + let room = row - (out.len() - start); if n >= 0 { let k = n as usize + 1; let s = src.get(i..i + k).ok_or(Error::Truncated)?; - out.extend_from_slice(s); + out.extend_from_slice(&s[..k.min(room)]); i += k; } else if n != -128 { let k = (1 - n as isize) as usize; let v = *src.get(i).ok_or(Error::Truncated)?; - out.extend(std::iter::repeat_n(v, k)); + out.extend(std::iter::repeat_n(v, k.min(room))); i += 1; } } diff --git a/crates/psd/src/tests.rs b/crates/psd/src/tests.rs index 074976510..fd041138a 100644 --- a/crates/psd/src/tests.rs +++ b/crates/psd/src/tests.rs @@ -493,3 +493,29 @@ fn rle_composite_matches_the_raw_one() { assert_eq!(rle.layer_pixels(2, true).unwrap().data, raw.layer_pixels(2, true).unwrap().data, "{depth} bit"); } } + +#[test] +fn packbits_packets_stop_at_the_row_end() { + // A 128-byte repeat and a 128-byte literal into 3-byte rows: the reserved row is never + // outgrown (capacity stays what was reserved). + for src in [vec![0x81u8, 7], { + let mut v = vec![0x7fu8]; + v.extend([1u8; 128]); + v + }] { + let mut out = Vec::with_capacity(3); + unpack_bits(&src, 3, &mut out).unwrap(); + assert_eq!(out.len(), 3); + assert_eq!(out.capacity(), 3); + } +} + +#[test] +fn a_failed_inflate_reports_its_status_not_its_output() { + // A stream that inflates a megabyte of zeros, then breaks its checksum. + let mut z = miniz_oxide::deflate::compress_to_vec_zlib(&vec![0u8; 1 << 20], 6); + let n = z.len(); + z[n - 1] ^= 0xff; + let Err(Error::Invalid(msg)) = crate::limits::inflate(&z, 1 << 21) else { panic!("the checksum fails") }; + assert!(msg.len() < 100, "{} bytes", msg.len()); +} diff --git a/crates/psd/src/tests/opt.rs b/crates/psd/src/tests/opt.rs index 908026d07..44bcec071 100644 --- a/crates/psd/src/tests/opt.rs +++ b/crates/psd/src/tests/opt.rs @@ -69,7 +69,7 @@ fn ref_decode_channel(p: &Psd, offset: usize, len: usize, w: usize, h: usize) -> } 2 | 3 => { let rest = r.bytes(r.remaining())?; - let mut v = miniz_oxide::inflate::decompress_to_vec_zlib_with_limit(rest, rowb * h).map_err(|e| Error::Invalid(format!("zip: {e:?}")))?; + let mut v = miniz_oxide::inflate::decompress_to_vec_zlib_with_limit(rest, rowb * h).map_err(|e| Error::Invalid(format!("zip: {:?}", e.status)))?; v.resize(rowb * h, 0); if comp == 3 { ref_unpredict(p.depth, &mut v, w, h); @@ -112,7 +112,8 @@ fn ref_composite(p: &Psd) -> Result { } 2 | 3 => { let rest = r.bytes(r.remaining())?; - let v = miniz_oxide::inflate::decompress_to_vec_zlib_with_limit(rest, want * rowb * h).map_err(|e| Error::Invalid(format!("zip: {e:?}")))?; + let v = + miniz_oxide::inflate::decompress_to_vec_zlib_with_limit(rest, want * rowb * h).map_err(|e| Error::Invalid(format!("zip: {:?}", e.status)))?; for (c, pl) in planes.iter_mut().enumerate().take(want) { let mut plane = v.get(c * rowb * h..(c + 1) * rowb * h).ok_or(Error::Truncated)?.to_vec(); if comp == 3 { @@ -448,7 +449,7 @@ fn unpredict_matches_the_reference_and_allocates_one_scratch_row() { let mut a: Vec = (0..w * h * bps).map(|i| (i * 31 % 251) as u8).collect(); let mut b = a.clone(); counters::take(&SCRATCH); - p.unpredict(&mut a, w, h); + p.unpredict(&mut a, w, h, &mut Budget::new(u64::MAX)).unwrap(); let scratch = counters::take(&SCRATCH); let old_allocs = ref_unpredict(depth, &mut b, w, h); assert_eq!(a, b, "{depth} bit"); @@ -473,11 +474,11 @@ fn raw_planes_are_borrowed_and_zip_planes_are_unpredicted_in_place() { let mut v: Vec = (0..3 * 40 * 30).map(|i| (i % 255) as u8).collect(); let range = v.as_ptr_range(); for comp in [2, 3] { - let got = q.zip_plane(comp, &mut v, 1, 40, 30).unwrap(); + let got = q.zip_plane(comp, &mut v, 1, 40, 30, &mut Budget::new(u64::MAX)).unwrap(); assert_eq!(got.len(), 1200); assert!(range.contains(&got.as_ptr()), "compression {comp}"); } - assert!(q.zip_plane(2, &mut v, 3, 40, 30).is_err()); + assert!(q.zip_plane(2, &mut v, 3, 40, 30, &mut Budget::new(u64::MAX)).is_err()); } #[test]