1.1

Cargo.toml

@@ -1,9 +1,8 @@
 [package]
 name = "image-proxy"
-version = "0.1.0"
+version = "0.1.1"
 edition = "2021"
 
-
 [profile.release]
 opt-level = 3
 debug = false
@@ -14,14 +13,47 @@ strip = true
 
 [dependencies]
 axum = "0.7"
+
 tokio = { version = "1", features = ["full"] }
-reqwest = { version = "0.12", default-features = false, features = ["rustls-tls", "gzip", "brotli", "deflate"] }
-image = { version = "0.25", default-features = false, features = ["jpeg", "png", "webp", "gif"] }
+
+reqwest = {
+    version = "0.12",
+    default-features = false,
+    features = ["rustls-tls", "gzip", "brotli", "deflate"]
+}
+
+image = {
+    version = "0.25",
+    default-features = false,
+    features = ["jpeg", "png", "webp", "gif"]
+}
+
+moka = {
+    version = "0.12",
+    features = ["future"]
+}
+
 mimalloc = "0.1"
+
 anyhow = "1"
-serde = { version = "1", features = ["derive"] }
+
+serde = {
+    version = "1",
+    features = ["derive"]
+}
+
 serde_json = "1"
+
 bytes = "1"
-tower-http = { version = "0.5", features = ["trace"] }
+
+tower-http = {
+    version = "0.5",
+    features = ["trace"]
+}
+
 tracing = "0.1"
-tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+
+tracing-subscriber = {
+    version = "0.3",
+    features = ["env-filter"]
+}

src/main.rs

@@ -1,6 +1,10 @@
-use std::{io::Cursor, net::SocketAddr, sync::Arc, time::Duration};
+use std::{
+    io::Cursor,
+    net::SocketAddr,
+    time::Duration,
+};
 
