Parse without temporary allocation (#119)

This commit is contained in:
Kornel 2018-07-15 14:11:26 +01:00 committed by Josh Holmer
parent ebd94934a5
commit efde5747c4
3 changed files with 45 additions and 76 deletions

View file

@ -6,6 +6,11 @@ use std::fmt;
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub enum PngError { pub enum PngError {
DeflatedDataTooLong(usize), DeflatedDataTooLong(usize),
NotPNG,
APNGNotSupported,
InvalidData,
TruncatedData,
ChunkMissing(&'static str),
Other(Box<str>), Other(Box<str>),
#[doc(hidden)] #[doc(hidden)]
_Nonexhaustive, _Nonexhaustive,
@ -23,6 +28,11 @@ impl fmt::Display for PngError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match *self { match *self {
PngError::DeflatedDataTooLong(_) => f.write_str("deflated data too long"), 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::Other(ref s) => f.write_str(s),
PngError::_Nonexhaustive => unreachable!(), PngError::_Nonexhaustive => unreachable!(),
} }

View file

@ -43,70 +43,38 @@ pub fn file_header_is_valid(bytes: &[u8]) -> bool {
*bytes == expected_header *bytes == expected_header
} }
pub fn parse_next_header( pub fn parse_next_header<'a>(
byte_data: &[u8], byte_data: &'a [u8],
byte_offset: &mut usize, byte_offset: &mut usize,
fix_errors: bool, fix_errors: bool,
) -> Result<Option<(String, Vec<u8>)>, PngError> { ) -> Result<Option<(&'a [u8], &'a [u8])>, PngError> {
let mut rdr = Cursor::new( let mut rdr = Cursor::new(byte_data.get(*byte_offset..*byte_offset + 4).ok_or(PngError::TruncatedData)?);
byte_data let length = rdr.read_u32::<BigEndian>().unwrap();
.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")),
};
*byte_offset += 4; *byte_offset += 4;
let mut header_bytes: Vec<u8> = byte_data let header_start = *byte_offset;
.iter() let chunk_name = byte_data.get(header_start..header_start + 4).ok_or(PngError::TruncatedData)?;
.skip(*byte_offset) if chunk_name == b"IEND" {
.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" {
// End of data // End of data
return Ok(None); return Ok(None);
} }
*byte_offset += 4; *byte_offset += 4;
let data: Vec<u8> = byte_data let data = byte_data.get(*byte_offset..*byte_offset + length as usize).ok_or(PngError::TruncatedData)?;
.iter()
.skip(*byte_offset)
.take(length as usize)
.cloned()
.collect();
*byte_offset += length as usize; *byte_offset += length as usize;
let mut rdr = Cursor::new( let mut rdr = Cursor::new(byte_data.get(*byte_offset..*byte_offset + 4).ok_or(PngError::TruncatedData)?);
byte_data let crc = rdr.read_u32::<BigEndian>().unwrap();
.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")),
};
*byte_offset += 4; *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!( return Err(PngError::new(&format!(
"CRC Mismatch in {} header; May be recoverable by using --fix", "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> { pub fn parse_ihdr_header(byte_data: &[u8]) -> Result<IhdrData, PngError> {

View file

@ -87,47 +87,39 @@ impl PngData {
pub fn from_slice(byte_data: &[u8], fix_errors: bool) -> Result<PngData, PngError> { pub fn from_slice(byte_data: &[u8], fix_errors: bool) -> Result<PngData, PngError> {
let mut byte_offset: usize = 0; let mut byte_offset: usize = 0;
// Test that png header is valid // Test that png header is valid
let header: Vec<u8> = byte_data.iter().take(8).cloned().collect(); let header = byte_data.get(0..8).ok_or(PngError::TruncatedData)?;
if !file_header_is_valid(header.as_ref()) { if !file_header_is_valid(header) {
return Err(PngError::new("Invalid PNG header detected")); return Err(PngError::NotPNG);
} }
byte_offset += 8; byte_offset += 8;
// Read the data headers // Read the data headers
let mut aux_headers: HashMap<String, Vec<u8>> = HashMap::new(); let mut aux_headers: HashMap<String, Vec<u8>> = HashMap::new();
let mut idat_headers: Vec<u8> = Vec::new(); let mut idat_headers: Vec<u8> = Vec::new();
loop { loop {
let header = parse_next_header(byte_data, &mut byte_offset, fix_errors); let header = parse_next_header(byte_data, &mut byte_offset, fix_errors)?;
let header = match header { let (name, data) = match header {
Ok(x) => x,
Err(x) => return Err(x),
};
let header = match header {
Some(x) => x, Some(x) => x,
None => break, None => break,
}; };
if header.0 == "IDAT" { match name {
idat_headers.extend(header.1); b"IDAT" => idat_headers.extend(data),
} else if header.0 == "acTL" { b"acTL" => return Err(PngError::APNGNotSupported),
return Err(PngError::new("APNG files are not (yet) supported")); _ => {
} else { let name = String::from_utf8(name.to_owned()).map_err(|_| PngError::InvalidData)?;
aux_headers.insert(header.0, header.1); aux_headers.insert(name, data.to_owned());
},
} }
} }
// Parse the headers into our PngData // Parse the headers into our PngData
if idat_headers.is_empty() { 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() { let ihdr = match aux_headers.remove("IHDR") {
return Err(PngError::new("Image header data was missing, skipping")); Some(ihdr) => ihdr,
} None => return Err(PngError::ChunkMissing("IHDR")),
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_header = parse_ihdr_header(&ihdr)?;
let raw_data = deflate::inflate(idat_headers.as_ref())?;
// Handle transparency header // Handle transparency header
let mut has_transparency_pixel = false; let mut has_transparency_pixel = false;
let mut has_transparency_palette = false; let mut has_transparency_palette = false;
@ -160,8 +152,7 @@ impl PngData {
Ok(png_data) Ok(png_data)
} }
#[doc(hidden)] pub(crate) fn reset_from_original(&mut self, original: &PngData) {
pub fn reset_from_original(&mut self, original: &PngData) {
self.idat_data = original.idat_data.clone(); self.idat_data = original.idat_data.clone();
self.ihdr_data = original.ihdr_data; self.ihdr_data = original.ihdr_data;
self.raw_data = original.raw_data.clone(); self.raw_data = original.raw_data.clone();