diff --git a/crates/codecs/src/psd.rs b/crates/codecs/src/psd.rs index 86cb2008..66dec826 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) @@ -255,25 +284,38 @@ pub(crate) fn header(bytes: &[u8]) -> Result<(u32, u32, u16)> { Ok((width as u32, height as u32, orientation)) } -/// 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) { @@ -307,11 +349,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:?}"); } }