diff --git a/Cargo.lock b/Cargo.lock index b126cdc0..74239300 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3544,9 +3544,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.153" +version = "0.2.159" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c198f91728a82281a64e1f4f9eeb25d82cb32a5de251c6bd1b5154d63a8e7bd" +checksum = "561d97a539a36e26a9a5fad1ea11a3039a67714694aaa379433e580854bc3dc5" [[package]] name = "libfuzzer-sys" @@ -5482,9 +5482,16 @@ dependencies = [ "aws-config", "aws-sdk-s3", "base64 0.22.1", + "ffprobe", + "image", + "imagesize", + "jxl-oxide", "revolt-config", "revolt-result", + "tempfile", + "tracing", "typenum", + "webp", ] [[package]] @@ -5498,10 +5505,12 @@ dependencies = [ "moka", "reqwest 0.12.4", "revolt-config", + "revolt-files", "revolt-models", "revolt-result", "serde", "serde_json", + "tempfile", "tokio 1.35.1", "tracing", "tracing-subscriber", @@ -5958,9 +5967,9 @@ dependencies = [ [[package]] name = "rustix" -version = "0.38.34" +version = "0.38.37" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70dc5ec042f7a43c4a73241207cecc9873a06d45debb38b329f8541d85c2730f" +checksum = "8acb788b847c24f28525660c4d7758620a7210875711f79e7f663cc152726811" dependencies = [ "bitflags 2.6.0", "errno", @@ -6825,9 +6834,9 @@ checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" [[package]] name = "tempfile" -version = "3.12.0" +version = "3.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04cbcdd0c794ebb0d4cf35e88edd2f7d2c4c3e9a5a6dab322839b321c6a87a64" +checksum = "f0f2c9fc62d0beef6951ccffd757e241266a2c833136efbe35af6cd2567dca5b" dependencies = [ "cfg-if", "fastrand 2.1.1", diff --git a/crates/core/files/Cargo.toml b/crates/core/files/Cargo.toml index a340894f..04292be9 100644 --- a/crates/core/files/Cargo.toml +++ b/crates/core/files/Cargo.toml @@ -7,6 +7,12 @@ authors = ["Paul Makles "] description = "Revolt Backend: S3 and encryption subroutines" [dependencies] +tracing = "0.1" + +ffprobe = "0.4.0" +imagesize = "0.13.0" +tempfile = "3.12.0" + base64 = "0.22.1" aes-gcm = "0.10.3" typenum = "1.17.0" @@ -18,3 +24,10 @@ revolt-config = { version = "0.7.16", path = "../config", features = [ "report-macros", ] } revolt-result = { version = "0.7.16", path = "../result" } + +# image processing +jxl-oxide = "0.8.1" +image = { version = "0.25.2" } + +# encoding +webp = "0.3.0" diff --git a/crates/core/files/src/lib.rs b/crates/core/files/src/lib.rs index 5d2bdbd6..a3028ab9 100644 --- a/crates/core/files/src/lib.rs +++ b/crates/core/files/src/lib.rs @@ -1,9 +1,10 @@ -use std::io::Write; +use std::io::{BufRead, Read, Seek, Write}; use aes_gcm::{ aead::{AeadCore, AeadMutInPlace, OsRng}, Aes256Gcm, Key, KeyInit, Nonce, }; +use image::DynamicImage; use revolt_config::{config, report_internal_error, FilesS3}; use revolt_result::{create_error, Result}; @@ -13,6 +14,7 @@ use aws_sdk_s3::{ }; use base64::prelude::*; +use tempfile::NamedTempFile; /// Size of the authentication tag in the buffer pub const AUTHENTICATION_TAG_SIZE_BYTES: usize = 16; @@ -108,3 +110,90 @@ pub async fn upload_to_s3(bucket_id: &str, path: &str, buf: &[u8]) -> Result Option<(usize, usize)> { + if let Ok(size) = imagesize::size(f.path()) + .inspect_err(|err| tracing::error!("Failed to generate image size! {err:?}")) + { + Some((size.width, size.height)) + } else { + None + } +} + +/// Determine size of video at temp file +pub fn video_size(f: &NamedTempFile) -> Option<(i64, i64)> { + if let Ok(data) = ffprobe::ffprobe(f.path()) + .inspect_err(|err| tracing::error!("Failed to ffprobe file! {err:?}")) + { + // Use first valid stream + for stream in data.streams { + if let (Some(w), Some(h)) = (stream.width, stream.height) { + return Some((w, h)); + } + } + + None + } else { + None + } +} + +/// Decode image from reader +pub fn decode_image( + reader: &mut R, + is_jxl: bool, +) -> Result { + if is_jxl { + Err(create_error!(LabelMe)) + } else { + // Check if we can read using image-rs crate + report_internal_error!(report_internal_error!( + image::ImageReader::new(reader).with_guessed_format() + )? + .decode()) + } +} + +/// Check whether given reader has a valid image +pub fn is_valid_image(reader: &mut R, is_jxl: bool) -> bool { + if is_jxl { + // Check if we can read using jxl-oxide crate + jxl_oxide::JxlImage::builder() + .read(reader) + .inspect_err(|err| tracing::error!("Failed to read JXL! {err:?}")) + .is_ok() + } else { + !matches!( + // Check if we can read using image-rs crate + image::ImageReader::new(reader) + .with_guessed_format() + .inspect_err(|err| tracing::error!("Failed to read image! {err:?}")) + .map(|f| f.decode()), + Err(_) | Ok(Err(_)) + ) + } +} + +/// Create thumbnail from given image +pub async fn create_thumbnail(image: DynamicImage, tag: &str) -> Vec { + // Load configuration + let config = config().await; + let [w, h] = config.files.preview.get(tag).unwrap(); + + // Create thumbnail + //.resize(width as u32, height as u32, image::imageops::FilterType::Gaussian) + // resize is about 2.5x slower, + // thumbnail doesn't have terrible quality + // so we use thumbnail + let image = image.thumbnail(image.width().min(*w as u32), image.height().min(*h as u32)); + + // Encode it into WEBP + let encoder = webp::Encoder::from_image(&image).expect("Could not create encoder."); + if config.files.webp_quality != 100.0 { + encoder.encode(config.files.webp_quality).to_vec() + } else { + encoder.encode_lossless().to_vec() + } +} diff --git a/crates/services/autumn/src/api.rs b/crates/services/autumn/src/api.rs index 3c61ee90..86b1dcc1 100644 --- a/crates/services/autumn/src/api.rs +++ b/crates/services/autumn/src/api.rs @@ -15,7 +15,9 @@ use image::ImageReader; use lazy_static::lazy_static; use revolt_config::{config, report_internal_error}; use revolt_database::{iso8601_timestamp::Timestamp, Database, FileHash, Metadata, User}; -use revolt_files::{fetch_from_s3, upload_to_s3, AUTHENTICATION_TAG_SIZE_BYTES}; +use revolt_files::{ + create_thumbnail, decode_image, fetch_from_s3, upload_to_s3, AUTHENTICATION_TAG_SIZE_BYTES, +}; use revolt_result::{create_error, Result}; use serde::{Deserialize, Serialize}; use sha2::Digest; @@ -352,30 +354,8 @@ async fn fetch_preview( // Original image data let data = retrieve_file_by_hash(&hash).await?; - // Dimensions we need to resize to - let config = config().await; - let [w, h] = config.files.preview.get(tag).unwrap(); - - // Read the image and resize it - // TODO: use jxl_oxide to process image/jxl files - let image = report_internal_error!(report_internal_error!(ImageReader::new(Cursor::new( - data - )) - .with_guessed_format())? - .decode())? - //.resize(width as u32, height as u32, image::imageops::FilterType::Gaussian) - // resize is about 2.5x slower, - // thumbnail doesn't have terrible quality - // so we use thumbnail - .thumbnail(*w as u32, *h as u32); - - // Encode it into WEBP - let encoder = webp::Encoder::from_image(&image).expect("Could not create encoder."); - let data = if config.files.webp_quality != 100.0 { - encoder.encode(config.files.webp_quality).to_vec() - } else { - encoder.encode_lossless().to_vec() - }; + // Read image and create thumbnail + let data = create_thumbnail(decode_image(&mut Cursor::new(data), false)?, tag).await; Ok(( [ diff --git a/crates/services/autumn/src/metadata.rs b/crates/services/autumn/src/metadata.rs index e7526a7c..72b610e9 100644 --- a/crates/services/autumn/src/metadata.rs +++ b/crates/services/autumn/src/metadata.rs @@ -1,6 +1,7 @@ use std::io::Cursor; use revolt_database::Metadata; +use revolt_files::{image_size, video_size}; use tempfile::NamedTempFile; /// Intersection of what infer can detect and what image-rs supports @@ -21,32 +22,19 @@ static SUPPORTED_IMAGE_MIME: [&str; 9] = [ /// Generate metadata from file, using mime type as a hint pub fn generate_metadata(f: &NamedTempFile, mime_type: &str) -> Metadata { if SUPPORTED_IMAGE_MIME.contains(&mime_type) { - if let Ok(size) = imagesize::size(f.path()) - .inspect_err(|err| tracing::error!("Failed to generate image size! {err:?}")) - { - if let (Ok(width), Ok(height)) = (size.width.try_into(), size.height.try_into()) { - return Metadata::Image { width, height }; - } - } - - Metadata::File + image_size(f) + .map(|(width, height)| Metadata::Image { + width: width as isize, + height: height as isize, + }) + .unwrap_or_default() } else if mime_type.starts_with("video/") { - if let Ok(data) = ffprobe::ffprobe(f.path()) - .inspect_err(|err| tracing::error!("Failed to ffprobe file! {err:?}")) - { - // Use first valid stream - for stream in data.streams { - if let (Some(w), Some(h)) = (stream.width, stream.height) { - if let (Ok(width), Ok(height)) = (w.try_into(), h.try_into()) { - return Metadata::Video { width, height }; - } - } - } - - Metadata::File - } else { - Metadata::File - } + video_size(f) + .map(|(width, height)| Metadata::Video { + width: width as isize, + height: height as isize, + }) + .unwrap_or_default() } else if mime_type.starts_with("audio/") { Metadata::Audio } else if mime_type == "plain/text" { diff --git a/crates/services/january/Cargo.toml b/crates/services/january/Cargo.toml index 242e4696..68f38ab5 100644 --- a/crates/services/january/Cargo.toml +++ b/crates/services/january/Cargo.toml @@ -6,6 +6,7 @@ edition = "2021" [dependencies] # Utility mime = "0.3.17" +tempfile = "3.13.0" lazy_static = "1.5.0" moka = { version = "0.12.8", features = ["future"] } @@ -30,6 +31,7 @@ revolt-result = { version = "0.7.16", path = "../../core/result", features = [ "utoipa", "axum", ] } +revolt-files = { version = "0.7.16", path = "../../core/files" } # Axum / web server axum = { version = "0.7.5" } diff --git a/crates/services/january/src/api.rs b/crates/services/january/src/api.rs index cba288f5..1d9c8557 100644 --- a/crates/services/january/src/api.rs +++ b/crates/services/january/src/api.rs @@ -1,4 +1,5 @@ -use axum::{body::Bytes, extract::Query, routing::get, Json, Router}; +use axum::{extract::Query, response::IntoResponse, routing::get, Json, Router}; +use reqwest::header; use revolt_models::v0::Embed; use revolt_result::Result; use serde::{Deserialize, Serialize}; @@ -11,6 +12,8 @@ use axum_extra::{ use crate::requests::Request; +pub static CACHE_CONTROL: &str = "public, max-age=600, immutable"; + pub async fn router() -> Router { Router::new() .route("/", get(root)) @@ -59,8 +62,17 @@ struct UrlQuery { ("url" = String, Query, description = "URL to fetch") ), )] -async fn proxy(Query(UrlQuery { url }): Query) -> Result { - Request::proxy_file(&url).await +async fn proxy(Query(UrlQuery { url }): Query) -> Result { + Request::proxy_file(&url).await.map(|(content_type, data)| { + ( + [ + (header::CONTENT_TYPE, content_type), + (header::CONTENT_DISPOSITION, "inline".to_owned()), + (header::CACHE_CONTROL, CACHE_CONTROL.to_owned()), + ], + data, + ) + }) } /// Generate embed for a given URL diff --git a/crates/services/january/src/requests.rs b/crates/services/january/src/requests.rs index a25c3918..5b8429d2 100644 --- a/crates/services/january/src/requests.rs +++ b/crates/services/january/src/requests.rs @@ -1,9 +1,13 @@ -use std::time::Duration; +use std::{ + io::{Cursor, Write}, + time::Duration, +}; -use axum::body::Bytes; use lazy_static::lazy_static; use mime::Mime; use reqwest::{header::CONTENT_TYPE, redirect, Client, Response}; +use revolt_config::report_internal_error; +use revolt_files::{create_thumbnail, decode_image, is_valid_image, video_size}; use revolt_models::v0::Embed; use revolt_result::{create_error, Result}; @@ -25,7 +29,7 @@ lazy_static! { .expect("reqwest Client"); /// Cache for proxy results - static ref PROXY_CACHE: moka::future::Cache> = moka::future::Cache::builder() + static ref PROXY_CACHE: moka::future::Cache)>> = moka::future::Cache::builder() .max_capacity(10_000) // TODO config .time_to_live(Duration::from_secs(60)) // TODO config .build(); @@ -45,11 +49,52 @@ pub struct Request { impl Request { /// Proxy a given URL - pub async fn proxy_file(url: &str) -> Result { + pub async fn proxy_file(url: &str) -> Result<(String, Vec)> { if let Some(hit) = PROXY_CACHE.get(url).await { hit } else { - todo!() + let Request { response, mime } = Request::new(url).await?; + + if matches!(mime.type_(), mime::IMAGE | mime::VIDEO) { + let bytes = response.bytes().await.map_err(|_| create_error!(LabelMe)); + + let result = match bytes { + Ok(bytes) => { + if matches!(mime.type_(), mime::IMAGE) { + let reader = &mut Cursor::new(&bytes); + + if matches!(mime.subtype(), mime::GIF) { + if is_valid_image(reader, false) { + Ok(("image/gif".to_owned(), bytes.to_vec())) + } else { + Err(create_error!(LabelMe)) + } + } else { + Ok(( + "image/webp".to_owned(), + create_thumbnail(decode_image(reader, false)?, "attachments") + .await, + )) + } + } else { + let mut file = report_internal_error!(tempfile::NamedTempFile::new())?; + report_internal_error!(file.write_all(&bytes))?; + if video_size(&file).is_some() { + Ok((mime.to_string(), bytes.to_vec())) + } else { + Err(create_error!(LabelMe)) + } + } + } + Err(err) => Err(err), + }; + + PROXY_CACHE.insert(url.to_owned(), result.clone()).await; + result + } else { + // Err(create_error!()) + todo!() // no proxy + } } }