From f74043df9aa59b406b5086c2e9fa2791a31aa41b Mon Sep 17 00:00:00 2001
From: Matthias Ahouansou <matthias@ahouansou.cz>
Date: Wed, 5 Jun 2024 20:17:32 +0100
Subject: [PATCH] fix: de-index pdus when redacted

---
 src/database/key_value/rooms/search.rs | 23 +++++++++++++++++++++++
 src/service/rooms/search/data.rs       |  7 +++++++
 src/service/rooms/timeline.rs          | 22 ++++++++++++++++++++--
 3 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/src/database/key_value/rooms/search.rs b/src/database/key_value/rooms/search.rs
index 8975c74d..5fb90462 100644
--- a/src/database/key_value/rooms/search.rs
+++ b/src/database/key_value/rooms/search.rs
@@ -33,6 +33,29 @@ impl service::rooms::search::Data for KeyValueDatabase {
         self.tokenids.insert_batch(&mut batch)
     }
 
+    #[tracing::instrument(skip(self))]
+    fn deindex_pdu(
+        &self,
+        shortroomid: u64,
+        pdu_id: &[u8],
+        message_body: &str,
+    ) -> Result<()> {
+        let batch = tokenize(message_body).map(|word| {
+            let mut key = shortroomid.to_be_bytes().to_vec();
+            key.extend_from_slice(word.as_bytes());
+            key.push(0xFF);
+            // TODO: currently we save the room id a second time here
+            key.extend_from_slice(pdu_id);
+            key
+        });
+
+        for token in batch {
+            self.tokenids.remove(&token)?;
+        }
+
+        Ok(())
+    }
+
     #[tracing::instrument(skip(self))]
     #[allow(clippy::type_complexity)]
     fn search_pdus<'a>(
diff --git a/src/service/rooms/search/data.rs b/src/service/rooms/search/data.rs
index 903a85e6..efc503fe 100644
--- a/src/service/rooms/search/data.rs
+++ b/src/service/rooms/search/data.rs
@@ -10,6 +10,13 @@ pub(crate) trait Data: Send + Sync {
         message_body: &str,
     ) -> Result<()>;
 
+    fn deindex_pdu(
+        &self,
+        shortroomid: u64,
+        pdu_id: &[u8],
+        message_body: &str,
+    ) -> Result<()>;
+
     #[allow(clippy::type_complexity)]
     fn search_pdus<'a>(
         &'a self,
diff --git a/src/service/rooms/timeline.rs b/src/service/rooms/timeline.rs
index 53926c3e..95f55819 100644
--- a/src/service/rooms/timeline.rs
+++ b/src/service/rooms/timeline.rs
@@ -414,7 +414,7 @@ impl Service {
                                 &pdu.room_id,
                                 false,
                             )? {
-                                self.redact_pdu(redact_id, pdu)?;
+                                self.redact_pdu(redact_id, pdu, shortroomid)?;
                             }
                         }
                     }
@@ -435,7 +435,7 @@ impl Service {
                                 &pdu.room_id,
                                 false,
                             )? {
-                                self.redact_pdu(redact_id, pdu)?;
+                                self.redact_pdu(redact_id, pdu, shortroomid)?;
                             }
                         }
                     }
@@ -1178,15 +1178,33 @@ impl Service {
         &self,
         event_id: &EventId,
         reason: &PduEvent,
+        shortroomid: u64,
     ) -> Result<()> {
         // TODO: Don't reserialize, keep original json
         if let Some(pdu_id) = self.get_pdu_id(event_id)? {
+            #[derive(Deserialize)]
+            struct ExtractBody {
+                body: String,
+            }
+
             let mut pdu = self.get_pdu_from_id(&pdu_id)?.ok_or_else(|| {
                 Error::bad_database("PDU ID points to invalid PDU.")
             })?;
+
+            if let Ok(content) =
+                serde_json::from_str::<ExtractBody>(pdu.content.get())
+            {
+                services().rooms.search.deindex_pdu(
+                    shortroomid,
+                    &pdu_id,
+                    &content.body,
+                )?;
+            }
+
             let room_version_id =
                 services().rooms.state.get_room_version(&pdu.room_id)?;
             pdu.redact(room_version_id, reason)?;
+
             self.replace_pdu(
                 &pdu_id,
                 &utils::to_canonical_object(&pdu).expect("PDU is an object"),