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)]
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<IhdrData> {

View file

@ -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::<Vec<String>>();
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();

View file

@ -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));
}
}

View file

@ -45,7 +45,7 @@ pub struct PngData {
/// A map of how transparent each color in the palette should be
pub transparency_palette: Option<Vec<u8>>,
/// 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 {
@ -93,16 +93,14 @@ impl PngData {
}
byte_offset += 8;
// 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();
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);

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,
// 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();
}

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);
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);
}

View file

@ -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();
}