HMC/handmade-revolt-backend
6
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

4
crates/core/config/Cargo.toml
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 }

10
crates/core/config/Revolt.toml
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 = ""

27
crates/core/config/src/lib.rs
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)]

4
crates/core/database/src/lib.rs
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)
}
}

39
crates/core/database/src/models/file_hashes/model.rs
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,
}
}
}

6
crates/core/database/src/models/file_hashes/ops.rs
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<()>;
}

29
crates/core/database/src/models/file_hashes/ops/mongodb.rs
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))
}
}

32
crates/core/database/src/models/file_hashes/ops/reference.rs
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))
}

19
crates/core/database/src/models/files/model.rs
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

2
crates/core/database/src/util/bridge/v0.rs
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,
}
}

4
crates/core/files/Cargo.toml
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" }

59
crates/core/files/src/lib.rs
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))
}