From ace6c30ba593665a9943ec520df73b397fb3340f Mon Sep 17 00:00:00 2001 From: Paul Makles Date: Wed, 11 Sep 2024 14:29:52 +0100 Subject: [PATCH] feat(services/autumn): file uploads feat(services/autumn): deduplicate uploads feat(services/autumn): ClamAV support --- Cargo.lock | 32 ++++- README.md | 12 ++ Revolt.toml | 6 +- compose.yml | 1 + crates/core/config/Cargo.toml | 4 + crates/core/config/Revolt.toml | 10 ++ crates/core/config/src/lib.rs | 27 +++++ crates/core/database/src/lib.rs | 4 +- .../database/src/models/file_hashes/model.rs | 39 ++++++ .../database/src/models/file_hashes/ops.rs | 6 + .../src/models/file_hashes/ops/mongodb.rs | 29 ++++- .../src/models/file_hashes/ops/reference.rs | 32 +++-- .../core/database/src/models/files/model.rs | 19 ++- crates/core/database/src/util/bridge/v0.rs | 2 +- crates/core/files/Cargo.toml | 4 +- crates/core/files/src/lib.rs | 59 ++++++--- crates/services/autumn/Cargo.toml | 9 +- crates/services/autumn/src/api.rs | 114 ++++++++++++------ crates/services/autumn/src/clamav.rs | 49 ++++++++ crates/services/autumn/src/exif.rs | 79 ++++++------ crates/services/autumn/src/main.rs | 6 +- crates/services/autumn/src/mime_type.rs | 17 ++- crates/services/january/src/main.rs | 2 +- 23 files changed, 427 insertions(+), 135 deletions(-) create mode 100644 crates/services/autumn/src/clamav.rs diff --git a/Cargo.lock b/Cargo.lock index 0ac9cf9c..00c8ad7c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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" diff --git a/README.md b/README.md index 6b67a0cb..48a770e7 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/Revolt.toml b/Revolt.toml index 0a2e2d1e..ff8827a8 100644 --- a/Revolt.toml +++ b/Revolt.toml @@ -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 diff --git a/compose.yml b/compose.yml index 1a78d35c..aeb1b2c4 100644 --- a/compose.yml +++ b/compose.yml @@ -24,6 +24,7 @@ services: - ./.data/minio:/data ports: - "14009:9000" + - "14010:9001" restart: always # Create buckets for minio. diff --git a/crates/core/config/Cargo.toml b/crates/core/config/Cargo.toml index 342ab9b6..81161841 100644 --- a/crates/core/config/Cargo.toml +++ b/crates/core/config/Cargo.toml @@ -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 } diff --git a/crates/core/config/Revolt.toml b/crates/core/config/Revolt.toml index 178ad36d..f25e7e3d 100644 --- a/crates/core/config/Revolt.toml +++ b/crates/core/config/Revolt.toml @@ -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 = "" diff --git a/crates/core/config/src/lib.rs b/crates/core/config/src/lib.rs index 33b8aeef..ff20fea9 100644 --- a/crates/core/config/src/lib.rs +++ b/crates/core/config/src/lib.rs @@ -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, pub clamd_host: String, + pub scan_mime_types: Vec, pub limit: FilesLimit, pub preview: HashMap, @@ -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)] diff --git a/crates/core/database/src/lib.rs b/crates/core/database/src/lib.rs index dd911f0d..e48821b4 100644 --- a/crates/core/database/src/lib.rs +++ b/crates/core/database/src/lib.rs @@ -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 { t != &Some(true) -} \ No newline at end of file +} diff --git a/crates/core/database/src/models/file_hashes/model.rs b/crates/core/database/src/models/file_hashes/model.rs index 703d17ea..c89951c5 100644 --- a/crates/core/database/src/models/file_hashes/model.rs +++ b/crates/core/database/src/models/file_hashes/model.rs @@ -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, + } + } +} diff --git a/crates/core/database/src/models/file_hashes/ops.rs b/crates/core/database/src/models/file_hashes/ops.rs index d91fc8a9..e7f62658 100644 --- a/crates/core/database/src/models/file_hashes/ops.rs +++ b/crates/core/database/src/models/file_hashes/ops.rs @@ -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; + + /// Update an attachment hash nonce value. + async fn set_attachment_hash_nonce(&self, hash: &str, nonce: &str) -> Result<()>; } diff --git a/crates/core/database/src/models/file_hashes/ops/mongodb.rs b/crates/core/database/src/models/file_hashes/ops/mongodb.rs index 46aef7d9..c55a0610 100644 --- a/crates/core/database/src/models/file_hashes/ops/mongodb.rs +++ b/crates/core/database/src/models/file_hashes/ops/mongodb.rs @@ -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 { 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::(COL) + .update_one( + doc! { + "_id": hash + }, + doc! { + "$set": { + "iv": nonce + } + }, + None, + ) + .await + .map(|_| ()) + .map_err(|_| create_database_error!("update_one", COL)) + } } diff --git a/crates/core/database/src/models/file_hashes/ops/reference.rs b/crates/core/database/src/models/file_hashes/ops/reference.rs index ef3795fb..86bbeee8 100644 --- a/crates/core/database/src/models/file_hashes/ops/reference.rs +++ b/crates/core/database/src/models/file_hashes/ops/reference.rs @@ -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 { + async fn fetch_attachment_hash(&self, hash_value: &str) -> Result { 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)) } diff --git a/crates/core/database/src/models/files/model.rs b/crates/core/database/src/models/files/model.rs index b04521a3..b1400eb4 100644 --- a/crates/core/database/src/models/files/model.rs +++ b/crates/core/database/src/models/files/model.rs @@ -19,7 +19,7 @@ auto_derived_partial!( /// When this file was uploaded pub uploaded_at: Option, // these are Option<>s to not break file uploads on legacy Autumn /// ID of user who uploaded this file - pub uploaded_id: Option, // these are Option<>s to not break file uploads on legacy Autumn + pub uploader_id: Option, // these are Option<>s to not break file uploads on legacy Autumn /// What the file was used for pub used_for: Option, @@ -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 diff --git a/crates/core/database/src/util/bridge/v0.rs b/crates/core/database/src/util/bridge/v0.rs index 8c0ecb3e..87b2f10d 100644 --- a/crates/core/database/src/util/bridge/v0.rs +++ b/crates/core/database/src/util/bridge/v0.rs @@ -421,7 +421,7 @@ impl From for crate::File { object_id: value.object_id, hash: None, uploaded_at: None, - uploaded_id: None, + uploader_id: None, used_for: None, } } diff --git a/crates/core/files/Cargo.toml b/crates/core/files/Cargo.toml index a2d421b9..a340894f 100644 --- a/crates/core/files/Cargo.toml +++ b/crates/core/files/Cargo.toml @@ -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" } diff --git a/crates/core/files/src/lib.rs b/crates/core/files/src/lib.rs index fb83871a..352ee0b0 100644 --- a/crates/core/files/src/lib.rs +++ b/crates/core/files/src/lib.rs @@ -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 = key.into(); Aes256Gcm::new(key) } @@ -46,23 +50,14 @@ pub async fn fetch_from_s3(bucket_id: &str, path: &str, nonce: &str) -> Result } @@ -78,3 +73,33 @@ pub async fn fetch_from_s3(bucket_id: &str, path: &str, nonce: &str) -> Result Result { + 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)) +} diff --git a/crates/services/autumn/Cargo.toml b/crates/services/autumn/Cargo.toml index c36354f2..7871b10f 100644 --- a/crates/services/autumn/Cargo.toml +++ b/crates/services/autumn/Cargo.toml @@ -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" diff --git a/crates/services/autumn/src/api.rs b/crates/services/autumn/src/api.rs index 455f5c2d..79fb413b 100644 --- a/crates/services/autumn/src/api.rs +++ b/crates/services/autumn/src/api.rs @@ -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 { } /// 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, TypedMultipart(UploadPayload { mut file }): TypedMultipart, -) -> axum::response::Result> { +) -> Result> { // 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::::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."); diff --git a/crates/services/autumn/src/clamav.rs b/crates/services/autumn/src/clamav.rs new file mode 100644 index 00000000..2994fdbd --- /dev/null +++ b/crates/services/autumn/src/clamav.rs @@ -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 { + 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) + } +} diff --git a/crates/services/autumn/src/exif.rs b/crates/services/autumn/src/exif.rs index aa991675..0220f73b 100644 --- a/crates/services/autumn/src/exif.rs +++ b/crates/services/autumn/src/exif.rs @@ -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::::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)) } diff --git a/crates/services/autumn/src/main.rs b/crates/services/autumn/src/main.rs index 7aff1c69..e8ed464a 100644 --- a/crates/services/autumn/src/main.rs +++ b/crates/services/autumn/src/main.rs @@ -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)] diff --git a/crates/services/autumn/src/mime_type.rs b/crates/services/autumn/src/mime_type.rs index 8bfd636a..19fb3cf7 100644 --- a/crates/services/autumn/src/mime_type.rs +++ b/crates/services/autumn/src/mime_type.rs @@ -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, 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, 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"; diff --git a/crates/services/january/src/main.rs b/crates/services/january/src/main.rs index 08020f21..fcf455ef 100644 --- a/crates/services/january/src/main.rs +++ b/crates/services/january/src/main.rs @@ -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)]