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

artifactize_tools/
image.rs
