use std::collections::HashMap;
pub(crate) const SIGNATURE_LEN: usize = 8;
const CHUNK_HEADER_LEN: usize = 8;
const CHUNK_CRC_LEN: usize = 4;
const MAX_PRESERVED_CHUNK_BYTES: usize = 4 * 1024 * 1024;
const MIN_IDAT_CHUNK_SIZE: usize = 512;
const MAX_IDAT_CHUNK_SIZE: usize = 32 * 1024 * 1024;
const DEFAULT_IDAT_CHUNK_SIZE: usize = 8192;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TechnicalChunk {
Gamma,
Srgb,
Chromaticities,
PhysicalDimensions,
IccProfile,
SignificantBits,
}
impl TechnicalChunk {
const ALL: [TechnicalChunk; 6] = [
TechnicalChunk::Gamma,
TechnicalChunk::Srgb,
TechnicalChunk::Chromaticities,
TechnicalChunk::PhysicalDimensions,
TechnicalChunk::IccProfile,
TechnicalChunk::SignificantBits,
];
pub fn type_code(self) -> [u8; 4] {
match self {
TechnicalChunk::Gamma => *b"gAMA",
TechnicalChunk::Srgb => *b"sRGB",
TechnicalChunk::Chromaticities => *b"cHRM",
TechnicalChunk::PhysicalDimensions => *b"pHYs",
TechnicalChunk::IccProfile => *b"iCCP",
TechnicalChunk::SignificantBits => *b"sBIT",
}
}
pub fn name(self) -> &'static str {
match self {
TechnicalChunk::Gamma => "gAMA",
TechnicalChunk::Srgb => "sRGB",
TechnicalChunk::Chromaticities => "cHRM",
TechnicalChunk::PhysicalDimensions => "pHYs",
TechnicalChunk::IccProfile => "iCCP",
TechnicalChunk::SignificantBits => "sBIT",
}
}
fn from_type_code(code: [u8; 4]) -> Option<Self> {
TechnicalChunk::ALL
.into_iter()
.find(|candidate| candidate.type_code() == code)
}
}
#[derive(Debug, Clone)]
pub struct PreservedChunk {
kind: TechnicalChunk,
data: Vec<u8>,
}
impl PreservedChunk {
pub fn kind(&self) -> TechnicalChunk {
self.kind
}
pub fn len(&self) -> usize {
self.data.len()
}
pub fn is_empty(&self) -> bool {
self.data.is_empty()
}
pub(crate) fn data(&self) -> &[u8] {
&self.data
}
}
#[derive(Debug, Clone)]
pub struct PngEnvelope {
chunks: Vec<PreservedChunk>,
idat_chunk_size: usize,
idat_chunk_count: usize,
discarded_chunks: usize,
}
impl PngEnvelope {
pub(crate) fn synthesised() -> Self {
let gamma = 45_455u32.to_be_bytes().to_vec();
let srgb = vec![0u8];
let chromaticities = [
31_270u32, 32_900, 64_000, 33_000, 30_000, 60_000, 15_000, 6_000,
]
.iter()
.flat_map(|value| value.to_be_bytes())
.collect();
let mut physical = Vec::with_capacity(9);
physical.extend_from_slice(&2835u32.to_be_bytes());
physical.extend_from_slice(&2835u32.to_be_bytes());
physical.push(1);
Self {
chunks: vec![
PreservedChunk {
kind: TechnicalChunk::Gamma,
data: gamma,
},
PreservedChunk {
kind: TechnicalChunk::Chromaticities,
data: chromaticities,
},
PreservedChunk {
kind: TechnicalChunk::Srgb,
data: srgb,
},
PreservedChunk {
kind: TechnicalChunk::PhysicalDimensions,
data: physical,
},
],
idat_chunk_size: DEFAULT_IDAT_CHUNK_SIZE,
idat_chunk_count: 0,
discarded_chunks: 0,
}
}
pub(crate) fn read(bytes: &[u8]) -> Self {
let mut envelope = Self {
chunks: Vec::new(),
idat_chunk_size: DEFAULT_IDAT_CHUNK_SIZE,
idat_chunk_count: 0,
discarded_chunks: 0,
};
let mut idat_lengths: Vec<usize> = Vec::new();
let mut offset = SIGNATURE_LEN;
while let Some((code, data, next)) = read_chunk(bytes, offset) {
offset = next;
match &code {
b"IEND" => break,
b"IDAT" => idat_lengths.push(data.len()),
_ => envelope.record(code, data, idat_lengths.is_empty()),
}
}
envelope.idat_chunk_count = idat_lengths.len();
if let Some(size) = dominant_length(&idat_lengths) {
if (MIN_IDAT_CHUNK_SIZE..=MAX_IDAT_CHUNK_SIZE).contains(&size) {
envelope.idat_chunk_size = size;
} else if size > MAX_IDAT_CHUNK_SIZE {
envelope.idat_chunk_size = MAX_IDAT_CHUNK_SIZE;
}
}
envelope
}
fn record(&mut self, code: [u8; 4], data: &[u8], before_pixels: bool) {
let preservable = TechnicalChunk::from_type_code(code)
.filter(|_| before_pixels)
.filter(|_| data.len() <= MAX_PRESERVED_CHUNK_BYTES)
.filter(|kind| !self.chunks.iter().any(|chunk| chunk.kind == *kind));
match preservable {
Some(kind) => self.chunks.push(PreservedChunk {
kind,
data: data.to_vec(),
}),
None => {
if code[0].is_ascii_lowercase() {
self.discarded_chunks += 1;
}
}
}
}
pub fn preserved_chunks(&self) -> &[PreservedChunk] {
&self.chunks
}
pub fn idat_chunk_size(&self) -> usize {
self.idat_chunk_size
}
pub fn idat_chunk_count(&self) -> usize {
self.idat_chunk_count
}
pub fn discarded_chunks(&self) -> usize {
self.discarded_chunks
}
}
pub(crate) fn read_chunk(bytes: &[u8], offset: usize) -> Option<([u8; 4], &[u8], usize)> {
let header = bytes.get(offset..offset.checked_add(CHUNK_HEADER_LEN)?)?;
let length = u32::from_be_bytes(header.get(..4)?.try_into().ok()?);
let code: [u8; 4] = header.get(4..)?.try_into().ok()?;
let length = usize::try_from(length).ok()?;
let start = offset.checked_add(CHUNK_HEADER_LEN)?;
let end = start.checked_add(length)?;
let next = end.checked_add(CHUNK_CRC_LEN)?;
if next > bytes.len() {
return None;
}
Some((code, bytes.get(start..end)?, next))
}
fn dominant_length(lengths: &[usize]) -> Option<usize> {
let mut counts: HashMap<usize, usize> = HashMap::new();
for &length in lengths {
*counts.entry(length).or_insert(0) += 1;
}
counts
.into_iter()
.max_by_key(|&(length, count)| (count, length))
.map(|(length, _)| length)
}
#[cfg(test)]
mod tests {
#![allow(clippy::expect_used)]
#![allow(clippy::panic)]
use super::*;
fn file(chunks: &[(&[u8; 4], Vec<u8>)]) -> Vec<u8> {
let mut bytes = vec![0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A];
for (code, data) in chunks {
bytes.extend_from_slice(&(data.len() as u32).to_be_bytes());
bytes.extend_from_slice(*code);
bytes.extend_from_slice(data);
bytes.extend_from_slice(&[0, 0, 0, 0]);
}
bytes
}
fn ihdr() -> (&'static [u8; 4], Vec<u8>) {
(b"IHDR", vec![0u8; 13])
}
#[test]
fn the_whitelist_is_copied_and_the_rest_is_dropped() {
let bytes = file(&[
ihdr(),
(b"sRGB", vec![0]),
(b"gAMA", 45_455u32.to_be_bytes().to_vec()),
(b"eXIf", vec![9; 400]),
(b"pHYs", vec![1; 9]),
(b"tEXt", b"Software\0Adobe".to_vec()),
(b"iTXt", vec![7; 20]),
(b"tIME", vec![0; 7]),
(b"zTXt", vec![3; 12]),
(b"IDAT", vec![0; 8192]),
(b"IDAT", vec![0; 100]),
(b"IEND", Vec::new()),
]);
let envelope = PngEnvelope::read(&bytes);
let kinds: Vec<TechnicalChunk> = envelope
.preserved_chunks()
.iter()
.map(PreservedChunk::kind)
.collect();
assert_eq!(
kinds,
vec![
TechnicalChunk::Srgb,
TechnicalChunk::Gamma,
TechnicalChunk::PhysicalDimensions,
]
);
assert_eq!(
envelope.preserved_chunks()[1].data(),
&45_455u32.to_be_bytes()
);
assert_eq!(envelope.preserved_chunks()[0].len(), 1);
assert!(!envelope.preserved_chunks()[0].is_empty());
assert_eq!(envelope.discarded_chunks(), 5);
assert_eq!(envelope.idat_chunk_count(), 2);
assert_eq!(envelope.idat_chunk_size(), 8192);
}
#[test]
fn no_identifying_chunk_can_be_recognised() {
for code in [b"eXIf", b"tEXt", b"iTXt", b"zTXt", b"tIME"] {
assert_eq!(
TechnicalChunk::from_type_code(*code),
None,
"{} must never be copyable",
String::from_utf8_lossy(code)
);
}
for kind in TechnicalChunk::ALL {
assert_eq!(
TechnicalChunk::from_type_code(kind.type_code()),
Some(kind),
"{} must round-trip through its type code",
kind.name()
);
assert_eq!(kind.name().as_bytes(), kind.type_code());
}
}
#[test]
fn a_chunk_behind_the_pixels_is_not_copied() {
let bytes = file(&[
ihdr(),
(b"IDAT", vec![0; 4096]),
(b"pHYs", vec![1; 9]),
(b"iTXt", vec![2; 30]),
(b"IEND", Vec::new()),
]);
let envelope = PngEnvelope::read(&bytes);
assert!(envelope.preserved_chunks().is_empty());
assert_eq!(envelope.discarded_chunks(), 2);
}
#[test]
fn a_duplicated_chunk_is_written_once() {
let bytes = file(&[
ihdr(),
(b"gAMA", vec![1, 1, 1, 1]),
(b"gAMA", vec![2, 2, 2, 2]),
(b"IDAT", vec![0; 512]),
(b"IEND", Vec::new()),
]);
let envelope = PngEnvelope::read(&bytes);
assert_eq!(envelope.preserved_chunks().len(), 1);
assert_eq!(envelope.preserved_chunks()[0].data(), &[1, 1, 1, 1]);
assert_eq!(envelope.discarded_chunks(), 1);
}
#[test]
fn an_implausible_profile_is_not_kept() {
let bytes = file(&[
ihdr(),
(b"iCCP", vec![0; MAX_PRESERVED_CHUNK_BYTES + 1]),
(b"IDAT", vec![0; 1024]),
(b"IEND", Vec::new()),
]);
let envelope = PngEnvelope::read(&bytes);
assert!(envelope.preserved_chunks().is_empty());
assert_eq!(envelope.discarded_chunks(), 1);
let bytes = file(&[
ihdr(),
(b"iCCP", vec![0; MAX_PRESERVED_CHUNK_BYTES]),
(b"IDAT", vec![0; 1024]),
(b"IEND", Vec::new()),
]);
assert_eq!(PngEnvelope::read(&bytes).preserved_chunks().len(), 1);
}
#[test]
fn an_absurd_length_ends_the_walk() {
let mut bytes = file(&[ihdr()]);
bytes.extend_from_slice(&u32::MAX.to_be_bytes());
bytes.extend_from_slice(b"iCCP");
bytes.extend_from_slice(&[0; 16]);
let envelope = PngEnvelope::read(&bytes);
assert!(envelope.preserved_chunks().is_empty());
assert_eq!(envelope.idat_chunk_count(), 0);
assert_eq!(envelope.idat_chunk_size(), DEFAULT_IDAT_CHUNK_SIZE);
}
#[test]
fn an_unreadable_file_yields_the_default_profile() {
for bytes in [
Vec::new(),
b"not a png at all".to_vec(),
vec![0x89, b'P', b'N', b'G'],
{
let mut truncated = file(&[ihdr(), (b"gAMA", vec![0; 4])]);
truncated.truncate(truncated.len() - 5);
truncated
},
] {
let envelope = PngEnvelope::read(&bytes);
assert_eq!(envelope.idat_chunk_size(), DEFAULT_IDAT_CHUNK_SIZE);
assert_eq!(envelope.idat_chunk_count(), 0);
}
}
#[test]
fn trailing_bytes_after_the_end_marker_are_ignored() {
let bytes = file(&[
ihdr(),
(b"gAMA", vec![0; 4]),
(b"IDAT", vec![0; 1024]),
(b"IEND", Vec::new()),
(b"pHYs", vec![1; 9]),
(b"IDAT", vec![0; 1024]),
]);
let envelope = PngEnvelope::read(&bytes);
assert_eq!(envelope.preserved_chunks().len(), 1);
assert_eq!(envelope.idat_chunk_count(), 1);
}
#[test]
fn the_idat_split_is_the_length_the_chunks_agree_on() {
let mut chunks = vec![ihdr(), (b"IDAT", vec![0; 65_445])];
chunks.extend((0..4).map(|_| (b"IDAT", vec![0; 65_524])));
chunks.push((b"IDAT", vec![0; 62_588]));
chunks.push((b"IEND", Vec::new()));
let envelope = PngEnvelope::read(&file(&chunks));
assert_eq!(envelope.idat_chunk_size(), 65_524);
assert_eq!(envelope.idat_chunk_count(), 6);
assert_eq!(dominant_length(&[8192, 300]), Some(8192));
assert_eq!(dominant_length(&[]), None);
}
#[test]
fn an_unreproducible_split_is_bounded() {
let tiny = file(&[
ihdr(),
(b"IDAT", vec![0; 4]),
(b"IDAT", vec![0; 4]),
(b"IEND", Vec::new()),
]);
assert_eq!(
PngEnvelope::read(&tiny).idat_chunk_size(),
DEFAULT_IDAT_CHUNK_SIZE
);
let mut huge = file(&[ihdr()]);
huge.extend_from_slice(&((MAX_IDAT_CHUNK_SIZE + 1) as u32).to_be_bytes());
huge.extend_from_slice(b"IDAT");
huge.resize(huge.len() + MAX_IDAT_CHUNK_SIZE + 1 + CHUNK_CRC_LEN, 0);
let envelope = PngEnvelope::read(&huge);
assert_eq!(envelope.idat_chunk_count(), 1);
assert_eq!(envelope.idat_chunk_size(), MAX_IDAT_CHUNK_SIZE);
}
#[test]
fn the_default_profile_looks_like_an_ordinary_export() {
let envelope = PngEnvelope::synthesised();
let kinds: Vec<&str> = envelope
.preserved_chunks()
.iter()
.map(|chunk| chunk.kind().name())
.collect();
assert_eq!(kinds, vec!["gAMA", "cHRM", "sRGB", "pHYs"]);
assert_eq!(envelope.idat_chunk_size(), DEFAULT_IDAT_CHUNK_SIZE);
assert_eq!(envelope.idat_chunk_count(), 0);
assert_eq!(envelope.discarded_chunks(), 0);
let lengths: Vec<usize> = envelope
.preserved_chunks()
.iter()
.map(PreservedChunk::len)
.collect();
assert_eq!(lengths, vec![4, 32, 1, 9]);
}
}