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

image_provider/
lib.rs
