Update src/main.rs

src/main.rs

@@ -1,4 +1,3 @@
-
 use std::{io::Cursor, net::SocketAddr, sync::Arc, time::Duration};
 
 use anyhow::Result;
@@ -9,9 +8,11 @@ use axum::{
     routing::get,
     Router,
 };
+use bytes::Bytes;
 use image::{
-    codecs::jpeg::JpegEncoder,
+    codecs::{gif::GifDecoder, jpeg::JpegEncoder},
     imageops::FilterType,
+    AnimationDecoder,
     GenericImageView,
 };
 use mimalloc::MiMalloc;
@@ -69,7 +70,6 @@ async fn main() -> Result<()> {
     info!("image backend listening on {}", addr);
 
     let listener = TcpListener::bind(addr).await?;
-
     axum::serve(listener, app).await?;
 
     Ok(())
@@ -87,12 +87,7 @@ async fn fetch(
         Ok(r) => r,
         Err(e) => {
             error!("fetch error: {:?}", e);
-
-            (
-                StatusCode::BAD_GATEWAY,
-                "failed to process image",
-            )
-                .into_response()
+            (StatusCode::BAD_GATEWAY, "failed to process image").into_response()
         }
     }
 }
@@ -105,12 +100,7 @@ async fn thumb(
         Ok(r) => r,
         Err(e) => {
             error!("thumb error: {:?}", e);
-
-            (
-                StatusCode::BAD_GATEWAY,
-                "failed to process thumbnail",
-            )
-                .into_response()
+            (StatusCode::BAD_GATEWAY, "failed to process thumbnail").into_response()
         }
     }
 }
@@ -123,7 +113,6 @@ fn validate_url(url: &str) -> Result<()> {
     }
 
     let host = parsed.host_str().unwrap_or("");
-
     if !host.eq_ignore_ascii_case("i.pinimg.com") {
         anyhow::bail!("invalid host");
     }
@@ -131,7 +120,7 @@ fn validate_url(url: &str) -> Result<()> {
     Ok(())
 }
 
-async fn download_image(client: &Client, url: &str) -> Result<bytes::Bytes> {
+async fn download_image(client: &Client, url: &str) -> Result<Bytes> {
     let resp = client.get(url).send().await?;
 
     if !resp.status().is_success() {
@@ -141,72 +130,99 @@ async fn download_image(client: &Client, url: &str) -> Result<bytes::Bytes> {
     Ok(resp.bytes().await?)
 }
 
-async fn process_fetch(
-    state: Arc<AppState>,
-    query: ImageQuery,
-) -> Result<Response> {
-    validate_url(&query.url)?;
+fn is_gif(data: &[u8]) -> bool {
+    data.starts_with(b"GIF87a") || data.starts_with(b"GIF89a")
+}
 
-    let data = download_image(&state.client, &query.url).await?;
+fn gif_passthrough(data: Bytes) -> Response {
+    let mut headers = HeaderMap::new();
 
-    let img = image::load_from_memory(&data)?;
+    headers.insert(
+        "content-type",
+        HeaderValue::from_static("image/gif"),
+    );
+    headers.insert(
+        "cache-control",
+        HeaderValue::from_static("public, max-age=86400, immutable"),
+    );
 
-    let rgb = img.to_rgb8();
-
-    let mut out = Vec::with_capacity(data.len() / 2);
-
-    {
-        let mut cursor = Cursor::new(&mut out);
-
-        let mut encoder =
-            JpegEncoder::new_with_quality(&mut cursor, 75);
-
-        encoder.encode_image(&rgb)?;
-    }
+    (headers, data).into_response()
+}
 
+fn jpeg_response(data: Vec<u8>) -> Response {
     let mut headers = HeaderMap::new();
 
     headers.insert(
         "content-type",
         HeaderValue::from_static("image/jpeg"),
     );
-
     headers.insert(
         "cache-control",
-        HeaderValue::from_static(
-            "public, max-age=86400, immutable",
-        ),
+        HeaderValue::from_static("public, max-age=86400, immutable"),
     );
 
-    Ok((headers, out).into_response())
+    (headers, data).into_response()
 }
 
-async fn process_thumb(
-    state: Arc<AppState>,
-    query: ThumbQuery,
-) -> Result<Response> {
+fn decode_first_gif_frame(data: &[u8]) -> Result<image::DynamicImage> {
+    let decoder = GifDecoder::new(Cursor::new(data))?;
+    let frames = decoder.into_frames().collect_frames()?;
+
+    let frame = frames
+        .into_iter()
+        .next()
+        .ok_or_else(|| anyhow::anyhow!("empty gif"))?;
+
+    Ok(image::DynamicImage::ImageRgba8(frame.into_buffer()))
+}
+
+async fn process_fetch(state: Arc<AppState>, query: ImageQuery) -> Result<Response> {
+    validate_url(&query.url)?;
+
+    let data = download_image(&state.client, &query.url).await?;
+
+    if is_gif(&data) {
+        return Ok(gif_passthrough(data));
+    }
+
+    let img = image::load_from_memory(&data)?;
+    let rgb = img.to_rgb8();
+
+    let mut out = Vec::with_capacity(data.len() / 2);
+
+    {
+        let mut cursor = Cursor::new(&mut out);
+        let mut encoder = JpegEncoder::new_with_quality(&mut cursor, 75);
+        encoder.encode_image(&rgb)?;
+    }
+
+    Ok(jpeg_response(out))
+}
+
+async fn process_thumb(state: Arc<AppState>, query: ThumbQuery) -> Result<Response> {
     validate_url(&query.url)?;
 
     let width = query.w.unwrap_or(260).clamp(64, 2048);
 
     let data = download_image(&state.client, &query.url).await?;
 
-    let img = image::load_from_memory(&data)?;
+    let img = if is_gif(&data) {
+        decode_first_gif_frame(&data)?
+    } else {
+        image::load_from_memory(&data)?
+    };
 
     let (src_w, src_h) = img.dimensions();
 
     if width >= src_w {
         return process_fetch(
             state,
-            ImageQuery {
-                url: query.url,
-            },
+            ImageQuery { url: query.url },
         )
         .await;
     }
 
-    let height =
-        ((src_h as f32 * width as f32) / src_w as f32) as u32;
+    let height = ((src_h as f32 * width as f32) / src_w as f32) as u32;
 
     let resized = img.resize_exact(
         width,
@@ -220,26 +236,9 @@ async fn process_thumb(
 
     {
         let mut cursor = Cursor::new(&mut out);
-
-        let mut encoder =
-            JpegEncoder::new_with_quality(&mut cursor, 72);
-
+        let mut encoder = JpegEncoder::new_with_quality(&mut cursor, 72);
         encoder.encode_image(&rgb)?;
     }
 
-    let mut headers = HeaderMap::new();
-
-    headers.insert(
-        "content-type",
-        HeaderValue::from_static("image/jpeg"),
-    );
-
-    headers.insert(
-        "cache-control",
-        HeaderValue::from_static(
-            "public, max-age=86400, immutable",
-        ),
-    );
-
-    Ok((headers, out).into_response())
+    Ok(jpeg_response(out))
 }