-use anyhow::Result;
+use anyhow::{anyhow, bail, Context, Result};
 use axum::{
     extract::{Query, State},
     http::{HeaderMap, HeaderValue, StatusCode},
@@ -8,34 +12,63 @@ use axum::{
     routing::get,
     Router,
 };
-use bytes::Bytes;
+use bytes::{Bytes, BytesMut};
 use image::{
-    codecs::{gif::GifDecoder, jpeg::JpegEncoder},
+    codecs::{
+        gif::GifDecoder,
+        jpeg::JpegEncoder,
+    },
     imageops::FilterType,
     AnimationDecoder,
     GenericImageView,
 };
 use mimalloc::MiMalloc;
-use reqwest::Client;
+use moka::future::Cache;
+use reqwest::{
+    redirect::{Attempt, Policy},
+    Client,
+};
 use serde::Deserialize;
-use tokio::net::TcpListener;
+use tokio::{
+    net::TcpListener,
+    task::spawn_blocking,
+};
 use tower_http::trace::TraceLayer;
 use tracing::{error, info};
 
 #[global_allocator]
 static GLOBAL: MiMalloc = MiMalloc;
 
+const DEFAULT_THUMB_WIDTH: u32 = 260;
+const MIN_THUMB_WIDTH: u32 = 64;
+const MAX_THUMB_WIDTH: u32 = 2048;
+
+const MAX_IMAGE_BYTES: usize = 20 * 1024 * 1024;
+const MAX_IMAGE_PIXELS: u64 = 40_000_000;
+const MAX_REDIRECTS: usize = 5;
+
+const CACHE_CAPACITY: u64 = 512;
+const CACHE_TTL_SECS: u64 = 3600;
+const CACHE_TTI_SECS: u64 = 900;
+
 #[derive(Clone)]
 struct AppState {
     client: Client,
+    cache: Cache<String, CachedResponse>,
 }
 
-#[derive(Deserialize)]
+#[derive(Clone)]
+struct CachedResponse {
+    content_type: &'static str,
+    body: Bytes,
+}
+
+#[derive(Debug, Deserialize)]
 struct ImageQuery {
     url: String,
 }
 
-#[derive(Deserialize)]
+#[derive(Debug, Deserialize)]
 struct ThumbQuery {
     url: String,
     w: Option<u32>,
@@ -43,20 +76,31 @@ struct ThumbQuery {
 
 #[tokio::main]
 async fn main() -> Result<()> {
+    let log_filter =
+        std::env::var("RUST_LOG").unwrap_or_else(|_| "info".to_string());
+
     tracing_subscriber::fmt()
-        .with_env_filter("info")
+        .with_env_filter(log_filter)
         .init();
 
     let client = Client::builder()
         .pool_idle_timeout(Duration::from_secs(60))
-        .pool_max_idle_per_host(8)
+        .pool_max_idle_per_host(32)
         .tcp_keepalive(Duration::from_secs(30))
         .connect_timeout(Duration::from_secs(10))
         .timeout(Duration::from_secs(20))
         .user_agent("PinataImageBackend/1.0")
-        .build()?;
+        .redirect(pinterest_redirect_policy())
+        .build()
+        .context("failed to build HTTP client")?;
 
-    let state = Arc::new(AppState { client });
+    let cache = Cache::builder()
+        .max_capacity(CACHE_CAPACITY)
+        .time_to_live(Duration::from_secs(CACHE_TTL_SECS))
+        .time_to_idle(Duration::from_secs(CACHE_TTI_SECS))
+        .build();
+
+    let state = AppState { client, cache };
 
     let app = Router::new()
         .route("/health", get(health))
@@ -70,6 +114,7 @@ async fn main() -> Result<()> {
     info!("image backend listening on {}", addr);
 
     let listener = TcpListener::bind(addr).await?;
+
     axum::serve(listener, app).await?;
 
     Ok(())
@@ -80,151 +125,335 @@ async fn health() -> &'static str {
 }
 
 async fn fetch(
-    State(state): State<Arc<AppState>>,
+    State(state): State<AppState>,
     Query(query): Query<ImageQuery>,
 ) -> Response {
-    match process_fetch(state, query).await {
-        Ok(r) => r,
-        Err(e) => {
-            error!("fetch error: {:?}", e);
-            (StatusCode::BAD_GATEWAY, "failed to process image").into_response()
+    if let Err(err) = validate_url(&query.url) {
+        return (
+            StatusCode::BAD_REQUEST,
+            err.to_string(),
+        )
+            .into_response();
+    }
+
+    match process_fetch_cached(&state, query.url).await {
+        Ok(response) => response,
+        Err(err) => {
+            error!(error = ?err, "fetch error");
+
+            (
+                StatusCode::BAD_GATEWAY,
+                "failed to process image",
+            )
+                .into_response()
         }
     }
 }
 
 async fn thumb(
-    State(state): State<Arc<AppState>>,
+    State(state): State<AppState>,
     Query(query): Query<ThumbQuery>,
 ) -> Response {
-    match process_thumb(state, query).await {
-        Ok(r) => r,
-        Err(e) => {
-            error!("thumb error: {:?}", e);
-            (StatusCode::BAD_GATEWAY, "failed to process thumbnail").into_response()
+    if let Err(err) = validate_url(&query.url) {
+        return (
+            StatusCode::BAD_REQUEST,
+            err.to_string(),
+        )
+            .into_response();
+    }
+
+    let width = query
+        .w
+        .unwrap_or(DEFAULT_THUMB_WIDTH)
+        .clamp(MIN_THUMB_WIDTH, MAX_THUMB_WIDTH);
+
+    match process_thumb_cached(&state, query.url, width).await {
+        Ok(response) => response,
+        Err(err) => {
+            error!(error = ?err, "thumbnail error");
+
+            (
+                StatusCode::BAD_GATEWAY,
+                "failed to process thumbnail",
+            )
+                .into_response()
         }
     }
 }
 
 fn validate_url(url: &str) -> Result<()> {
-    let parsed = reqwest::Url::parse(url)?;
+    let parsed = reqwest::Url::parse(url)
+        .context("invalid URL")?;
 
     if parsed.scheme() != "https" {
-        anyhow::bail!("https required");
+        bail!("https required");
     }
 
     let host = parsed.host_str().unwrap_or("");
+
     if !host.eq_ignore_ascii_case("i.pinimg.com") {
-        anyhow::bail!("invalid host");
+        bail!("invalid host");
+    }
+
+    if !parsed.username().is_empty() || parsed.password().is_some() {
+        bail!("credentials are not allowed");
+    }
+
+    if parsed.port().is_some() {
+        bail!("custom ports are not allowed");
     }
 
     Ok(())
 }
 
-async fn download_image(client: &Client, url: &str) -> Result<Bytes> {
-    let resp = client.get(url).send().await?;
-
-    if !resp.status().is_success() {
-        anyhow::bail!("upstream status {}", resp.status());
+fn pinterest_redirect_policy() -> Policy {
+    Policy::custom(|attempt: Attempt<'_>| {
+        if attempt.previous().len() >= MAX_REDIRECTS {
+            return attempt.error(anyhow!("too many redirects"));
         }
 
-    Ok(resp.bytes().await?)
+        if validate_url(attempt.url().as_str()).is_ok() {
+            attempt.follow()
+        } else {
+            attempt.error(anyhow!(
+                "redirected outside i.pinimg.com"
+            ))
+        }
+    })
+}
+
+async fn download_image(
+    client: &Client,
+    url: &str,
+) -> Result<Bytes> {
+    let response = client
+        .get(url)
+        .send()
+        .await
+        .context("upstream request failed")?;
+
+    let status = response.status();
+
+    if !status.is_success() {
+        bail!("upstream status {}", status);
+    }
+
+    if let Some(content_length) = response.content_length() {
+        if content_length > MAX_IMAGE_BYTES as u64 {
+            bail!("image exceeds maximum allowed size");
+        }
+    }
+
+    let capacity = response
+        .content_length()
+        .map(|size| {
+            size.min(MAX_IMAGE_BYTES as u64) as usize
+        })
+        .unwrap_or(64 * 1024);
+
+    let mut body = BytesMut::with_capacity(capacity);
+    let mut response = response;
+
+    while let Some(chunk) = response
+        .chunk()
+        .await
+        .context("failed reading upstream body")?
+    {
+        if body.len() + chunk.len() > MAX_IMAGE_BYTES {
+            bail!("image exceeds maximum allowed size");
+        }
+
+        body.extend_from_slice(&chunk);
+    }
+
+    if body.is_empty() {
+        bail!("upstream returned an empty body");
+    }
+
+    Ok(body.freeze())
 }
 
 fn is_gif(data: &[u8]) -> bool {
     data.starts_with(b"GIF87a") || data.starts_with(b"GIF89a")
 }
 
-fn gif_passthrough(data: Bytes) -> Response {
+fn is_jpeg(data: &[u8]) -> bool {
+    data.len() >= 3
+        && data[0] == 0xFF
+        && data[1] == 0xD8
+        && data[2] == 0xFF
+}
+
+fn is_png(data: &[u8]) -> bool {
+    data.starts_with(b"\x89PNG\r\n\x1a\n")
+}
+
+fn is_webp(data: &[u8]) -> bool {
+    data.len() >= 12
+        && data.starts_with(b"RIFF")
+        && &data[8..12] == b"WEBP"
+}
+
+fn is_supported_image(data: &[u8]) -> bool {
+    is_gif(data)
+        || is_jpeg(data)
+        || is_png(data)
+        || is_webp(data)
+}
+
+fn response_from_cached(
+    cached: CachedResponse,
+) -> Response {
     let mut headers = HeaderMap::new();
 
     headers.insert(
         "content-type",
-        HeaderValue::from_static("image/gif"),
+        HeaderValue::from_static(cached.content_type),
     );
+
     headers.insert(
         "cache-control",
-        HeaderValue::from_static("public, max-age=86400, immutable"),
+        HeaderValue::from_static(
+            "public, max-age=86400, immutable",
+        ),
     );
 
-    (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"),
+        "x-content-type-options",
+        HeaderValue::from_static("nosniff"),
     );
 
-    (headers, data).into_response()
+    (headers, cached.body).into_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()))
+fn cached_jpeg(data: Vec<u8>) -> CachedResponse {
+    CachedResponse {
+        content_type: "image/jpeg",
+        body: Bytes::from(data),
+    }
 }
 
-async fn process_fetch(state: Arc<AppState>, query: ImageQuery) -> Result<Response> {
-    validate_url(&query.url)?;
+fn cached_gif(data: Bytes) -> CachedResponse {
+    CachedResponse {
+        content_type: "image/gif",
+        body: data,
+    }
+}
 
-    let data = download_image(&state.client, &query.url).await?;
+fn cached_original_jpeg(data: Bytes) -> CachedResponse {
+    CachedResponse {
+        content_type: "image/jpeg",
+        body: data,
+    }
+}
+
+async fn process_fetch_cached(
+    state: &AppState,
+    url: String,
+) -> Result<Response> {
+    let key = format!("fetch:{url}");
+
+    let cache = state.cache.clone();
+    let client = state.client.clone();
+
+    let cached = cache
+        .try_get_with(key, async move {
+            build_fetch_response(&client, &url)
+                .await
+                .map_err(|err| err.to_string())
+        })
+        .await
+        .map_err(|err| anyhow!(err.to_string()))?;
+
+    Ok(response_from_cached(cached))
+}
+
+async fn build_fetch_response(
+    client: &Client,
+    url: &str,
+) -> Result<CachedResponse> {
+    let data = download_image(client, url).await?;
+
+    if !is_supported_image(&data) {
+        bail!("unsupported image format");
+    }
+
+    if is_jpeg(&data) {
+        return Ok(cached_original_jpeg(data));
+    }
 
     if is_gif(&data) {
-        return Ok(gif_passthrough(data));
+        return Ok(cached_gif(data));
     }
 
-    let img = image::load_from_memory(&data)?;
-    let rgb = img.to_rgb8();
+    let jpeg = spawn_blocking(move || {
+        transcode_to_jpeg(data, 75)
+    })
+    .await
+    .context("image processing task failed")??;
 
-    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))
+    Ok(cached_jpeg(jpeg))
 }
 
-async fn process_thumb(state: Arc<AppState>, query: ThumbQuery) -> Result<Response> {
-    validate_url(&query.url)?;
+async fn process_thumb_cached(
+    state: &AppState,
+    url: String,
+    width: u32,
+) -> Result<Response> {
+    let key = format!("thumb:{width}:{url}");
 
-    let width = query.w.unwrap_or(260).clamp(64, 2048);
+    let cache = state.cache.clone();
+    let client = state.client.clone();
 
-    let data = download_image(&state.client, &query.url).await?;
+    let cached = cache
+        .try_get_with(key, async move {
+            build_thumb_response(&client, &url, width)
+                .await
+                .map_err(|err| err.to_string())
+        })
+        .await
+        .map_err(|err| anyhow!(err.to_string()))?;
 
-    let img = if is_gif(&data) {
-        decode_first_gif_frame(&data)?
-    } else {
-        image::load_from_memory(&data)?
-    };
+    Ok(response_from_cached(cached))
+}
+
+async fn build_thumb_response(
+    client: &Client,
+    url: &str,
+    width: u32,
+) -> Result<CachedResponse> {
+    let data = download_image(client, url).await?;
+
+    if !is_supported_image(&data) {
+        bail!("unsupported image format");
+    }
+
+    spawn_blocking(move || {
+        build_thumbnail(data, width)
+    })
+    .await
+    .context("thumbnail processing task failed")?
+}
+
+fn build_thumbnail(
+    data: Bytes,
+    width: u32,
+) -> Result<CachedResponse> {
+    if is_gif(&data) {
+        let img = decode_first_gif_frame(&data)?;
+
+        validate_dimensions(&img)?;
 
         let (src_w, src_h) = img.dimensions();
 
         if width >= src_w {
-        return process_fetch(
-            state,
-            ImageQuery { url: query.url },
-        )
-        .await;
+            return Ok(cached_gif(data));
         }
 
-    let height = ((src_h as f32 * width as f32) / src_w as f32) as u32;
+        let height =
+            calculate_height(src_w, src_h, width);
 
-    let resized = img.resize_exact(
+        let resized = img.resize(
             width,
             height.max(1),
             FilterType::Triangle,
@@ -232,13 +461,161 @@ async fn process_thumb(state: Arc<AppState>, query: ThumbQuery) -> Result<Respon
 
         let rgb = resized.to_rgb8();
 
-    let mut out = Vec::new();
+        let jpeg = encode_jpeg(&rgb, 72)?;
 
-    {
-        let mut cursor = Cursor::new(&mut out);
-        let mut encoder = JpegEncoder::new_with_quality(&mut cursor, 72);
-        encoder.encode_image(&rgb)?;
+        return Ok(cached_jpeg(jpeg));
     }
 
-    Ok(jpeg_response(out))
+    if is_jpeg(&data) {
+        let img = image::load_from_memory(&data)
+            .context("failed to decode JPEG")?;
+
+        validate_dimensions(&img)?;
+
+        let (src_w, src_h) = img.dimensions();
+
+        if width >= src_w {
+            return Ok(cached_original_jpeg(data));
+        }
+
+        let height =
+            calculate_height(src_w, src_h, width);
+
+        let resized = img.resize(
+            width,
+            height.max(1),
+            FilterType::Triangle,
+        );
+
+        let rgb = resized.to_rgb8();
+
+        let jpeg = encode_jpeg(&rgb, 72)?;
+
+        return Ok(cached_jpeg(jpeg));
+    }
+
+    let img = image::load_from_memory(&data)
+        .context("failed to decode image")?;
+
+    validate_dimensions(&img)?;
+
+    let (src_w, src_h) = img.dimensions();
+
+    if width >= src_w {
+        let rgb = img.to_rgb8();
+
+        let jpeg = encode_jpeg(&rgb, 75)?;
+
+        return Ok(cached_jpeg(jpeg));
+    }
+
+    let height =
+        calculate_height(src_w, src_h, width);
+
+    let resized = img.resize(
+        width,
+        height.max(1),
+        FilterType::Triangle,
+    );
+
+    let rgb = resized.to_rgb8();
+
+    let jpeg = encode_jpeg(&rgb, 72)?;
+
+    Ok(cached_jpeg(jpeg))
+}
+
+fn validate_dimensions(
+    img: &image::DynamicImage,
+) -> Result<()> {
+    let (width, height) = img.dimensions();
+
+    if width == 0 || height == 0 {
+        bail!("invalid image dimensions");
+    }
+
+    let pixels = (width as u64)
+        .checked_mul(height as u64)
+        .ok_or_else(|| anyhow!("image dimensions overflow"))?;
+
+    if pixels > MAX_IMAGE_PIXELS {
+        bail!(
+            "image dimensions too large: {}x{}",
+            width,
+            height
+        );
+    }
+
+    Ok(())
+}
+
+fn decode_first_gif_frame(
+    data: &[u8],
+) -> Result<image::DynamicImage> {
+    let decoder = GifDecoder::new(Cursor::new(data))
+        .context("failed to create GIF decoder")?;
+
+    let mut frames = decoder.into_frames();
+
+    let frame = frames
+        .next()
+        .ok_or_else(|| anyhow!("empty GIF"))?
+        .context("failed to decode GIF frame")?;
+
+    Ok(image::DynamicImage::ImageRgba8(
+        frame.into_buffer(),
+    ))
+}
+
+fn encode_jpeg(
+    rgb: &image::RgbImage,
+    quality: u8,
+) -> Result<Vec<u8>> {
+    let pixel_count = (rgb.width() as usize)
+        .saturating_mul(rgb.height() as usize);
+
+    let capacity = (pixel_count / 3)
+        .clamp(4 * 1024, 4 * 1024 * 1024);
+
+    let mut output = Vec::with_capacity(capacity);
+
+    {
+        let mut encoder =
+            JpegEncoder::new_with_quality(
+                &mut output,
+                quality,
+            );
+
+        encoder
+            .encode_image(rgb)
+            .context("failed to encode JPEG")?;
+    }
+
+    Ok(output)
+}
+
+fn transcode_to_jpeg(
+    data: Bytes,
+    quality: u8,
+) -> Result<Vec<u8>> {
+    let img = image::load_from_memory(&data)
+        .context("failed to decode image")?;
+
+    validate_dimensions(&img)?;
+
+    let rgb = img.to_rgb8();
+
+    encode_jpeg(&rgb, quality)
+}
+
+fn calculate_height(
+    src_width: u32,
+    src_height: u32,
+    target_width: u32,
+) -> u32 {
+    (
+        (src_height as u64)
+            .saturating_mul(target_width as u64)
+            / src_width as u64
+    ) as u32
 }