diff --git a/src/headers.rs b/src/headers.rs index 0695bd0e..7a328e7b 100644 --- a/src/headers.rs +++ b/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 { diff --git a/src/interlace.rs b/src/interlace.rs index 2d391934..aacf2514 100644 --- a/src/interlace.rs +++ b/src/interlace.rs @@ -5,7 +5,7 @@ use bit_vec::BitVec; #[must_use] pub fn interlace_image(png: &PngImage) -> PngImage { let mut passes: Vec = 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 diff --git a/src/png/mod.rs b/src/png/mod.rs index f2aa30b9..d599a871 100644 --- a/src/png/mod.rs +++ b/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"), diff --git a/src/png/scan_lines.rs b/src/png/scan_lines.rs index dc923cca..90b92589 100644 --- a/src/png/scan_lines.rs +++ b/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)) } }