Merge pull request #161 from kornelski/depanic

Avoid panic on truncated palette
image-rs · Jan 10, 2024 · c33079b · c33079b
2 parents ff242b0 + fee622d
commit c33079b
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 16 deletions.
diff --git a/src/common.rs b/src/common.rs
@@ -239,10 +239,10 @@ impl Frame<'static> {
         let mut colors_vec: Vec<(u8, u8, u8, u8)> = colors.into_iter().collect();
         colors_vec.sort_unstable();
         let palette = colors_vec.iter().flat_map(|&(r, g, b, _a)| [r, g, b]).collect();
-        let colors_lookup: HashMap<(u8, u8, u8, u8), u8> =  colors_vec.into_iter().zip(0..=255).collect();
+        let colors_lookup: HashMap<(u8, u8, u8, u8), u8> = colors_vec.into_iter().zip(0..=255).collect();
 
         let index_of = | pixel: &[u8] |
-            *colors_lookup.get(&(pixel[0], pixel[1], pixel[2], pixel[3])).unwrap();
+            colors_lookup.get(&(pixel[0], pixel[1], pixel[2], pixel[3])).copied().unwrap_or(0);
 
         return Frame {
             width,

diff --git a/src/reader/decoder.rs b/src/reader/decoder.rs
@@ -403,8 +403,9 @@ impl StreamingDecoder {
         self.current.as_mut().unwrap()
     }
 
-    #[inline(always)]
     /// Current frame info as a ref.
+    #[inline(always)]
+    #[track_caller]
     pub fn current_frame<'a>(&'a self) -> &'a Frame<'static> {
         self.current.as_ref().unwrap()
     }
@@ -486,9 +487,9 @@ impl StreamingDecoder {
                         goto!(Byte(ByteValue::GlobalFlags))
                     },
                     (Delay, delay) => {
+                        self.current_frame_mut().delay = delay;
                         self.ext.data.push(value as u8);
                         self.ext.data.push(b);
-                        self.current_frame_mut().delay = delay;
                         goto!(Byte(ByteValue::TransparentIdx))
                     },
                     (ImageLeft, left) => {
@@ -534,14 +535,15 @@ impl StreamingDecoder {
                     },
                     ControlFlags => {
                         self.ext.data.push(b);
+                        let frame = self.current_frame_mut();
                         let control_flags = b;
                         if control_flags & 1 != 0 {
                             // Set to Some(...), gets overwritten later
-                            self.current_frame_mut().transparent = Some(0)
+                            frame.transparent = Some(0)
                         }
-                        self.current_frame_mut().needs_user_input =
+                        frame.needs_user_input =
                             control_flags & 0b10 != 0;
-                        self.current_frame_mut().dispose = match DisposalMethod::from_u8(
+                        frame.dispose = match DisposalMethod::from_u8(
                             (control_flags & 0b11100) >> 2
                         ) {
                             Some(method) => method,
@@ -560,13 +562,14 @@ impl StreamingDecoder {
                         let local_table = (b & 0b1000_0000) != 0;
                         let interlaced   = (b & 0b0100_0000) != 0;
                         let table_size  =  b & 0b0000_0111;
+                        let check_frame_consistency = self.check_frame_consistency;
+                        let (width, height) = (self.width, self.height);
 
-                        self.current_frame_mut().interlaced = interlaced;
+                        let frame = self.current_frame_mut();
 
-                        if self.check_frame_consistency {
+                        frame.interlaced = interlaced;
+                        if check_frame_consistency {
                             // Consistency checks.
-                            let (width, height) = (self.width, self.height);
-                            let frame = self.current_frame_mut();
                             if width.checked_sub(frame.width) < Some(frame.left)
                                 || height.checked_sub(frame.height) < Some(frame.top)
                             {
@@ -578,7 +581,7 @@ impl StreamingDecoder {
                             let entries = PLTE_CHANNELS * (1 << (table_size + 1));
                             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);
+                            frame.palette = Some(pal);
                             goto!(LocalPalette(entries))
                         } else {
                             goto!(Byte(CodeSize))
@@ -594,7 +597,7 @@ impl StreamingDecoder {
                     goto!(n, GlobalPalette(left - n))
                 } else {
                     let idx = self.background_color[0];
-                    match self.global_color_table.chunks(PLTE_CHANNELS).nth(idx as usize) {
+                    match self.global_color_table.chunks_exact(PLTE_CHANNELS).nth(idx as usize) {
                         Some(chunk) => self.background_color[..PLTE_CHANNELS]
                             .copy_from_slice(&chunk[..PLTE_CHANNELS]),
                         None => self.background_color[0] = 0
@@ -675,9 +678,13 @@ impl StreamingDecoder {
             LocalPalette(left) => {
                 let n = cmp::min(left, buf.len());
                 if left > 0 {
-
-                    self.current_frame_mut().palette
-                        .as_mut().unwrap().extend(buf[..n].iter().copied());
+                    let src = &buf[..n];
+                    if let Some(pal) = self.current_frame_mut().palette.as_mut() {
+                        // capacity has already been reserved in ImageFlags
+                        if pal.capacity() - pal.len() >= src.len() {
+                            pal.extend_from_slice(src);
+                        }
+                    }
                     goto!(n, LocalPalette(left - n))
                 } else {
                     goto!(LzwInit(b))