From 757427ebbc3778f6bfc122cee5851d68def61ff9 Mon Sep 17 00:00:00 2001 From: gigirassy Date: Thu, 28 May 2026 18:16:25 +0200 Subject: [PATCH] Update src/main.rs --- src/main.rs | 139 ++++++++++++++++++++++++++-------------------------- 1 file changed, 69 insertions(+), 70 deletions(-) diff --git a/src/main.rs b/src/main.rs index dc4eada..b68d3b8 100644 --- a/src/main.rs +++ b/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 { +async fn download_image(client: &Client, url: &str) -> Result { 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 { Ok(resp.bytes().await?) } -async fn process_fetch( - state: Arc, - query: ImageQuery, -) -> Result { - 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) -> 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, - query: ThumbQuery, -) -> Result { +fn decode_first_gif_frame(data: &[u8]) -> Result { + 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, query: ImageQuery) -> Result { + 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, query: ThumbQuery) -> Result { 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)) } \ No newline at end of file