use std::fmt;
use std::io::Cursor;
use std::path::Path;
use image::{codecs::png::PngDecoder, ColorType, DynamicImage, ImageDecoder};
use crate::image_io::buffer::{ColorSpace, ImageBuffer};
use crate::image_io::jpeg_detect;
const MIN_DIMENSION: u32 = 2000;
const MAX_PIXELS: u64 = 128 * 1024 * 1024;
const MIN_HEADER_LEN: usize = 12;
const PNG_HEADER_LEN: usize = 33;
const IHDR_WIDTH_OFFSET: usize = 16;
const IHDR_CHUNK_LENGTH: u32 = 13;
const PNG_MAGIC: [u8; 8] = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
#[derive(Debug)]
pub enum ValidationError {
IoError(std::io::Error),
JpegDetected,
WebpDetected,
NotPng,
UnsupportedColorSpace {
found: String,
},
ImageTooSmall {
width: u32,
height: u32,
min: u32,
},
ImageTooLarge {
width: u32,
height: u32,
pixels: u64,
max: u64,
},
DecodingError(String),
JpegArtifactsDetected {
ratio: f32,
},
}
impl fmt::Display for ValidationError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
ValidationError::IoError(err) => write!(f, "failed to read the image file: {err}"),
ValidationError::JpegDetected => {
write!(
f,
"the file is a JPEG; only lossless PNG containers are supported"
)
}
ValidationError::WebpDetected => {
write!(
f,
"the file is a WebP; only lossless PNG containers are supported"
)
}
ValidationError::NotPng => write!(f, "the file is not a PNG image"),
ValidationError::UnsupportedColorSpace { found } => {
write!(f, "unsupported pixel layout: {found}")
}
ValidationError::ImageTooSmall { width, height, min } => write!(
f,
"image is {width}x{height}; both sides must be at least {min} pixels"
),
ValidationError::ImageTooLarge {
width,
height,
pixels,
max,
} => write!(
f,
"image is {width}x{height}, which is {} megapixels; analysing it would need more \
memory than this limit allows, so it is refused rather than attempted. The \
maximum is {} megapixels",
pixels / (1024 * 1024),
max / (1024 * 1024)
),
ValidationError::DecodingError(message) => {
write!(f, "failed to decode the PNG stream: {message}")
}
ValidationError::JpegArtifactsDetected { ratio } => write!(
f,
"image shows an 8x8 JPEG block structure (blocking ratio {ratio:.2}, a clean \
image scores about 1.00); use a container that was never JPEG-compressed"
),
}
}
}
impl std::error::Error for ValidationError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
ValidationError::IoError(err) => Some(err),
_ => None,
}
}
}
impl From<std::io::Error> for ValidationError {
fn from(err: std::io::Error) -> Self {
ValidationError::IoError(err)
}
}
struct RawBytes(Vec<u8>);
struct VerifiedPngFile(Vec<u8>);
struct DecodedPng {
pixels: Vec<u8>,
width: u32,
height: u32,
color_space: ColorSpace,
}
fn validate_magic_bytes(raw: RawBytes) -> Result<VerifiedPngFile, ValidationError> {
let bytes = raw.0;
if bytes.len() < MIN_HEADER_LEN {
return Err(ValidationError::NotPng);
}
if bytes[0..3] == [0xFF, 0xD8, 0xFF] {
return Err(ValidationError::JpegDetected);
}
if &bytes[0..4] == b"RIFF" && &bytes[8..12] == b"WEBP" {
return Err(ValidationError::WebpDetected);
}
if bytes[0..8] != PNG_MAGIC {
return Err(ValidationError::NotPng);
}
Ok(VerifiedPngFile(bytes))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ImageGeometry {
pub width: u32,
pub height: u32,
}
impl ImageGeometry {
pub fn pixel_count(&self) -> u64 {
u64::from(self.width) * u64::from(self.height)
}
}
fn check_dimensions(geometry: ImageGeometry) -> Result<ImageGeometry, ValidationError> {
let ImageGeometry { width, height } = geometry;
if width < MIN_DIMENSION || height < MIN_DIMENSION {
return Err(ValidationError::ImageTooSmall {
width,
height,
min: MIN_DIMENSION,
});
}
let pixels = geometry.pixel_count();
if pixels > MAX_PIXELS {
return Err(ValidationError::ImageTooLarge {
width,
height,
pixels,
max: MAX_PIXELS,
});
}
Ok(geometry)
}
pub fn probe_geometry(path: &Path) -> Result<ImageGeometry, ValidationError> {
use std::io::Read;
let mut file = std::fs::File::open(path)?;
let mut header = [0u8; PNG_HEADER_LEN];
let mut filled = 0usize;
loop {
match file.read(&mut header[filled..]) {
Ok(0) => break,
Ok(read) => filled += read,
Err(err) if err.kind() == std::io::ErrorKind::Interrupted => continue,
Err(err) => return Err(ValidationError::IoError(err)),
}
if filled == PNG_HEADER_LEN {
break;
}
}
validate_magic_bytes(RawBytes(header[..filled].to_vec()))?;
if filled < PNG_HEADER_LEN {
return Err(ValidationError::NotPng);
}
let declares_ihdr = header.get(8..16).is_some_and(|chunk| {
chunk[..4] == IHDR_CHUNK_LENGTH.to_be_bytes() && &chunk[4..] == b"IHDR"
});
if !declares_ihdr {
return Err(ValidationError::DecodingError(
"the file begins with a PNG signature but no IHDR chunk follows it".to_owned(),
));
}
let Some(fields) = header.get(IHDR_WIDTH_OFFSET..IHDR_WIDTH_OFFSET + 8) else {
return Err(ValidationError::NotPng);
};
let (width_bytes, height_bytes) = fields.split_at(4);
let width = u32::from_be_bytes(width_bytes.try_into().unwrap_or([0; 4]));
let height = u32::from_be_bytes(height_bytes.try_into().unwrap_or([0; 4]));
check_dimensions(ImageGeometry { width, height })
}
fn decode_png(file: VerifiedPngFile) -> Result<DecodedPng, ValidationError> {
let decoder = PngDecoder::new(Cursor::new(file.0))
.map_err(|err| ValidationError::DecodingError(err.to_string()))?;
let (width, height) = decoder.dimensions();
check_dimensions(ImageGeometry { width, height })?;
let color_type = decoder.color_type();
let color_space = match color_type {
ColorType::Rgb8 => ColorSpace::Rgb8,
ColorType::Rgb16 => ColorSpace::Rgb16,
ColorType::Rgba8 => ColorSpace::Rgba8,
ColorType::L8 => ColorSpace::Luma8,
other => {
return Err(ValidationError::UnsupportedColorSpace {
found: format!("{other:?}"),
});
}
};
let decoded = DynamicImage::from_decoder(decoder)
.map_err(|err| ValidationError::DecodingError(err.to_string()))?;
let pixels = match color_space {
ColorSpace::Rgb8 => decoded.into_rgb8().into_raw(),
ColorSpace::Rgba8 => decoded.into_rgba8().into_raw(),
ColorSpace::Luma8 => decoded.into_luma8().into_raw(),
ColorSpace::Rgb16 => decoded
.into_rgb16()
.into_raw()
.into_iter()
.flat_map(u16::to_le_bytes)
.collect(),
};
Ok(DecodedPng {
pixels,
width,
height,
color_space,
})
}
fn validate_no_jpeg_artifacts(decoded: DecodedPng) -> Result<ImageBuffer, ValidationError> {
let DecodedPng {
pixels,
width,
height,
color_space,
} = decoded;
match jpeg_detect::detect_jpeg_artifacts(&pixels, width, height, color_space) {
Some(ratio) => Err(ValidationError::JpegArtifactsDetected { ratio }),
None => Ok(ImageBuffer::new(pixels, width, height, color_space)),
}
}
pub fn load_and_validate(path: &Path) -> Result<ImageBuffer, ValidationError> {
probe_geometry(path)?;
let raw = RawBytes(std::fs::read(path)?);
let verified = validate_magic_bytes(raw)?;
let decoded = decode_png(verified)?;
validate_no_jpeg_artifacts(decoded)
}
#[cfg(test)]
mod tests {
#![allow(clippy::expect_used)]
#![allow(clippy::panic)]
use super::*;
use image::{GrayAlphaImage, GrayImage, ImageFormat, Rgb, RgbImage, RgbaImage};
use tempfile::NamedTempFile;
use crate::image_io::buffer::CoverSource;
const SIDE: u32 = MIN_DIMENSION;
fn header(prefix: &[u8]) -> RawBytes {
let mut bytes = prefix.to_vec();
bytes.resize(MIN_HEADER_LEN.max(prefix.len()), 0);
RawBytes(bytes)
}
fn flat_png(color_space: ColorSpace) -> NamedTempFile {
let file = NamedTempFile::new().expect("temporary png file");
let saved = match color_space {
ColorSpace::Rgb8 => RgbImage::from_pixel(SIDE, SIDE, Rgb([90, 110, 130]))
.save_with_format(file.path(), ImageFormat::Png),
ColorSpace::Rgba8 => {
RgbaImage::from_pixel(SIDE, SIDE, image::Rgba([90, 110, 130, 255]))
.save_with_format(file.path(), ImageFormat::Png)
}
ColorSpace::Luma8 => GrayImage::from_pixel(SIDE, SIDE, image::Luma([110]))
.save_with_format(file.path(), ImageFormat::Png),
ColorSpace::Rgb16 => image::ImageBuffer::<Rgb<u16>, Vec<u16>>::from_pixel(
SIDE,
SIDE,
Rgb([23_000, 28_000, 33_000]),
)
.save_with_format(file.path(), ImageFormat::Png),
};
saved.expect("a flat png must be writable");
file
}
#[test]
fn a_truncated_header_is_not_a_png() {
let error = validate_magic_bytes(RawBytes(vec![0x89, 0x50, 0x4E]))
.map(|_| ())
.expect_err("three bytes cannot identify a format");
assert!(matches!(error, ValidationError::NotPng), "got: {error:?}");
}
#[test]
fn the_magic_number_decides_the_format() {
let jpeg = validate_magic_bytes(header(&[0xFF, 0xD8, 0xFF]))
.map(|_| ())
.expect_err("a jpeg must be named as such");
assert!(
matches!(jpeg, ValidationError::JpegDetected),
"got: {jpeg:?}"
);
let mut webp = b"RIFF".to_vec();
webp.extend_from_slice(&[0, 0, 0, 0]);
webp.extend_from_slice(b"WEBP");
let webp = validate_magic_bytes(RawBytes(webp))
.map(|_| ())
.expect_err("a webp must be named as such");
assert!(
matches!(webp, ValidationError::WebpDetected),
"got: {webp:?}"
);
let unknown = validate_magic_bytes(header(b"GIF89a"))
.map(|_| ())
.expect_err("an unknown format must be refused");
assert!(
matches!(unknown, ValidationError::NotPng),
"got: {unknown:?}"
);
assert!(validate_magic_bytes(header(&PNG_MAGIC)).is_ok());
}
#[test]
fn a_missing_file_is_an_io_error() {
let error = load_and_validate(Path::new("no-such-container-image.png"))
.map(|_| ())
.expect_err("a path that does not exist must be refused");
assert!(
matches!(error, ValidationError::IoError(_)),
"got: {error:?}"
);
assert!(std::error::Error::source(&error).is_some());
}
#[test]
fn the_probe_reads_the_geometry_from_the_header() {
let file = flat_png(ColorSpace::Rgb8);
match probe_geometry(file.path()) {
Ok(geometry) => {
assert_eq!(geometry.width, SIDE);
assert_eq!(geometry.height, SIDE);
assert_eq!(geometry.pixel_count(), u64::from(SIDE) * u64::from(SIDE));
}
Err(error) => panic!("a flat container must probe: {error}"),
}
}
#[test]
fn the_probe_refuses_what_the_loader_refuses() {
let scratch = NamedTempFile::new().expect("temporary file");
std::fs::write(
scratch.path(),
[0xFF, 0xD8, 0xFF, 0, 0, 0, 0, 0, 0, 0, 0, 0],
)
.expect("a jpeg header must be writable");
assert!(matches!(
probe_geometry(scratch.path()),
Err(ValidationError::JpegDetected)
));
std::fs::write(scratch.path(), b"not an image at all").expect("writable");
assert!(matches!(
probe_geometry(scratch.path()),
Err(ValidationError::NotPng)
));
std::fs::write(scratch.path(), PNG_MAGIC).expect("writable");
assert!(matches!(
probe_geometry(scratch.path()),
Err(ValidationError::NotPng)
));
assert!(matches!(
probe_geometry(Path::new("no-such-container.png")),
Err(ValidationError::IoError(_))
));
}
#[test]
fn the_size_gates_bound_the_geometry_from_both_ends() {
let accepted = ImageGeometry {
width: MIN_DIMENSION,
height: MIN_DIMENSION,
};
assert!(check_dimensions(accepted).is_ok());
assert!(matches!(
check_dimensions(ImageGeometry {
width: MIN_DIMENSION - 1,
height: MIN_DIMENSION,
}),
Err(ValidationError::ImageTooSmall { .. })
));
let enormous = ImageGeometry {
width: 32_767,
height: 32_767,
};
match check_dimensions(enormous) {
Err(ValidationError::ImageTooLarge { pixels, max, .. }) => {
assert_eq!(pixels, 32_767 * 32_767);
assert_eq!(max, MAX_PIXELS);
assert!(pixels > max);
}
other => panic!("a billion-pixel container must be refused, got: {other:?}"),
}
let (width, height) = (16_384u32, 8_192u32);
assert_eq!(u64::from(width) * u64::from(height), MAX_PIXELS);
assert!(check_dimensions(ImageGeometry { width, height }).is_ok());
assert!(matches!(
check_dimensions(ImageGeometry {
width,
height: height + 1,
}),
Err(ValidationError::ImageTooLarge { .. })
));
}
#[test]
fn the_pixel_count_does_not_wrap() {
let geometry = ImageGeometry {
width: u32::MAX,
height: u32::MAX,
};
assert_eq!(
geometry.pixel_count(),
u64::from(u32::MAX) * u64::from(u32::MAX)
);
assert!(geometry.pixel_count() > MAX_PIXELS);
assert!(matches!(
check_dimensions(geometry),
Err(ValidationError::ImageTooLarge { .. })
));
}
#[test]
fn a_corrupt_png_stream_is_a_decoding_error() {
let mut bytes = PNG_MAGIC.to_vec();
bytes.extend_from_slice(&[0x13; 64]);
let file = NamedTempFile::new().expect("temporary png file");
std::fs::write(file.path(), &bytes).expect("the corrupt file must be writable");
let error = load_and_validate(file.path())
.map(|_| ())
.expect_err("a malformed png stream must be refused");
assert!(
matches!(error, ValidationError::DecodingError(_)),
"got: {error:?}"
);
}
#[test]
fn an_undersized_container_is_refused() {
let file = NamedTempFile::new().expect("temporary png file");
RgbImage::from_pixel(100, 100, Rgb([10, 20, 30]))
.save_with_format(file.path(), ImageFormat::Png)
.expect("a small png must be writable");
let error = load_and_validate(file.path())
.map(|_| ())
.expect_err("a 100x100 container must be refused");
assert!(
matches!(
error,
ValidationError::ImageTooSmall {
width: 100,
height: 100,
min: MIN_DIMENSION,
}
),
"got: {error:?}"
);
}
#[test]
fn an_unsupported_layout_is_refused_by_name() {
let file = NamedTempFile::new().expect("temporary png file");
GrayAlphaImage::from_pixel(SIDE, SIDE, image::LumaA([110, 255]))
.save_with_format(file.path(), ImageFormat::Png)
.expect("a grayscale-alpha png must be writable");
let error = load_and_validate(file.path())
.map(|_| ())
.expect_err("grayscale with alpha must be refused");
match error {
ValidationError::UnsupportedColorSpace { found } => {
assert!(found.contains("La8"), "the layout must be named: {found}");
}
other => panic!("expected an unsupported layout, got: {other:?}"),
}
}
#[test]
fn every_supported_layout_decodes_to_its_own_stride() {
for expected in [
ColorSpace::Rgb8,
ColorSpace::Rgba8,
ColorSpace::Luma8,
ColorSpace::Rgb16,
] {
let file = flat_png(expected);
let image = match load_and_validate(file.path()) {
Ok(image) => image,
Err(error) => panic!("a flat {expected:?} container must load: {error}"),
};
assert_eq!(image.color_space(), expected);
assert_eq!(image.dimensions(), (SIDE, SIDE));
assert_eq!(
image.pixels().len(),
image.pixel_count() * expected.bytes_per_pixel()
);
}
}
#[test]
fn every_rejection_explains_itself() {
let messages = [
ValidationError::IoError(std::io::Error::other("disk on fire")).to_string(),
ValidationError::JpegDetected.to_string(),
ValidationError::WebpDetected.to_string(),
ValidationError::NotPng.to_string(),
ValidationError::UnsupportedColorSpace {
found: "Rgba16".to_owned(),
}
.to_string(),
ValidationError::ImageTooSmall {
width: 10,
height: 20,
min: MIN_DIMENSION,
}
.to_string(),
ValidationError::ImageTooLarge {
width: 32_767,
height: 32_767,
pixels: 32_767 * 32_767,
max: MAX_PIXELS,
}
.to_string(),
ValidationError::DecodingError("truncated".to_owned()).to_string(),
ValidationError::JpegArtifactsDetected { ratio: 3.25 }.to_string(),
];
for message in &messages {
assert!(!message.is_empty());
}
assert!(messages[1].contains("JPEG"));
assert!(messages[2].contains("WebP"));
assert!(messages[4].contains("Rgba16"));
assert!(messages[5].contains("10x20"));
assert!(messages[6].contains("32767x32767"));
assert!(messages[6].contains("1023 megapixels"));
assert!(messages[6].contains("128 megapixels"));
assert!(messages[8].contains("3.25"));
assert!(std::error::Error::source(&ValidationError::NotPng).is_none());
}
#[test]
fn an_io_error_converts_into_a_validation_error() {
let converted =
ValidationError::from(std::io::Error::new(std::io::ErrorKind::NotFound, "gone"));
assert!(matches!(converted, ValidationError::IoError(_)));
}
}