Merge pull request #218 from image-rs/ext

StreamingDecoder streaming ext data

Author: 197g

src/reader/converter.rs

@@ -28,25 +28,26 @@ pub(crate) type FillBufferCallback<'a> =
 
 /// Deinterlaces and expands to RGBA if needed
 pub(crate) struct PixelConverter {
-    memory_limit: MemoryLimit,
     color_output: ColorOutput,
     buffer: Vec<u8>,
     global_palette: Option<Vec<u8>>,
 }
 
 impl PixelConverter {
-    pub(crate) const fn new(color_output: ColorOutput, memory_limit: MemoryLimit) -> Self {
+    pub(crate) const fn new(color_output: ColorOutput) -> Self {
         Self {
-            memory_limit,
             color_output,
             buffer: Vec::new(),
             global_palette: None,
         }
     }
 
-    pub(crate) fn check_buffer_size(&self, frame: &Frame<'_>) -> Result<usize, DecodingError> {
-        let pixel_bytes = self
-            .memory_limit
+    pub(crate) fn check_buffer_size(
+        &self,
+        frame: &Frame<'_>,
+        memory_limit: &MemoryLimit,
+    ) -> Result<usize, DecodingError> {
+        let pixel_bytes = memory_limit
             .buffer_size(self.color_output, frame.width, frame.height)
             .ok_or_else(|| DecodingError::OutOfMemory)?;
 
@@ -63,8 +64,9 @@ impl PixelConverter {
         &mut self,
         frame: &mut Frame<'_>,
         data_callback: FillBufferCallback<'_>,
+        memory_limit: &MemoryLimit,
     ) -> Result<(), DecodingError> {
-        let pixel_bytes = self.check_buffer_size(frame)?;
+        let pixel_bytes = self.check_buffer_size(frame, memory_limit)?;
         let mut vec = match mem::replace(&mut frame.buffer, Cow::Borrowed(&[])) {
             // reuse buffer if possible without reallocating
             Cow::Owned(mut vec) if vec.capacity() >= pixel_bytes => {

src/reader/decoder.rs

@@ -13,16 +13,11 @@ use std::io;
 use crate::common::{AnyExtension, Block, DisposalMethod, Extension, Frame};
 use crate::reader::DecodeOptions;
 use crate::MemoryLimit;
-use crate::Repeat;
 
 use weezl::{decode::Decoder as LzwDecoder, BitOrder, LzwError, LzwStatus};
 
 /// GIF palettes are RGB
 pub const PLTE_CHANNELS: usize = 3;
-/// Headers for supported extensions.
-const EXT_NAME_NETSCAPE: &[u8] = b"NETSCAPE2.0\x01";
-const EXT_NAME_XMP: &[u8] = b"XMP DataXMP";
-const EXT_NAME_ICC: &[u8] = b"ICCRGBG1012";
 
 /// An error returned in the case of the image not being formatted properly.
 #[derive(Debug)]
@@ -50,6 +45,8 @@ impl error::Error for DecodingFormatError {
 pub enum DecodingError {
     /// Failed to internally allocate a buffer of sufficient size.
     OutOfMemory,
+    /// Allocation exceeded set memory limit
+    MemoryLimit,
     /// Expected a decoder but none found.
     DecoderNotFound,
     /// Expected an end-code, but none found.
@@ -78,6 +75,7 @@ impl fmt::Display for DecodingError {
     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
         match *self {
             Self::OutOfMemory => fmt.write_str("Out of Memory"),
+            Self::MemoryLimit => fmt.write_str("Memory limit reached"),
             Self::DecoderNotFound => fmt.write_str("Decoder Not Found"),
             Self::EndCodeNotFound => fmt.write_str("End-Code Not Found"),
             Self::UnexpectedEof => fmt.write_str("Unexpected End of File"),
@@ -93,6 +91,7 @@ impl error::Error for DecodingError {
     fn source(&self) -> Option<&(dyn error::Error + 'static)> {
         match *self {
             Self::OutOfMemory => None,
+            Self::MemoryLimit => None,
             Self::DecoderNotFound => None,
             Self::EndCodeNotFound => None,
             Self::UnexpectedEof => None,
@@ -146,28 +145,21 @@ pub enum Decoded {
     GlobalPalette(Box<[u8]>),
     /// Index of the background color in the global palette.
     BackgroundColor(u8),
-    /// Loop count is known
-    Repetitions(Repeat),
     /// Palette and optional `Application` extension have been parsed,
     /// reached frame data.
     HeaderEnd,
     /// The start of a block.
     /// `BlockStart(Block::Trailer)` is the very last decode event
     BlockStart(Block),
-    /// Decoded a sub-block. More sub-block are available.
+    /// Decoded a sub-block.
     ///
-    /// Indicates the label of the extension which might be unknown. A label of `0` is used when
-    /// the sub block does not belong to an extension.
-    ///
-    /// Call `last_ext()` to get the data
-    SubBlockFinished(AnyExtension),
-    /// Decoded the last (or only) sub-block of a block.
-    ///
-    /// Indicates the label of the extension which might be unknown. A label of `0` is used when
-    /// the sub block does not belong to an extension.
-    ///
-    /// Call `last_ext()` to get the data
-    BlockFinished(AnyExtension),
+    /// Call `last_ext_sub_block()` to get the sub-block data. It won't be available after this event.
+    SubBlock {
+        /// An ext label of `0` is used when the sub block does not belong to an extension.
+        ext: AnyExtension,
+        /// if true, then no more sub-blocks are available in this block.
+        is_last: bool,
+    },
     /// Decoded all information of the next frame, except the image data.
     ///
     /// The returned frame does **not** contain any owned image data.
@@ -192,9 +184,11 @@ enum State {
     BlockStart(u8),
     BlockEnd,
     ExtensionBlockStart,
+    /// Resets ext.data
+    ExtensionDataSubBlockStart(usize),
     /// Collects data in ext.data
-    ExtensionDataBlock(usize),
-    ApplicationExtension,
+    ExtensionDataSubBlock(usize),
+    ExtensionBlockEnd,
     LocalPalette(usize),
     LzwInit(u8),
     /// Decompresses LZW
@@ -212,6 +206,7 @@ use super::converter::PixelConverter;
 pub struct FrameDecoder {
     lzw_reader: LzwReader,
     pixel_converter: PixelConverter,
+    memory_limit: MemoryLimit,
 }
 
 impl FrameDecoder {
@@ -221,7 +216,8 @@ impl FrameDecoder {
     pub fn new(options: DecodeOptions) -> Self {
         Self {
             lzw_reader: LzwReader::new(options.check_for_end_code),
-            pixel_converter: PixelConverter::new(options.color_output, options.memory_limit),
+            pixel_converter: PixelConverter::new(options.color_output),
+            memory_limit: options.memory_limit.clone(),
         }
     }
 
@@ -236,7 +232,9 @@ impl FrameDecoder {
     /// If you get an error about invalid min code size, the buffer was probably pixels, not compressed data.
     #[inline]
     pub fn decode_lzw_encoded_frame(&mut self, frame: &mut Frame<'_>) -> Result<(), DecodingError> {
-        let pixel_bytes = self.pixel_converter.check_buffer_size(frame)?;
+        let pixel_bytes = self
+            .pixel_converter
+            .check_buffer_size(frame, &self.memory_limit)?;
         let mut vec = vec![0; pixel_bytes];
         self.decode_lzw_encoded_frame_into_buffer(frame, &mut vec)?;
         frame.buffer = Cow::Owned(vec);
@@ -382,10 +380,6 @@ pub struct StreamingDecoder {
     current: Option<Frame<'static>>,
     /// Needs to emit `HeaderEnd` once
     header_end_reached: bool,
-    /// XMP metadata bytes.
-    xmp_metadata: Option<Vec<u8>>,
-    /// ICC profile bytes.
-    icc_profile: Option<Vec<u8>>,
 }
 
 /// One version number of the GIF standard.
@@ -400,8 +394,6 @@ pub enum Version {
 struct ExtensionData {
     id: AnyExtension,
     data: Vec<u8>,
-    sub_block_lens: Vec<u8>,
-    is_block_end: bool,
 }
 
 /// Destination to write to for `StreamingDecoder::update`
@@ -429,9 +421,7 @@ impl OutputBuffer<'_> {
             OutputBuffer::Vec(vec) => {
                 let vec: &mut Vec<u8> = vec;
                 let len = buf.len();
-                memory_limit.check_size(vec.len() + len)?;
-                vec.try_reserve(len)
-                    .map_err(|_| DecodingError::OutOfMemory)?;
+                memory_limit.try_reserve(vec, len)?;
                 if vec.capacity() - vec.len() >= len {
                     vec.extend_from_slice(buf);
                 }
@@ -470,13 +460,9 @@ impl StreamingDecoder {
             ext: ExtensionData {
                 id: AnyExtension(0),
                 data: Vec::with_capacity(256), // 0xFF + 1 byte length
-                sub_block_lens: Vec::new(),
-                is_block_end: true,
             },
             current: None,
             header_end_reached: false,
-            xmp_metadata: None,
-            icc_profile: None,
         }
     }
 
@@ -503,10 +489,11 @@ impl StreamingDecoder {
         Ok((len - buf.len(), Decoded::Nothing))
     }
 
-    /// Returns the data of the last extension that has been decoded.
+    /// Data of the last extension sub block that has been decoded.
+    /// You need to concatenate all subblocks together to get the overall block content.
     #[must_use]
-    pub fn last_ext(&self) -> (AnyExtension, &[u8], bool) {
-        (self.ext.id, &self.ext.data, self.ext.is_block_end)
+    pub fn last_ext_sub_block(&mut self) -> &[u8] {
+        &self.ext.data
     }
 
     /// Current frame info as a mutable ref.
@@ -543,18 +530,6 @@ impl StreamingDecoder {
         self.height
     }
 
-    /// XMP metadata stored in the image.
-    #[must_use]
-    pub fn xmp_metadata(&self) -> Option<&[u8]> {
-        self.xmp_metadata.as_deref()
-    }
-
-    /// ICC profile stored in the image.
-    #[must_use]
-    pub fn icc_profile(&self) -> Option<&[u8]> {
-        self.icc_profile.as_deref()
-    }
-
     /// The version number of the GIF standard used in this image.
     ///
     /// We suppose a minimum of `V87a` compatibility. This value will be reported until we have
@@ -739,6 +714,13 @@ impl StreamingDecoder {
                     }
                 }
             }
+            ExtensionBlockStart => {
+                goto!(ExtensionDataSubBlockStart(b as usize), emit Decoded::BlockStart(Block::Extension))
+            }
+            ExtensionBlockEnd => {
+                self.ext.data.clear();
+                goto!(0, BlockEnd)
+            }
             BlockEnd => {
                 if b == Block::Trailer as u8 {
                     // can't consume yet, because the trailer is not a real block,
@@ -748,49 +730,27 @@ impl StreamingDecoder {
                     goto!(BlockStart(b))
                 }
             }
-            ExtensionBlockStart => {
+            ExtensionDataSubBlockStart(sub_block_len) => {
                 self.ext.data.clear();
-                self.ext.sub_block_lens.clear();
-                self.ext.is_block_end = false;
-                self.ext.sub_block_lens.push(b);
-                goto!(ExtensionDataBlock(b as usize), emit Decoded::BlockStart(Block::Extension))
+                goto!(0, ExtensionDataSubBlock(sub_block_len))
             }
-            ExtensionDataBlock(left) => {
+            ExtensionDataSubBlock(left) => {
                 if left > 0 {
                     let n = cmp::min(left, buf.len());
-                    self.memory_limit.check_size(self.ext.data.len() + n)?;
-                    self.ext
-                        .data
-                        .try_reserve(n)
-                        .map_err(|_| DecodingError::OutOfMemory)?;
+                    let needs_to_grow =
+                        n > self.ext.data.capacity().wrapping_sub(self.ext.data.len());
+                    if needs_to_grow {
+                        return Err(DecodingError::OutOfMemory);
+                    }
                     self.ext.data.extend_from_slice(&buf[..n]);
-                    goto!(n, ExtensionDataBlock(left - n))
+                    goto!(n, ExtensionDataSubBlock(left - n))
                 } else if b == 0 {
-                    self.ext.is_block_end = true;
-                    match self.ext.id.into_known() {
-                        Some(Extension::Application) => {
-                            goto!(0, ApplicationExtension, emit Decoded::BlockFinished(self.ext.id))
-                        }
-                        Some(Extension::Control) => {
-                            self.read_control_extension()?;
-                            goto!(BlockEnd, emit Decoded::BlockFinished(self.ext.id))
-                        }
-                        _ => {
-                            goto!(BlockEnd, emit Decoded::BlockFinished(self.ext.id))
-                        }
+                    if self.ext.id.into_known() == Some(Extension::Control) {
+                        self.read_control_extension()?;
                     }
+                    goto!(ExtensionBlockEnd, emit Decoded::SubBlock { ext: self.ext.id, is_last: true })
                 } else {
-                    self.ext.sub_block_lens.push(b);
-                    self.ext.is_block_end = false;
-                    goto!(ExtensionDataBlock(b as usize), emit Decoded::SubBlockFinished(self.ext.id))
-                }
-            }
-            ApplicationExtension => {
-                debug_assert_eq!(0, b);
-                if let Some(decoded) = self.read_application_extension() {
-                    goto!(BlockEnd, emit decoded)
-                } else {
-                    goto!(BlockEnd)
+                    goto!(ExtensionDataSubBlockStart(b as usize), emit Decoded::SubBlock { ext: self.ext.id, is_last: false })
                 }
             }
             LocalPalette(left) => {
@@ -897,48 +857,6 @@ impl StreamingDecoder {
         Ok(())
     }
 
-    fn read_application_extension(&mut self) -> Option<Decoded> {
-        if let Some(&[first, second]) = self.ext.data.strip_prefix(EXT_NAME_NETSCAPE) {
-            let repeat = u16::from(first) | u16::from(second) << 8;
-            return Some(Decoded::Repetitions(if repeat == 0 {
-                Repeat::Infinite
-            } else {
-                Repeat::Finite(repeat)
-            }));
-        } else if let Some(mut rest) = self.ext.data.strip_prefix(EXT_NAME_XMP) {
-            let (&id_len, data_lens) = self.ext.sub_block_lens.split_first()?;
-            if id_len as usize != EXT_NAME_XMP.len() {
-                return None;
-            }
-            // XMP is not written as a valid "pascal-string", so we need to stitch together
-            // the text from our collected sublock-lengths.
-            let mut xmp_metadata = Vec::with_capacity(data_lens.len() + self.ext.data.len());
-            for &len in data_lens {
-                xmp_metadata.push(len);
-                let (sub_block, tail) = rest.split_at(len as usize);
-                xmp_metadata.extend_from_slice(sub_block);
-                rest = tail;
-            }
-
-            // XMP adds a "ramp" of 257 bytes to the end of the metadata to let the "pascal-strings"
-            // parser converge to the null byte. The ramp looks like "0x01, 0xff, .., 0x01, 0x00".
-            // For convenience and to allow consumers to not be bothered with this implementation detail,
-            // we cut the ramp.
-            const RAMP_SIZE: usize = 257;
-            if xmp_metadata.len() >= RAMP_SIZE
-                && xmp_metadata.ends_with(&[0x03, 0x02, 0x01, 0x00])
-                && xmp_metadata[xmp_metadata.len() - RAMP_SIZE..].starts_with(&[0x01, 0x0ff])
-            {
-                xmp_metadata.truncate(xmp_metadata.len() - RAMP_SIZE);
-            }
-
-            self.xmp_metadata = Some(xmp_metadata);
-        } else if let Some(rest) = self.ext.data.strip_prefix(EXT_NAME_ICC) {
-            self.icc_profile = Some(rest.to_vec());
-        }
-        None
-    }
-
     fn add_frame(&mut self) {
         if self.current.is_none() {
             self.current = Some(Frame::default());