Parse without temporary allocation (#119)

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<str>),
     #[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!(),
         }

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<Option<(String, Vec<u8>)>, PngError> {
-    let mut rdr = Cursor::new(
-        byte_data
-            .iter()
-            .skip(*byte_offset)
-            .take(4)
-            .cloned()
-            .collect::<Vec<u8>>(),
-    );
-    let length: u32 = match rdr.read_u32::<BigEndian>() {
-        Ok(x) => x,
-        Err(_) => return Err(PngError::new("Invalid data found; unable to read PNG file")),
-    };
+    ) -> Result<Option<(&'a [u8], &'a [u8])>, PngError> {
+    let mut rdr = Cursor::new(byte_data.get(*byte_offset..*byte_offset + 4).ok_or(PngError::TruncatedData)?);
+    let length = rdr.read_u32::<BigEndian>().unwrap();
     *byte_offset += 4;
 
-    let mut header_bytes: Vec<u8> = 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<u8> = 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::<Vec<u8>>(),
-    );
-    let crc: u32 = match rdr.read_u32::<BigEndian>() {
-        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::<BigEndian>().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<IhdrData, PngError> {

src/png/mod.rs

@@ -87,47 +87,39 @@ impl PngData {
     pub fn from_slice(byte_data: &[u8], fix_errors: bool) -> Result<PngData, PngError> {
         let mut byte_offset: usize = 0;
         // Test that png header is valid
-        let header: Vec<u8> = 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<String, Vec<u8>> = HashMap::new();
         let mut idat_headers: Vec<u8> = 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();