Merge pull request #1153 from storytold/fix/psd-head-view

psd: expose bounded borrowed header resources

crates/psd/src/head.rs

@@ -2,19 +2,67 @@
 //! without the layer and mask section or any pixels.
 //!
 //! This is what a file browser or a metadata reader needs (size, depth, colour mode, ICC profile,
-//! EXIF, XMP, resolution, thumbnail): a few kilobytes at the start of a file that may be
-//! gigabytes long. [`Head::from_source`] reads exactly those bytes; [`Head::from_bytes`] accepts a
-//! prefix of the file and reports how many more bytes it needs when the prefix is too short.
+//! EXIF, XMP, resolution, thumbnail): a few bytes at the start of a file that may be gigabytes
+//! long. [`Head::from_source`] reads exactly those bytes; [`Head::from_bytes`] accepts a prefix of
+//! the file and reports how many more bytes it needs when the prefix is too short. Both return
+//! owned data. For untrusted files with large resource sections, [`HeadView`] borrows the input
+//! and parses resource records lazily.
 
 use crate::error::Result;
 use crate::header::Header;
 use crate::io::Reader;
-use crate::resources::{self, ImageResource};
+use crate::resources::{self, ImageResource, ResourceIter};
 use crate::source::Source;
 
 /// Length of the file header.
 const HEADER_LEN: u64 = 26;
 
+/// A zero-copy view of a PSD/PSB head.
+///
+/// The header and section lengths are validated eagerly. Colour-mode data and resource records
+/// remain borrowed from the input; [`HeadView::resources`] validates one resource at a time and
+/// stops at the first malformed record. This avoids cloning unknown or attacker-sized resource
+/// payloads and lets callers keep valid records that precede a damaged one.
+#[derive(Debug, Clone)]
+pub struct HeadView<'a> {
+    /// The file header.
+    pub header: Header,
+    /// The colour mode data section (the palette of an indexed file, duotone or HDR toning data).
+    pub color_mode_data: &'a [u8],
+    resources: &'a [u8],
+    /// Offset in the file of the layer and mask information section (of its length field).
+    pub layer_section_offset: u64,
+}
+
+impl<'a> HeadView<'a> {
+    /// Parses the fixed header and section bounds from `data` without copying colour data or
+    /// resources. Individual resource records are validated by the iterator returned from
+    /// [`HeadView::resources`].
+    pub fn from_bytes(data: &'a [u8]) -> Result<Self> {
+        let header = Header::from_bytes(data)?;
+        let mut r = Reader::new(data);
+        r.skip(HEADER_LEN as usize)?;
+        let color_mode_len = r.u32()?;
+        let color_mode_data = r.bytes_u64(u64::from(color_mode_len))?;
+        let resource_len = r.u32()?;
+        let resources = r.bytes_u64(u64::from(resource_len))?;
+        Ok(HeadView { header, color_mode_data, resources, layer_section_offset: r.pos() as u64 })
+    }
+
+    /// Iterates over resource records without allocating. If a record is malformed, its error is
+    /// returned once and iteration stops; earlier records remain available to the caller.
+    pub fn resources(&self) -> ResourceIter<'a> {
+        ResourceIter::new(self.resources)
+    }
+
+    /// Iterates over at most `max_records` resources. Returns one limit error if additional bytes
+    /// remain after that many records, then stops. Use this for untrusted files where record
+    /// processing itself needs a caller-defined work bound.
+    pub fn resources_limited(&self, max_records: usize) -> ResourceIter<'a> {
+        ResourceIter::with_limit(self.resources, max_records)
+    }
+}
+
 /// The first three sections of a file. The same values [`crate::PsdFile`] holds in its fields of
 /// the same names.
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -132,6 +180,72 @@ mod tests {
         }
     }
 
