Merge pull request #155 from kornelski/skip-decoding

Option to skip decoding, while keeping LZW-compressed data

Author: kornelski

fuzz/Cargo.toml

@@ -1,10 +1,9 @@
-
 [package]
 name = "gif-fuzz"
 version = "0.0.0"
 authors = ["Automatically generated"]
 publish = false
-edition = "2018"
+edition = "2021"
 
 [package.metadata]
 cargo-fuzz = true

gif-afl/Cargo.toml

@@ -2,11 +2,12 @@
 name = "gif-afl"
 version = "0.1.0"
 authors = ["nwin <[email protected]>"]
+edition = "2021"
 
 [dependencies.gif]
 version = "*"
 path = "../"
 
 [dependencies]
 afl  = "*"
-afl-plugin  = "*"
+afl-plugin  = "*"

gif-c/Cargo.toml

@@ -4,10 +4,9 @@ license = "MIT OR Apache-2.0"
 version = "0.0.0"
 description = "C bindings to the gif library"
 authors = ["HeroicKatora <[email protected]>"]
-
 homepage = "https://github.com/image-rs/image-gif"
 repository = "https://github.com/image-rs/image-gif"
-edition = "2018"
+edition = "2021"
 
 [lib]
 crate-type = ["staticlib", "cdylib"]

src/common.rs

@@ -325,7 +325,8 @@ impl Frame<'static> {
     #[cfg(feature = "color_quant")]
     pub fn from_rgb_speed(width: u16, height: u16, pixels: &[u8], speed: i32) -> Frame<'static> {
         assert_eq!(width as usize * height as usize * 3, pixels.len(), "Too much or too little pixel data for the given width and height to create a GIF Frame");
-        let mut vec: Vec<u8> = Vec::with_capacity(pixels.len() + width as usize * height as usize);
+        let mut vec: Vec<u8> = Vec::new();
+        vec.try_reserve_exact(pixels.len() + width as usize * height as usize).expect("OOM");
         for v in pixels.chunks_exact(3) {
             vec.extend_from_slice(&[v[0], v[1], v[2], 0xFF])
         }

src/encoder.rs

