From b079b9471514360b236f1243914e87f438efd2a3 Mon Sep 17 00:00:00 2001 From: strawberry Date: Sun, 17 Mar 2024 01:42:30 -0400 Subject: [PATCH] track media uploads by user Signed-off-by: strawberry --- src/api/client_server/media.rs | 11 ++++++++++- src/database/key_value/media.rs | 27 ++++++++++++++++++++++----- src/database/mod.rs | 2 ++ src/service/media/data.rs | 3 ++- src/service/media/mod.rs | 24 ++++++++++++++++++------ 5 files changed, 54 insertions(+), 13 deletions(-) diff --git a/src/api/client_server/media.rs b/src/api/client_server/media.rs index 380d4e3c..bb98814b 100644 --- a/src/api/client_server/media.rs +++ b/src/api/client_server/media.rs @@ -138,6 +138,8 @@ pub async fn get_media_preview_v1_route( /// - Some metadata will be saved in the database /// - Media will be saved in the media/ directory pub async fn create_content_route(body: Ruma) -> Result { + let sender_user = body.sender_user.as_ref().expect("user is authenticated"); + let mxc = format!( "mxc://{}/{}", services().globals.server_name(), @@ -147,6 +149,7 @@ pub async fn create_content_route(body: Ruma) -> Re services() .media .create( + Some(sender_user.clone()), mxc.clone(), body.filename.as_ref().map(|filename| "inline; filename=".to_owned() + filename).as_deref(), body.content_type.as_deref(), @@ -175,6 +178,8 @@ pub async fn create_content_route(body: Ruma) -> Re pub async fn create_content_v1_route( body: Ruma, ) -> Result> { + let sender_user = body.sender_user.as_ref().expect("user is authenticated"); + let mxc = format!( "mxc://{}/{}", services().globals.server_name(), @@ -184,6 +189,7 @@ pub async fn create_content_v1_route( services() .media .create( + Some(sender_user.clone()), mxc.clone(), body.filename.as_ref().map(|filename| "inline; filename=".to_owned() + filename).as_deref(), body.content_type.as_deref(), @@ -231,6 +237,7 @@ pub async fn get_remote_content( services() .media .create( + None, mxc.to_owned(), content_response.content_disposition.as_deref(), content_response.content_type.as_deref(), @@ -484,6 +491,7 @@ pub async fn get_content_thumbnail_route( services() .media .upload_thumbnail( + None, mxc, None, get_thumbnail_response.content_type.as_deref(), @@ -566,6 +574,7 @@ pub async fn get_content_thumbnail_v1_route( services() .media .upload_thumbnail( + None, mxc, None, get_thumbnail_response.content_type.as_deref(), @@ -589,7 +598,7 @@ async fn download_image(client: &reqwest::Client, url: &str) -> Result (None, None), diff --git a/src/database/key_value/media.rs b/src/database/key_value/media.rs index f00f6b55..af7a883a 100644 --- a/src/database/key_value/media.rs +++ b/src/database/key_value/media.rs @@ -4,12 +4,14 @@ use tracing::debug; use crate::{ database::KeyValueDatabase, service::{self, media::UrlPreviewData}, - utils, Error, Result, + utils::string_from_bytes, + Error, Result, }; impl service::media::Data for KeyValueDatabase { fn create_file_metadata( - &self, mxc: String, width: u32, height: u32, content_disposition: Option<&str>, content_type: Option<&str>, + &self, sender_user: Option<&str>, mxc: String, width: u32, height: u32, content_disposition: Option<&str>, + content_type: Option<&str>, ) -> Result> { let mut key = mxc.as_bytes().to_vec(); key.push(0xFF); @@ -22,6 +24,12 @@ impl service::media::Data for KeyValueDatabase { self.mediaid_file.insert(&key, &[])?; + if let Some(user) = sender_user { + let key = mxc.as_bytes().to_vec(); + let user = user.as_bytes().to_vec(); + self.mediaid_user.insert(&key, &user)?; + } + Ok(key) } @@ -31,13 +39,22 @@ impl service::media::Data for KeyValueDatabase { let mut prefix = mxc.as_bytes().to_vec(); prefix.push(0xFF); - debug!("MXC db prefix: {:?}", prefix); + debug!("MXC db prefix: {prefix:?}"); for (key, _) in self.mediaid_file.scan_prefix(prefix) { debug!("Deleting key: {:?}", key); self.mediaid_file.remove(&key)?; } + for (key, value) in self.mediaid_user.scan_prefix(mxc.as_bytes().to_vec()) { + if key == mxc.as_bytes().to_vec() { + let user = string_from_bytes(&value).unwrap_or_default(); + + debug!("Deleting key \"{key:?}\" which was uploaded by user {user}"); + self.mediaid_user.remove(&key)?; + } + } + Ok(()) } @@ -85,7 +102,7 @@ impl service::media::Data for KeyValueDatabase { let content_type = parts .next() .map(|bytes| { - utils::string_from_bytes(bytes) + string_from_bytes(bytes) .map_err(|_| Error::bad_database("Content type in mediaid_file is invalid unicode.")) }) .transpose()?; @@ -97,7 +114,7 @@ impl service::media::Data for KeyValueDatabase { None } else { Some( - utils::string_from_bytes(content_disposition_bytes) + string_from_bytes(content_disposition_bytes) .map_err(|_| Error::bad_database("Content Disposition in mediaid_file is invalid unicode."))?, ) }; diff --git a/src/database/mod.rs b/src/database/mod.rs index 994b9273..e2b25bf0 100644 --- a/src/database/mod.rs +++ b/src/database/mod.rs @@ -157,6 +157,7 @@ pub struct KeyValueDatabase { //pub media: media::Media, pub(super) mediaid_file: Arc, // MediaId = MXC + WidthHeight + ContentDisposition + ContentType pub(super) url_previews: Arc, + pub(super) mediaid_user: Arc, //pub key_backups: key_backups::KeyBackups, pub(super) backupid_algorithm: Arc, // BackupId = UserId + Version(Count) pub(super) backupid_etag: Arc, // BackupId = UserId + Version(Count) @@ -365,6 +366,7 @@ impl KeyValueDatabase { roomusertype_roomuserdataid: builder.open_tree("roomusertype_roomuserdataid")?, mediaid_file: builder.open_tree("mediaid_file")?, url_previews: builder.open_tree("url_previews")?, + mediaid_user: builder.open_tree("mediaid_user")?, backupid_algorithm: builder.open_tree("backupid_algorithm")?, backupid_etag: builder.open_tree("backupid_etag")?, backupkeyid_backup: builder.open_tree("backupkeyid_backup")?, diff --git a/src/service/media/data.rs b/src/service/media/data.rs index 9da50860..7cbde755 100644 --- a/src/service/media/data.rs +++ b/src/service/media/data.rs @@ -2,7 +2,8 @@ use crate::Result; pub trait Data: Send + Sync { fn create_file_metadata( - &self, mxc: String, width: u32, height: u32, content_disposition: Option<&str>, content_type: Option<&str>, + &self, sender_user: Option<&str>, mxc: String, width: u32, height: u32, content_disposition: Option<&str>, + content_type: Option<&str>, ) -> Result>; fn delete_file_mxc(&self, mxc: String) -> Result<()>; diff --git a/src/service/media/mod.rs b/src/service/media/mod.rs index a4e78378..696fa9f0 100644 --- a/src/service/media/mod.rs +++ b/src/service/media/mod.rs @@ -3,7 +3,7 @@ use std::{collections::HashMap, io::Cursor, sync::Arc, time::SystemTime}; pub(crate) use data::Data; use image::imageops::FilterType; -use ruma::OwnedMxcUri; +use ruma::{OwnedMxcUri, OwnedUserId}; use serde::Serialize; use tokio::{ fs::{self, File}, @@ -45,10 +45,15 @@ pub struct Service { impl Service { /// Uploads a file. pub async fn create( - &self, mxc: String, content_disposition: Option<&str>, content_type: Option<&str>, file: &[u8], + &self, sender_user: Option, mxc: String, content_disposition: Option<&str>, + content_type: Option<&str>, file: &[u8], ) -> Result<()> { // Width, Height = 0 if it's not a thumbnail - let key = self.db.create_file_metadata(mxc, 0, 0, content_disposition, content_type)?; + let key = if let Some(user) = sender_user { + self.db.create_file_metadata(Some(user.as_str()), mxc, 0, 0, content_disposition, content_type)? + } else { + self.db.create_file_metadata(None, mxc, 0, 0, content_disposition, content_type)? + }; let path; @@ -106,11 +111,17 @@ impl Service { } /// Uploads or replaces a file thumbnail. + #[allow(clippy::too_many_arguments)] pub async fn upload_thumbnail( - &self, mxc: String, content_disposition: Option<&str>, content_type: Option<&str>, width: u32, height: u32, - file: &[u8], + &self, sender_user: Option, mxc: String, content_disposition: Option<&str>, + content_type: Option<&str>, width: u32, height: u32, file: &[u8], ) -> Result<()> { - let key = self.db.create_file_metadata(mxc, width, height, content_disposition, content_type)?; + let key = if let Some(user) = sender_user { + self.db.create_file_metadata(Some(user.as_str()), mxc, width, height, content_disposition, content_type)? + } else { + self.db.create_file_metadata(None, mxc, width, height, content_disposition, content_type)? + }; + let path; #[allow(clippy::unnecessary_operation)] // error[E0658]: attributes on expressions are experimental @@ -403,6 +414,7 @@ impl Service { // Save thumbnail in database so we don't have to generate it again next time let thumbnail_key = self.db.create_file_metadata( + None, mxc, width, height,