feat(services/autumn): file uploads

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

32
Cargo.lock generated
View File

@@ -2515,8 +2515,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "94b22e06ecb0110981051723910cbf0b5f5e09a2062dd7663334ee79a9d1286c"
dependencies = [
"cfg-if",
"js-sys",
"libc",
"wasi",
"wasm-bindgen",
]
[[package]]
@@ -5308,10 +5310,12 @@ dependencies = [
"kamadak-exif",
"lazy_static",
"moka",
"nanoid",
"revolt-config",
"revolt-database",
"revolt-files",
"revolt-result",
"revolt_clamav-client",
"serde",
"serde_json",
"sha2",
@@ -5321,6 +5325,7 @@ dependencies = [
"tokio 1.35.1",
"tracing",
"tracing-subscriber",
"ulid 1.1.3",
"utoipa",
"utoipa-scalar",
"webp",
@@ -5368,6 +5373,7 @@ dependencies = [
"log",
"once_cell",
"pretty_env_logger",
"revolt-result",
"sentry",
"serde",
]
@@ -5413,7 +5419,7 @@ dependencies = [
"schemars",
"serde",
"serde_json",
"ulid 1.0.0",
"ulid 1.1.3",
"unicode-segmentation",
"url-escape",
"validator 0.16.0",
@@ -5597,6 +5603,12 @@ dependencies = [
"tokio 1.35.1",
]
[[package]]
name = "revolt_clamav-client"
version = "0.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3d719e9fe861e2e05cbf01fd82d310fdeb2160c74fee21689a5c223eeb65cea8"
[[package]]
name = "revolt_okapi"
version = "0.9.1"
@@ -7325,11 +7337,13 @@ dependencies = [
[[package]]
name = "ulid"
version = "1.0.0"
version = "1.1.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "13a3aaa69b04e5b66cc27309710a569ea23593612387d67daaf102e73aa974fd"
checksum = "04f903f293d11f31c0c29e4148f6dc0d033a7f80cebc0282bea147611667d289"
dependencies = [
"getrandom",
"rand 0.8.5",
"web-time",
]
[[package]]
@@ -7510,7 +7524,7 @@ dependencies = [
"quote 1.0.37",
"regex",
"syn 2.0.76",
"ulid 1.0.0",
"ulid 1.1.3",
]
[[package]]
@@ -7792,6 +7806,16 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "web-time"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb"
dependencies = [
"js-sys",
"wasm-bindgen",
]
[[package]]
name = "webp"
version = "0.3.0"

View File

@@ -77,6 +77,18 @@ A default configuration `Revolt.toml` is present in this project that is suited
If you'd like to change anything, create a `Revolt.overrides.toml` file and specify relevant variables.
> [!TIP]
> Use Sentry to catch unexpected service errors:
>
> ```toml
> # Revolt.overrides.toml
> [sentry]
> api = "https://abc@your.sentry/1"
> events = "https://abc@your.sentry/1"
> files = "https://abc@your.sentry/1"
> proxy = "https://abc@your.sentry/1"
> ```
Then continue:
```bash

View File

@@ -4,10 +4,10 @@
[database]
# MongoDB connection URL
# Defaults to the container name specified in self-hosted
mongodb = "mongodb://localhost:14017"
mongodb = "mongodb://127.0.0.1:14017"
# Redis connection URL
# Defaults to the container name specified in self-hosted
redis = "redis://localhost:14079/"
redis = "redis://127.0.0.1:14079/"
[hosts]
# Web locations of various services
@@ -38,7 +38,7 @@ use_tls = false
[files.s3]
# S3 protocol endpoint
endpoint = "http://localhost:14009"
endpoint = "http://127.0.0.1:14009"
# S3 region name
region = "minio"
# S3 protocol key ID

View File

@@ -24,6 +24,7 @@ services:
- ./.data/minio:/data
ports:
- "14009:9000"
- "14010:9001"
restart: always
# Create buckets for minio.

View File

@@ -9,6 +9,7 @@ description = "Revolt Backend: Configuration"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[features]
report-macros = ["revolt-result"]
test = ["async-std"]
default = ["test"]
@@ -32,3 +33,6 @@ pretty_env_logger = "0.4.0"
# Sentry
sentry = "0.31.5"
# Core
revolt-result = { version = "0.7.16", path = "../result", optional = true }

View File

@@ -98,6 +98,14 @@ blocked_mime_types = []
# ClamAV service
# hostname:port
clamd_host = ""
# Mime types that should be virus scanned
#
# Leave empty to scan all file types
scan_mime_types = [
"application/vnd.microsoft.portable-executable",
"application/vnd.android.package-archive",
"application/zip",
]
[files.limit]
# Minimum file size (in bytes)
@@ -220,3 +228,5 @@ emojis = 500_000
# Configuration for Sentry error reporting
api = ""
events = ""
files = ""
proxy = ""

View File

@@ -8,6 +8,30 @@ use serde::Deserialize;
pub use sentry::capture_error;
#[cfg(feature = "report-macros")]
#[macro_export]
macro_rules! report_error {
( $expr: expr, $error: ident $( $tt:tt )? ) => {
$expr
.inspect_err(|err| {
$crate::capture_error(err);
})
.map_err(|_| ::revolt_result::create_error!($error))
};
}
#[cfg(feature = "report-macros")]
#[macro_export]
macro_rules! report_internal_error {
( $expr: expr ) => {
$expr
.inspect_err(|err| {
$crate::capture_error(err);
})
.map_err(|_| ::revolt_result::create_error!(InternalError))
};
}
/// Paths to search for configuration
static CONFIG_SEARCH_PATHS: [&str; 3] = [
// current working directory
@@ -157,6 +181,7 @@ pub struct Files {
pub webp_quality: f32,
pub blocked_mime_types: Vec<String>,
pub clamd_host: String,
pub scan_mime_types: Vec<String>,
pub limit: FilesLimit,
pub preview: HashMap<String, [usize; 2]>,
@@ -211,6 +236,8 @@ pub struct Features {
pub struct Sentry {
pub api: String,
pub events: String,
pub files: String,
pub proxy: String,
}
#[derive(Deserialize, Debug, Clone)]

View File

@@ -16,6 +16,8 @@ extern crate revolt_optional_struct;
#[macro_use]
extern crate revolt_result;
pub use iso8601_timestamp;
#[cfg(feature = "mongodb")]
pub use mongodb;
@@ -91,4 +93,4 @@ pub fn if_false(t: &bool) -> bool {
/// Utility function to check if an option doesnt contain true
pub fn if_option_false(t: &Option<bool>) -> bool {
t != &Some(true)
}
}

View File

@@ -1,5 +1,7 @@
use iso8601_timestamp::Timestamp;
use crate::File;
auto_derived_partial!(
/// File hash
pub struct FileHash {
@@ -51,3 +53,40 @@ auto_derived!(
Audio,
}
);
impl FileHash {
/// Create a file from a file hash
pub fn into_file(
&self,
id: String,
tag: String,
filename: String,
uploader_id: String,
) -> File {
File {
id,
tag,
filename,
hash: Some(self.id.clone()),
uploaded_at: Some(Timestamp::now_utc()),
uploader_id: Some(uploader_id),
used_for: None,
deleted: None,
reported: None,
// TODO: remove this data
metadata: self.metadata.clone(),
content_type: self.content_type.clone(),
size: self.size,
// TODO: superseded by "used_for"
message_id: None,
object_id: None,
server_id: None,
user_id: None,
}
}
}

View File

@@ -7,6 +7,12 @@ mod reference;
#[async_trait]
pub trait AbstractAttachmentHashes: Sync + Send {
/// Insert a new attachment hash into the database.
async fn insert_attachment_hash(&self, hash: &FileHash) -> Result<()>;
/// Fetch an attachment hash entry by sha256 hash.
async fn fetch_attachment_hash(&self, hash: &str) -> Result<FileHash>;
/// Update an attachment hash nonce value.
async fn set_attachment_hash_nonce(&self, hash: &str, nonce: &str) -> Result<()>;
}

View File

@@ -9,6 +9,11 @@ static COL: &str = "attachment_hashes";
#[async_trait]
impl AbstractAttachmentHashes for MongoDb {
/// Insert a new attachment hash into the database.
async fn insert_attachment_hash(&self, hash: &FileHash) -> Result<()> {
query!(self, insert_one, COL, &hash).map(|_| ())
}
/// Fetch an attachment hash entry by sha256 hash.
async fn fetch_attachment_hash(&self, hash: &str) -> Result<FileHash> {
query!(
@@ -16,9 +21,31 @@ impl AbstractAttachmentHashes for MongoDb {
find_one,
COL,
doc! {
"_id": hash
"$or": [
{"_id": hash},
{"processed_hash": hash}
]
}
)?
.ok_or_else(|| create_error!(NotFound))
}
/// Update an attachment hash nonce value.
async fn set_attachment_hash_nonce(&self, hash: &str, nonce: &str) -> Result<()> {
self.col::<FileHash>(COL)
.update_one(
doc! {
"_id": hash
},
doc! {
"$set": {
"iv": nonce
}
},
None,
)
.await
.map(|_| ())
.map_err(|_| create_database_error!("update_one", COL))
}
}

View File

@@ -7,15 +7,33 @@ use super::AbstractAttachmentHashes;
#[async_trait]
impl AbstractAttachmentHashes for ReferenceDb {
/// Insert a new attachment hash into the database.
async fn insert_attachment_hash(&self, hash: &FileHash) -> Result<()> {
let mut hashes = self.file_hashes.lock().await;
if hashes.contains_key(&hash.id) {
Err(create_database_error!("insert", "attachment"))
} else {
hashes.insert(hash.id.to_string(), hash.clone());
Ok(())
}
}
/// Fetch an attachment hash entry by sha256 hash.
async fn fetch_attachment_hash(&self, hash: &str) -> Result<FileHash> {
async fn fetch_attachment_hash(&self, hash_value: &str) -> Result<FileHash> {
let hashes = self.file_hashes.lock().await;
if let Some(file) = hashes.get(hash) {
if file.id == hash {
Ok(file.clone())
} else {
Err(create_error!(NotFound))
}
hashes
.values()
.cloned()
.find(|hash| hash.id == hash_value || hash.processed_hash == hash_value)
.ok_or(create_error!(NotFound))
}
/// Update an attachment hash nonce value.
async fn set_attachment_hash_nonce(&self, hash: &str, nonce: &str) -> Result<()> {
let mut hashes = self.file_hashes.lock().await;
if let Some(file) = hashes.get_mut(hash) {
file.iv = nonce.to_owned();
Ok(())
} else {
Err(create_error!(NotFound))
}

View File

@@ -19,7 +19,7 @@ auto_derived_partial!(
/// When this file was uploaded
pub uploaded_at: Option<Timestamp>, // these are Option<>s to not break file uploads on legacy Autumn
/// ID of user who uploaded this file
pub uploaded_id: Option<String>, // these are Option<>s to not break file uploads on legacy Autumn
pub uploader_id: Option<String>, // these are Option<>s to not break file uploads on legacy Autumn
/// What the file was used for
pub used_for: Option<FileUsedFor>,
@@ -59,15 +59,14 @@ auto_derived_partial!(
auto_derived!(
/// Type of object file was used for
pub enum FileUsedForType {
// TODO: changing this requires a db migration
message,
serverBanner,
emoji,
userAvatar,
userProfileBackground,
legacyGroupIcon,
channelIcon,
serverIcon,
Message,
ServerBanner,
Emoji,
UserAvatar,
UserProfileBackground,
LegacyGroupIcon,
ChannelIcon,
ServerIcon,
}
/// Information about what the file was used for

View File

@@ -421,7 +421,7 @@ impl From<File> for crate::File {
object_id: value.object_id,
hash: None,
uploaded_at: None,
uploaded_id: None,
uploader_id: None,
used_for: None,
}
}

View File

@@ -14,5 +14,7 @@ typenum = "1.17.0"
aws-config = "1.5.5"
aws-sdk-s3 = { version = "1.46.0", features = ["behavior-version-latest"] }
revolt-config = { version = "0.7.16", path = "../config" }
revolt-config = { version = "0.7.16", path = "../config", features = [
"report-macros",
] }
revolt-result = { version = "0.7.16", path = "../result" }

View File

@@ -1,7 +1,10 @@
use std::io::Write;
use aes_gcm::{aead::AeadMutInPlace, Aes256Gcm, Key, KeyInit, Nonce};
use revolt_config::{config, FilesS3};
use aes_gcm::{
aead::{AeadCore, AeadMutInPlace, OsRng},
Aes256Gcm, Key, KeyInit, Nonce,
};
use revolt_config::{config, report_internal_error, FilesS3};
use revolt_result::{create_error, Result};
use aws_sdk_s3::{
@@ -11,7 +14,8 @@ use aws_sdk_s3::{
use base64::prelude::*;
pub static AUTHENTICATION_TAG_SIZE_BYTES: usize = 16;
/// Size of the authentication tag in the buffer
pub const AUTHENTICATION_TAG_SIZE_BYTES: usize = 16;
/// Create an S3 client
pub fn create_client(s3_config: FilesS3) -> Client {
@@ -35,7 +39,7 @@ pub fn create_client(s3_config: FilesS3) -> Client {
/// Create an AES-256-GCM cipher
pub fn create_cipher(key: &str) -> Aes256Gcm {
let key = &BASE64_STANDARD.decode(key).unwrap()[..];
let key = &BASE64_STANDARD.decode(key).expect("valid base64 string")[..];
let key: &Key<Aes256Gcm> = key.into();
Aes256Gcm::new(key)
}
@@ -46,23 +50,14 @@ pub async fn fetch_from_s3(bucket_id: &str, path: &str, nonce: &str) -> Result<V
let client = create_client(config.files.s3);
// Send a request for the file
let mut obj = client
.get_object()
.bucket(bucket_id)
.key(path)
.send()
.await
.inspect_err(|err| {
revolt_config::capture_error(err);
})
.map_err(|_| create_error!(InternalError))?;
let mut obj =
report_internal_error!(client.get_object().bucket(bucket_id).key(path).send().await)?;
// Read the file from remote
let mut buf = vec![];
while let Some(bytes) = obj.body.next().await {
let data = bytes.map_err(|_| create_error!(InternalError))?;
buf.write_all(&data)
.map_err(|_| create_error!(InternalError))?;
let data = report_internal_error!(bytes)?;
report_internal_error!(buf.write_all(&data))?;
// is there a more efficient way to do this?
// we just want the Vec<u8>
}
@@ -78,3 +73,33 @@ pub async fn fetch_from_s3(bucket_id: &str, path: &str, nonce: &str) -> Result<V
Ok(buf)
}
/// Encrypt and upload a file to S3 (returning its nonce/IV)
pub async fn upload_to_s3(bucket_id: &str, path: &str, buf: &[u8]) -> Result<String> {
let config = config().await;
let client = create_client(config.files.s3);
// Generate a nonce
let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
// Extend the buffer for in-place encryption
let mut buf = [buf, &[0; AUTHENTICATION_TAG_SIZE_BYTES]].concat();
// Encrypt the file in place
create_cipher(&config.files.encryption_key)
.encrypt_in_place(&nonce, b"", &mut buf)
.map_err(|_| create_error!(InternalError))?;
// Upload the file to remote
report_internal_error!(
client
.put_object()
.bucket(bucket_id)
.key(path)
.body(buf.into())
.send()
.await
)?;
Ok(BASE64_STANDARD.encode(nonce))
}

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"

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.");

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)
}
}

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))
}

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)]

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";

View File

@@ -15,7 +15,7 @@ pub mod requests;
#[tokio::main]
async fn main() -> Result<(), std::io::Error> {
// Configure logging and environment
revolt_config::configure!(api);
revolt_config::configure!(proxy);
// Configure API schema
#[derive(OpenApi)]