@@ -141,12 +141,13 @@ impl<W: Write> Encoder<W> {
     /// if no global palette shall be used an empty slice may be supplied.
     pub fn new(w: W, width: u16, height: u16, global_palette: &[u8]) -> Result<Self, EncodingError> {
         let buffer_size = (width as usize) * (height as usize);
+        let mut buffer = Vec::new();
+        buffer.try_reserve_exact(buffer_size).map_err(|_| io::Error::from(io::ErrorKind::OutOfMemory))?;
         Encoder {
             w: Some(w),
             global_palette: false,
-            width: width,
-            height: height,
-            buffer: Vec::with_capacity(buffer_size)
+            width, height,
+            buffer,
         }.write_global_palette(global_palette)
     }
 

src/reader/decoder.rs

@@ -110,6 +110,12 @@ pub enum Extensions {
     Skip
 }
 
+#[derive(Debug, Copy, Clone)]
+pub enum FrameDataType {
+    Pixels,
+    Lzw { min_code_size: u8 },
+}
+
 /// Indicates whether a certain object has been decoded
 #[derive(Debug)]
 pub enum Decoded<'a> {
@@ -133,15 +139,16 @@ pub enum Decoded<'a> {
     /// 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.
     BlockFinished(AnyExtension, &'a [u8]),
-    /// Decoded all information of the next frame.
+    /// Decoded all information of the next frame, except the image data.
     ///
     /// The returned frame does **not** contain any owned image data.
-    Frame(&'a Frame<'static>),
+    FrameMetadata(&'a Frame<'static>, FrameDataType),
     /// Decoded some data of the current frame.
     BytesDecoded(usize),
+    /// Copied compressed data of the current frame.
+    LzwData(&'a [u8]),
     /// No more data available the current frame.
     DataEnd,
-
 }
 
 /// Internal state of the GIF decoder
@@ -159,7 +166,10 @@ enum State {
     SkipBlock(usize),
     LocalPalette(usize),
     LzwInit(u8),
+    /// Decompresses LZW
     DecodeSubBlock(usize),
+    /// Keeps LZW compressed
+    CopySubBlock(usize),
     FrameDecoded,
     Trailer
 }
@@ -258,6 +268,7 @@ pub struct StreamingDecoder {
     state: Option<State>,
     lzw_reader: LzwReader,
     skip_extensions: bool,
+    skip_frame_decoding: bool,
     check_frame_consistency: bool,
     allow_unknown_blocks: bool,
     version: Version,
@@ -299,6 +310,7 @@ impl StreamingDecoder {
             state: Some(Magic(0, [0; 6])),
             lzw_reader: LzwReader::new(options.check_for_end_code),
             skip_extensions: true,
+            skip_frame_decoding: options.skip_frame_decoding,
             check_frame_consistency: options.check_frame_consistency,
             allow_unknown_blocks: options.allow_unknown_blocks,
             version: Version::V87a,
@@ -405,7 +417,7 @@ impl StreamingDecoder {
         }
     }
 
-    fn next_state<'a>(&'a mut self, buf: &[u8], decode_bytes_into: Option<&mut [u8]>) -> Result<(usize, Decoded<'a>), DecodingError> {
+    fn next_state<'a>(&'a mut self, buf: &'a [u8], decode_bytes_into: Option<&mut [u8]>) -> Result<(usize, Decoded<'a>), DecodingError> {
         macro_rules! goto (
             ($n:expr, $state:expr) => ({
                 self.state = Some($state); 
@@ -495,7 +507,7 @@ impl StreamingDecoder {
                         let global_table = b & 0x80 != 0;
                         let entries = if global_table {
                             let entries = PLTE_CHANNELS*(1 << ((b & 0b111) + 1) as usize);
-                            self.global_color_table.reserve_exact(entries);
+                            self.global_color_table.try_reserve_exact(entries).map_err(|_| io::Error::from(io::ErrorKind::OutOfMemory))?;
                             entries
                         } else {
                             0usize
@@ -556,9 +568,9 @@ impl StreamingDecoder {
 
                         if local_table {
                             let entries = PLTE_CHANNELS * (1 << (table_size + 1));
-                            
-                            self.current_frame_mut().palette =
-                                Some(Vec::with_capacity(entries));
+                            let mut pal = Vec::new();
+                            pal.try_reserve_exact(entries).map_err(|_| io::Error::from(io::ErrorKind::OutOfMemory))?;
+                            self.current_frame_mut().palette = Some(pal);
                             goto!(LocalPalette(entries))
                         } else {
                             goto!(Byte(CodeSize))
@@ -663,12 +675,29 @@ impl StreamingDecoder {
                 }
             }
             LzwInit(min_code_size) => {
-                // Reset validates the min code size
-                self.lzw_reader.reset(min_code_size)?;
-
-                goto!(DecodeSubBlock(b as usize), emit Decoded::Frame(self.current_frame_mut()))
+                if !self.skip_frame_decoding {
+                    // Reset validates the min code size
+                    self.lzw_reader.reset(min_code_size)?;
+                    goto!(DecodeSubBlock(b as usize), emit Decoded::FrameMetadata(self.current_frame_mut(), FrameDataType::Pixels))
+                } else {
+                    goto!(CopySubBlock(b as usize), emit Decoded::FrameMetadata(self.current_frame_mut(), FrameDataType::Lzw { min_code_size }))
+                }
+            }
+            CopySubBlock(left) => {
+                debug_assert!(self.skip_frame_decoding);
+                let n = cmp::min(left, buf.len());
+                if left > 0 {
+                    goto!(n, CopySubBlock(left - n), emit Decoded::LzwData(&buf[..n]))
+                } else if b != 0 {
+                    goto!(CopySubBlock(b as usize))
+                } else {
+                    // end of image data reached
+                    self.current = None;
+                    goto!(0, FrameDecoded, emit Decoded::DataEnd)
+                }
             }
             DecodeSubBlock(left) => {
+                debug_assert!(!self.skip_frame_decoding);
                 if left > 0 {
                     let n = cmp::min(left, buf.len());
                     if self.lzw_reader.has_ended() || decode_bytes_into.is_none() {

src/reader/mod.rs

@@ -11,7 +11,7 @@ use crate::common::{Block, Frame};
 mod decoder;
 pub use self::decoder::{
     PLTE_CHANNELS, StreamingDecoder, Decoded, DecodingError, DecodingFormatError, Extensions,
-    Version
+    Version, FrameDataType
 };
 
 const N_CHANNELS: usize = 4;
@@ -89,6 +89,7 @@ pub struct DecodeOptions {
     memory_limit: MemoryLimit,
     color_output: ColorOutput,
     check_frame_consistency: bool,
+    skip_frame_decoding: bool,
     check_for_end_code: bool,
     allow_unknown_blocks: bool,
 }
@@ -100,6 +101,7 @@ impl DecodeOptions {
             memory_limit: MemoryLimit::Bytes(50_000_000.try_into().unwrap()), // 50 MB
             color_output: ColorOutput::Indexed,
             check_frame_consistency: false,
+            skip_frame_decoding: false,
             check_for_end_code: false,
             allow_unknown_blocks: false,
         }
@@ -129,6 +131,18 @@ impl DecodeOptions {
         self.check_frame_consistency = check;
     }
 
+    /// Configure whether to skip decoding frames.
+    ///
+    /// The default is false.
+    ///
+    /// When turned on, LZW decoding is skipped. `Decoder::read_next_frame` will return
+    /// compressed LZW bytes in frame's data.
+    /// `Decoder::next_frame_info` will return the metadata of the next frame as usual.
+    /// This is useful to count frames without incurring the overhead of decoding.
+    pub fn skip_frame_decoding(&mut self, skip: bool) {
+        self.skip_frame_decoding = skip;
+    }
+
     /// Configure if LZW encoded blocks must end with a marker end code.
     ///
     /// The default is `false`.
@@ -210,6 +224,7 @@ pub struct Decoder<R: Read> {
     bg_color: Option<u8>,
     global_palette: Option<Vec<u8>>,
     current_frame: Frame<'static>,
+    current_frame_data_type: FrameDataType,
     buffer: Vec<u8>,
 }
 
@@ -237,6 +252,7 @@ impl<R> Decoder<R> where R: Read {
             color_output: options.color_output,
             memory_limit: options.memory_limit,
             current_frame: Frame::default(),
+            current_frame_data_type: FrameDataType::Pixels,
         }
     }
 
@@ -277,8 +293,9 @@ impl<R> Decoder<R> where R: Read {
     pub fn next_frame_info(&mut self) -> Result<Option<&Frame<'static>>, DecodingError> {
         loop {
             match self.decoder.decode_next(None)? {
-                Some(Decoded::Frame(frame)) => {
+                Some(Decoded::FrameMetadata(frame, frame_data_type)) => {
                     self.current_frame = frame.clone();
+                    self.current_frame_data_type = frame_data_type;
                     if frame.palette.is_none() && self.global_palette.is_none() {
                         return Err(DecodingError::format(
                             "no color table available for current frame"
@@ -311,11 +328,21 @@ impl<R> Decoder<R> where R: Read {
                 pixel_bytes, self.buffer_size(),
                 "Checked computation diverges from required buffer size"
             );
-
-            let mut vec = vec![0; pixel_bytes];
-            self.read_into_buffer(&mut vec)?;
-            self.current_frame.buffer = Cow::Owned(vec);
-            self.current_frame.interlaced = false;
+            match self.current_frame_data_type {
+                FrameDataType::Pixels => {
+                    let mut vec = vec![0; pixel_bytes];
+                    self.read_into_buffer(&mut vec)?;
+                    self.current_frame.buffer = Cow::Owned(vec);
+                    self.current_frame.interlaced = false;
+                }
+                FrameDataType::Lzw { min_code_size } => {
+                    let mut vec = Vec::new();
+                    // Guesstimate 2bpp
+                    vec.try_reserve(pixel_bytes/4).map_err(|_| io::Error::from(io::ErrorKind::OutOfMemory))?;
+                    self.copy_lzw_into_buffer(min_code_size, &mut vec)?;
+                    self.current_frame.buffer = Cow::Owned(vec);
+                },
+            }
             Ok(Some(&self.current_frame))
         } else {
             Ok(None)
@@ -348,6 +375,21 @@ impl<R> Decoder<R> where R: Read {
         Ok(())
     }
 
+    fn copy_lzw_into_buffer(&mut self, min_code_size: u8, buf: &mut Vec<u8>) -> Result<(), DecodingError> {
+        // `write_lzw_pre_encoded_frame` smuggles `min_code_size` in the first byte.
+        buf.push(min_code_size);
+        loop {
+            match self.decoder.decode_next(None)? {
+                Some(Decoded::LzwData(data)) => {
+                    buf.try_reserve(data.len()).map_err(|_| io::Error::from(io::ErrorKind::OutOfMemory))?;
+                    buf.extend_from_slice(data);
+                },
+                Some(Decoded::DataEnd) => return Ok(()),
+                _ => return Err(DecodingError::format("unexpected data")),
+            }
+        }
+    }
+
     /// Reads data of the current frame into a pre-allocated buffer until the buffer has been
     /// filled completely.
     ///

tests/decode.rs

@@ -82,3 +82,34 @@ fn check_for_end_code_is_configurable() {
         assert!(decoder.read_next_frame().is_err());
     }
 }
+
+#[test]
+fn check_skip_frame_data() {
+    let image: &[u8] = include_bytes!("samples/moon_impact.gif");
+
+    let mut options = DecodeOptions::new();
+    options.skip_frame_decoding(true);
+    let mut decoder = options.clone().read_info(image).unwrap();
+
+    for _ in 0..14 {
+        assert!(matches!(decoder.next_frame_info().unwrap(), Some(_)));
+    }
+
+    assert!(matches!(decoder.next_frame_info(), Ok(None)));
+}
+
+#[test]
+fn check_skip_frame_data_decode_frame_error() {
+    let image: &[u8] = include_bytes!("samples/moon_impact.gif");
+
+    let mut options = DecodeOptions::new();
+    options.skip_frame_decoding(true);
+    let mut skipping_decoder = options.read_info(image).unwrap();
+    let mut normal_decoder = DecodeOptions::new().read_info(image).unwrap();
+
+    while let Ok(Some(_normal_frame)) = normal_decoder.read_next_frame() {
+        let _compressed_frame = skipping_decoder.read_next_frame().unwrap().unwrap();
+    }
+    assert!(skipping_decoder.read_next_frame().unwrap().is_none());
+
+}