From 5e680ccf7ac92c51468201a6d31d016cee30f3ea Mon Sep 17 00:00:00 2001
From: Jonathan Behrens <fintelia@gmail.com>
Date: Mon, 1 Jan 2024 01:51:58 -0500
Subject: [PATCH] Validate size of uncompressed alpha chunks and other fixes
 (#40)

---
 fuzz/Cargo.lock | 16 ----------------
 src/decoder.rs  | 10 +++++++---
 src/extended.rs | 23 ++++++-----------------
 src/lossless.rs |  2 +-
 src/vp8.rs      |  2 +-
 5 files changed, 15 insertions(+), 38 deletions(-)

diff --git a/fuzz/Cargo.lock b/fuzz/Cargo.lock
index 81abac1..b8716fd 100644
--- a/fuzz/Cargo.lock
+++ b/fuzz/Cargo.lock
@@ -8,12 +8,6 @@ version = "1.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7d5a26814d8dcb93b0e5a0ff3c6d80a8843bafb21b39e8e18a6f05471870e110"
 
-[[package]]
-name = "autocfg"
-version = "1.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
-
 [[package]]
 name = "byteorder"
 version = "1.5.0"
@@ -56,15 +50,6 @@ dependencies = [
  "once_cell",
 ]
 
-[[package]]
-name = "num-traits"
-version = "0.2.17"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "39e3200413f237f41ab11ad6d161bc7239c84dcb631773ccd7de3dfe4b5c267c"
-dependencies = [
- "autocfg",
-]
-
 [[package]]
 name = "once_cell"
 version = "1.19.0"
@@ -131,7 +116,6 @@ name = "webp"
 version = "0.1.0"
 dependencies = [
  "byteorder",
- "num-traits",
  "thiserror",
 ]
 
diff --git a/src/decoder.rs b/src/decoder.rs
index f29f8f7..fcc1d22 100644
--- a/src/decoder.rs
+++ b/src/decoder.rs
@@ -571,8 +571,8 @@ impl<R: Read + Seek> WebPDecoder<R> {
                     .clone();
                 let alpha_chunk = read_alpha_chunk(
                     &mut range_reader(&mut self.r, range.start..range.end)?,
-                    self.width,
-                    self.height,
+                    self.width as u16,
+                    self.height as u16,
                 )?;
 
                 for y in 0..frame.height {
@@ -632,6 +632,9 @@ impl<R: Read + Seek> WebPDecoder<R> {
         let frame_y = extended::read_3_bytes(&mut self.r)? * 2;
         let frame_width = extended::read_3_bytes(&mut self.r)? + 1;
         let frame_height = extended::read_3_bytes(&mut self.r)? + 1;
+        if frame_width > 16384 || frame_height > 16384 {
+            return Err(DecodingError::ImageTooLarge);
+        }
         if frame_x + frame_width > self.width || frame_y + frame_height > self.height {
             return Err(DecodingError::FrameOutsideImage);
         }
@@ -688,7 +691,8 @@ impl<R: Read + Seek> WebPDecoder<R> {
                 // read alpha
                 let next_chunk_start = self.r.stream_position()? + chunk_size_rounded as u64;
                 let mut reader = (&mut self.r).take(chunk_size as u64);
-                let alpha_chunk = read_alpha_chunk(&mut reader, frame_width, frame_height)?;
+                let alpha_chunk =
+                    read_alpha_chunk(&mut reader, frame_width as u16, frame_height as u16)?;
 
                 // read opaque
                 self.r.seek(io::SeekFrom::Start(next_chunk_start))?;
diff --git a/src/extended.rs b/src/extended.rs
index ea9ee74..bcc5563 100644
--- a/src/extended.rs
+++ b/src/extended.rs
@@ -2,7 +2,7 @@ use super::lossless::LosslessDecoder;
 use crate::decoder::DecodingError;
 use byteorder::ReadBytesExt;
 use std::convert::TryInto;
-use std::io::{self, Read};
+use std::io::Read;
 
 #[derive(Debug, Clone)]
 pub(crate) struct WebPExtendedInfo {
@@ -264,8 +264,8 @@ pub(crate) enum FilteringMethod {
 
 pub(crate) fn read_alpha_chunk<R: Read>(
     reader: &mut R,
-    width: u32,
-    height: u32,
+    width: u16,
+    height: u16,
 ) -> Result<AlphaChunk, DecodingError> {
     let info_byte = reader.read_u8()?;
 
@@ -298,27 +298,16 @@ pub(crate) fn read_alpha_chunk<R: Read>(
         _ => return Err(DecodingError::InvalidCompressionMethod),
     };
 
-    let mut framedata = Vec::new();
-    reader.read_to_end(&mut framedata)?;
-
     let data = if lossless_compression {
-        let cursor = io::Cursor::new(framedata);
-
-        let mut decoder = LosslessDecoder::new(cursor);
-        //this is a potential problem for large images; would require rewriting lossless decoder to
-        //use u32 for width and height
-        let width: u16 = width.try_into().map_err(|_| DecodingError::ImageTooLarge)?;
-        let height: u16 = height
-            .try_into()
-            .map_err(|_| DecodingError::ImageTooLarge)?;
+        let mut decoder = LosslessDecoder::new(reader);
         let frame = decoder.decode_frame(Some((width, height)))?;
 
         let mut data = vec![0u8; usize::from(width) * usize::from(height)];
-
         frame.fill_green(&mut data);
-
         data
     } else {
+        let mut framedata = vec![0; width as usize * height as usize];
+        reader.read_exact(&mut framedata)?;
         framedata
     };
 
diff --git a/src/lossless.rs b/src/lossless.rs
index d76d445..489ebcd 100644
--- a/src/lossless.rs
+++ b/src/lossless.rs
@@ -236,7 +236,7 @@ impl<R: Read> LosslessDecoder<R> {
     }
 
     /// Adjusts the color map since it's subtraction coded
-    fn adjust_color_map(color_map: &mut Vec<u32>) {
+    fn adjust_color_map(color_map: &mut [u32]) {
         for i in 1..color_map.len() {
             color_map[i] = add_pixels(color_map[i], color_map[i - 1]);
         }
diff --git a/src/vp8.rs b/src/vp8.rs
index fba5e63..71c0af6 100644
--- a/src/vp8.rs
+++ b/src/vp8.rs
@@ -1307,7 +1307,7 @@ impl<R: Read> Vp8Decoder<R> {
 
             self.top = init_top_macroblocks(self.frame.width as usize);
             // Almost always the first macro block, except when non exists (i.e. `width == 0`)
-            self.left = self.top.get(0).cloned().unwrap_or_default();
+            self.left = self.top.first().cloned().unwrap_or_default();
 
             self.mbwidth = (self.frame.width + 15) / 16;
             self.mbheight = (self.frame.height + 15) / 16;