track media uploads by user

Signed-off-by: strawberry <strawberry@puppygock.gay>

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<create_content::v3::Request>) -> Result<create_content::v3::Response> {
+	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<create_content::v3::Request>) -> 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<create_content::v3::Request>) -> Re
 pub async fn create_content_v1_route(
 	body: Ruma<create_content::v3::Request>,
 ) -> Result<RumaResponse<create_content::v3::Response>> {
+	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<UrlPrevie
 		utils::random_string(MXC_LENGTH)
 	);
 
-	services().media.create(mxc.clone(), None, None, &image).await?;
+	services().media.create(None, mxc.clone(), None, None, &image).await?;
 
 	let (width, height) = match ImgReader::new(Cursor::new(&image)).with_guessed_format() {
 		Err(_) => (None, None),

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<Vec<u8>> {
 		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."))?,
 			)
 		};

src/database/mod.rs

@@ -157,6 +157,7 @@ pub struct KeyValueDatabase {
 	//pub media: media::Media,
 	pub(super) mediaid_file: Arc<dyn KvTree>, // MediaId = MXC + WidthHeight + ContentDisposition + ContentType
 	pub(super) url_previews: Arc<dyn KvTree>,
+	pub(super) mediaid_user: Arc<dyn KvTree>,
 	//pub key_backups: key_backups::KeyBackups,
 	pub(super) backupid_algorithm: Arc<dyn KvTree>, // BackupId = UserId + Version(Count)
 	pub(super) backupid_etag: Arc<dyn KvTree>,      // 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")?,

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<Vec<u8>>;
 
 	fn delete_file_mxc(&self, mxc: String) -> Result<()>;

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<OwnedUserId>, 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,
+		&self, sender_user: Option<OwnedUserId>, mxc: String, content_disposition: Option<&str>,
-		file: &[u8],
+		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,