HMC/handmade-revolt-backend
6
0
Fork 0
forked from jmug/stoatchat
Code Pull Requests Actions Packages Releases 2 Activity

feat(services/january): proxy images/videos

Browse Source
This commit is contained in:
Paul Makles
2024-10-01 19:02:39 +01:00
parent 8fc791f81a
commit 21335b3297
8 changed files with 203 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
Cargo.lock generated
View File

@@ -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",

13
crates/core/files/Cargo.toml
View File

@@ -7,6 +7,12 @@ authors = ["Paul Makles <me@insrt.uk>"]
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"

91
crates/core/files/src/lib.rs
View File

@@ -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<Str
Ok(BASE64_STANDARD.encode(nonce))
}
/// Determine size of image at temp file
pub fn image_size(f: &NamedTempFile) -> 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<R: Read + BufRead + Seek>(
reader: &mut R,
is_jxl: bool,
) -> Result<DynamicImage> {
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<R: Read + BufRead + Seek>(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<u8> {
// 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()
}
}

30
crates/services/autumn/src/api.rs
View File

@@ -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((
[

38
crates/services/autumn/src/metadata.rs
View File

@@ -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" {

2
crates/services/january/Cargo.toml
View File

@@ -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" }

18
crates/services/january/src/api.rs
View File

@@ -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<UrlQuery>) -> Result<Bytes> {
Request::proxy_file(&url).await
async fn proxy(Query(UrlQuery { url }): Query<UrlQuery>) -> Result<impl IntoResponse> {
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

55
crates/services/january/src/requests.rs
View File

@@ -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<String, Result<Bytes>> = moka::future::Cache::builder()
static ref PROXY_CACHE: moka::future::Cache<String, Result<(String, Vec<u8>)>> = 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<Bytes> {
pub async fn proxy_file(url: &str) -> Result<(String, Vec<u8>)> {
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
}
}
}