diff --git a/Cargo.toml b/Cargo.toml index eec8d27..8c36195 100644 --- a/Cargo.toml +++ b/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"] } \ No newline at end of file + +tracing-subscriber = { + version = "0.3", + features = ["env-filter"] +} diff --git a/src/main.rs b/src/main.rs index b68d3b8..de918e2 100644 --- a/src/main.rs +++ b/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, } -#[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, @@ -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,394 @@ async fn health() -> &'static str { } async fn fetch( - State(state): State>, + State(state): State, Query(query): Query, ) -> 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>, + State(state): State, Query(query): Query, ) -> 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 { - let resp = client.get(url).send().await?; +fn pinterest_redirect_policy() -> Policy { + Policy::custom(|attempt: Attempt<'_>| { + if attempt.previous().len() >= MAX_REDIRECTS { + return attempt.error(anyhow!("too many redirects")); + } - if !resp.status().is_success() { - anyhow::bail!("upstream status {}", resp.status()); + 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 { + let response = client + .get(url) + .send() + .await + .context("upstream request failed")?; + + let status = response.status(); + + if !status.is_success() { + bail!("upstream status {}", status); } - Ok(resp.bytes().await?) + 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) -> 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 { - 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) -> CachedResponse { + CachedResponse { + content_type: "image/jpeg", + body: Bytes::from(data), + } } -async fn process_fetch(state: Arc, query: ImageQuery) -> Result { - 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 { + 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 { + 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, query: ThumbQuery) -> Result { - validate_url(&query.url)?; +async fn process_thumb_cached( + state: &AppState, + url: String, + width: u32, +) -> Result { + 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 { + 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 { + 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 Ok(cached_gif(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)); + } + + 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 { - return process_fetch( - state, - ImageQuery { url: query.url }, - ) - .await; + let rgb = img.to_rgb8(); + + let jpeg = encode_jpeg(&rgb, 75)?; + + return Ok(cached_jpeg(jpeg)); } - 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 +520,102 @@ async fn process_thumb(state: Arc, query: ThumbQuery) -> Result Result<()> { + let (width, height) = img.dimensions(); + + if width == 0 || height == 0 { + bail!("invalid image dimensions"); } - Ok(jpeg_response(out)) -} \ No newline at end of file + 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 { + 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> { + 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> { + 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 +}