diff --git a/src/headers.rs b/src/headers.rs index 1ac462b7..3f46768f 100644 --- a/src/headers.rs +++ b/src/headers.rs @@ -49,7 +49,7 @@ pub fn file_header_is_valid(bytes: &[u8]) -> bool { #[derive(Debug, Clone, Copy)] pub struct RawHeader<'a> { - pub name: &'a [u8], + pub name: [u8; 4], pub data: &'a [u8], } @@ -98,10 +98,9 @@ pub fn parse_next_header<'a>( ))); } - Ok(Some(RawHeader { - name: chunk_name, - data, - })) + let mut name = [0u8; 4]; + name.copy_from_slice(chunk_name); + Ok(Some(RawHeader {name, data})) } pub fn parse_ihdr_header(byte_data: &[u8]) -> PngResult { diff --git a/src/lib.rs b/src/lib.rs index ce9a1764..4d8240da 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -828,21 +828,16 @@ fn perform_strip(png: &mut PngData, opts: &Options) { // Strip headers Headers::None => (), 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 { - png.aux_headers.remove(hdr); + png.aux_headers.remove(hdr.as_bytes()); }, Headers::Safe => { - const PRESERVED_HEADERS: [&str; 9] = [ - "cHRM", "gAMA", "iCCP", "sBIT", "sRGB", "bKGD", "hIST", "pHYs", "sPLT", + const PRESERVED_HEADERS: [[u8; 4]; 9] = [ + *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::>(); - for hdr in hdrs { - if !PRESERVED_HEADERS.contains(&hdr.as_ref()) { - png.aux_headers.remove(&hdr); - } - } + png.aux_headers.retain(|hdr, _| PRESERVED_HEADERS.contains(hdr)); } Headers::All => { png.aux_headers = HashMap::new(); diff --git a/src/main.rs b/src/main.rs index 6c9ebf39..48b9a4bc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -453,9 +453,9 @@ fn parse_opts_into_struct( opts.strip = Headers::All; } } 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 { - 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)); } } diff --git a/src/png/mod.rs b/src/png/mod.rs index eb1cbdcc..56c682e1 100644 --- a/src/png/mod.rs +++ b/src/png/mod.rs @@ -45,7 +45,7 @@ pub struct PngData { /// A map of how transparent each color in the palette should be pub transparency_palette: Option>, /// All non-critical headers from the PNG are stored here - pub aux_headers: HashMap>, + pub aux_headers: HashMap<[u8; 4], Vec>, } impl PngData { @@ -93,16 +93,14 @@ impl PngData { } byte_offset += 8; // Read the data headers - let mut aux_headers: HashMap> = HashMap::new(); + let mut aux_headers: HashMap<[u8; 4], Vec> = HashMap::new(); let mut idat_headers: Vec = Vec::new(); 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"acTL" => return Err(PngError::APNGNotSupported), _ => { - let name = String::from_utf8(header.name.to_owned()) - .map_err(|_| PngError::InvalidData)?; - aux_headers.insert(name, header.data.to_owned()); + aux_headers.insert(header.name, header.data.to_owned()); } } } @@ -110,7 +108,7 @@ impl PngData { if idat_headers.is_empty() { return Err(PngError::ChunkMissing("IDAT")); } - let ihdr = match aux_headers.remove("IHDR") { + let ihdr = match aux_headers.remove(b"IHDR") { Some(ihdr) => ihdr, None => return Err(PngError::ChunkMissing("IHDR")), }; @@ -119,7 +117,7 @@ impl PngData { // Handle transparency header let mut has_transparency_pixel = 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 { has_transparency_palette = true; } else { @@ -130,14 +128,14 @@ impl PngData { idat_data: idat_headers, ihdr_data: ihdr_header, raw_data, - palette: aux_headers.remove("PLTE"), + palette: aux_headers.remove(b"PLTE"), transparency_pixel: if has_transparency_pixel { - aux_headers.remove("tRNS") + aux_headers.remove(b"tRNS") } else { None }, transparency_palette: if has_transparency_palette { - aux_headers.remove("tRNS") + aux_headers.remove(b"tRNS") } else { None }, @@ -182,9 +180,9 @@ impl PngData { for (key, header) in self .aux_headers .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 if let Some(ref palette) = self.palette { @@ -201,9 +199,9 @@ impl PngData { for (key, header) in self .aux_headers .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 write_png_block(b"IDAT", &self.idat_data, &mut output); diff --git a/src/reduction/alpha.rs b/src/reduction/alpha.rs index ebe062d8..658b8c28 100644 --- a/src/reduction/alpha.rs +++ b/src/reduction/alpha.rs @@ -28,7 +28,7 @@ pub fn reduce_alpha_channel(png: &mut PngData, channels: u8) -> Option> // sBIT contains information about alpha channel's original depth, // 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); sbit_header.pop(); } diff --git a/src/reduction/color.rs b/src/reduction/color.rs index a1be2552..02712baa 100644 --- a/src/reduction/color.rs +++ b/src/reduction/color.rs @@ -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); 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); 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( - palette.len() * 3 + if png.aux_headers.contains_key(&"bKGD".to_string()) { + palette.len() * 3 + if png.aux_headers.contains_key(b"bKGD") { 6 } else { 0 @@ -123,7 +123,7 @@ pub fn reduce_rgba_to_palette(png: &mut PngData) -> bool { 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); let header_pixels = bkgd_header .iter() @@ -144,7 +144,7 @@ pub fn reduce_rgba_to_palette(png: &mut PngData) -> bool { *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); sbit_header.pop(); } @@ -209,7 +209,7 @@ pub fn reduce_rgb_to_palette(png: &mut PngData) -> bool { 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); let header_pixels = bkgd_header .iter() @@ -281,11 +281,11 @@ pub fn reduce_rgb_to_grayscale(png: &mut PngData) -> bool { } *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); 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); bkgd_header.truncate(2); } diff --git a/tests/flags.rs b/tests/flags.rs index 049e88b7..0270a9c2 100644 --- a/tests/flags.rs +++ b/tests/flags.rs @@ -81,9 +81,9 @@ fn strip_headers_list() { let png = PngData::new(&input, opts.fix_errors).unwrap(); - assert!(png.aux_headers.contains_key("tEXt")); - assert!(png.aux_headers.contains_key("iTXt")); - assert!(png.aux_headers.contains_key("iCCP")); + assert!(png.aux_headers.contains_key(b"tEXt")); + assert!(png.aux_headers.contains_key(b"iTXt")); + assert!(png.aux_headers.contains_key(b"iCCP")); match oxipng::optimize(&InFile::Path(input), &output, &opts) { Ok(_) => (), @@ -100,9 +100,9 @@ fn strip_headers_list() { } }; - assert!(!png.aux_headers.contains_key("tEXt")); - assert!(png.aux_headers.contains_key("iTXt")); - assert!(!png.aux_headers.contains_key("iCCP")); + assert!(!png.aux_headers.contains_key(b"tEXt")); + assert!(png.aux_headers.contains_key(b"iTXt")); + assert!(!png.aux_headers.contains_key(b"iCCP")); remove_file(output).ok(); } @@ -115,9 +115,9 @@ fn strip_headers_safe() { let png = PngData::new(&input, opts.fix_errors).unwrap(); - assert!(png.aux_headers.contains_key("tEXt")); - assert!(png.aux_headers.contains_key("iTXt")); - assert!(png.aux_headers.contains_key("iCCP")); + assert!(png.aux_headers.contains_key(b"tEXt")); + assert!(png.aux_headers.contains_key(b"iTXt")); + assert!(png.aux_headers.contains_key(b"iCCP")); match oxipng::optimize(&InFile::Path(input), &output, &opts) { Ok(_) => (), @@ -134,9 +134,9 @@ fn strip_headers_safe() { } }; - assert!(!png.aux_headers.contains_key("tEXt")); - assert!(!png.aux_headers.contains_key("iTXt")); - assert!(png.aux_headers.contains_key("iCCP")); + assert!(!png.aux_headers.contains_key(b"tEXt")); + assert!(!png.aux_headers.contains_key(b"iTXt")); + assert!(png.aux_headers.contains_key(b"iCCP")); remove_file(output).ok(); } @@ -149,9 +149,9 @@ fn strip_headers_all() { let png = PngData::new(&input, opts.fix_errors).unwrap(); - assert!(png.aux_headers.contains_key("tEXt")); - assert!(png.aux_headers.contains_key("iTXt")); - assert!(png.aux_headers.contains_key("iCCP")); + assert!(png.aux_headers.contains_key(b"tEXt")); + assert!(png.aux_headers.contains_key(b"iTXt")); + assert!(png.aux_headers.contains_key(b"iCCP")); match oxipng::optimize(&InFile::Path(input), &output, &opts) { Ok(_) => (), @@ -168,9 +168,9 @@ fn strip_headers_all() { } }; - assert!(!png.aux_headers.contains_key("tEXt")); - assert!(!png.aux_headers.contains_key("iTXt")); - assert!(!png.aux_headers.contains_key("iCCP")); + assert!(!png.aux_headers.contains_key(b"tEXt")); + assert!(!png.aux_headers.contains_key(b"iTXt")); + assert!(!png.aux_headers.contains_key(b"iCCP")); remove_file(output).ok(); } @@ -183,9 +183,9 @@ fn strip_headers_none() { let png = PngData::new(&input, opts.fix_errors).unwrap(); - assert!(png.aux_headers.contains_key("tEXt")); - assert!(png.aux_headers.contains_key("iTXt")); - assert!(png.aux_headers.contains_key("iCCP")); + assert!(png.aux_headers.contains_key(b"tEXt")); + assert!(png.aux_headers.contains_key(b"iTXt")); + assert!(png.aux_headers.contains_key(b"iCCP")); match oxipng::optimize(&InFile::Path(input), &output, &opts) { Ok(_) => (), @@ -202,9 +202,9 @@ fn strip_headers_none() { } }; - assert!(png.aux_headers.contains_key("tEXt")); - assert!(png.aux_headers.contains_key("iTXt")); - assert!(png.aux_headers.contains_key("iCCP")); + assert!(png.aux_headers.contains_key(b"tEXt")); + assert!(png.aux_headers.contains_key(b"iTXt")); + assert!(png.aux_headers.contains_key(b"iCCP")); remove_file(output).ok(); }