From 6aa599f869d9724ec9f6c63298a675a9f067e659 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 8 Oct 2026 22:19:42 +0000 Subject: [PATCH] Fix: PSD RLE decode bounds its allocations by the input A tiny hostile PSD could abort the process on out-of-memory in `psd::decode` (allocation failure doesn't unwind, so catch_unwind can't save it): 1. The RLE row byte-count table (channels x rows entries, from the header) was reserved with Vec::with_capacity before checking that the file holds it: 56 channels x 2^30 rows asked for ~480 GB from a ~40-byte file. The table is now taken as a slice of the input with checked size math and read lazily, so nothing is allocated from it. 2. RLE rows were zero-padded / run-expanded to the declared width with no tie to the input size, so all-zero row counts on a 1 Gpx (the DecodeOptions::max_pixels limit, already enforced by check_size) 32-bit image expanded to 4 GiB per channel. A PackBits row expands at most 64x (2-byte repeat run -> 128 bytes), so each decoded row must now declare at least ceil(row_bytes / 128) * 2 encoded bytes, and the decoded channels' data must be present in the file, before any plane is allocated. Decoded planes are thus bounded by 64x the file size as well as by max_pixels. PackBits now decodes into a pre-sized, zeroed row and a run that would write past the declared width is an error instead of being silently clipped. Also: size math (width x height x bpc, samples) is checked; image buffers are allocated with try_reserve_exact (TooLarge on failure); only the colour (+ alpha) planes are decoded, extra channels are skipped; the PSB layer-section length no longer truncates on 32-bit targets; interleaving uses iterators instead of indexing. Valid PSDs (raw and RLE, 8/16/32-bit) decode identically; regression tests cover a huge RLE header with a truncated table (aborted with a 481 GB allocation before), all-empty RLE rows, counts larger than the file, a repeat run past the row width (was Ok before) and a 16-bit gray RLE file with an extra channel. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01912WaK6Rmew3wxZuPyX7E8 --- crates/codecs/src/psd.rs | 262 ++++++++++++++++++++++++++++----------- 1 file changed, 191 insertions(+), 71 deletions(-) diff --git a/crates/codecs/src/psd.rs b/crates/codecs/src/psd.rs index de668915..d49a3885 100644 --- a/crates/codecs/src/psd.rs +++ b/crates/codecs/src/psd.rs @@ -69,8 +69,8 @@ pub(crate) fn decode(bytes: &[u8], opts: &DecodeOptions) -> Result { let (icc, exif, xmp) = resources(res); // Layer and mask information: only the layer count sign matters (merged transparency). - let lm_len = if psb { r.u64()? as usize } else { r.u32()? as usize }; - let lm = r.take(lm_len)?; + let lm_len = if psb { r.u64()? } else { u64::from(r.u32()?) }; + let lm = r.take(usize::try_from(lm_len).map_err(|_| err("truncated"))?)?; let merged_transparency = if lm.len() >= if psb { 10 } else { 6 } { let off = if psb { 8 } else { 4 }; let count = i16::from_be_bytes([lm[off], lm[off + 1]]); @@ -88,33 +88,6 @@ pub(crate) fn decode(bytes: &[u8], opts: &DecodeOptions) -> Result { 1 => return Err(Error::Unsupported(F, "1-bit bitmap PSD")), _ => return Err(err("bad depth")), }; - let plane = width * height * bpc; - let row_bytes = width * bpc; - let mut planes: Vec> = Vec::with_capacity(channels); - match compression { - 0 => { - for _ in 0..channels { - planes.push(r.take(plane)?.to_vec()); - } - } - 1 => { - let counts_len = channels * height; - let mut counts = Vec::with_capacity(counts_len); - for _ in 0..counts_len { - counts.push(if psb { r.u32()? as usize } else { r.u16()? as usize }); - } - for c in 0..channels { - let mut out = Vec::with_capacity(plane); - for y in 0..height { - let src = r.take(counts[c * height + y])?; - unpackbits(src, row_bytes, &mut out); - } - planes.push(out); - } - } - _ => return Err(Error::Unsupported(F, "ZIP-compressed merged image data")), - } - let (model, n_color) = match mode { 1 | 8 => (Model::Gray, 1), 2 => (Model::Rgb, 1), @@ -128,50 +101,106 @@ pub(crate) fn decode(bytes: &[u8], opts: &DecodeOptions) -> Result { return Err(err("too few channels")); } let alpha = merged_transparency && channels > n_color && mode != 2; + // Only the colour (+ alpha) planes are decoded; extra channels (spot colours, masks) are + // stored after them and never touched. let used = n_color + alpha as usize; - let n = width * height; + let too_large = || Error::TooLarge(width as u64, height as u64); + let n = width.checked_mul(height).ok_or_else(too_large)?; + let row_bytes = width.checked_mul(bpc).ok_or_else(too_large)?; + let plane = n.checked_mul(bpc).ok_or_else(too_large)?; + let mut planes: Vec> = Vec::with_capacity(used); + match compression { + 0 => { + // Raw planes are borrowed from the input, so their size is bounded by the file. + for _ in 0..used { + planes.push(r.take(plane)?.to_vec()); + } + } + 1 => { + // Row byte-count table: one entry per row per channel (all channels, since the data + // follows it). Taken from the input before anything is allocated from it. + let entry = if psb { 4 } else { 2 }; + let table_len = channels.checked_mul(height).and_then(|c| c.checked_mul(entry)).ok_or_else(|| err("truncated"))?; + let table = r.take(table_len)?; + let count = |i: usize| -> usize { + let at = i.saturating_mul(entry); + match table.get(at..at.saturating_add(entry)) { + Some(&[a, b]) => usize::from(u16::from_be_bytes([a, b])), + Some(&[a, b, c, d]) => u32::from_be_bytes([a, b, c, d]) as usize, + _ => 0, + } + }; + // PackBits expands at most 64× (a 2-byte repeat run yields 128 bytes), so a row needs + // at least this many encoded bytes. Requiring it (and the data being present) before + // allocating bounds the decoded planes by the file size: a tiny file can no longer + // declare gigabytes of zero-padded or repeat-run rows. + let min_row = row_bytes.div_ceil(128).saturating_mul(2); + let mut data_len = 0usize; + for i in 0..used.saturating_mul(height) { + let c = count(i); + if c < min_row { + return Err(err("RLE row too short for its width")); + } + data_len = data_len.saturating_add(c); + } + if data_len > bytes.len().saturating_sub(r.p) { + return Err(err("truncated")); + } + for c in 0..used { + let mut out = zeroed::(plane, too_large)?; + for (y, row) in out.chunks_exact_mut(row_bytes).enumerate() { + let src = r.take(count(c * height + y))?; + unpackbits(src, row)?; + } + planes.push(out); + } + } + _ => return Err(Error::Unsupported(F, "ZIP-compressed merged image data")), + } + let samples = n.checked_mul(used).ok_or_else(too_large)?; // Interleave to device samples. let buf = if mode == 2 { // Indexed → 8-bit RGB via the 768-byte palette. - if color_mode_data.len() < 768 || depth != 8 { + let (Some(pal), true) = (color_mode_data.get(..768), depth == 8) else { return Err(err("bad indexed palette")); - } - let mut v = Vec::with_capacity(n * 3); - for &i in &planes[0] { + }; + let idx = planes.first().ok_or_else(|| err("missing plane"))?; + let mut v = zeroed::(n.checked_mul(3).ok_or_else(too_large)?, too_large)?; + for (o, &i) in v.chunks_exact_mut(3).zip(idx) { let i = i as usize; - v.extend_from_slice(&[color_mode_data[i], color_mode_data[256 + i], color_mode_data[512 + i]]); + o.copy_from_slice(&[pal[i], pal[256 + i], pal[512 + i]]); } Buf::U8(v) } else { + let cmyk = model == Model::Cmyk; match bpc { 1 => { - let mut v = vec![0u8; n * used]; - for (c, p) in planes.iter().take(used).enumerate() { - for i in 0..n { - let x = p[i]; + let mut v = zeroed::(samples, too_large)?; + for (c, p) in planes.iter().enumerate() { + for (o, &x) in v.iter_mut().skip(c).step_by(used).zip(p) { // PSD CMYK stores 255 = no ink; our CMYK convention is 0 = no ink. - v[i * used + c] = if model == Model::Cmyk && c < 4 { 255 - x } else { x }; + *o = if cmyk && c < 4 { 255 - x } else { x }; } } Buf::U8(v) } 2 => { - let mut v = vec![0u16; n * used]; - for (c, p) in planes.iter().take(used).enumerate() { - for i in 0..n { - let x = u16::from_be_bytes([p[2 * i], p[2 * i + 1]]); - v[i * used + c] = if model == Model::Cmyk && c < 4 { 65535 - x } else { x }; + let mut v = zeroed::(samples, too_large)?; + for (c, p) in planes.iter().enumerate() { + for (o, s) in v.iter_mut().skip(c).step_by(used).zip(p.chunks_exact(2)) { + let x = u16::from_be_bytes([s[0], s[1]]); + *o = if cmyk && c < 4 { 65535 - x } else { x }; } } Buf::U16(v) } _ => { - let mut v = vec![0f32; n * used]; - for (c, p) in planes.iter().take(used).enumerate() { - for i in 0..n { - let x = f32::from_be_bytes([p[4 * i], p[4 * i + 1], p[4 * i + 2], p[4 * i + 3]]); - v[i * used + c] = if model == Model::Cmyk && c < 4 { 1.0 - x } else { x }; + let mut v = zeroed::(samples, too_large)?; + for (c, p) in planes.iter().enumerate() { + for (o, s) in v.iter_mut().skip(c).step_by(used).zip(p.chunks_exact(4)) { + let x = f32::from_be_bytes([s[0], s[1], s[2], s[3]]); + *o = if cmyk && c < 4 { 1.0 - x } else { x }; } } Buf::F32(v) @@ -182,25 +211,38 @@ pub(crate) fn decode(bytes: &[u8], opts: &DecodeOptions) -> Result { finish(F, raw, Meta { icc, exif, xmp, ..Default::default() }, (width as u32, height as u32), opts) } -/// PackBits: append exactly `row_bytes` bytes (zero-padded if the run data is short). -fn unpackbits(mut src: &[u8], row_bytes: usize, out: &mut Vec) { - let target = out.len() + row_bytes; - while out.len() < target && !src.is_empty() { - let h = src[0] as i8; - src = &src[1..]; +/// A zero-filled buffer whose allocation failure is an error, not an abort. +fn zeroed(len: usize, too_large: impl Fn() -> Error) -> Result> { + let mut v = Vec::new(); + v.try_reserve_exact(len).map_err(|_| too_large())?; + v.resize(len, T::default()); + Ok(v) +} + +/// PackBits: decode one row into `row` (pre-zeroed, so short run data leaves zero padding). +/// A run that would write past the end of the row is an error; trailing bytes are ignored. +fn unpackbits(mut src: &[u8], row: &mut [u8]) -> Result<()> { + let mut at = 0; + while at < row.len() { + let Some((&h, rest)) = src.split_first() else { break }; + src = rest; + let h = h as i8; if h >= 0 { - let n = (h as usize + 1).min(src.len()).min(target - out.len()); - out.extend_from_slice(&src[..n]); - src = &src[(h as usize + 1).min(src.len())..]; + let n = h as usize + 1; + let lit = src.get(..n).unwrap_or(src); + let dst = row.get_mut(at..at + lit.len()).ok_or_else(|| err("RLE row overruns its width"))?; + dst.copy_from_slice(lit); + at += lit.len(); + src = src.get(n..).unwrap_or(&[]); } else if h != -128 { - let n = ((1 - h as isize) as usize).min(target - out.len()); - if let Some(&b) = src.first() { - out.extend(std::iter::repeat_n(b, n)); - src = &src[1..]; - } + let n = (1 - h as isize) as usize; + let Some((&b, rest)) = src.split_first() else { break }; + src = rest; + row.get_mut(at..at + n).ok_or_else(|| err("RLE row overruns its width"))?.fill(b); + at += n; } } - out.resize(target, 0); + Ok(()) } fn resources(mut b: &[u8]) -> (Option>, Option>, Option) { @@ -234,11 +276,89 @@ mod tests { #[test] fn packbits() { - let mut out = Vec::new(); - unpackbits(&[0xFE, 0xAA, 0x02, 0x80, 0x00, 0x2A], 6, &mut out); - assert_eq!(out, vec![0xAA, 0xAA, 0xAA, 0x80, 0x00, 0x2A]); - let mut out = Vec::new(); - unpackbits(&[0x05, 1], 4, &mut out); - assert_eq!(out, vec![1, 0, 0, 0]); + let mut out = [0u8; 6]; + unpackbits(&[0xFE, 0xAA, 0x02, 0x80, 0x00, 0x2A], &mut out).unwrap(); + assert_eq!(out, [0xAA, 0xAA, 0xAA, 0x80, 0x00, 0x2A]); + // Short run data is zero-padded. + let mut out = [0u8; 4]; + unpackbits(&[0x01, 1], &mut out).unwrap(); + assert_eq!(out, [1, 0, 0, 0]); + // Truncated literal run: copies what is there. + let mut out = [0u8; 4]; + unpackbits(&[0x03, 7, 8], &mut out).unwrap(); + assert_eq!(out, [7, 8, 0, 0]); + } + + #[test] + fn packbits_overrun_is_an_error() { + // Repeat run of 128 into a 4-byte row. + let mut out = [0u8; 4]; + assert!(unpackbits(&[0x81, 0xAA], &mut out).is_err()); + // Literal run of 6 into a 4-byte row. + assert!(unpackbits(&[0x05, 1, 2, 3, 4, 5, 6], &mut out).is_err()); + } + + /// PSD header + empty colour-mode/resources/layer sections + compression id. + fn header(channels: u16, w: u32, h: u32, depth: u16, mode: u16, compression: u16) -> Vec { + let mut b = b"8BPS".to_vec(); + b.extend_from_slice(&1u16.to_be_bytes()); + b.extend_from_slice(&[0; 6]); + b.extend_from_slice(&channels.to_be_bytes()); + b.extend_from_slice(&h.to_be_bytes()); + b.extend_from_slice(&w.to_be_bytes()); + b.extend_from_slice(&depth.to_be_bytes()); + b.extend_from_slice(&mode.to_be_bytes()); + b.extend_from_slice(&[0; 12]); // colour mode data, resources, layer & mask: all empty + b.extend_from_slice(&compression.to_be_bytes()); + b + } + + #[test] + fn huge_rle_with_truncated_count_table_is_rejected_before_allocating() { + // 1 × 2^30 px, 56 channels: the row table alone would be 56 Gi entries; the file has 4 bytes. + let mut b = header(56, 1, 1 << 30, 32, 3, 1); + b.extend_from_slice(&[0, 0, 0, 0]); + let e = decode(&b, &DecodeOptions::default()).unwrap_err(); + assert!(matches!(e, Error::Malformed(_, ref m) if m == "truncated"), "{e:?}"); + } + + #[test] + fn huge_rle_of_empty_rows_is_rejected_before_allocating() { + // 2^30 × 1 px RGB at 32 bits: every row count is 0, which used to zero-pad 3 × 4 GiB. + let mut b = header(3, 1 << 30, 1, 32, 3, 1); + b.extend_from_slice(&[0; 6]); + let e = decode(&b, &DecodeOptions::default()).unwrap_err(); + assert!(matches!(e, Error::Malformed(_, ref m) if m.contains("too short")), "{e:?}"); + // 1 Gpx RGB 8-bit whose (plausible) row counts claim more data than the file holds. + let mut b = header(3, 1 << 20, 1 << 10, 8, 3, 1); + b.extend_from_slice(&[0xFF; 3 * 1024 * 2]); + let e = decode(&b, &DecodeOptions::default()).unwrap_err(); + assert!(matches!(e, Error::Malformed(_, ref m) if m == "truncated"), "{e:?}"); + } + + #[test] + fn rle_repeat_run_past_row_width_is_an_error() { + // 4 × 1 gray 8-bit; the single row is a 128-byte repeat run. + let mut b = header(1, 4, 1, 8, 1, 1); + b.extend_from_slice(&2u16.to_be_bytes()); + b.extend_from_slice(&[0x81, 0x55]); + let e = decode(&b, &DecodeOptions::default()).unwrap_err(); + assert!(matches!(e, Error::Malformed(_, ref m) if m.contains("overruns")), "{e:?}"); + } + + #[test] + fn rle_gray16_and_extra_channels_decode() { + // 2 × 2 gray 16-bit with an extra (unused) channel whose data is absent: only the used + // plane is decoded. Row: repeat run of 4 bytes 0x80 -> samples 0x8080. + let mut b = header(2, 2, 2, 16, 1, 1); + for _ in 0..4 { + b.extend_from_slice(&2u16.to_be_bytes()); + } + b.extend_from_slice(&[0xFD, 0x80, 0xFD, 0x80]); + let d = decode(&b, &DecodeOptions::default()).unwrap(); + assert_eq!((d.width, d.height), (2, 2)); + let p = d.image.get(1, 1); + let q = d.image.get(0, 0); + assert!(p[0] > 0.0 && (p[0] - q[0]).abs() < 1e-6, "{p:?}"); } }