+    #[test]
+    fn head_view_borrows_sections_and_records() {
+        for (file, bytes) in cases() {
+            let view = HeadView::from_bytes(&bytes).unwrap();
+            let owned = Head::from_bytes(&bytes).unwrap();
+            assert_eq!(view.header, owned.header);
+            assert_eq!(view.color_mode_data, owned.color_mode_data);
+            assert_eq!(view.layer_section_offset, owned.layer_section_offset);
+            let bytes_start = bytes.as_ptr() as usize;
+            let bytes_end = bytes_start + bytes.len();
+            let color_start = view.color_mode_data.as_ptr() as usize;
+            assert!(bytes_start <= color_start && color_start + view.color_mode_data.len() <= bytes_end);
+
+            let borrowed = view.resources().collect::<Result<Vec<_>>>().unwrap();
+            assert_eq!(borrowed.iter().map(|r| (*r).into_owned()).collect::<Vec<_>>(), file.resources);
+            let nonempty = borrowed.iter().find(|r| !r.data.is_empty()).expect("test resource has bytes");
+            let data_start = nonempty.data.as_ptr() as usize;
+            assert!(bytes_start <= data_start && data_start + nonempty.data.len() <= bytes_end);
+        }
+    }
+
+    #[test]
+    fn head_view_keeps_valid_resources_before_a_malformed_record() {
+        let (_, mut bytes) = cases().remove(0);
+        let color_len = u32::from_be_bytes(bytes[26..30].try_into().unwrap()) as usize;
+        let resource_len_at = 30 + color_len;
+        let resource_start = resource_len_at + 4;
+        let name_len = usize::from(bytes[resource_start + 6]);
+        let name_total = (1 + name_len + 1) & !1;
+        let data_len_at = resource_start + 6 + name_total;
+        let data_len = u32::from_be_bytes(bytes[data_len_at..data_len_at + 4].try_into().unwrap()) as usize;
+        let first_end = data_len_at + 4 + data_len + data_len % 2;
+        let resource_len = u32::from_be_bytes(bytes[resource_len_at..resource_len_at + 4].try_into().unwrap());
+        let damaged_record = b"JUNK";
+        bytes[resource_len_at..resource_len_at + 4].copy_from_slice(&(resource_len + damaged_record.len() as u32).to_be_bytes());
+        bytes.splice(first_end..first_end, damaged_record.iter().copied());
+
+        let view = HeadView::from_bytes(&bytes).unwrap();
+        let mut resources = view.resources();
+        let first = resources.next().expect("first valid resource").unwrap();
+        assert_eq!(first.id, 1005);
+        assert!(resources.next().expect("malformed second resource").is_err());
+        assert!(resources.next().is_none(), "iteration stops after the first malformed record");
+        assert!(Head::from_bytes(&bytes).is_err(), "the strict owned API still rejects malformed resources");
+    }
+
+    #[test]
+    fn head_view_resource_limit_is_enforced_after_valid_records() {
+        let (_, mut bytes) = cases().remove(0);
+        let color_len = u32::from_be_bytes(bytes[26..30].try_into().unwrap()) as usize;
+        let resource_len_at = 30 + color_len;
+        let resource_start = resource_len_at + 4;
+        let resource_len = u32::from_be_bytes(bytes[resource_len_at..resource_start].try_into().unwrap()) as usize;
+        let duplicate = bytes[resource_start..resource_start + resource_len].to_vec();
+        bytes[resource_len_at..resource_start].copy_from_slice(&((resource_len * 2) as u32).to_be_bytes());
+        bytes.splice(resource_start + resource_len..resource_start + resource_len, duplicate);
+        let view = HeadView::from_bytes(&bytes).unwrap();
+        let all = view.resources().collect::<Result<Vec<_>>>().unwrap();
+        assert!(all.len() > 1);
+        let mut limited = view.resources_limited(1);
+        assert!(limited.next().unwrap().is_ok());
+        assert!(limited.size_hint().1.unwrap() >= 1, "limit error is included in the upper bound");
+        assert!(limited.next().unwrap().is_err());
+        assert!(limited.next().is_none());
+    }
+
     #[test]
     fn a_prefix_ending_with_the_resources_is_enough_and_one_byte_less_is_not() {
         for (_, bytes) in cases() {

crates/psd/src/lib.rs

@@ -105,13 +105,13 @@ pub use descriptor::{Descriptor, VersionedDescriptor};
 pub use error::{PsdError, Result};
 pub use file::{GlobalLayerMask, LayerInfoPlacement, PsdFile};
 pub use float::FloatImage;
-pub use head::Head;
+pub use head::{Head, HeadView};
 pub use header::{ColorMode, Header, Version};
 pub use image_data::ImageData;
 pub use layer::{BlendingRanges, ChannelData, LayerFlags, LayerInfo, LayerMask, LayerRecord, MaskData, MaskParameters, RealMask, Rect};
 pub use pixels::{GrayImage, Layer, RgbaImage};
 pub use preview::Preview;
-pub use resources::{ImageResource, ResolutionInfo, ResourceData};
+pub use resources::{ImageResource, ImageResourceView, ResolutionInfo, ResourceData, ResourceIter};
 pub use rows::RowDecoder;
 #[cfg(any(unix, windows))]
 pub use source::FileSource;

crates/psd/src/resources.rs

@@ -2,7 +2,7 @@
 
 use crate::descriptor::UnicodeString;
 use crate::error::{PsdError, Result};
-use crate::io::{Reader, WriteExt, read_pascal, write_pascal};
+use crate::io::{Reader, WriteExt, write_pascal};
 
 /// Well-known resource ids.
 pub mod ids {
@@ -48,6 +48,124 @@ pub struct ImageResource {
     pub data: Vec<u8>,
 }
 
+/// An image-resource record that borrows its name and data from the PSD/PSB input.
+///
+/// Use [`ImageResourceView::into_owned`] only when the data must outlive the input bytes. Parsing a
+/// view does not allocate proportional to the resource's declared size.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct ImageResourceView<'a> {
+    /// Signature, normally `8BIM` (also `MeSa`, `AgHg`, `PHUT`, `DCSR` seen in the wild).
+    pub signature: [u8; 4],
+    /// Resource id.
+    pub id: u16,
+    /// Pascal name bytes (usually empty).
+    pub name: &'a [u8],
+    /// Resource data (without the even padding).
+    pub data: &'a [u8],
+}
+
+impl ImageResourceView<'_> {
+    /// Copies this view into an owned resource.
+    pub fn into_owned(self) -> ImageResource {
+        ImageResource { signature: self.signature, id: self.id, name: self.name.to_vec(), data: self.data.to_vec() }
+    }
+}
+
+impl<'a> ImageResourceView<'a> {
+    pub(crate) fn read(r: &mut Reader<'a>) -> Result<Self> {
+        let signature = r.array::<4>()?;
+        if !matches!(&signature, b"8BIM" | b"MeSa" | b"AgHg" | b"PHUT" | b"DCSR") {
+            return Err(PsdError::InvalidSignature { expected: "8BIM resource", found: signature });
+        }
+        let id = r.u16()?;
+        let name_len = usize::from(r.u8()?);
+        let name = r.bytes(name_len)?;
+        let name_total = 1 + name_len;
+        let name_pad = (2 - name_total % 2) % 2;
+        r.skip(name_pad)?;
+        let len = r.u32()?;
+        let data = r.bytes_u64(u64::from(len))?;
+        if len % 2 == 1 {
+            r.skip(1)?;
+        }
+        Ok(ImageResourceView { signature, id, name, data })
+    }
+}
+
+/// Lazily parses image-resource records from a borrowed resource section.
+///
+/// Records are yielded one at a time. If a record is malformed, the iterator yields its error
+/// once and then stops; records yielded before it remain usable.
+pub struct ResourceIter<'a> {
+    reader: Reader<'a>,
+    remaining_records: usize,
+    failed: bool,
+}
+
+impl<'a> ResourceIter<'a> {
+    pub(crate) fn new(bytes: &'a [u8]) -> Self {
+        ResourceIter::with_limit(bytes, usize::MAX)
+    }
+
+    pub(crate) fn with_limit(bytes: &'a [u8], max_records: usize) -> Self {
+        ResourceIter { reader: Reader::new(bytes), remaining_records: max_records, failed: false }
+    }
+}
+
+impl<'a> Iterator for ResourceIter<'a> {
+    type Item = Result<ImageResourceView<'a>>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.failed || self.reader.is_empty() {
+            return None;
+        }
+        if self.remaining_records == 0 {
+            self.failed = true;
+            return Some(Err(PsdError::LimitExceeded("image resource record count")));
+        }
+        match ImageResourceView::read(&mut self.reader) {
+            Ok(resource) => {
+                self.remaining_records -= 1;
+                Some(Ok(resource))
+            }
+            Err(error) => {
+                self.failed = true;
+                Some(Err(error))
+            }
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        // A valid record consumes at least 12 bytes. A truncated record or exhausted limit can
+        // produce one final error without consuming a full record.
+        let upper =
+            if self.failed || self.reader.is_empty() { 0 } else { (self.reader.remaining() / 12).min(self.remaining_records).saturating_add(1) };
+        (0, Some(upper))
+    }
+}
+
+impl std::iter::FusedIterator for ResourceIter<'_> {}
+
+#[cfg(test)]
+mod view_tests {
+    use super::ResourceIter;
+
+    #[test]
+    fn resource_size_hint_respects_the_record_limit() {
+        let data = vec![0; 12_000];
+        let iter = ResourceIter::with_limit(&data, 3);
+        assert_eq!(iter.size_hint(), (0, Some(4)));
+    }
+
+    #[test]
+    fn truncated_resource_size_hint_includes_error_item() {
+        let mut iter = ResourceIter::new(b"JUNK");
+        assert!(iter.size_hint().1.unwrap() >= 1);
+        assert!(iter.next().unwrap().is_err());
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+    }
+}
+
 /// Fixed-point resolution info (resource 1005).
 ///
 /// `h_res_fixed` and `v_res_fixed` are always pixels per inch (16.16 fixed).
@@ -292,18 +410,7 @@ impl ImageResource {
     }
 
     pub(crate) fn read(r: &mut Reader<'_>) -> Result<Self> {
-        let signature = r.array::<4>()?;
-        if !matches!(&signature, b"8BIM" | b"MeSa" | b"AgHg" | b"PHUT" | b"DCSR") {
-            return Err(PsdError::InvalidSignature { expected: "8BIM resource", found: signature });
-        }
-        let id = r.u16()?;
-        let name = read_pascal(r, 2)?;
-        let len = r.u32()?;
-        let data = r.bytes_u64(u64::from(len))?.to_vec();
-        if len % 2 == 1 {
-            r.skip(1)?;
-        }
-        Ok(ImageResource { signature, id, name, data })
+        Ok(ImageResourceView::read(r)?.into_owned())
     }
 
     pub(crate) fn write(&self, out: &mut Vec<u8>) -> Result<()> {