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

feat(services/autumn): file uploads

Browse Source 
feat(services/autumn): deduplicate uploads
feat(services/autumn): ClamAV support
This commit is contained in:
Paul Makles
2024-09-11 14:29:52 +01:00
parent f4104612b2
commit ace6c30ba5
23 changed files with 427 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
crates/services/autumn/Cargo.toml
View File

@@ -4,14 +4,21 @@ version = "0.7.14"
edition = "2021"
[dependencies]
# ID generation
ulid = "1.1.3"
nanoid = "0.4.0"
# Media processing
webp = "0.3.0"
sha2 = "0.10.8"
jxl-oxide = "0.8.1"
kamadak-exif = "0.5.4"
# revolt_little_exif = "0.4.0"
image = { version = "0.25.2" } # avif encode requires dav1d system library: features = ["avif-native"]
# File processing
revolt_clamav-client = { version = "0.1.5" }
simdutf8 = { version = "0.1.4", features = ["aarch64_neon"] }
image = { version = "0.25.2" } # requires dav1d system library: features = ["avif-native"]
# Content type processing
infer = "0.16.0"

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

@@ -5,7 +5,7 @@ use std::{
use axum::{
extract::{DefaultBodyLimit, Path, State},
http::{header, HeaderMap},
http::header,
response::{IntoResponse, Redirect, Response},
routing::{get, post},
Json, Router,
@@ -13,9 +13,9 @@ use axum::{
use axum_typed_multipart::{FieldData, TryFromMultipart, TypedMultipart};
use image::ImageReader;
use lazy_static::lazy_static;
use revolt_config::config;
use revolt_database::{Database, FileHash, Metadata, User};
use revolt_files::{fetch_from_s3, AUTHENTICATION_TAG_SIZE_BYTES};
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_result::{create_error, Result};
use serde::{Deserialize, Serialize};
use sha2::Digest;
@@ -87,7 +87,7 @@ async fn root() -> Json<RootResponse> {
}
/// Available tags to upload to
#[derive(Deserialize, Debug, ToSchema, strum_macros::IntoStaticStr)]
#[derive(Clone, Deserialize, Debug, ToSchema, strum_macros::IntoStaticStr)]
#[allow(non_camel_case_types)]
pub enum Tag {
attachments,
@@ -111,7 +111,7 @@ pub struct UploadPayload {
#[derive(Serialize, Debug, ToSchema)]
pub struct UploadResponse {
/// ID to attach uploaded file to object
id: &'static str,
id: String,
}
/// Upload a file
@@ -146,7 +146,7 @@ async fn upload_file(
user: User,
Path(tag): Path<Tag>,
TypedMultipart(UploadPayload { mut file }): TypedMultipart<UploadPayload>,
) -> axum::response::Result<Json<UploadResponse>> {
) -> Result<Json<UploadResponse>> {
// Fetch configuration
let config = config().await;
@@ -154,31 +154,29 @@ async fn upload_file(
let now = Instant::now();
// Extract the filename, or give it a generic name
let file_name = file.metadata.file_name.unwrap_or("unnamed-file".to_owned());
let filename = file.metadata.file_name.unwrap_or("unnamed-file".to_owned());
// Load file to memory
let mut buf = Vec::<u8>::new();
file.contents
.read_to_end(&mut buf)
.map_err(|_| create_error!(InternalError))?;
report_internal_error!(file.contents.read_to_end(&mut buf))?;
// Take note of original file size
let original_file_size = buf.len();
// Ensure the file is not empty
if original_file_size < config.files.limit.min_file_size {
return Err(create_error!(FileTooSmall).into());
return Err(create_error!(FileTooSmall));
}
// Get user's file upload limits
let limits = user.limits().await;
let size_limit = *limits
.file_upload_size_limit
.get(tag.into())
.get(tag.clone().into())
.expect("size limit");
if original_file_size > size_limit {
return Err(create_error!(FileTooLarge { max: size_limit }).into());
return Err(create_error!(FileTooLarge { max: size_limit }));
}
// Generate sha256 hash
@@ -188,18 +186,27 @@ async fn upload_file(
hasher.finalize()
};
// Generate an ID for this file
let id = if matches!(tag, Tag::emojis) {
ulid::Ulid::new().to_string()
} else {
nanoid::nanoid!(42)
};
// Find an existing hash and use that if possible
if let Ok(file_hash) = db
.fetch_attachment_hash(&format!("{original_hash:02x}"))
// or match processed (add to query TODO)
.await
{
unimplemented!() // TODO
// then: create attachment and return UploadResponse { id }
let tag: &'static str = tag.into();
db.insert_attachment(&file_hash.into_file(id.clone(), tag.to_owned(), filename, user.id))
.await?;
return Ok(Json(UploadResponse { id }));
}
// Determine the mime type for the file
let mime_type = determine_mime_type(&mut file.contents, &buf, &file_name);
let mime_type = determine_mime_type(&mut file.contents, &buf, &filename);
// Check blocklist for mime type
if config
@@ -208,7 +215,7 @@ async fn upload_file(
.iter()
.any(|m| m == mime_type)
{
return Err(create_error!(FileTypeNotAllowed).into());
return Err(create_error!(FileTypeNotAllowed));
}
// Determine metadata for the file
@@ -217,7 +224,14 @@ async fn upload_file(
// Strip metadata
let (buf, metadata) = strip_metadata(file.contents, buf, metadata, mime_type).await?;
// TODO: Virus scan Metadata::File
// Virus scan files if ClamAV is configured
if matches!(metadata, Metadata::File)
&& (config.files.scan_mime_types.is_empty()
|| config.files.scan_mime_types.iter().any(|v| v == mime_type))
&& crate::clamav::is_malware(&buf).await?
{
return Err(create_error!(InternalError));
}
// Print file information for debug purposes
let new_file_size = buf.len() + AUTHENTICATION_TAG_SIZE_BYTES;
@@ -229,15 +243,41 @@ async fn upload_file(
let process_ratio = new_file_size as f32 / original_file_size as f32;
let time_to_process = Instant::now() - now;
tracing::info!("Received file {file_name}\nOriginal hash: {original_hash:02x}\nOriginal size: {original_file_size} bytes\nMime type: {mime_type}\nMetadata: {metadata:?}\nProcessed file size: {new_file_size} bytes ({:.2}%).\nProcessed hash: {processed_hash:02x}\nProcessing took {time_to_process:?}", process_ratio * 100.0);
tracing::info!("Received file {filename}\nOriginal hash: {original_hash:02x}\nOriginal size: {original_file_size} bytes\nMime type: {mime_type}\nMetadata: {metadata:?}\nProcessed file size: {new_file_size} bytes ({:.2}%).\nProcessed hash: {processed_hash:02x}\nProcessing took {time_to_process:?}", process_ratio * 100.0);
// TODO: encrypt the file and generate FileHash
// note: file_size is processed file's size (incl. authTag in encrypted buffer)
// TODO: insert FileHash
// TODO: upload to S3
// TODO: create attachment
// Create hash entry in database
let file_hash = FileHash {
id: format!("{original_hash:02x}"),
processed_hash: format!("{processed_hash:02x}"),
Ok(Json(UploadResponse { id: "aaa" }))
created_at: Timestamp::now_utc(),
bucket_id: config.files.s3.default_bucket,
path: format!("{original_hash:02x}"),
iv: String::new(), // indicates file is not uploaded yet
metadata,
content_type: mime_type.to_owned(),
size: new_file_size as isize,
};
db.insert_attachment_hash(&file_hash).await?;
// Upload the file to S3 and commit nonce to database
let upload_start = Instant::now();
let nonce = upload_to_s3(&file_hash.bucket_id, &file_hash.id, &buf).await?;
db.set_attachment_hash_nonce(&file_hash.id, &nonce).await?;
// Debug information
let time_to_upload = Instant::now() - upload_start;
tracing::info!("Took {time_to_upload:?} to upload {new_file_size} bytes to S3.");
// Finally, create the file and return its ID
let tag: &'static str = tag.into();
db.insert_attachment(&file_hash.into_file(id.clone(), tag.to_owned(), filename, user.id))
.await?;
Ok(Json(UploadResponse { id }))
}
/// Header value used for cache control
@@ -306,16 +346,16 @@ async fn fetch_preview(
// Read the image and resize it
// TODO: use jxl_oxide to process image/jxl files
let image = ImageReader::new(Cursor::new(data))
.with_guessed_format()
.map_err(|_| create_error!(InternalError))?
.decode()
.map_err(|_| create_error!(InternalError))?
//.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);
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.");

49
crates/services/autumn/src/clamav.rs Normal file
View File

@@ -0,0 +1,49 @@
use std::time::Duration;
use revolt_config::{config, report_internal_error};
use revolt_result::Result;
/// Initialise ClamAV
pub async fn init() {
let config = config().await;
if !config.files.clamd_host.is_empty() {
tracing::info!("Waiting for clamd to be ready...");
loop {
let clamd_available =
match revolt_clamav_client::ping_tcp(config.files.clamd_host.clone()) {
Ok(ping_response) => ping_response == b"PONG\0",
Err(_) => false,
};
if clamd_available {
tracing::info!("clamd is ready, virus protection enabled!");
break;
} else {
tracing::error!(
"Could not ping clamd host at {}, retrying in 10 seconds...",
config.files.clamd_host
);
std::thread::sleep(Duration::from_secs(10));
}
}
}
}
/// Scan for malware
pub async fn is_malware(buf: &[u8]) -> Result<bool> {
let config = config().await;
if config.files.clamd_host.is_empty() {
Ok(false)
} else {
let scan_response = report_internal_error!(revolt_clamav_client::scan_buffer_tcp(
buf,
config.files.clamd_host,
None
))?;
report_internal_error!(revolt_clamav_client::clean(&scan_response)).map(|v| !v)
}
}

79
crates/services/autumn/src/exif.rs
View File

@@ -2,6 +2,7 @@ use std::io::{Cursor, Read};
use exif::Reader;
use image::{ImageFormat, ImageReader};
use revolt_config::report_internal_error;
use revolt_database::Metadata;
use revolt_result::{create_error, Result};
use tempfile::NamedTempFile;
@@ -39,11 +40,11 @@ pub async fn strip_metadata(
let mut cursor = Cursor::new(buf);
// Decode the image
let image = ImageReader::new(&mut cursor)
.with_guessed_format()
.map_err(|_| create_error!(InternalError))?
.decode()
.map_err(|_| create_error!(InternalError));
let image = report_internal_error!(report_internal_error!(ImageReader::new(
&mut cursor
)
.with_guessed_format())?
.decode());
// Reset read position
cursor.set_position(0);
@@ -64,7 +65,7 @@ pub async fn strip_metadata(
// Apply the EXIF rotation
// See https://jdhao.github.io/2019/07/31/image_rotation_exif_info/
match &rotation {
report_internal_error!(match &rotation {
2 => image?.fliph(),
3 => image?.rotate180(),
4 => image?.rotate180().fliph(),
@@ -83,8 +84,7 @@ pub async fn strip_metadata(
"image/tiff" => ImageFormat::Tiff,
_ => todo!(),
},
)
.map_err(|_| create_error!(InternalError))?;
))?;
// Calculate dimensions after rotation.
let (width, height) = match &rotation {
@@ -112,47 +112,44 @@ pub async fn strip_metadata(
};
// Temporary output file
let mut out_file =
NamedTempFile::new().map_err(|_| create_error!(InternalError))?;
let mut out_file = report_internal_error!(NamedTempFile::new())?;
// Process the file with ffmpeg
Command::new("ffmpeg")
.args([
// Overwrite the temporary file
"-y",
// Read original uploaded file
"-i",
file.path().to_str().ok_or(create_error!(InternalError))?,
// Strip any metadata
"-map_metadata",
"-1",
// Copy video / audio data to new file
"-c:v",
"copy",
"-c:a",
"copy",
// Select correct file format
"-f",
ext,
// Save to new temporary file
out_file
.path()
.to_str()
.ok_or(create_error!(InternalError))?,
])
.output()
.await
.inspect_err(|err| tracing::error!("{err:?}"))
.map_err(|_| create_error!(InternalError))?;
report_internal_error!(
Command::new("ffmpeg")
.args([
// Overwrite the temporary file
"-y",
// Read original uploaded file
"-i",
file.path().to_str().ok_or(create_error!(InternalError))?,
// Strip any metadata
"-map_metadata",
"-1",
// Copy video / audio data to new file
"-c:v",
"copy",
"-c:a",
"copy",
// Select correct file format
"-f",
ext,
// Save to new temporary file
out_file
.path()
.to_str()
.ok_or(create_error!(InternalError))?,
])
.output()
.await
)?;
// Probe the file again
let metadata = crate::metadata::generate_metadata(&out_file, mime);
// Read the file from disk
let mut buf = Vec::<u8>::new();
out_file
.read_to_end(&mut buf)
.map_err(|_| create_error!(InternalError))?;
report_internal_error!(out_file.read_to_end(&mut buf))?;
Ok((buf, metadata))
}

6
crates/services/autumn/src/main.rs
View File

@@ -11,6 +11,7 @@ use utoipa::{
use utoipa_scalar::{Scalar, Servable as ScalarServable};
mod api;
pub mod clamav;
pub mod exif;
pub mod metadata;
pub mod mime_type;
@@ -18,7 +19,10 @@ pub mod mime_type;
#[tokio::main]
async fn main() -> Result<(), std::io::Error> {
// Configure logging and environment
revolt_config::configure!(api);
revolt_config::configure!(files);
// Wait for ClamAV
clamav::init().await;
// Configure API schema
#[derive(OpenApi)]

17
crates/services/autumn/src/mime_type.rs
View File

@@ -1,7 +1,14 @@
use tempfile::NamedTempFile;
/// Determine the mime type of the given temporary file and filename
pub fn determine_mime_type(f: &mut NamedTempFile, buf: &Vec<u8>, file_name: &str) -> &'static str {
pub fn determine_mime_type(f: &mut NamedTempFile, buf: &[u8], file_name: &str) -> &'static str {
// Force certain extensions into particular mime types
if file_name.to_lowercase().ends_with(".apk") {
return "application/vnd.android.package-archive";
} else if file_name.to_lowercase().ends_with(".exe") {
return "application/vnd.microsoft.portable-executable";
}
// Use magic signatures to determine mime type
let kind = infer::get_from_path(f.path()).expect("file read successfully");
let mime_type = if let Some(kind) = kind {
@@ -10,14 +17,6 @@ pub fn determine_mime_type(f: &mut NamedTempFile, buf: &Vec<u8>, file_name: &str
"application/octet-stream"
};
// Map any known conflicts where appropriate
let mime_type = if mime_type == "application/zip" && file_name.to_lowercase().ends_with(".apk")
{
"application/vnd.android.package-archive"
} else {
mime_type
};
// See if the file is actually just plain Unicode/ASCII text
if mime_type == "application/octet-stream" && simdutf8::basic::from_utf8(buf).is_ok() {
return "plain/text";