From d7087c66bb53f789fa32312bdc9705daa10c3fd2 Mon Sep 17 00:00:00 2001
From: Benjamin Lee <benjamin@computer.surgery>
Date: Sat, 14 Sep 2024 20:20:02 -0700
Subject: [PATCH] add admin command to delete individual media files

---
 src/database/key_value/media.rs | 21 +++++++++++++++++++++
 src/service/admin.rs            | 16 ++++++++++++++--
 src/service/media.rs            | 33 ++++++++++++++++++++++++++++++++-
 src/service/media/data.rs       | 11 +++++++++++
 4 files changed, 78 insertions(+), 3 deletions(-)

diff --git a/src/database/key_value/media.rs b/src/database/key_value/media.rs
index 5ecb0388..10a55554 100644
--- a/src/database/key_value/media.rs
+++ b/src/database/key_value/media.rs
@@ -140,4 +140,25 @@ impl service::media::Data for KeyValueDatabase {
         let parts = MediaFileKeyParts::try_from(&key)?;
         Ok((parts.meta, key))
     }
+
+    fn delete_file_metadata(&self, key: MediaFileKey) -> Result<()> {
+        self.mediaid_file.remove(key.as_bytes())
+    }
+
+    fn search_thumbnails_metadata(
+        &self,
+        mxc: OwnedMxcUri,
+    ) -> Result<Vec<(FileMeta, MediaFileKey)>> {
+        let mut prefix = mxc.as_bytes().to_vec();
+        prefix.push(0xFF);
+
+        self.mediaid_file
+            .scan_prefix(prefix)
+            .map(|(key, _)| {
+                let key = MediaFileKey::new(key);
+                let parts = MediaFileKeyParts::try_from(&key)?;
+                Ok((parts.meta, key))
+            })
+            .collect()
+    }
 }
diff --git a/src/service/admin.rs b/src/service/admin.rs
index b3f098a5..5bb56c4f 100644
--- a/src/service/admin.rs
+++ b/src/service/admin.rs
@@ -23,8 +23,8 @@ use ruma::{
         TimelineEventType,
     },
     signatures::verify_json,
-    EventId, MilliSecondsSinceUnixEpoch, OwnedRoomId, RoomId, RoomVersionId,
-    ServerName, UserId,
+    EventId, MilliSecondsSinceUnixEpoch, OwnedMxcUri, OwnedRoomId, RoomId,
+    RoomVersionId, ServerName, UserId,
 };
 use serde_json::value::to_raw_value;
 use tokio::sync::{mpsc, Mutex, RwLock};
@@ -179,6 +179,12 @@ enum AdminCommand {
         room_id: Box<RoomId>,
     },
 
+    /// Delete media and all associated thumbnails.
+    DeleteMedia {
+        /// mxc:// URI of the media to delete
+        mxc: OwnedMxcUri,
+    },
+
     /// Verify json signatures
     /// [commandbody]()
     /// # ```
@@ -791,6 +797,12 @@ impl Service {
                 services().rooms.metadata.disable_room(&room_id, false)?;
                 RoomMessageEventContent::text_plain("Room enabled.")
             }
+            AdminCommand::DeleteMedia {
+                mxc,
+            } => {
+                services().media.delete(mxc).await?;
+                RoomMessageEventContent::text_plain("Media deleted.")
+            }
             AdminCommand::DeactivateUser {
                 leave_rooms,
                 user_id,
diff --git a/src/service/media.rs b/src/service/media.rs
index fa626d45..830a43e4 100644
--- a/src/service/media.rs
+++ b/src/service/media.rs
@@ -3,7 +3,7 @@ use std::io::Cursor;
 use image::imageops::FilterType;
 use ruma::{http_headers::ContentDisposition, OwnedMxcUri};
 use tokio::{
-    fs::File,
+    fs::{self, File},
     io::{AsyncReadExt, AsyncWriteExt},
 };
 use tracing::{debug, warn};
@@ -110,6 +110,37 @@ impl Service {
         }
     }
 
+    /// Deletes a media object and all associated thumbnails.
+    #[tracing::instrument(skip(self))]
+    pub(crate) async fn delete(&self, mxc: OwnedMxcUri) -> Result<()> {
+        let (_, key) = self.db.search_file_metadata(mxc.clone(), 0, 0)?;
+
+        let thumbnails = self.db.search_thumbnails_metadata(mxc)?;
+        for (_, thumbnail_key) in thumbnails {
+            self.delete_by_key(thumbnail_key).await?;
+        }
+
+        self.delete_by_key(key).await?;
+
+        Ok(())
+    }
+
+    /// Deletes a specific media key, which may or may not be a thumbnail.
+    ///
+    /// When deleting a non-thumbnail key with this method, the associated
+    /// thumbnails are not deleted.
+    async fn delete_by_key(&self, key: MediaFileKey) -> Result<()> {
+        let path = services().globals.get_media_file(&key);
+        match fs::remove_file(path).await {
+            Ok(()) => (),
+            // The file in the fs may already have been deleted by hand
+            Err(e) if e.kind() == std::io::ErrorKind::NotFound => (),
+            other_error => other_error?,
+        }
+        self.db.delete_file_metadata(key)?;
+        Ok(())
+    }
+
     /// Returns width, height of the thumbnail and whether it should be cropped.
     /// Returns None when the server should send the original file.
     fn thumbnail_properties(
diff --git a/src/service/media/data.rs b/src/service/media/data.rs
index e5fe4767..d7734845 100644
--- a/src/service/media/data.rs
+++ b/src/service/media/data.rs
@@ -18,4 +18,15 @@ pub(crate) trait Data: Send + Sync {
         width: u32,
         height: u32,
     ) -> Result<(FileMeta, MediaFileKey)>;
+
+    fn delete_file_metadata(&self, key: MediaFileKey) -> Result<()>;
+
+    /// Return all thumbnail keys/metadata associated with a MXC.
+    ///
+    /// The original file is not returned. To fetch the key/metadata of the
+    /// original file, use [`Data::search_file_metadata`].
+    fn search_thumbnails_metadata(
+        &self,
+        mxc: OwnedMxcUri,
+    ) -> Result<Vec<(FileMeta, MediaFileKey)>>;
 }