codex_utils_image/
lib.rs

1use std::num::NonZeroUsize;
2use std::path::Path;
3use std::sync::LazyLock;
4
5use crate::error::ImageProcessingError;
6use base64::Engine;
7use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
8use codex_utils_cache::BlockingLruCache;
9use codex_utils_cache::sha1_digest;
10use image::ColorType;
11use image::DynamicImage;
12use image::GenericImageView;
13use image::ImageEncoder;
14use image::ImageFormat;
15use image::codecs::jpeg::JpegEncoder;
16use image::codecs::png::PngEncoder;
17use image::imageops::FilterType;
18/// Maximum width used when resizing images before uploading.
19pub const MAX_WIDTH: u32 = 2048;
20/// Maximum height used when resizing images before uploading.
21pub const MAX_HEIGHT: u32 = 768;
22
23pub mod error;
24
25#[derive(Debug, Clone)]
26pub struct EncodedImage {
27    pub bytes: Vec<u8>,
28    pub mime: String,
29    pub width: u32,
30    pub height: u32,
31}
32
33impl EncodedImage {
34    pub fn into_data_url(self) -> String {
35        let encoded = BASE64_STANDARD.encode(&self.bytes);
36        format!("data:{};base64,{}", self.mime, encoded)
37    }
38}
39
40static IMAGE_CACHE: LazyLock<BlockingLruCache<[u8; 20], EncodedImage>> =
41    LazyLock::new(|| BlockingLruCache::new(NonZeroUsize::new(32).unwrap_or(NonZeroUsize::MIN)));
42
43pub fn load_and_resize_to_fit(path: &Path) -> Result<EncodedImage, ImageProcessingError> {
44    let path_buf = path.to_path_buf();
45
46    let file_bytes = read_file_bytes(path, &path_buf)?;
47
48    let key = sha1_digest(&file_bytes);
49
50    IMAGE_CACHE.get_or_try_insert_with(key, move || {
51        let format = match image::guess_format(&file_bytes) {
52            Ok(ImageFormat::Png) => Some(ImageFormat::Png),
53            Ok(ImageFormat::Jpeg) => Some(ImageFormat::Jpeg),
54            _ => None,
55        };
56
57        let dynamic = image::load_from_memory(&file_bytes).map_err(|source| {
58            ImageProcessingError::Decode {
59                path: path_buf.clone(),
60                source,
61            }
62        })?;
63
64        let (width, height) = dynamic.dimensions();
65
66        let encoded = if width <= MAX_WIDTH && height <= MAX_HEIGHT {
67            if let Some(format) = format {
68                let mime = format_to_mime(format);
69                EncodedImage {
70                    bytes: file_bytes,
71                    mime,
72                    width,
73                    height,
74                }
75            } else {
76                let (bytes, output_format) = encode_image(&dynamic, ImageFormat::Png)?;
77                let mime = format_to_mime(output_format);
78                EncodedImage {
79                    bytes,
80                    mime,
81                    width,
82                    height,
83                }
84            }
85        } else {
86            let resized = dynamic.resize(MAX_WIDTH, MAX_HEIGHT, FilterType::Triangle);
87            let target_format = format.unwrap_or(ImageFormat::Png);
88            let (bytes, output_format) = encode_image(&resized, target_format)?;
89            let mime = format_to_mime(output_format);
90            EncodedImage {
91                bytes,
92                mime,
93                width: resized.width(),
94                height: resized.height(),
95            }
96        };
97
98        Ok(encoded)
99    })
100}
101
102fn read_file_bytes(path: &Path, path_for_error: &Path) -> Result<Vec<u8>, ImageProcessingError> {
103    match tokio::runtime::Handle::try_current() {
104        // If we're inside a Tokio runtime, avoid block_on (it panics on worker threads).
105        // Use block_in_place and do a standard blocking read safely.
106        Ok(_) => tokio::task::block_in_place(|| std::fs::read(path)).map_err(|source| {
107            ImageProcessingError::Read {
108                path: path_for_error.to_path_buf(),
109                source,
110            }
111        }),
112        // Outside a runtime, just read synchronously.
113        Err(_) => std::fs::read(path).map_err(|source| ImageProcessingError::Read {
114            path: path_for_error.to_path_buf(),
115            source,
116        }),
117    }
118}
119
120fn encode_image(
121    image: &DynamicImage,
122    preferred_format: ImageFormat,
123) -> Result<(Vec<u8>, ImageFormat), ImageProcessingError> {
124    let target_format = match preferred_format {
125        ImageFormat::Jpeg => ImageFormat::Jpeg,
126        _ => ImageFormat::Png,
127    };
128
129    let mut buffer = Vec::new();
130
131    match target_format {
132        ImageFormat::Png => {
133            let rgba = image.to_rgba8();
134            let encoder = PngEncoder::new(&mut buffer);
135            encoder
136                .write_image(
137                    rgba.as_raw(),
138                    image.width(),
139                    image.height(),
140                    ColorType::Rgba8.into(),
141                )
142                .map_err(|source| ImageProcessingError::Encode {
143                    format: target_format,
144                    source,
145                })?;
146        }
147        ImageFormat::Jpeg => {
148            let mut encoder = JpegEncoder::new_with_quality(&mut buffer, 85);
149            encoder
150                .encode_image(image)
151                .map_err(|source| ImageProcessingError::Encode {
152                    format: target_format,
153                    source,
154                })?;
155        }
156        _ => unreachable!("unsupported target_format should have been handled earlier"),
157    }
158
159    Ok((buffer, target_format))
160}
161
162fn format_to_mime(format: ImageFormat) -> String {
163    match format {
164        ImageFormat::Jpeg => "image/jpeg".to_string(),
165        _ => "image/png".to_string(),
166    }
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172    use image::GenericImageView;
173    use image::ImageBuffer;
174    use image::Rgba;
175    use tempfile::NamedTempFile;
176
177    #[tokio::test(flavor = "multi_thread")]
178    async fn returns_original_image_when_within_bounds() {
179        let temp_file = NamedTempFile::new().expect("temp file");
180        let image = ImageBuffer::from_pixel(64, 32, Rgba([10u8, 20, 30, 255]));
181        image
182            .save_with_format(temp_file.path(), ImageFormat::Png)
183            .expect("write png to temp file");
184
185        let original_bytes = std::fs::read(temp_file.path()).expect("read written image");
186
187        let encoded = load_and_resize_to_fit(temp_file.path()).expect("process image");
188
189        assert_eq!(encoded.width, 64);
190        assert_eq!(encoded.height, 32);
191        assert_eq!(encoded.mime, "image/png");
192        assert_eq!(encoded.bytes, original_bytes);
193    }
194
195    #[tokio::test(flavor = "multi_thread")]
196    async fn downscales_large_image() {
197        let temp_file = NamedTempFile::new().expect("temp file");
198        let image = ImageBuffer::from_pixel(4096, 2048, Rgba([200u8, 10, 10, 255]));
199        image
200            .save_with_format(temp_file.path(), ImageFormat::Png)
201            .expect("write png to temp file");
202
203        let processed = load_and_resize_to_fit(temp_file.path()).expect("process image");
204
205        assert!(processed.width <= MAX_WIDTH);
206        assert!(processed.height <= MAX_HEIGHT);
207
208        let loaded =
209            image::load_from_memory(&processed.bytes).expect("read resized bytes back into image");
210        assert_eq!(loaded.dimensions(), (processed.width, processed.height));
211    }
212
213    #[tokio::test(flavor = "multi_thread")]
214    async fn fails_cleanly_for_invalid_images() {
215        let temp_file = NamedTempFile::new().expect("temp file");
216        std::fs::write(temp_file.path(), b"not an image").expect("write bytes");
217
218        let err = load_and_resize_to_fit(temp_file.path()).expect_err("invalid image should fail");
219        match err {
220            ImageProcessingError::Decode { .. } => {}
221            _ => panic!("unexpected error variant"),
222        }
223    }
224
225    #[tokio::test(flavor = "multi_thread")]
226    async fn reprocesses_updated_file_contents() {
227        {
228            IMAGE_CACHE.clear();
229        }
230
231        let temp_file = NamedTempFile::new().expect("temp file");
232        let first_image = ImageBuffer::from_pixel(32, 16, Rgba([20u8, 120, 220, 255]));
233        first_image
234            .save_with_format(temp_file.path(), ImageFormat::Png)
235            .expect("write initial image");
236
237        let first = load_and_resize_to_fit(temp_file.path()).expect("process first image");
238
239        let second_image = ImageBuffer::from_pixel(96, 48, Rgba([50u8, 60, 70, 255]));
240        second_image
241            .save_with_format(temp_file.path(), ImageFormat::Png)
242            .expect("write updated image");
243
244        let second = load_and_resize_to_fit(temp_file.path()).expect("process updated image");
245
246        assert_eq!(first.width, 32);
247        assert_eq!(first.height, 16);
248        assert_eq!(second.width, 96);
249        assert_eq!(second.height, 48);
250        assert_ne!(second.bytes, first.bytes);
251    }
252}