Skip to main content

artifactize_tools/
image.rs

1use std::{fs::File, io::Read, path::Path};
2
3use base64::{
4    Engine, alphabet,
5    engine::{DecodePaddingMode, GeneralPurpose, GeneralPurposeConfig, general_purpose::STANDARD},
6};
7
8use crate::scope;
9
10use super::Content;
11
12/// Bound decoded image memory and base64 provider payloads while allowing review screenshots.
13pub const IMAGE_LIMIT: usize = 4 * 1024 * 1024;
14/// The eight bytes every PNG file starts with.
15const PNG_SIGNATURE: &[u8] = b"\x89PNG\r\n\x1a\n";
16/// The first chunk header of a PNG: IHDR with its fixed 13-byte length.
17const PNG_IHDR_HEADER: &[u8] = b"\0\0\0\rIHDR";
18/// A PNG chunk header: a 4-byte big-endian data length and a 4-byte chunk type.
19const PNG_CHUNK_HEADER_BYTES: usize = 8;
20/// PNG chunk data lengths are encoded as one big-endian u32.
21const PNG_LENGTH_BYTES: usize = std::mem::size_of::<u32>();
22/// Base64 encodes each group of three bytes as four ASCII characters.
23const BASE64_INPUT_BYTES: usize = 3;
24const BASE64_OUTPUT_BYTES: usize = 4;
25/// A PNG chunk's bytes besides its data: the header and a 4-byte CRC.
26const PNG_CHUNK_OVERHEAD_BYTES: usize = 12;
27/// The chunk type that marks an animated PNG, which reviews do not accept.
28const PNG_ANIMATION_CONTROL: &[u8] = b"acTL";
29/// The chunk type that ends a PNG.
30const PNG_END: &[u8] = b"IEND";
31/// A JPEG start-of-image marker followed by the first byte of the next marker.
32const JPEG_START: &[u8] = b"\xff\xd8\xff";
33/// The marker code after `JPEG_START` that means JPEG-LS, which is not a baseline JPEG.
34const JPEG_LS_MARKER: u8 = 0xf7;
35/// A WebP file is a RIFF container whose form type, at bytes 8 to 12, is `WEBP`.
36const RIFF_SIGNATURE: &[u8] = b"RIFF";
37const WEBP_FORM: &[u8] = b"WEBP";
38const RIFF_FORM_TYPE: std::ops::Range<usize> = 8..12;
39
40pub fn from_base64(data: &str, mime_type: &str) -> Result<Content, String> {
41    if data.len() > IMAGE_LIMIT.div_ceil(BASE64_INPUT_BYTES) * BASE64_OUTPUT_BYTES {
42        return Err("Image exceeds the 4 MiB limit.".into());
43    }
44    let decoder = GeneralPurpose::new(
45        &alphabet::STANDARD,
46        GeneralPurposeConfig::new().with_decode_padding_mode(DecodePaddingMode::Indifferent),
47    );
48    let bytes = decoder
49        .decode(data)
50        .map_err(|_| "Image requires valid base64 data.")?;
51    normalize(&bytes, Some(mime_type))
52}
53
54pub fn from_output(root: &Path, path: &str, mime_type: &str) -> Result<Content, String> {
55    let relative = crate::files::logical_below(root, path)
56        .ok_or("Image path must be a UTF-8 path below the tool output directory.")?;
57    let file = scope::open_scoped(root, &relative).map_err(|e| format!("Image {relative}: {e}"))?;
58    normalize(&read(file)?, Some(mime_type))
59}
60
61pub fn read(file: File) -> Result<Vec<u8>, String> {
62    let metadata = file.metadata().map_err(|_| "Cannot inspect image file.")?;
63    if !metadata.is_file() {
64        return Err("Viewing an image requires a regular file.".into());
65    }
66    if metadata.len() == 0 || metadata.len() > IMAGE_LIMIT as u64 {
67        return Err("Image exceeds the 4 MiB limit or contains no data.".into());
68    }
69    let mut bytes = Vec::new();
70    file.take((IMAGE_LIMIT + 1) as u64)
71        .read_to_end(&mut bytes)
72        .map_err(|_| "Cannot read image file.")?;
73    Ok(bytes)
74}
75
76pub fn normalize(bytes: &[u8], declared: Option<&str>) -> Result<Content, String> {
77    if bytes.is_empty() || bytes.len() > IMAGE_LIMIT {
78        return Err("Image exceeds the 4 MiB limit or contains no data.".into());
79    }
80    let mime_type = mime_type(bytes)
81        .ok_or("Image must be PNG, JPEG or WebP; GIF, BMP and animated PNG are not supported.")?;
82    if declared.is_some_and(|declared| declared != mime_type) {
83        return Err("Image bytes do not match declared MIME type.".into());
84    }
85    Ok(Content::Image {
86        data: STANDARD.encode(bytes),
87        mime_type: mime_type.into(),
88    })
89}
90
91fn mime_type(bytes: &[u8]) -> Option<&'static str> {
92    if bytes.starts_with(JPEG_START) && bytes.get(JPEG_START.len()) != Some(&JPEG_LS_MARKER) {
93        Some("image/jpeg")
94    } else if bytes.starts_with(PNG_SIGNATURE) {
95        let ihdr = PNG_SIGNATURE.len()..PNG_SIGNATURE.len() + PNG_IHDR_HEADER.len();
96        if bytes.get(ihdr) != Some(PNG_IHDR_HEADER) {
97            return None;
98        }
99        // Scan chunk boundaries, not image payloads, for animation control.
100        let mut offset = PNG_SIGNATURE.len();
101        while offset < bytes.len() {
102            let header = bytes.get(offset..offset + PNG_CHUNK_HEADER_BYTES)?;
103            let (length, kind) = header.split_at(PNG_LENGTH_BYTES);
104            let length = u32::from_be_bytes(length.try_into().ok()?) as usize;
105            let end = offset
106                .checked_add(PNG_CHUNK_OVERHEAD_BYTES)?
107                .checked_add(length)?;
108            if end > bytes.len() || kind == PNG_ANIMATION_CONTROL {
109                return None;
110            }
111            if kind == PNG_END {
112                return Some("image/png");
113            }
114            offset = end;
115        }
116        None
117    } else if bytes.starts_with(RIFF_SIGNATURE) && bytes.get(RIFF_FORM_TYPE) == Some(WEBP_FORM) {
118        Some("image/webp")
119    } else {
120        None
121    }
122}