Skip to main content

oxidize_html/
image.rs

1use std::collections::HashMap;
2
3use base64::Engine;
4
5#[derive(Debug, Clone, PartialEq)]
6/// The format of an image, used to determine how to decode the image data.
7pub enum ImageFormat {
8    Png,
9    Jpeg,
10    Gif,
11    Webp,
12    Unknown,
13}
14
15#[derive(Debug, Clone, PartialEq)]
16/// The source of an image, which can be a local file path, a remote URL, a data URI, or a CID.
17pub enum ImageSource {
18    DataUri(Vec<u8>, ImageFormat),
19    Remote(String),
20    Cid(String),
21    LocalPath(String),
22    Invalid,
23}
24
25/// The image data, including the image format and the raw image data.
26#[derive(Debug, Clone, PartialEq)]
27pub struct ImageData {
28    pub bytes: Vec<u8>,
29    pub format: ImageFormat,
30}
31
32/// Takes an image source string and determines the type of source it is, returning an `ImageSource` enum.
33pub fn parse_source(src: &str) -> ImageSource {
34    let src = src.trim();
35    if src.starts_with("data:") {
36        return resolve_data_uri(src).unwrap_or(ImageSource::Invalid);
37    }
38    if src.starts_with("cid:") {
39        return ImageSource::Cid(src.trim_start_matches("cid:").to_string());
40    }
41    if src.starts_with("http://") || src.starts_with("https://") {
42        return ImageSource::Remote(src.to_string());
43    }
44    if !src.is_empty() {
45        return ImageSource::LocalPath(src.to_string());
46    }
47    ImageSource::Invalid
48}
49
50/// Takes an image source string and a map of MIME types to image data and attempts to resolve the image source to an [`ImageData`] struct.
51pub fn resolve_image(src: &str, mime_parts: &HashMap<String, Vec<u8>>) -> Option<ImageData> {
52    match parse_source(src) {
53        ImageSource::DataUri(bytes, format) => Some(ImageData { bytes, format }),
54        ImageSource::Cid(id) => mime_parts.get(&id).map(|bytes| ImageData {
55            bytes: bytes.clone(),
56            format: detect_image_format(bytes),
57        }),
58        ImageSource::LocalPath(path) => std::fs::read(path).ok().map(|bytes| ImageData {
59            format: detect_image_format(&bytes),
60            bytes,
61        }),
62        ImageSource::Remote(_) | ImageSource::Invalid => None,
63    }
64}
65
66/// Returns the dimensions of the image source if it can be determined, otherwise returns `None`.
67pub fn source_dimensions(source: &ImageSource) -> Option<(u32, u32)> {
68    match source {
69        ImageSource::DataUri(bytes, _) => image::load_from_memory(bytes)
70            .ok()
71            .map(|img| (img.width(), img.height())),
72        ImageSource::LocalPath(path) => image::image_dimensions(path).ok(),
73        _ => None,
74    }
75}
76
77/// Takes a data URI string and decodes it into an [`ImageData`] struct.
78fn resolve_data_uri(src: &str) -> Option<ImageSource> {
79    let payload = src.strip_prefix("data:")?;
80    let (meta, data) = payload.split_once(',')?;
81    if !meta.contains(";base64") {
82        return None;
83    }
84    let mime = meta.split(';').next().unwrap_or_default();
85    let bytes = base64::engine::general_purpose::STANDARD
86        .decode(data)
87        .ok()?;
88    Some(ImageSource::DataUri(bytes, format_from_mime(mime)))
89}
90///
91///
92/// Determines the image format from a MIME type string, returning an [`ImageFormat`] enum.
93///
94fn format_from_mime(mime: &str) -> ImageFormat {
95    match mime {
96        "image/png" => ImageFormat::Png,
97        "image/jpeg" | "image/jpg" => ImageFormat::Jpeg,
98        "image/gif" => ImageFormat::Gif,
99        "image/webp" => ImageFormat::Webp,
100        _ => ImageFormat::Unknown,
101    }
102}
103
104/// Supports Png, Jpeg, GIF, and Webp formats.
105fn detect_image_format(bytes: &[u8]) -> ImageFormat {
106    if bytes.starts_with(&[0x89, b'P', b'N', b'G']) {
107        ImageFormat::Png
108    } else if bytes.starts_with(&[0xFF, 0xD8]) {
109        ImageFormat::Jpeg
110    } else if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") {
111        ImageFormat::Gif
112    } else if bytes.starts_with(b"RIFF") && bytes.len() >= 12 && &bytes[8..12] == b"WEBP" {
113        ImageFormat::Webp
114    } else {
115        ImageFormat::Unknown
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use super::{ImageSource, parse_source, resolve_image};
122    use std::collections::HashMap;
123
124    #[test]
125    fn parses_data_uri_source() {
126        let src = "data:image/png;base64,aGVsbG8=";
127        let parsed = parse_source(src);
128        match parsed {
129            ImageSource::DataUri(bytes, _) => assert_eq!(bytes, b"hello"),
130            _ => panic!("expected data uri"),
131        }
132    }
133
134    #[test]
135    fn resolves_cid_from_map() {
136        let mut map = HashMap::new();
137        map.insert("logo".to_string(), vec![0x89, b'P', b'N', b'G']);
138        let data = resolve_image("cid:logo", &map).expect("cid image");
139        assert_eq!(data.bytes.len(), 4);
140    }
141}