fix: check both hash and processed hash when deleting files. (#1076)

Change-Id: I99a8042440a3732ee0d88c67a92e51a459be1ae3

Signed-off-by: IAmTomahawkx <iamtomahawkx@gmail.com>

crates/core/database/src/models/file_hashes/model.rs

@@ -91,3 +91,49 @@ impl FileHash {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use iso8601_timestamp::Timestamp;
+    use lapin::{ExchangeKind, options::ExchangeDeclareOptions, types::FieldTable};
+
+    use crate::{FileHash, Metadata};
+
+    #[tokio::test]
+    #[allow(clippy::disallowed_methods)]
+    async fn count_references_by_original_and_processed_hash() {
+        database_test!(|db| async move {
+            let hash = FileHash {
+                id: "original".to_string(),
+                processed_hash: "processed".to_string(),
+                created_at: Timestamp::now_utc(),
+                bucket_id: "bucket".to_string(),
+                path: "path".to_string(),
+                iv: String::new(),
+                metadata: Metadata::File,
+                content_type: "application/octet-stream".to_string(),
+                size: 0,
+            };
+
+            let current = hash.into_file(
+                "current".to_string(),
+                "attachments".to_string(),
+                "current.bin".to_string(),
+                "user".to_string(),
+            );
+
+            let mut legacy = hash.into_file(
+                "legacy".to_string(),
+                "attachments".to_string(),
+                "legacy.bin".to_string(),
+                "user".to_string(),
+            );
+            legacy.hash = Some(hash.processed_hash.clone());
+
+            db.insert_attachment(&current).await.unwrap();
+            db.insert_attachment(&legacy).await.unwrap();
+
+            assert_eq!(db.count_file_hash_references(&hash).await.unwrap(), 2);
+        });
+    }
+}

crates/core/database/src/models/files/ops.rs

@@ -1,6 +1,7 @@
 use revolt_result::Result;
 
 use crate::File;
+use crate::FileHash;
 
 use super::FileUsedFor;
 
@@ -22,8 +23,8 @@ pub trait AbstractAttachments: Sync + Send {
     /// Fetch all dangling attachments.
     async fn fetch_dangling_files(&self) -> Result<Vec<File>>;
 
-    /// Count references to a given hash.
-    async fn count_file_hash_references(&self, hash: &str) -> Result<usize>;
+    /// Count references to a given hash, by either its original or processed hash.
+    async fn count_file_hash_references(&self, hash: &FileHash) -> Result<usize>;
 
     /// Find an attachment by its details and mark it as used by a given parent.
     async fn find_and_use_attachment(

crates/core/database/src/models/files/ops/mongodb.rs

@@ -4,6 +4,7 @@ use revolt_config::report_internal_error;
 use revolt_result::Result;
 
 use crate::File;
+use crate::FileHash;
 use crate::FileUsedFor;
 use crate::MongoDb;
 
@@ -64,14 +65,16 @@ impl AbstractAttachments for MongoDb {
         )
     }
 
-    /// Count references to a given hash.
-    async fn count_file_hash_references(&self, hash: &str) -> Result<usize> {
+    /// Count references to a given hash, by either its original or processed hash.
+    async fn count_file_hash_references(&self, hash: &FileHash) -> Result<usize> {
         query!(
             self,
             count_documents,
             COL,
             doc! {
-                "hash": hash
+                "hash": {
+                    "$in": [&hash.id, &hash.processed_hash]
+                }
             }
         )
         .map(|count| count as usize)

crates/core/database/src/models/files/ops/reference.rs

@@ -1,6 +1,7 @@
 use revolt_result::Result;
 
 use crate::File;
+use crate::FileHash;
 use crate::FileUsedFor;
 use crate::ReferenceDb;
 
@@ -58,13 +59,16 @@ impl AbstractAttachments for ReferenceDb {
             .collect())
     }
 
-    /// Count references to a given hash.
-    async fn count_file_hash_references(&self, hash: &str) -> Result<usize> {
+    /// Count references to a given hash, by either its original or processed hash.
+    async fn count_file_hash_references(&self, hash: &FileHash) -> Result<usize> {
         let files = self.files.lock().await;
         Ok(files
             .values()
-            .filter(|file| file.hash.as_ref().is_some_and(|h| h == hash))
-            .cloned()
+            .filter(|file| {
+                file.hash
+                    .as_ref()
+                    .is_some_and(|h| h == &hash.id || h == &hash.processed_hash)
+            })
             .count())
     }
 

crates/daemons/crond/src/tasks/file_deletion.rs

@@ -1,9 +1,9 @@
 use std::time::Duration;
 
-use log::info;
+use log::{info, warn};
 use revolt_database::Database;
 use revolt_files::delete_from_s3;
-use revolt_result::Result;
+use revolt_result::{ErrorType, Result};
 use tokio::time::sleep;
 
 pub async fn task(db: Database, _: revolt_database::AMQP) -> Result<()> {
@@ -12,18 +12,25 @@ pub async fn task(db: Database, _: revolt_database::AMQP) -> Result<()> {
 
         for file in files {
             if let Some(hash) = &file.hash {
-                let count = db.count_file_hash_references(hash).await?;
+                match db.fetch_attachment_hash(hash).await {
+                    Ok(file_hash) => {
+                        let count = db.count_file_hash_references(&file_hash).await?;
 
-                // No other files reference this file on disk anymore
-                if count <= 1 {
-                    let file_hash = db.fetch_attachment_hash(hash).await?;
+                        // No other files reference this file on disk anymore
+                        if count <= 1 {
+                            // Delete from S3
+                            delete_from_s3(&file_hash.bucket_id, &file_hash.path).await?;
 
-                    // Delete from S3
-                    delete_from_s3(&file_hash.bucket_id, &file_hash.path).await?;
-
-                    // Delete the hash
-                    db.delete_attachment_hash(&file_hash.id).await?;
-                    info!("Deleted file hash {}", file_hash.id);
+                            // Delete the hash
+                            db.delete_attachment_hash(&file_hash.id).await?;
+                            info!("Deleted file hash {}", file_hash.id);
+                        }
+                    }
+                    // Hash is already gone
+                    Err(error) if matches!(error.error_type, ErrorType::NotFound) => {
+                        warn!("File {} references missing hash {hash}", file.id);
+                    }
+                    Err(error) => return Err(error),
                 }
             }