Fix writing palettes with non-power-of-two sizes

src/encoder.rs

@@ -151,18 +151,15 @@ impl<W: Write> Encoder<W> {
     fn write_global_palette(mut self, palette: &[u8]) -> Result<Self, EncodingError> {
         let mut flags = 0;
         flags |= 0b1000_0000;
-        let num_colors = palette.len() / 3;
-        if num_colors > 256 {
-            return Err(EncodingError::from(EncodingFormatError::TooManyColors));
-        }
-        self.global_palette = num_colors > 0;
+        let (palette, padding, table_size) = Self::check_color_table(palette)?;
+        self.global_palette = !palette.is_empty();
         // Size of global color table.
-        flags |= flag_size(num_colors);
+        flags |= table_size;
         // Color resolution .. FIXME. This is mostly ignored (by ImageMagick at least) but hey, we
         // should use some sensible value here or even allow configuring it?
-        flags |= flag_size(num_colors) << 4; // wtf flag
+        flags |= table_size << 4; // wtf flag
         self.write_screen_desc(flags)?;
-        self.write_color_table(palette)?;
+        Self::write_color_table(self.writer()?, palette, padding)?;
         Ok(self)
     }
 
@@ -192,12 +189,9 @@ impl<W: Write> Encoder<W> {
         let palette = match frame.palette {
             Some(ref palette) => {
                 flags |= 0b1000_0000;
-                let num_colors = palette.len() / 3;
-                if num_colors > 256 {
-                    return Err(EncodingError::from(EncodingFormatError::TooManyColors));
-                }
-                flags |= flag_size(num_colors);
-                Some(palette)
+                let (palette, padding, table_size) = Self::check_color_table(&palette)?;
+                flags |= table_size;
+                Some((palette, padding))
             },
             None if self.global_palette => None,
             _ => return Err(EncodingError::from(EncodingFormatError::MissingColorPalette))
@@ -211,8 +205,8 @@ impl<W: Write> Encoder<W> {
         tmp.write_le(flags)?;
         let writer = self.writer()?;
         tmp.finish(&mut *writer)?;
-        if let Some(palette) = palette {
-            writer.write_all(palette)?;
+        if let Some((palette, padding)) = palette {
+            Self::write_color_table(writer, palette, padding)?;
         }
         Ok(())
     }
@@ -246,21 +240,26 @@ impl<W: Write> Encoder<W> {
         writer.write_le(0u8).map_err(Into::into)
     }
 
-    fn write_color_table(&mut self, table: &[u8]) -> Result<(), EncodingError> {
-        let writer = self.writer()?;
-        let num_colors = table.len() / 3;
-        if num_colors > 256 {
-            return Err(EncodingError::from(EncodingFormatError::TooManyColors));
-        }
-        writer.write_all(&table[..num_colors * 3])?;
-        let size = flag_size(num_colors);
+    fn write_color_table(writer: &mut W, table: &[u8], padding: usize) -> Result<(), EncodingError> {
+        writer.write_all(&table)?;
         // Waste some space as of gif spec
-        for _ in 0..((2 << size) - num_colors) {
+        for _ in 0..padding {
             writer.write_all(&[0, 0, 0])?;
         }
         Ok(())
     }
 
+    /// returns rounded palette size, number of missing colors, and table size flag
+    fn check_color_table(table: &[u8]) -> Result<(&[u8], usize, u8), EncodingError> {
+        let num_colors = table.len() / 3;
+        if num_colors > 256 {
+            return Err(EncodingError::from(EncodingFormatError::TooManyColors));
+        }
+        let table_size = flag_size(num_colors);
+        let padding = (2 << table_size) - num_colors;
+        Ok((&table[..num_colors * 3], padding, table_size))
+    }
+
     /// Writes an extension to the image.
     ///
     /// It is normally not necessary to call this method manually.

src/reader/decoder.rs

@@ -290,12 +290,17 @@ impl LzwReader {
         }
     }
 
-    pub fn reset(&mut self, min_code_size: u8) -> Result<(), DecodingError> {
+    pub fn check_code_size(min_code_size: u8) -> Result<(), DecodingError> {
         // LZW spec: max 12 bits per code. This check helps catch confusion
         // between LZW-compressed buffers and raw pixel data
         if min_code_size > 11 || min_code_size < 1 {
             return Err(DecodingError::format("invalid minimal code size"));
         }
+        Ok(())
+    }
+
+    pub fn reset(&mut self, min_code_size: u8) -> Result<(), DecodingError> {
+        Self::check_code_size(min_code_size)?;
 
         // The decoder can be reused if the code size stayed the same
         if self.min_code_size != min_code_size || self.decoder.is_none() {
@@ -752,6 +757,7 @@ impl StreamingDecoder {
                     self.lzw_reader.reset(min_code_size)?;
                     goto!(DecodeSubBlock(b as usize), emit Decoded::FrameMetadata(FrameDataType::Pixels))
                 } else {
+                    LzwReader::check_code_size(min_code_size)?;
                     goto!(CopySubBlock(b as usize), emit Decoded::FrameMetadata(FrameDataType::Lzw { min_code_size }))
                 }
             }

tests/roundtrip.rs

@@ -1,6 +1,6 @@
 #![cfg(feature = "std")]
 
-use gif::{ColorOutput, Decoder, Encoder, Frame, AnyExtension};
+use gif::{ColorOutput, Decoder, Encoder, Frame, AnyExtension, DecodeOptions};
 
 #[test]
 fn round_trip() {
@@ -152,6 +152,49 @@ fn encode_roundtrip_few_colors() {
 }
 
 
+#[test]
+fn palette_sizes() {
+    let global_pal = (0..=255u8).flat_map(|i| [i, i/2, i.wrapping_add(13)]).collect::<Vec<_>>();
+    let local_pal = (0..=255u8).flat_map(|i| [i^0x55, i, i.wrapping_add(7)]).collect::<Vec<_>>();
+
+    for size in 1..=256 {
+        let global = &global_pal[..size * 3];
+        let local = &local_pal[..size * 3];
+
+        let mut encoder = Encoder::new(vec![], 1, 1, global).unwrap();
+        let mut f = Frame::default();
+        f.width = 1;
+        f.height = 1;
+        f.buffer = [1][..].into();
+        encoder.write_frame(&f).unwrap();
+
+        let mut f = Frame::default();
+        f.width = 1;
+        f.height = 1;
+        f.buffer = [1][..].into();
+        f.palette = Some(local.to_vec());
+        encoder.write_frame(&f).unwrap();
+        let gif = encoder.into_inner().unwrap();
+        let gif = &mut gif.as_slice();
+
+        let mut d = DecodeOptions::new().read_info(gif).unwrap();
+        let (decoded_global_pal, padding) = d.global_palette().unwrap().split_at(global.len());
+        assert_eq!(global, decoded_global_pal);
+        assert_eq!(padding.len(), 3 * (size.max(2).next_power_of_two() - size));
+        assert!(padding.iter().all(|&b| b == 0));
+
+        assert!(d.read_next_frame().unwrap().unwrap().palette.is_none());
+        let decoded_local_pal = d.read_next_frame().unwrap().unwrap().palette.as_deref().unwrap();
+        let (decoded_local_pal, padding) = decoded_local_pal.split_at(local.len());
+        assert_eq!(local, decoded_local_pal);
+        assert_eq!(padding.len(), 3 * (size.max(2).next_power_of_two() - size));
+        assert!(padding.iter().all(|&b| b == 0));
+
+        assert!(d.read_next_frame().unwrap().is_none());
+    }
+}
+
+
 #[test]
 fn palette_fail() {
     let mut encoder = Encoder::new(vec![], 0xFFFF, 0xFFFF, &[]).unwrap();