From 88ecf61d49945a3e82307b9f940cdc5e39eafb70 Mon Sep 17 00:00:00 2001
From: Jade Ellis <jade@ellis.link>
Date: Wed, 28 May 2025 01:04:00 +0100
Subject: [PATCH] feat: Store the original content of redacted PDUs

---
 src/database/maps.rs               |  9 +++++++++
 src/service/rooms/timeline/data.rs | 13 +++++++++++++
 src/service/rooms/timeline/mod.rs  | 28 ++++++++++++++++++++++++++--
 3 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/src/database/maps.rs b/src/database/maps.rs
index 19f9ced4..c72ed414 100644
--- a/src/database/maps.rs
+++ b/src/database/maps.rs
@@ -121,6 +121,15 @@ pub(super) static MAPS: &[Descriptor] = &[
 		index_size: 512,
 		..descriptor::SEQUENTIAL
 	},
+	Descriptor {
+		name: "pduid_originalcontent",
+		cache_disp: CacheDisp::SharedWith("pduid_pdu"),
+		key_size_hint: Some(16),
+		val_size_hint: Some(1520),
+		block_size: 2048,
+		index_size: 512,
+		..descriptor::RANDOM
+	},
 	Descriptor {
 		name: "publicroomids",
 		..descriptor::RANDOM_SMALL
diff --git a/src/service/rooms/timeline/data.rs b/src/service/rooms/timeline/data.rs
index 94c78bb0..a0b407d6 100644
--- a/src/service/rooms/timeline/data.rs
+++ b/src/service/rooms/timeline/data.rs
@@ -19,6 +19,8 @@ pub(super) struct Data {
 	pduid_pdu: Arc<Map>,
 	userroomid_highlightcount: Arc<Map>,
 	userroomid_notificationcount: Arc<Map>,
+	/// Stores the original content of redacted PDUs.
+	pduid_originalcontent: Arc<Map>,
 	pub(super) db: Arc<Database>,
 	services: Services,
 }
@@ -38,6 +40,7 @@ impl Data {
 			pduid_pdu: db["pduid_pdu"].clone(),
 			userroomid_highlightcount: db["userroomid_highlightcount"].clone(),
 			userroomid_notificationcount: db["userroomid_notificationcount"].clone(),
+			pduid_originalcontent: db["pduid_originalcontent"].clone(), // Initialize new table
 			db: args.db.clone(),
 			services: Services {
 				short: args.depend::<rooms::short::Service>("rooms::short"),
@@ -177,6 +180,16 @@ impl Data {
 		self.pduid_pdu.get(pdu_id).await.deserialized()
 	}
 
+	/// Stores the original content of a PDU that is about to be redacted.
+	pub(super) async fn store_redacted_pdu_content(
+		&self,
+		pdu_id: &RawPduId,
+		pdu_json: &CanonicalJsonObject,
+	) -> Result<()> {
+		self.pduid_originalcontent.raw_put(pdu_id, Json(pdu_json));
+		Ok(())
+	}
+
 	pub(super) async fn append_pdu(
 		&self,
 		pdu_id: &RawPduId,
diff --git a/src/service/rooms/timeline/mod.rs b/src/service/rooms/timeline/mod.rs
index 4b2f3cb2..dee12a41 100644
--- a/src/service/rooms/timeline/mod.rs
+++ b/src/service/rooms/timeline/mod.rs
@@ -260,6 +260,16 @@ impl Service {
 		self.db.replace_pdu(pdu_id, pdu_json, pdu).await
 	}
 
+	/// Stores the content of a to-be redacted pdu.
+	#[tracing::instrument(skip(self), level = "debug")]
+	pub async fn store_redacted_pdu_content(
+		&self,
+		pdu_id: &RawPduId,
+		pdu_json: &CanonicalJsonObject,
+	) -> Result<()> {
+		self.db.store_redacted_pdu_content(pdu_id, pdu_json).await
+	}
+
 	/// Creates a new persisted data unit and adds it to a room.
 	///
 	/// By this point the incoming event should be fully authenticated, no auth
@@ -472,7 +482,7 @@ impl Service {
 								.user_can_redact(redact_id, &pdu.sender, &pdu.room_id, false)
 								.await?
 							{
-								self.redact_pdu(redact_id, pdu, shortroomid).await?;
+								self.redact_pdu(redact_id, pdu, shortroomid, true).await?;
 							}
 						}
 					},
@@ -485,7 +495,7 @@ impl Service {
 								.user_can_redact(redact_id, &pdu.sender, &pdu.room_id, false)
 								.await?
 							{
-								self.redact_pdu(redact_id, pdu, shortroomid).await?;
+								self.redact_pdu(redact_id, pdu, shortroomid, true).await?;
 							}
 						}
 					},
@@ -1033,6 +1043,7 @@ impl Service {
 		event_id: &EventId,
 		reason: &PduEvent,
 		shortroomid: ShortRoomId,
+		keep_original_content: bool,
 	) -> Result {
 		// TODO: Don't reserialize, keep original json
 		let Ok(pdu_id) = self.get_pdu_id(event_id).await else {
@@ -1054,6 +1065,19 @@ impl Service {
 
 		let room_version_id = self.services.state.get_room_version(&pdu.room_id).await?;
 
+		if keep_original_content {
+			let original_pdu_json = utils::to_canonical_object(&pdu).map_err(|e| {
+				err!(Database(error!(
+					?event_id,
+					?e,
+					"Failed to convert PDU to canonical JSON for original content storage"
+				)))
+			})?;
+			self.db
+				.store_redacted_pdu_content(&pdu_id, &original_pdu_json)
+				.await?;
+		}
+
 		pdu.redact(&room_version_id, reason)?;
 
 		let obj = utils::to_canonical_object(&pdu).map_err(|e| {