diff --git a/src/error.rs b/src/error.rs index 6dcde42b..157bd7b2 100644 --- a/src/error.rs +++ b/src/error.rs @@ -6,6 +6,11 @@ use std::fmt; #[derive(Debug, Clone)] pub enum PngError { DeflatedDataTooLong(usize), + NotPNG, + APNGNotSupported, + InvalidData, + TruncatedData, + ChunkMissing(&'static str), Other(Box), #[doc(hidden)] _Nonexhaustive, @@ -23,6 +28,11 @@ impl fmt::Display for PngError { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match *self { PngError::DeflatedDataTooLong(_) => f.write_str("deflated data too long"), + PngError::NotPNG => f.write_str("Invalid header detected; Not a PNG file"), + PngError::InvalidData => f.write_str("Invalid data found; unable to read PNG file"), + PngError::TruncatedData => f.write_str("Missing data in the file; the file is truncated"), + PngError::APNGNotSupported => f.write_str("APNG files are not (yet) supported"), + PngError::ChunkMissing(s) => write!(f, "Chunk {} missing or empty", s), PngError::Other(ref s) => f.write_str(s), PngError::_Nonexhaustive => unreachable!(), } diff --git a/src/headers.rs b/src/headers.rs index fff2847d..d08c4a11 100644 --- a/src/headers.rs +++ b/src/headers.rs @@ -43,70 +43,38 @@ pub fn file_header_is_valid(bytes: &[u8]) -> bool { *bytes == expected_header } -pub fn parse_next_header( - byte_data: &[u8], +pub fn parse_next_header<'a>( + byte_data: &'a [u8], byte_offset: &mut usize, fix_errors: bool, -) -> Result)>, PngError> { - let mut rdr = Cursor::new( - byte_data - .iter() - .skip(*byte_offset) - .take(4) - .cloned() - .collect::>(), - ); - let length: u32 = match rdr.read_u32::() { - Ok(x) => x, - Err(_) => return Err(PngError::new("Invalid data found; unable to read PNG file")), - }; + ) -> Result, PngError> { + let mut rdr = Cursor::new(byte_data.get(*byte_offset..*byte_offset + 4).ok_or(PngError::TruncatedData)?); + let length = rdr.read_u32::().unwrap(); *byte_offset += 4; - let mut header_bytes: Vec = byte_data - .iter() - .skip(*byte_offset) - .take(4) - .cloned() - .collect(); - let header = match String::from_utf8(header_bytes.clone()) { - Ok(x) => x, - Err(_) => return Err(PngError::new("Invalid data found; unable to read PNG file")), - }; - if header == "IEND" { + let header_start = *byte_offset; + let chunk_name = byte_data.get(header_start..header_start + 4).ok_or(PngError::TruncatedData)?; + if chunk_name == b"IEND" { // End of data return Ok(None); } *byte_offset += 4; - let data: Vec = byte_data - .iter() - .skip(*byte_offset) - .take(length as usize) - .cloned() - .collect(); + let data = byte_data.get(*byte_offset..*byte_offset + length as usize).ok_or(PngError::TruncatedData)?; *byte_offset += length as usize; - let mut rdr = Cursor::new( - byte_data - .iter() - .skip(*byte_offset) - .take(4) - .cloned() - .collect::>(), - ); - let crc: u32 = match rdr.read_u32::() { - Ok(x) => x, - Err(_) => return Err(PngError::new("Invalid data found; unable to read PNG file")), - }; + let mut rdr = Cursor::new(byte_data.get(*byte_offset..*byte_offset + 4).ok_or(PngError::TruncatedData)?); + let crc = rdr.read_u32::().unwrap(); *byte_offset += 4; - header_bytes.extend_from_slice(&data); - if !fix_errors && crc32::checksum_ieee(header_bytes.as_ref()) != crc { + + let header_bytes = byte_data.get(header_start..header_start + 4 + length as usize).ok_or(PngError::TruncatedData)?; + if !fix_errors && crc32::checksum_ieee(header_bytes) != crc { return Err(PngError::new(&format!( "CRC Mismatch in {} header; May be recoverable by using --fix", - header + String::from_utf8_lossy(chunk_name) ))); } - Ok(Some((header, data))) + Ok(Some((chunk_name, data))) } pub fn parse_ihdr_header(byte_data: &[u8]) -> Result { diff --git a/src/png/mod.rs b/src/png/mod.rs index 0c1e4c58..d8afa106 100644 --- a/src/png/mod.rs +++ b/src/png/mod.rs @@ -87,47 +87,39 @@ impl PngData { pub fn from_slice(byte_data: &[u8], fix_errors: bool) -> Result { let mut byte_offset: usize = 0; // Test that png header is valid - let header: Vec = byte_data.iter().take(8).cloned().collect(); - if !file_header_is_valid(header.as_ref()) { - return Err(PngError::new("Invalid PNG header detected")); + let header = byte_data.get(0..8).ok_or(PngError::TruncatedData)?; + if !file_header_is_valid(header) { + return Err(PngError::NotPNG); } byte_offset += 8; // Read the data headers let mut aux_headers: HashMap> = HashMap::new(); let mut idat_headers: Vec = Vec::new(); loop { - let header = parse_next_header(byte_data, &mut byte_offset, fix_errors); - let header = match header { - Ok(x) => x, - Err(x) => return Err(x), - }; - let header = match header { + let header = parse_next_header(byte_data, &mut byte_offset, fix_errors)?; + let (name, data) = match header { Some(x) => x, None => break, }; - if header.0 == "IDAT" { - idat_headers.extend(header.1); - } else if header.0 == "acTL" { - return Err(PngError::new("APNG files are not (yet) supported")); - } else { - aux_headers.insert(header.0, header.1); + match name { + b"IDAT" => idat_headers.extend(data), + b"acTL" => return Err(PngError::APNGNotSupported), + _ => { + let name = String::from_utf8(name.to_owned()).map_err(|_| PngError::InvalidData)?; + aux_headers.insert(name, data.to_owned()); + }, } } // Parse the headers into our PngData if idat_headers.is_empty() { - return Err(PngError::new("Image data was empty, skipping")); + return Err(PngError::ChunkMissing("IDAT")); } - if aux_headers.get("IHDR").is_none() { - return Err(PngError::new("Image header data was missing, skipping")); - } - let ihdr_header = match parse_ihdr_header(aux_headers.remove("IHDR").unwrap().as_ref()) { - Ok(x) => x, - Err(x) => return Err(x), - }; - let raw_data = match deflate::inflate(idat_headers.as_ref()) { - Ok(x) => x, - Err(x) => return Err(x), + let ihdr = match aux_headers.remove("IHDR") { + Some(ihdr) => ihdr, + None => return Err(PngError::ChunkMissing("IHDR")), }; + let ihdr_header = parse_ihdr_header(&ihdr)?; + let raw_data = deflate::inflate(idat_headers.as_ref())?; // Handle transparency header let mut has_transparency_pixel = false; let mut has_transparency_palette = false; @@ -160,8 +152,7 @@ impl PngData { Ok(png_data) } - #[doc(hidden)] - pub fn reset_from_original(&mut self, original: &PngData) { + pub(crate) fn reset_from_original(&mut self, original: &PngData) { self.idat_data = original.idat_data.clone(); self.ihdr_data = original.ihdr_data; self.raw_data = original.raw_data.clone();