use indexmap::IndexSet;
use log::warn;
use rgb::{RGB16, RGBA8};
use crate::{
colors::{BitDepth, ColorType},
deflate::{crc32, inflate},
display_chunks::DISPLAY_CHUNKS,
error::PngError,
interlace::Interlacing,
AtomicMin, Deflaters, PngResult,
};
#[derive(Debug, Clone)]
pub struct IhdrData {
pub width: u32,
pub height: u32,
pub color_type: ColorType,
pub bit_depth: BitDepth,
pub interlaced: Interlacing,
}
impl IhdrData {
#[must_use]
#[inline]
pub fn bpp(&self) -> usize {
self.bit_depth as usize * self.color_type.channels_per_pixel() as usize
}
#[must_use]
pub fn raw_data_size(&self) -> usize {
let w = self.width as usize;
let h = self.height as usize;
let bpp = self.bpp();
fn bitmap_size(bpp: usize, w: usize, h: usize) -> usize {
((w * bpp + 7) / 8) * h
}
if self.interlaced == Interlacing::None {
bitmap_size(bpp, w, h) + h
} else {
let mut size = bitmap_size(bpp, (w + 7) >> 3, (h + 7) >> 3) + ((h + 7) >> 3);
if w > 4 {
size += bitmap_size(bpp, (w + 3) >> 3, (h + 7) >> 3) + ((h + 7) >> 3);
}
size += bitmap_size(bpp, (w + 3) >> 2, (h + 3) >> 3) + ((h + 3) >> 3);
if w > 2 {
size += bitmap_size(bpp, (w + 1) >> 2, (h + 3) >> 2) + ((h + 3) >> 2);
}
size += bitmap_size(bpp, (w + 1) >> 1, (h + 1) >> 2) + ((h + 1) >> 2);
if w > 1 {
size += bitmap_size(bpp, w >> 1, (h + 1) >> 1) + ((h + 1) >> 1);
}
size + bitmap_size(bpp, w, h >> 1) + (h >> 1)
}
}
}
#[derive(Debug, Clone)]
pub struct Chunk {
pub name: [u8; 4],
pub data: Vec<u8>,
}
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum StripChunks {
None,
Strip(IndexSet<[u8; 4]>),
Safe,
Keep(IndexSet<[u8; 4]>),
All,
}
impl StripChunks {
pub(crate) fn keep(&self, name: &[u8; 4]) -> bool {
match &self {
StripChunks::None => true,
StripChunks::Keep(names) => names.contains(name),
StripChunks::Strip(names) => !names.contains(name),
StripChunks::Safe => DISPLAY_CHUNKS.contains(name),
StripChunks::All => false,
}
}
}
#[inline]
pub fn file_header_is_valid(bytes: &[u8]) -> bool {
let expected_header: [u8; 8] = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
*bytes == expected_header
}
#[derive(Debug, Clone, Copy)]
pub struct RawChunk<'a> {
pub name: [u8; 4],
pub data: &'a [u8],
}
pub fn parse_next_chunk<'a>(
byte_data: &'a [u8],
byte_offset: &mut usize,
fix_errors: bool,
) -> PngResult<Option<RawChunk<'a>>> {
let length = read_be_u32(
byte_data
.get(*byte_offset..*byte_offset + 4)
.ok_or(PngError::TruncatedData)?,
);
*byte_offset += 4;
let chunk_start = *byte_offset;
let chunk_name = byte_data
.get(chunk_start..chunk_start + 4)
.ok_or(PngError::TruncatedData)?;
if chunk_name == b"IEND" {
return Ok(None);
}
*byte_offset += 4;
let data = byte_data
.get(*byte_offset..*byte_offset + length as usize)
.ok_or(PngError::TruncatedData)?;
*byte_offset += length as usize;
let crc = read_be_u32(
byte_data
.get(*byte_offset..*byte_offset + 4)
.ok_or(PngError::TruncatedData)?,
);
*byte_offset += 4;
let chunk_bytes = byte_data
.get(chunk_start..chunk_start + 4 + length as usize)
.ok_or(PngError::TruncatedData)?;
if !fix_errors && crc32(chunk_bytes) != crc {
return Err(PngError::new(&format!(
"CRC Mismatch in {} chunk; May be recoverable by using --fix",
String::from_utf8_lossy(chunk_name)
)));
}
let name: [u8; 4] = chunk_name.try_into().unwrap();
Ok(Some(RawChunk { name, data }))
}
pub fn parse_ihdr_chunk(
byte_data: &[u8],
palette_data: Option<Vec<u8>>,
trns_data: Option<Vec<u8>>,
) -> PngResult<IhdrData> {
let interlaced = byte_data.get(12).copied().ok_or(PngError::TruncatedData)?;
Ok(IhdrData {
color_type: match byte_data[9] {
0 => ColorType::Grayscale {
transparent_shade: trns_data
.filter(|t| t.len() >= 2)
.map(|t| u16::from_be_bytes([t[0], t[1]])),
},
2 => ColorType::RGB {
transparent_color: trns_data.filter(|t| t.len() >= 6).map(|t| RGB16 {
r: u16::from_be_bytes([t[0], t[1]]),
g: u16::from_be_bytes([t[2], t[3]]),
b: u16::from_be_bytes([t[4], t[5]]),
}),
},
3 => ColorType::Indexed {
palette: palette_to_rgba(palette_data, trns_data).unwrap_or_default(),
},
4 => ColorType::GrayscaleAlpha,
6 => ColorType::RGBA,
_ => return Err(PngError::new("Unexpected color type in header")),
},
bit_depth: byte_data[8].try_into()?,
width: read_be_u32(&byte_data[0..4]),
height: read_be_u32(&byte_data[4..8]),
interlaced: interlaced.try_into()?,
})
}
fn palette_to_rgba(
palette_data: Option<Vec<u8>>,
trns_data: Option<Vec<u8>>,
) -> Result<Vec<RGBA8>, PngError> {
let palette_data = palette_data.ok_or_else(|| PngError::new("no palette in indexed image"))?;
let mut palette: Vec<_> = palette_data
.chunks(3)
.map(|color| RGBA8::new(color[0], color[1], color[2], 255))
.collect();
if let Some(trns_data) = trns_data {
for (color, trns) in palette.iter_mut().zip(trns_data) {
color.a = trns;
}
}
Ok(palette)
}
#[inline]
fn read_be_u32(bytes: &[u8]) -> u32 {
u32::from_be_bytes(bytes.try_into().unwrap())
}
pub fn extract_icc(iccp: &Chunk) -> Option<Vec<u8>> {
let mut data = iccp.data.as_slice();
loop {
let (&n, rest) = data.split_first()?;
data = rest;
if n == 0 {
break;
}
}
let (&compression_method, compressed_data) = data.split_first()?;
if compression_method != 0 {
return None; }
let max_size = compressed_data.len() * 2 + 1000;
match inflate(compressed_data, max_size) {
Ok(icc) => Some(icc),
Err(e) => {
warn!("Failed to decompress icc: {}", e);
None
}
}
}
pub fn construct_iccp(icc: &[u8], deflater: Deflaters) -> PngResult<Chunk> {
let mut compressed = deflater.deflate(icc, &AtomicMin::new(None))?;
let mut data = Vec::with_capacity(compressed.len() + 5);
data.extend(b"icc"); data.extend([0, 0]); data.append(&mut compressed);
Ok(Chunk {
name: *b"iCCP",
data,
})
}
pub fn srgb_rendering_intent(icc_data: &[u8]) -> Option<u8> {
let rendering_intent = *icc_data.get(67)?;
match icc_data.get(84..100)? {
b"\x29\xf8\x3d\xde\xaf\xf2\x55\xae\x78\x42\xfa\xe4\xca\x83\x39\x0d"
| b"\xc9\x5b\xd6\x37\xe9\x5d\x8a\x3b\x0d\xf3\x8f\x99\xc1\x32\x03\x89"
| b"\xfc\x66\x33\x78\x37\xe2\x88\x6b\xfd\x72\xe9\x83\x82\x28\xf1\xb8"
| b"\x34\x56\x2a\xbf\x99\x4c\xcd\x06\x6d\x2c\x57\x21\xd0\xd6\x8c\x5d" => {
Some(rendering_intent)
}
b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" => {
match (crc32(icc_data), icc_data.len()) {
(0x5d51_29ce, 3024) | (0x182e_a552, 3144) | (0xf29e_526d, 3144) => {
Some(rendering_intent)
}
_ => None,
}
}
_ => None,
}
}