Reject truncated image data instead of panicking

src/headers.rs

@@ -25,6 +25,42 @@ pub struct IhdrData {
     pub interlaced: u8,
 }
 
+impl IhdrData {
+    /// Bits per pixel
+    pub fn bpp(&self) -> u8 {
+        self.bit_depth.as_u8() * self.color_type.channels_per_pixel()
+    }
+
+    /// Byte length of IDAT that is correct for this IHDR
+    pub fn raw_data_size(&self) -> usize {
+        let w = self.width as usize;
+        let h = self.height as usize;
+        let bpp = self.bpp();
+
+        fn bitmap_size(bpp: u8, w: usize, h: usize) -> usize {
+            (((w / 8) * bpp as usize) + ((w & 7) * bpp as usize + 7) / 8) * h
+        }
+
+        if self.interlaced == 0 {
+            bitmap_size(bpp, w, h) + h
+        } else {
+            let mut size = bitmap_size(bpp, (w + 7) >> 3, (h + 7) >> 3) + ((h + 7) >> 3);
+            if w > 4 {
+                size += bitmap_size(bpp, (w + 3) >> 3, (h + 7) >> 3) + ((h + 7) >> 3);
+            }
+            size += bitmap_size(bpp, (w + 3) >> 2, (h + 3) >> 3) + ((h + 3) >> 3);
+            if w > 2 {
+                size += bitmap_size(bpp, (w + 1) >> 2, (h + 3) >> 2) + ((h + 3) >> 2);
+            }
+            size += bitmap_size(bpp, (w + 1) >> 1, (h + 1) >> 2) + ((h + 1) >> 2);
+            if w > 1 {
+                size += bitmap_size(bpp, w >> 1, (h + 1) >> 1) + ((h + 1) >> 1);
+            }
+            size + bitmap_size(bpp, w, h >> 1) + (h >> 1)
+        }
+    }
+}
+
 #[derive(Debug, PartialEq, Clone)]
 /// Options to use for performing operations on headers (such as stripping)
 pub enum Headers {

src/interlace.rs

@@ -5,7 +5,7 @@ use bit_vec::BitVec;
 #[must_use]
 pub fn interlace_image(png: &PngImage) -> PngImage {
     let mut passes: Vec<BitVec> = vec![BitVec::new(); 7];
-    let bits_per_pixel = png.ihdr.bit_depth.as_u8() * png.channels_per_pixel();
+    let bits_per_pixel = png.ihdr.bpp();
     for (index, line) in png.scan_lines().enumerate() {
         match index % 8 {
             // Add filter bytes to passes that will be in the output image
@@ -95,7 +95,7 @@ pub fn interlace_image(png: &PngImage) -> PngImage {
 }
 
 pub fn deinterlace_image(png: &PngImage) -> PngImage {
-    let bits_per_pixel = png.ihdr.bit_depth.as_u8() * png.channels_per_pixel();
+    let bits_per_pixel = png.ihdr.bpp();
     let bits_per_line = 8 + bits_per_pixel as usize * png.ihdr.width as usize;
     // Initialize each output line with a starting filter byte of 0
     // as well as some blank data

src/png/mod.rs

@@ -118,6 +118,11 @@ impl PngData {
         let ihdr_header = parse_ihdr_header(&ihdr)?;
         let raw_data = deflate::inflate(idat_headers.as_ref())?;
 
+        // Reject files with incorrect width/height or truncated data
+        if raw_data.len() != ihdr_header.raw_data_size() {
+            return Err(PngError::TruncatedData);
+        }
+
         let (palette, transparency_pixel) = Self::palette_to_rgba(
             ihdr_header.color_type,
             aux_headers.remove(b"PLTE"),

src/png/scan_lines.rs

@@ -167,8 +167,7 @@ impl Iterator for ScanLineRanges {
         let bits_per_line = pixels_per_line * u32::from(self.bits_per_pixel);
         let bytes_per_line = ((bits_per_line + 7) / 8) as usize;
         let len = bytes_per_line + 1;
-        assert!(self.left >= len);
-        self.left -= len;
+        self.left = self.left.checked_sub(len)?;
         Some((len, current_pass))
     }
 }