1use std::fmt;
19use std::io::Cursor;
20use std::path::Path;
21
22use image::{codecs::png::PngDecoder, ColorType, DynamicImage, ImageDecoder};
23
24use crate::image_io::buffer::{ColorSpace, ImageBuffer};
25use crate::image_io::jpeg_detect;
26
27const MIN_DIMENSION: u32 = 2000;
32
33const MAX_PIXELS: u64 = 128 * 1024 * 1024;
51
52const MIN_HEADER_LEN: usize = 12;
54
55const PNG_HEADER_LEN: usize = 33;
61
62const IHDR_WIDTH_OFFSET: usize = 16;
66
67const IHDR_CHUNK_LENGTH: u32 = 13;
72
73const PNG_MAGIC: [u8; 8] = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
75
76#[derive(Debug)]
78pub enum ValidationError {
79 IoError(std::io::Error),
81 JpegDetected,
83 WebpDetected,
85 NotPng,
87 UnsupportedColorSpace {
89 found: String,
91 },
92 ImageTooSmall {
94 width: u32,
96 height: u32,
98 min: u32,
100 },
101 ImageTooLarge {
106 width: u32,
108 height: u32,
110 pixels: u64,
112 max: u64,
114 },
115 DecodingError(String),
117 JpegArtifactsDetected {
119 ratio: f32,
123 },
124}
125
126impl fmt::Display for ValidationError {
127 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
128 match self {
129 ValidationError::IoError(err) => write!(f, "failed to read the image file: {err}"),
130 ValidationError::JpegDetected => {
131 write!(
132 f,
133 "the file is a JPEG; only lossless PNG containers are supported"
134 )
135 }
136 ValidationError::WebpDetected => {
137 write!(
138 f,
139 "the file is a WebP; only lossless PNG containers are supported"
140 )
141 }
142 ValidationError::NotPng => write!(f, "the file is not a PNG image"),
143 ValidationError::UnsupportedColorSpace { found } => {
144 write!(f, "unsupported pixel layout: {found}")
145 }
146 ValidationError::ImageTooSmall { width, height, min } => write!(
147 f,
148 "image is {width}x{height}; both sides must be at least {min} pixels"
149 ),
150 ValidationError::ImageTooLarge {
151 width,
152 height,
153 pixels,
154 max,
155 } => write!(
156 f,
157 "image is {width}x{height}, which is {} megapixels; analysing it would need more \
158 memory than this limit allows, so it is refused rather than attempted. The \
159 maximum is {} megapixels",
160 pixels / (1024 * 1024),
161 max / (1024 * 1024)
162 ),
163 ValidationError::DecodingError(message) => {
164 write!(f, "failed to decode the PNG stream: {message}")
165 }
166 ValidationError::JpegArtifactsDetected { ratio } => write!(
167 f,
168 "image shows an 8x8 JPEG block structure (blocking ratio {ratio:.2}, a clean \
169 image scores about 1.00); use a container that was never JPEG-compressed"
170 ),
171 }
172 }
173}
174
175impl std::error::Error for ValidationError {
176 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
177 match self {
178 ValidationError::IoError(err) => Some(err),
179 _ => None,
180 }
181 }
182}
183
184impl From<std::io::Error> for ValidationError {
185 fn from(err: std::io::Error) -> Self {
186 ValidationError::IoError(err)
187 }
188}
189
190struct RawBytes(Vec<u8>);
192
193struct VerifiedPngFile(Vec<u8>);
195
196struct DecodedPng {
198 pixels: Vec<u8>,
199 width: u32,
200 height: u32,
201 color_space: ColorSpace,
202}
203
204fn validate_magic_bytes(raw: RawBytes) -> Result<VerifiedPngFile, ValidationError> {
210 let bytes = raw.0;
211
212 if bytes.len() < MIN_HEADER_LEN {
213 return Err(ValidationError::NotPng);
214 }
215
216 if bytes[0..3] == [0xFF, 0xD8, 0xFF] {
217 return Err(ValidationError::JpegDetected);
218 }
219
220 if &bytes[0..4] == b"RIFF" && &bytes[8..12] == b"WEBP" {
221 return Err(ValidationError::WebpDetected);
222 }
223
224 if bytes[0..8] != PNG_MAGIC {
225 return Err(ValidationError::NotPng);
226 }
227
228 Ok(VerifiedPngFile(bytes))
229}
230
231#[derive(Debug, Clone, Copy, PartialEq, Eq)]
233pub struct ImageGeometry {
234 pub width: u32,
236 pub height: u32,
238}
239
240impl ImageGeometry {
241 pub fn pixel_count(&self) -> u64 {
247 u64::from(self.width) * u64::from(self.height)
248 }
249}
250
251fn check_dimensions(geometry: ImageGeometry) -> Result<ImageGeometry, ValidationError> {
256 let ImageGeometry { width, height } = geometry;
257
258 if width < MIN_DIMENSION || height < MIN_DIMENSION {
259 return Err(ValidationError::ImageTooSmall {
260 width,
261 height,
262 min: MIN_DIMENSION,
263 });
264 }
265
266 let pixels = geometry.pixel_count();
267 if pixels > MAX_PIXELS {
268 return Err(ValidationError::ImageTooLarge {
269 width,
270 height,
271 pixels,
272 max: MAX_PIXELS,
273 });
274 }
275
276 Ok(geometry)
277}
278
279pub fn probe_geometry(path: &Path) -> Result<ImageGeometry, ValidationError> {
305 use std::io::Read;
306
307 let mut file = std::fs::File::open(path)?;
308 let mut header = [0u8; PNG_HEADER_LEN];
309
310 let mut filled = 0usize;
313 loop {
314 match file.read(&mut header[filled..]) {
315 Ok(0) => break,
316 Ok(read) => filled += read,
317 Err(err) if err.kind() == std::io::ErrorKind::Interrupted => continue,
318 Err(err) => return Err(ValidationError::IoError(err)),
319 }
320
321 if filled == PNG_HEADER_LEN {
322 break;
323 }
324 }
325
326 validate_magic_bytes(RawBytes(header[..filled].to_vec()))?;
329
330 if filled < PNG_HEADER_LEN {
331 return Err(ValidationError::NotPng);
332 }
333
334 let declares_ihdr = header.get(8..16).is_some_and(|chunk| {
341 chunk[..4] == IHDR_CHUNK_LENGTH.to_be_bytes() && &chunk[4..] == b"IHDR"
342 });
343 if !declares_ihdr {
344 return Err(ValidationError::DecodingError(
345 "the file begins with a PNG signature but no IHDR chunk follows it".to_owned(),
346 ));
347 }
348
349 let Some(fields) = header.get(IHDR_WIDTH_OFFSET..IHDR_WIDTH_OFFSET + 8) else {
350 return Err(ValidationError::NotPng);
351 };
352
353 let (width_bytes, height_bytes) = fields.split_at(4);
356 let width = u32::from_be_bytes(width_bytes.try_into().unwrap_or([0; 4]));
357 let height = u32::from_be_bytes(height_bytes.try_into().unwrap_or([0; 4]));
358
359 check_dimensions(ImageGeometry { width, height })
360}
361
362fn decode_png(file: VerifiedPngFile) -> Result<DecodedPng, ValidationError> {
368 let decoder = PngDecoder::new(Cursor::new(file.0))
369 .map_err(|err| ValidationError::DecodingError(err.to_string()))?;
370
371 let (width, height) = decoder.dimensions();
372 check_dimensions(ImageGeometry { width, height })?;
373
374 let color_type = decoder.color_type();
375 let color_space = match color_type {
376 ColorType::Rgb8 => ColorSpace::Rgb8,
377 ColorType::Rgb16 => ColorSpace::Rgb16,
378 ColorType::Rgba8 => ColorSpace::Rgba8,
379 ColorType::L8 => ColorSpace::Luma8,
380 other => {
381 return Err(ValidationError::UnsupportedColorSpace {
382 found: format!("{other:?}"),
383 });
384 }
385 };
386
387 let decoded = DynamicImage::from_decoder(decoder)
388 .map_err(|err| ValidationError::DecodingError(err.to_string()))?;
389
390 let pixels = match color_space {
393 ColorSpace::Rgb8 => decoded.into_rgb8().into_raw(),
394 ColorSpace::Rgba8 => decoded.into_rgba8().into_raw(),
395 ColorSpace::Luma8 => decoded.into_luma8().into_raw(),
396 ColorSpace::Rgb16 => decoded
401 .into_rgb16()
402 .into_raw()
403 .into_iter()
404 .flat_map(u16::to_le_bytes)
405 .collect(),
406 };
407
408 Ok(DecodedPng {
409 pixels,
410 width,
411 height,
412 color_space,
413 })
414}
415
416fn validate_no_jpeg_artifacts(decoded: DecodedPng) -> Result<ImageBuffer, ValidationError> {
422 let DecodedPng {
423 pixels,
424 width,
425 height,
426 color_space,
427 } = decoded;
428
429 match jpeg_detect::detect_jpeg_artifacts(&pixels, width, height, color_space) {
430 Some(ratio) => Err(ValidationError::JpegArtifactsDetected { ratio }),
431 None => Ok(ImageBuffer::new(pixels, width, height, color_space)),
432 }
433}
434
435pub fn load_and_validate(path: &Path) -> Result<ImageBuffer, ValidationError> {
446 probe_geometry(path)?;
452
453 let raw = RawBytes(std::fs::read(path)?);
454 let verified = validate_magic_bytes(raw)?;
455 let decoded = decode_png(verified)?;
456 validate_no_jpeg_artifacts(decoded)
457}
458
459#[cfg(test)]
460mod tests {
461 #![allow(clippy::expect_used)]
465 #![allow(clippy::panic)]
466
467 use super::*;
468
469 use image::{GrayAlphaImage, GrayImage, ImageFormat, Rgb, RgbImage, RgbaImage};
470 use tempfile::NamedTempFile;
471
472 use crate::image_io::buffer::CoverSource;
473
474 const SIDE: u32 = MIN_DIMENSION;
479
480 fn header(prefix: &[u8]) -> RawBytes {
483 let mut bytes = prefix.to_vec();
484 bytes.resize(MIN_HEADER_LEN.max(prefix.len()), 0);
485
486 RawBytes(bytes)
487 }
488
489 fn flat_png(color_space: ColorSpace) -> NamedTempFile {
497 let file = NamedTempFile::new().expect("temporary png file");
498
499 let saved = match color_space {
500 ColorSpace::Rgb8 => RgbImage::from_pixel(SIDE, SIDE, Rgb([90, 110, 130]))
501 .save_with_format(file.path(), ImageFormat::Png),
502 ColorSpace::Rgba8 => {
503 RgbaImage::from_pixel(SIDE, SIDE, image::Rgba([90, 110, 130, 255]))
504 .save_with_format(file.path(), ImageFormat::Png)
505 }
506 ColorSpace::Luma8 => GrayImage::from_pixel(SIDE, SIDE, image::Luma([110]))
507 .save_with_format(file.path(), ImageFormat::Png),
508 ColorSpace::Rgb16 => image::ImageBuffer::<Rgb<u16>, Vec<u16>>::from_pixel(
509 SIDE,
510 SIDE,
511 Rgb([23_000, 28_000, 33_000]),
512 )
513 .save_with_format(file.path(), ImageFormat::Png),
514 };
515 saved.expect("a flat png must be writable");
516
517 file
518 }
519
520 #[test]
522 fn a_truncated_header_is_not_a_png() {
523 let error = validate_magic_bytes(RawBytes(vec![0x89, 0x50, 0x4E]))
524 .map(|_| ())
525 .expect_err("three bytes cannot identify a format");
526
527 assert!(matches!(error, ValidationError::NotPng), "got: {error:?}");
528 }
529
530 #[test]
532 fn the_magic_number_decides_the_format() {
533 let jpeg = validate_magic_bytes(header(&[0xFF, 0xD8, 0xFF]))
534 .map(|_| ())
535 .expect_err("a jpeg must be named as such");
536 assert!(
537 matches!(jpeg, ValidationError::JpegDetected),
538 "got: {jpeg:?}"
539 );
540
541 let mut webp = b"RIFF".to_vec();
542 webp.extend_from_slice(&[0, 0, 0, 0]);
543 webp.extend_from_slice(b"WEBP");
544 let webp = validate_magic_bytes(RawBytes(webp))
545 .map(|_| ())
546 .expect_err("a webp must be named as such");
547 assert!(
548 matches!(webp, ValidationError::WebpDetected),
549 "got: {webp:?}"
550 );
551
552 let unknown = validate_magic_bytes(header(b"GIF89a"))
553 .map(|_| ())
554 .expect_err("an unknown format must be refused");
555 assert!(
556 matches!(unknown, ValidationError::NotPng),
557 "got: {unknown:?}"
558 );
559
560 assert!(validate_magic_bytes(header(&PNG_MAGIC)).is_ok());
561 }
562
563 #[test]
565 fn a_missing_file_is_an_io_error() {
566 let error = load_and_validate(Path::new("no-such-container-image.png"))
567 .map(|_| ())
568 .expect_err("a path that does not exist must be refused");
569
570 assert!(
571 matches!(error, ValidationError::IoError(_)),
572 "got: {error:?}"
573 );
574
575 assert!(std::error::Error::source(&error).is_some());
578 }
579
580 #[test]
582 fn the_probe_reads_the_geometry_from_the_header() {
583 let file = flat_png(ColorSpace::Rgb8);
584
585 match probe_geometry(file.path()) {
586 Ok(geometry) => {
587 assert_eq!(geometry.width, SIDE);
588 assert_eq!(geometry.height, SIDE);
589 assert_eq!(geometry.pixel_count(), u64::from(SIDE) * u64::from(SIDE));
590 }
591 Err(error) => panic!("a flat container must probe: {error}"),
592 }
593 }
594
595 #[test]
597 fn the_probe_refuses_what_the_loader_refuses() {
598 let scratch = NamedTempFile::new().expect("temporary file");
599
600 std::fs::write(
601 scratch.path(),
602 [0xFF, 0xD8, 0xFF, 0, 0, 0, 0, 0, 0, 0, 0, 0],
603 )
604 .expect("a jpeg header must be writable");
605 assert!(matches!(
606 probe_geometry(scratch.path()),
607 Err(ValidationError::JpegDetected)
608 ));
609
610 std::fs::write(scratch.path(), b"not an image at all").expect("writable");
611 assert!(matches!(
612 probe_geometry(scratch.path()),
613 Err(ValidationError::NotPng)
614 ));
615
616 std::fs::write(scratch.path(), PNG_MAGIC).expect("writable");
619 assert!(matches!(
620 probe_geometry(scratch.path()),
621 Err(ValidationError::NotPng)
622 ));
623
624 assert!(matches!(
625 probe_geometry(Path::new("no-such-container.png")),
626 Err(ValidationError::IoError(_))
627 ));
628 }
629
630 #[test]
637 fn the_size_gates_bound_the_geometry_from_both_ends() {
638 let accepted = ImageGeometry {
639 width: MIN_DIMENSION,
640 height: MIN_DIMENSION,
641 };
642 assert!(check_dimensions(accepted).is_ok());
643
644 assert!(matches!(
645 check_dimensions(ImageGeometry {
646 width: MIN_DIMENSION - 1,
647 height: MIN_DIMENSION,
648 }),
649 Err(ValidationError::ImageTooSmall { .. })
650 ));
651
652 let enormous = ImageGeometry {
655 width: 32_767,
656 height: 32_767,
657 };
658 match check_dimensions(enormous) {
659 Err(ValidationError::ImageTooLarge { pixels, max, .. }) => {
660 assert_eq!(pixels, 32_767 * 32_767);
661 assert_eq!(max, MAX_PIXELS);
662 assert!(pixels > max);
663 }
664 other => panic!("a billion-pixel container must be refused, got: {other:?}"),
665 }
666
667 let (width, height) = (16_384u32, 8_192u32);
671 assert_eq!(u64::from(width) * u64::from(height), MAX_PIXELS);
672 assert!(check_dimensions(ImageGeometry { width, height }).is_ok());
673 assert!(matches!(
674 check_dimensions(ImageGeometry {
675 width,
676 height: height + 1,
677 }),
678 Err(ValidationError::ImageTooLarge { .. })
679 ));
680 }
681
682 #[test]
689 fn the_pixel_count_does_not_wrap() {
690 let geometry = ImageGeometry {
691 width: u32::MAX,
692 height: u32::MAX,
693 };
694
695 assert_eq!(
696 geometry.pixel_count(),
697 u64::from(u32::MAX) * u64::from(u32::MAX)
698 );
699 assert!(geometry.pixel_count() > MAX_PIXELS);
700 assert!(matches!(
701 check_dimensions(geometry),
702 Err(ValidationError::ImageTooLarge { .. })
703 ));
704 }
705
706 #[test]
711 fn a_corrupt_png_stream_is_a_decoding_error() {
712 let mut bytes = PNG_MAGIC.to_vec();
713 bytes.extend_from_slice(&[0x13; 64]);
714
715 let file = NamedTempFile::new().expect("temporary png file");
716 std::fs::write(file.path(), &bytes).expect("the corrupt file must be writable");
717
718 let error = load_and_validate(file.path())
719 .map(|_| ())
720 .expect_err("a malformed png stream must be refused");
721
722 assert!(
723 matches!(error, ValidationError::DecodingError(_)),
724 "got: {error:?}"
725 );
726 }
727
728 #[test]
730 fn an_undersized_container_is_refused() {
731 let file = NamedTempFile::new().expect("temporary png file");
732 RgbImage::from_pixel(100, 100, Rgb([10, 20, 30]))
733 .save_with_format(file.path(), ImageFormat::Png)
734 .expect("a small png must be writable");
735
736 let error = load_and_validate(file.path())
737 .map(|_| ())
738 .expect_err("a 100x100 container must be refused");
739
740 assert!(
741 matches!(
742 error,
743 ValidationError::ImageTooSmall {
744 width: 100,
745 height: 100,
746 min: MIN_DIMENSION,
747 }
748 ),
749 "got: {error:?}"
750 );
751 }
752
753 #[test]
755 fn an_unsupported_layout_is_refused_by_name() {
756 let file = NamedTempFile::new().expect("temporary png file");
757 GrayAlphaImage::from_pixel(SIDE, SIDE, image::LumaA([110, 255]))
758 .save_with_format(file.path(), ImageFormat::Png)
759 .expect("a grayscale-alpha png must be writable");
760
761 let error = load_and_validate(file.path())
762 .map(|_| ())
763 .expect_err("grayscale with alpha must be refused");
764
765 match error {
766 ValidationError::UnsupportedColorSpace { found } => {
767 assert!(found.contains("La8"), "the layout must be named: {found}");
768 }
769 other => panic!("expected an unsupported layout, got: {other:?}"),
770 }
771 }
772
773 #[test]
781 fn every_supported_layout_decodes_to_its_own_stride() {
782 for expected in [
783 ColorSpace::Rgb8,
784 ColorSpace::Rgba8,
785 ColorSpace::Luma8,
786 ColorSpace::Rgb16,
787 ] {
788 let file = flat_png(expected);
789 let image = match load_and_validate(file.path()) {
790 Ok(image) => image,
791 Err(error) => panic!("a flat {expected:?} container must load: {error}"),
792 };
793
794 assert_eq!(image.color_space(), expected);
795 assert_eq!(image.dimensions(), (SIDE, SIDE));
796 assert_eq!(
797 image.pixels().len(),
798 image.pixel_count() * expected.bytes_per_pixel()
799 );
800 }
801 }
802
803 #[test]
805 fn every_rejection_explains_itself() {
806 let messages = [
807 ValidationError::IoError(std::io::Error::other("disk on fire")).to_string(),
808 ValidationError::JpegDetected.to_string(),
809 ValidationError::WebpDetected.to_string(),
810 ValidationError::NotPng.to_string(),
811 ValidationError::UnsupportedColorSpace {
812 found: "Rgba16".to_owned(),
813 }
814 .to_string(),
815 ValidationError::ImageTooSmall {
816 width: 10,
817 height: 20,
818 min: MIN_DIMENSION,
819 }
820 .to_string(),
821 ValidationError::ImageTooLarge {
822 width: 32_767,
823 height: 32_767,
824 pixels: 32_767 * 32_767,
825 max: MAX_PIXELS,
826 }
827 .to_string(),
828 ValidationError::DecodingError("truncated".to_owned()).to_string(),
829 ValidationError::JpegArtifactsDetected { ratio: 3.25 }.to_string(),
830 ];
831
832 for message in &messages {
833 assert!(!message.is_empty());
834 }
835
836 assert!(messages[1].contains("JPEG"));
837 assert!(messages[2].contains("WebP"));
838 assert!(messages[4].contains("Rgba16"));
839 assert!(messages[5].contains("10x20"));
840 assert!(messages[6].contains("32767x32767"));
843 assert!(messages[6].contains("1023 megapixels"));
847 assert!(messages[6].contains("128 megapixels"));
848 assert!(messages[8].contains("3.25"));
849
850 assert!(std::error::Error::source(&ValidationError::NotPng).is_none());
852 }
853
854 #[test]
856 fn an_io_error_converts_into_a_validation_error() {
857 let converted =
858 ValidationError::from(std::io::Error::new(std::io::ErrorKind::NotFound, "gone"));
859
860 assert!(matches!(converted, ValidationError::IoError(_)));
861 }
862}