Chunk name as 4 bytes (#135)

This commit is contained in:
Kornel 2018-09-16 21:09:55 +01:00 committed by Josh Holmer
parent edd4966cea
commit 3c466df80e
7 changed files with 57 additions and 65 deletions

View file

@ -49,7 +49,7 @@ pub fn file_header_is_valid(bytes: &[u8]) -> bool {
#[derive(Debug, Clone, Copy)] #[derive(Debug, Clone, Copy)]
pub struct RawHeader<'a> { pub struct RawHeader<'a> {
pub name: &'a [u8], pub name: [u8; 4],
pub data: &'a [u8], pub data: &'a [u8],
} }
@ -98,10 +98,9 @@ pub fn parse_next_header<'a>(
))); )));
} }
Ok(Some(RawHeader { let mut name = [0u8; 4];
name: chunk_name, name.copy_from_slice(chunk_name);
data, Ok(Some(RawHeader {name, data}))
}))
} }
pub fn parse_ihdr_header(byte_data: &[u8]) -> PngResult<IhdrData> { pub fn parse_ihdr_header(byte_data: &[u8]) -> PngResult<IhdrData> {

View file

@ -828,21 +828,16 @@ fn perform_strip(png: &mut PngData, opts: &Options) {
// Strip headers // Strip headers
Headers::None => (), Headers::None => (),
Headers::Keep(ref hdrs) => { Headers::Keep(ref hdrs) => {
png.aux_headers.retain(|chunk, _| hdrs.contains(chunk)); png.aux_headers.retain(|chunk, _| std::str::from_utf8(chunk).ok().map_or(false, |name| hdrs.contains(name)));
} }
Headers::Strip(ref hdrs) => for hdr in hdrs { Headers::Strip(ref hdrs) => for hdr in hdrs {
png.aux_headers.remove(hdr); png.aux_headers.remove(hdr.as_bytes());
}, },
Headers::Safe => { Headers::Safe => {
const PRESERVED_HEADERS: [&str; 9] = [ const PRESERVED_HEADERS: [[u8; 4]; 9] = [
"cHRM", "gAMA", "iCCP", "sBIT", "sRGB", "bKGD", "hIST", "pHYs", "sPLT", *b"cHRM", *b"gAMA", *b"iCCP", *b"sBIT", *b"sRGB", *b"bKGD", *b"hIST", *b"pHYs", *b"sPLT",
]; ];
let hdrs = png.aux_headers.keys().cloned().collect::<Vec<String>>(); png.aux_headers.retain(|hdr, _| PRESERVED_HEADERS.contains(hdr));
for hdr in hdrs {
if !PRESERVED_HEADERS.contains(&hdr.as_ref()) {
png.aux_headers.remove(&hdr);
}
}
} }
Headers::All => { Headers::All => {
png.aux_headers = HashMap::new(); png.aux_headers = HashMap::new();

View file

@ -453,9 +453,9 @@ fn parse_opts_into_struct(
opts.strip = Headers::All; opts.strip = Headers::All;
} }
} else { } else {
const FORBIDDEN_CHUNKS: [&str; 5] = ["IHDR", "IDAT", "tRNS", "PLTE", "IEND"]; const FORBIDDEN_CHUNKS: [[u8; 4]; 5] = [*b"IHDR", *b"IDAT", *b"tRNS", *b"PLTE", *b"IEND"];
for i in &hdrs { for i in &hdrs {
if FORBIDDEN_CHUNKS.contains(&i.as_ref()) { if FORBIDDEN_CHUNKS.iter().any(|chunk| chunk == i.as_bytes()) {
return Err(format!("{} chunk is not allowed to be stripped", i)); return Err(format!("{} chunk is not allowed to be stripped", i));
} }
} }

View file

@ -45,7 +45,7 @@ pub struct PngData {
/// A map of how transparent each color in the palette should be /// A map of how transparent each color in the palette should be
pub transparency_palette: Option<Vec<u8>>, pub transparency_palette: Option<Vec<u8>>,
/// All non-critical headers from the PNG are stored here /// All non-critical headers from the PNG are stored here
pub aux_headers: HashMap<String, Vec<u8>>, pub aux_headers: HashMap<[u8; 4], Vec<u8>>,
} }
impl PngData { impl PngData {
@ -93,16 +93,14 @@ impl PngData {
} }
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<[u8; 4], Vec<u8>> = HashMap::new();
let mut idat_headers: Vec<u8> = Vec::new(); let mut idat_headers: Vec<u8> = Vec::new();
while let Some(header) = parse_next_header(byte_data, &mut byte_offset, fix_errors)? { while let Some(header) = parse_next_header(byte_data, &mut byte_offset, fix_errors)? {
match header.name { match &header.name {
b"IDAT" => idat_headers.extend(header.data), b"IDAT" => idat_headers.extend(header.data),
b"acTL" => return Err(PngError::APNGNotSupported), b"acTL" => return Err(PngError::APNGNotSupported),
_ => { _ => {
let name = String::from_utf8(header.name.to_owned()) aux_headers.insert(header.name, header.data.to_owned());
.map_err(|_| PngError::InvalidData)?;
aux_headers.insert(name, header.data.to_owned());
} }
} }
} }
@ -110,7 +108,7 @@ impl PngData {
if idat_headers.is_empty() { if idat_headers.is_empty() {
return Err(PngError::ChunkMissing("IDAT")); return Err(PngError::ChunkMissing("IDAT"));
} }
let ihdr = match aux_headers.remove("IHDR") { let ihdr = match aux_headers.remove(b"IHDR") {
Some(ihdr) => ihdr, Some(ihdr) => ihdr,
None => return Err(PngError::ChunkMissing("IHDR")), None => return Err(PngError::ChunkMissing("IHDR")),
}; };
@ -119,7 +117,7 @@ impl PngData {
// 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;
if aux_headers.contains_key("tRNS") { if aux_headers.contains_key(b"tRNS") {
if ihdr_header.color_type == ColorType::Indexed { if ihdr_header.color_type == ColorType::Indexed {
has_transparency_palette = true; has_transparency_palette = true;
} else { } else {
@ -130,14 +128,14 @@ impl PngData {
idat_data: idat_headers, idat_data: idat_headers,
ihdr_data: ihdr_header, ihdr_data: ihdr_header,
raw_data, raw_data,
palette: aux_headers.remove("PLTE"), palette: aux_headers.remove(b"PLTE"),
transparency_pixel: if has_transparency_pixel { transparency_pixel: if has_transparency_pixel {
aux_headers.remove("tRNS") aux_headers.remove(b"tRNS")
} else { } else {
None None
}, },
transparency_palette: if has_transparency_palette { transparency_palette: if has_transparency_palette {
aux_headers.remove("tRNS") aux_headers.remove(b"tRNS")
} else { } else {
None None
}, },
@ -182,9 +180,9 @@ impl PngData {
for (key, header) in self for (key, header) in self
.aux_headers .aux_headers
.iter() .iter()
.filter(|&(key, _)| !(*key == "bKGD" || *key == "hIST" || *key == "tRNS")) .filter(|&(key, _)| !(key == b"bKGD" || key == b"hIST" || key == b"tRNS"))
{ {
write_png_block(key.as_bytes(), header, &mut output); write_png_block(key, header, &mut output);
} }
// Palette // Palette
if let Some(ref palette) = self.palette { if let Some(ref palette) = self.palette {
@ -201,9 +199,9 @@ impl PngData {
for (key, header) in self for (key, header) in self
.aux_headers .aux_headers
.iter() .iter()
.filter(|&(key, _)| *key == "bKGD" || *key == "hIST" || *key == "tRNS") .filter(|&(key, _)| key == b"bKGD" || key == b"hIST" || key == b"tRNS")
{ {
write_png_block(key.as_bytes(), header, &mut output); write_png_block(key, header, &mut output);
} }
// IDAT data // IDAT data
write_png_block(b"IDAT", &self.idat_data, &mut output); write_png_block(b"IDAT", &self.idat_data, &mut output);

View file

@ -28,7 +28,7 @@ pub fn reduce_alpha_channel(png: &mut PngData, channels: u8) -> Option<Vec<u8>>
// sBIT contains information about alpha channel's original depth, // sBIT contains information about alpha channel's original depth,
// and alpha has just been removed // and alpha has just been removed
if let Some(sbit_header) = png.aux_headers.get_mut(&"sBIT".to_string()) { if let Some(sbit_header) = png.aux_headers.get_mut(b"sBIT") {
assert_eq!(sbit_header.len(), channels as usize); assert_eq!(sbit_header.len(), channels as usize);
sbit_header.pop(); sbit_header.pop();
} }

View file

@ -56,12 +56,12 @@ pub fn reduce_rgba_to_grayscale_alpha(png: &mut PngData) -> bool {
} }
} }
if let Some(sbit_header) = png.aux_headers.get_mut(&"sBIT".to_string()) { if let Some(sbit_header) = png.aux_headers.get_mut(b"sBIT") {
assert_eq!(sbit_header.len(), 4); assert_eq!(sbit_header.len(), 4);
sbit_header.remove(1); sbit_header.remove(1);
sbit_header.remove(1); sbit_header.remove(1);
} }
if let Some(bkgd_header) = png.aux_headers.get_mut(&"bKGD".to_string()) { if let Some(bkgd_header) = png.aux_headers.get_mut(b"bKGD") {
assert_eq!(bkgd_header.len(), 6); assert_eq!(bkgd_header.len(), 6);
bkgd_header.truncate(2); bkgd_header.truncate(2);
} }
@ -100,7 +100,7 @@ pub fn reduce_rgba_to_palette(png: &mut PngData) -> bool {
} }
let mut color_palette = Vec::with_capacity( let mut color_palette = Vec::with_capacity(
palette.len() * 3 + if png.aux_headers.contains_key(&"bKGD".to_string()) { palette.len() * 3 + if png.aux_headers.contains_key(b"bKGD") {
6 6
} else { } else {
0 0
@ -123,7 +123,7 @@ pub fn reduce_rgba_to_palette(png: &mut PngData) -> bool {
return false; return false;
} }
if let Some(bkgd_header) = png.aux_headers.get_mut(&"bKGD".to_string()) { if let Some(bkgd_header) = png.aux_headers.get_mut(b"bKGD") {
assert_eq!(bkgd_header.len(), 6); assert_eq!(bkgd_header.len(), 6);
let header_pixels = bkgd_header let header_pixels = bkgd_header
.iter() .iter()
@ -144,7 +144,7 @@ pub fn reduce_rgba_to_palette(png: &mut PngData) -> bool {
*bkgd_header = vec![entry as u8]; *bkgd_header = vec![entry as u8];
} }
} }
if let Some(sbit_header) = png.aux_headers.get_mut(&"sBIT".to_string()) { if let Some(sbit_header) = png.aux_headers.get_mut(b"sBIT") {
assert_eq!(sbit_header.len(), 4); assert_eq!(sbit_header.len(), 4);
sbit_header.pop(); sbit_header.pop();
} }
@ -209,7 +209,7 @@ pub fn reduce_rgb_to_palette(png: &mut PngData) -> bool {
return false; return false;
} }
if let Some(bkgd_header) = png.aux_headers.get_mut(&"bKGD".to_string()) { if let Some(bkgd_header) = png.aux_headers.get_mut(b"bKGD") {
assert_eq!(bkgd_header.len(), 6); assert_eq!(bkgd_header.len(), 6);
let header_pixels = bkgd_header let header_pixels = bkgd_header
.iter() .iter()
@ -281,11 +281,11 @@ pub fn reduce_rgb_to_grayscale(png: &mut PngData) -> bool {
} }
*trns = trns[0..2].to_owned(); *trns = trns[0..2].to_owned();
} }
if let Some(sbit_header) = png.aux_headers.get_mut(&"sBIT".to_string()) { if let Some(sbit_header) = png.aux_headers.get_mut(b"sBIT") {
assert_eq!(sbit_header.len(), 3); assert_eq!(sbit_header.len(), 3);
sbit_header.truncate(1); sbit_header.truncate(1);
} }
if let Some(bkgd_header) = png.aux_headers.get_mut(&"bKGD".to_string()) { if let Some(bkgd_header) = png.aux_headers.get_mut(b"bKGD") {
assert_eq!(bkgd_header.len(), 6); assert_eq!(bkgd_header.len(), 6);
bkgd_header.truncate(2); bkgd_header.truncate(2);
} }

View file

@ -81,9 +81,9 @@ fn strip_headers_list() {
let png = PngData::new(&input, opts.fix_errors).unwrap(); let png = PngData::new(&input, opts.fix_errors).unwrap();
assert!(png.aux_headers.contains_key("tEXt")); assert!(png.aux_headers.contains_key(b"tEXt"));
assert!(png.aux_headers.contains_key("iTXt")); assert!(png.aux_headers.contains_key(b"iTXt"));
assert!(png.aux_headers.contains_key("iCCP")); assert!(png.aux_headers.contains_key(b"iCCP"));
match oxipng::optimize(&InFile::Path(input), &output, &opts) { match oxipng::optimize(&InFile::Path(input), &output, &opts) {
Ok(_) => (), Ok(_) => (),
@ -100,9 +100,9 @@ fn strip_headers_list() {
} }
}; };
assert!(!png.aux_headers.contains_key("tEXt")); assert!(!png.aux_headers.contains_key(b"tEXt"));
assert!(png.aux_headers.contains_key("iTXt")); assert!(png.aux_headers.contains_key(b"iTXt"));
assert!(!png.aux_headers.contains_key("iCCP")); assert!(!png.aux_headers.contains_key(b"iCCP"));
remove_file(output).ok(); remove_file(output).ok();
} }
@ -115,9 +115,9 @@ fn strip_headers_safe() {
let png = PngData::new(&input, opts.fix_errors).unwrap(); let png = PngData::new(&input, opts.fix_errors).unwrap();
assert!(png.aux_headers.contains_key("tEXt")); assert!(png.aux_headers.contains_key(b"tEXt"));
assert!(png.aux_headers.contains_key("iTXt")); assert!(png.aux_headers.contains_key(b"iTXt"));
assert!(png.aux_headers.contains_key("iCCP")); assert!(png.aux_headers.contains_key(b"iCCP"));
match oxipng::optimize(&InFile::Path(input), &output, &opts) { match oxipng::optimize(&InFile::Path(input), &output, &opts) {
Ok(_) => (), Ok(_) => (),
@ -134,9 +134,9 @@ fn strip_headers_safe() {
} }
}; };
assert!(!png.aux_headers.contains_key("tEXt")); assert!(!png.aux_headers.contains_key(b"tEXt"));
assert!(!png.aux_headers.contains_key("iTXt")); assert!(!png.aux_headers.contains_key(b"iTXt"));
assert!(png.aux_headers.contains_key("iCCP")); assert!(png.aux_headers.contains_key(b"iCCP"));
remove_file(output).ok(); remove_file(output).ok();
} }
@ -149,9 +149,9 @@ fn strip_headers_all() {
let png = PngData::new(&input, opts.fix_errors).unwrap(); let png = PngData::new(&input, opts.fix_errors).unwrap();
assert!(png.aux_headers.contains_key("tEXt")); assert!(png.aux_headers.contains_key(b"tEXt"));
assert!(png.aux_headers.contains_key("iTXt")); assert!(png.aux_headers.contains_key(b"iTXt"));
assert!(png.aux_headers.contains_key("iCCP")); assert!(png.aux_headers.contains_key(b"iCCP"));
match oxipng::optimize(&InFile::Path(input), &output, &opts) { match oxipng::optimize(&InFile::Path(input), &output, &opts) {
Ok(_) => (), Ok(_) => (),
@ -168,9 +168,9 @@ fn strip_headers_all() {
} }
}; };
assert!(!png.aux_headers.contains_key("tEXt")); assert!(!png.aux_headers.contains_key(b"tEXt"));
assert!(!png.aux_headers.contains_key("iTXt")); assert!(!png.aux_headers.contains_key(b"iTXt"));
assert!(!png.aux_headers.contains_key("iCCP")); assert!(!png.aux_headers.contains_key(b"iCCP"));
remove_file(output).ok(); remove_file(output).ok();
} }
@ -183,9 +183,9 @@ fn strip_headers_none() {
let png = PngData::new(&input, opts.fix_errors).unwrap(); let png = PngData::new(&input, opts.fix_errors).unwrap();
assert!(png.aux_headers.contains_key("tEXt")); assert!(png.aux_headers.contains_key(b"tEXt"));
assert!(png.aux_headers.contains_key("iTXt")); assert!(png.aux_headers.contains_key(b"iTXt"));
assert!(png.aux_headers.contains_key("iCCP")); assert!(png.aux_headers.contains_key(b"iCCP"));
match oxipng::optimize(&InFile::Path(input), &output, &opts) { match oxipng::optimize(&InFile::Path(input), &output, &opts) {
Ok(_) => (), Ok(_) => (),
@ -202,9 +202,9 @@ fn strip_headers_none() {
} }
}; };
assert!(png.aux_headers.contains_key("tEXt")); assert!(png.aux_headers.contains_key(b"tEXt"));
assert!(png.aux_headers.contains_key("iTXt")); assert!(png.aux_headers.contains_key(b"iTXt"));
assert!(png.aux_headers.contains_key("iCCP")); assert!(png.aux_headers.contains_key(b"iCCP"));
remove_file(output).ok(); remove_file(output).ok();
} }