image_provider/
lib.rs

1use std::{io::Cursor, path::PathBuf, sync::Arc};
2
3use axum::{
4    extract::{Path, Query, State},
5    http::{
6        header::{CACHE_CONTROL, CONTENT_TYPE, X_CONTENT_TYPE_OPTIONS},
7        HeaderMap, HeaderValue, StatusCode,
8    },
9    response::{IntoResponse, Response},
10    routing::get,
11    Router,
12};
13use axum_extra::{headers::Range, TypedHeader};
14use axum_range::{KnownSize, Ranged};
15use bytes::Bytes;
16use cached::{Cached, TimedSizedCache};
17use fast_image_resize::{images::Image, IntoImageView, ResizeAlg, ResizeOptions, Resizer};
18use image::{
19    codecs::{jpeg::JpegEncoder, png::PngEncoder, webp::WebPEncoder},
20    load_from_memory, ColorType, DynamicImage, ImageEncoder, ImageFormat,
21};
22use log::{debug, info};
23use mime_guess::MimeGuess;
24use serde::Deserialize;
25use tokio::{fs::File, io::AsyncReadExt, sync::Mutex};
26
27pub mod config;
28
29pub use config::*;
30
31pub fn get_images_router(root: PathBuf, config: ResizeConfig) -> Router {
32    const CACHE_LIFESPAN: u64 = 24 * 60 * 60; // 1 days in seconds
33    let cache = TimedSizedCache::with_size_and_lifespan_and_refresh(
34        config.cache_size,
35        CACHE_LIFESPAN,
36        true,
37    );
38    let cache = Arc::new(Mutex::new(cache));
39
40    Router::new()
41        .route("/{*path}", get(provide_images))
42        .route(
43            "/",
44            get(|| async { (StatusCode::NOT_FOUND, "File not found".to_string()) }),
45        )
46        .with_state(ImageState {
47            root,
48            config,
49            cache,
50        })
51}
52
53type Error = (StatusCode, String);
54type Result<T, E = Error> = std::result::Result<T, E>;
55
56#[derive(Clone)]
57struct ImageState {
58    root: PathBuf,
59    config: ResizeConfig,
60    cache: Arc<Mutex<TimedSizedCache<(PathBuf, ImageQuery), Bytes>>>,
61}
62
63#[derive(Debug, Clone, Deserialize, PartialEq, Eq, Hash)]
64pub struct ImageQuery {
65    pub output: Option<String>,
66    pub dpr: Option<u32>,
67    pub w: Option<u32>,
68    pub h: Option<u32>,
69}
70
71impl ImageQuery {
72    fn output(&self) -> Result<Option<ImageFormat>> {
73        self.output
74            .as_ref()
75            .map(|ext| {
76                find_image_mime(MimeGuess::from_ext(ext)).ok_or((
77                    StatusCode::BAD_REQUEST,
78                    format!("Unsupported output format: {ext}"),
79                ))
80            })
81            .transpose()
82    }
83
84    fn size(&self) -> (Option<u32>, Option<u32>) {
85        (self.w, self.h)
86    }
87
88    fn dpr(&self) -> u32 {
89        self.dpr.unwrap_or(1).clamp(1, 3)
90    }
91}
92
93async fn provide_images(
94    State(ImageState {
95        root,
96        config,
97        cache,
98    }): State<ImageState>,
99    Query(query): Query<ImageQuery>,
100    Path(path): Path<PathBuf>,
101    range: Option<TypedHeader<Range>>,
102) -> Result<Response> {
103    let (path, raw_mime) = get_path_and_mime(root, path)?;
104    let dst_mime = query.output()?.unwrap_or(raw_mime);
105    let (dst_width, dst_height) = query.size();
106    let dpr = query.dpr();
107    info!(
108        "Processing image: {path:?} to mime: {dst_mime:?}, size: {dst_width:?}x{dst_height:?}, dpr: {dpr}"
109    );
110
111    let range = range.map(|TypedHeader(range)| range);
112    let headers = get_response_headers(&dst_mime);
113
114    // If no resizing is needed, serve the original file directly
115    let eq_raw = dst_width.is_none() && dst_height.is_none() && dpr == 1 && raw_mime == dst_mime;
116    let exclude = matches!(raw_mime, image::ImageFormat::Ico | image::ImageFormat::Gif);
117    if eq_raw || exclude {
118        debug!("Serving original image: {path:?}");
119        let file = load_file(&path).await?;
120        let body = KnownSize::file(file).await.unwrap();
121        let ranged = Ranged::new(range, body);
122        return Ok((headers, ranged).into_response());
123    }
124
125    if let Some(cached) = cache.lock().await.cache_get(&(path.clone(), query.clone())) {
126        debug!(
127            "Serving cached image: {path:?} (mime: {dst_mime:?}, size {dst_width:?}x{dst_height:?}, dpr: {dpr})"
128        );
129        let body = KnownSize::seek(Cursor::new(cached.clone())).await.unwrap();
130        return Ok((headers, Ranged::new(range, body)).into_response());
131    }
132
133    let file = load_file(&path).await?;
134    let src_image = load_image(file).await?;
135
136    let (dst_width, dst_height) = get_output_size(
137        (src_image.width(), src_image.height()),
138        (dst_width, dst_height),
139        dpr,
140    );
141
142    let mut dst_image = Image::new(dst_width, dst_height, src_image.pixel_type().unwrap());
143    resize_image(&config, &src_image, &mut dst_image)?;
144
145    let bytes = encode_image(dst_mime, &dst_image, src_image.color())?;
146
147    // Cache the processed image
148    cache
149        .lock()
150        .await
151        .cache_set((path.clone(), query), bytes.clone());
152    debug!(
153        "Cached processed image: {:?} (mime: {:?}, size {:?}x{:?}, dpr: {})",
154        &path, dst_mime, dst_width, dst_height, dpr
155    );
156
157    let body = KnownSize::seek(Cursor::new(bytes)).await.unwrap();
158
159    debug!(
160        "Serving processed image: {path:?} (mime: {dst_mime:?}, size {dst_width:?}x{dst_height:?}, dpr: {dpr})"
161    );
162    Ok((headers, Ranged::new(range, body)).into_response())
163}
164
165fn get_path_and_mime(root: PathBuf, rel_path: PathBuf) -> Result<(PathBuf, ImageFormat)> {
166    let path = path_clean::clean(rel_path);
167    let path = path.strip_prefix("/").unwrap_or(&path);
168    let path = root.join(path);
169
170    if !path.exists() || !path.is_file() {
171        return Err((StatusCode::NOT_FOUND, "File not found".to_string()));
172    }
173
174    match find_image_mime(MimeGuess::from_path(&path)) {
175        Some(mime) => Ok((path.clone(), mime)),
176        None => Err((StatusCode::BAD_REQUEST, "Unsupported file type".to_string())),
177    }
178}
179
180fn find_image_mime(mime: MimeGuess) -> Option<ImageFormat> {
181    mime.into_iter()
182        .flat_map(|m| ImageFormat::from_mime_type(&m))
183        .next()
184}
185
186fn get_response_headers(image_format: &ImageFormat) -> HeaderMap {
187    debug!("Setting response headers for format: {image_format:?}");
188    let mut headers = HeaderMap::new();
189    for (name, value) in [
190        (CONTENT_TYPE, image_format.to_mime_type()),
191        (CACHE_CONTROL, "public, max-age=31536000"),
192        (X_CONTENT_TYPE_OPTIONS, "nosniff"),
193    ] {
194        debug!("Setting header: {name}: {value}");
195        headers.insert(name, HeaderValue::from_static(value));
196    }
197    debug!("Response headers set: {headers:?}");
198    headers
199}
200
201async fn load_file(path: &PathBuf) -> Result<File> {
202    debug!("Loading file: {path:?}");
203    File::open(&path).await.map_err(|_| {
204        (
205            StatusCode::INTERNAL_SERVER_ERROR,
206            "Failed to read image".to_string(),
207        )
208    })
209}
210
211fn get_output_size(src: (u32, u32), dst: (Option<u32>, Option<u32>), dpr: u32) -> (u32, u32) {
212    let (src_width, src_height) = src;
213    let (dst_width, dst_height) = dst;
214    let aspect_ratio = src_width as f32 / src_height as f32;
215
216    let (mut width, mut height) = match (dst_width, dst_height) {
217        (Some(w), Some(h)) => (w, h),
218        (Some(w), None) => (w, (w as f32 / aspect_ratio).round() as u32),
219        (None, Some(h)) => ((h as f32 * aspect_ratio).round() as u32, h),
220        (None, None) => (src_width, src_height),
221    };
222
223    width *= dpr;
224    height *= dpr;
225
226    (width, height)
227}
228
229async fn load_image(mut file: File) -> Result<DynamicImage> {
230    let mut buffer = Vec::new();
231    file.read_to_end(&mut buffer).await.unwrap();
232    load_from_memory(&buffer).map_err(|e| {
233        (
234            StatusCode::INTERNAL_SERVER_ERROR,
235            format!("Failed to decode image: {e}"),
236        )
237    })
238}
239
240fn resize_image(
241    config: &ResizeConfig,
242    src_image: &DynamicImage,
243    dst_image: &mut Image<'_>,
244) -> Result<()> {
245    let mut resizer = Resizer::new();
246
247    let algorithm = if cfg!(debug_assertions) {
248        ResizeAlg::Nearest
249    } else {
250        config.resize_algorithm()
251    };
252    let options = ResizeOptions::new()
253        .resize_alg(algorithm)
254        .fit_into_destination(Some((0.5, 0.5)));
255
256    resizer
257        .resize(src_image, dst_image, Some(&options))
258        .map_err(|_| {
259            (
260                StatusCode::INTERNAL_SERVER_ERROR,
261                "Failed to resize image".to_string(),
262            )
263        })
264}
265
266fn encode_image(format: ImageFormat, image: &Image<'_>, color: ColorType) -> Result<Bytes> {
267    macro_rules! match_format {
268        ($format: expr , $( $target: pat => $encoder: expr, )+ ) => {
269            match $format {$(
270                $target => $encoder
271                    .write_image(image.buffer(), image.width(), image.height(), color.into())
272                    .map_err(|e| {
273                        (
274                            StatusCode::INTERNAL_SERVER_ERROR,
275                            format!("Failed to encode image: {e}"),
276                        )
277                    }),
278            )+
279                _ => Err((StatusCode::BAD_REQUEST, "Unsupported output format".to_string())),
280            }
281        };
282    }
283
284    let mut bytes = vec![];
285    match_format! {
286        format,
287        ImageFormat::WebP => WebPEncoder::new_lossless(&mut bytes),
288        ImageFormat::Png => PngEncoder::new(&mut bytes),
289        ImageFormat::Jpeg => JpegEncoder::new(&mut bytes),
290    }?;
291
292    Ok(Bytes::from(bytes))
293}
image_provider/lib.rs

image_provider/
lib.rs
