1use std::collections::HashMap;
2
3use base64::Engine;
4
5#[derive(Debug, Clone, PartialEq)]
6pub enum ImageFormat {
8 Png,
9 Jpeg,
10 Gif,
11 Webp,
12 Unknown,
13}
14
15#[derive(Debug, Clone, PartialEq)]
16pub enum ImageSource {
18 DataUri(Vec<u8>, ImageFormat),
19 Remote(String),
20 Cid(String),
21 LocalPath(String),
22 Invalid,
23}
24
25#[derive(Debug, Clone, PartialEq)]
27pub struct ImageData {
28 pub bytes: Vec<u8>,
29 pub format: ImageFormat,
30}
31
32pub 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
50pub 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
66pub 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
77fn 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}
90fn 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
104fn 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}