From 234e0be3376ee466a4050a7cd55cbbdcda432597 Mon Sep 17 00:00:00 2001
From: Benjamin Bouvier <benjamin@bouvier.cc>
Date: Tue, 22 Apr 2025 20:45:06 +0200
Subject: [PATCH] refactor(timeline): reuse the same push context for all the
 events we're trying to re-decrypt

---
 .../controller/decryption_retry_task.rs       | 13 ++------
 .../src/timeline/controller/state.rs          |  9 ++----
 crates/matrix-sdk-ui/src/timeline/traits.rs   | 30 ++++++++++++++-----
 3 files changed, 27 insertions(+), 25 deletions(-)

diff --git a/crates/matrix-sdk-ui/src/timeline/controller/decryption_retry_task.rs b/crates/matrix-sdk-ui/src/timeline/controller/decryption_retry_task.rs
index 0b2f89704..bea7e4b77 100644
--- a/crates/matrix-sdk-ui/src/timeline/controller/decryption_retry_task.rs
+++ b/crates/matrix-sdk-ui/src/timeline/controller/decryption_retry_task.rs
@@ -256,6 +256,7 @@ async fn decrypt_by_index<D: Decryptor>(
     retry_indices: Vec<usize>,
 ) {
     let push_context = room_data_provider.push_context().await;
+    let push_context = push_context.as_ref();
     let unable_to_decrypt_hook = state.meta.unable_to_decrypt_hook.clone();
 
     let retry_one = |item: Arc<TimelineItem>| {
@@ -290,7 +291,7 @@ async fn decrypt_by_index<D: Decryptor>(
                 return None;
             };
 
-            match decryptor.decrypt_event_impl(original_json).await {
+            match decryptor.decrypt_event_impl(original_json, push_context).await {
                 Ok(event) => {
                     if let SdkTimelineEventKind::UnableToDecrypt { utd_info, .. } = event.kind {
                         info!(
@@ -320,15 +321,7 @@ async fn decrypt_by_index<D: Decryptor>(
         ))
     };
 
-    state
-        .retry_event_decryption(
-            retry_one,
-            retry_indices,
-            push_context,
-            room_data_provider,
-            settings,
-        )
-        .await;
+    state.retry_event_decryption(retry_one, retry_indices, room_data_provider, settings).await;
 }
 
 #[cfg(test)]
diff --git a/crates/matrix-sdk-ui/src/timeline/controller/state.rs b/crates/matrix-sdk-ui/src/timeline/controller/state.rs
index d6e88028c..945b0312f 100644
--- a/crates/matrix-sdk-ui/src/timeline/controller/state.rs
+++ b/crates/matrix-sdk-ui/src/timeline/controller/state.rs
@@ -15,9 +15,7 @@
 use std::{future::Future, sync::Arc};
 
 use eyeball_im::VectorDiff;
-use matrix_sdk::{
-    deserialized_responses::TimelineEvent, room::PushContext, send_queue::SendHandle,
-};
+use matrix_sdk::{deserialized_responses::TimelineEvent, send_queue::SendHandle};
 #[cfg(test)]
 use ruma::events::receipt::ReceiptEventContent;
 use ruma::{
@@ -180,7 +178,6 @@ impl TimelineState {
         &mut self,
         retry_one: impl Fn(Arc<TimelineItem>) -> Fut,
         retry_indices: Vec<usize>,
-        push_context: Option<PushContext>,
         room_data_provider: &P,
         settings: &TimelineSettings,
     ) where
@@ -197,12 +194,10 @@ impl TimelineState {
         let mut offset = 0;
         for idx in retry_indices {
             let idx = idx - offset;
-            let Some(mut event) = retry_one(txn.items[idx].clone()).await else {
+            let Some(event) = retry_one(txn.items[idx].clone()).await else {
                 continue;
             };
 
-            event.push_actions = push_context.as_ref().map(|ctx| ctx.for_event(event.raw()));
-
             let handle_one_res = txn
                 .handle_remote_event(
                     event,
diff --git a/crates/matrix-sdk-ui/src/timeline/traits.rs b/crates/matrix-sdk-ui/src/timeline/traits.rs
index 8dddf2eb0..e3a899972 100644
--- a/crates/matrix-sdk-ui/src/timeline/traits.rs
+++ b/crates/matrix-sdk-ui/src/timeline/traits.rs
@@ -282,30 +282,44 @@ pub(super) trait Decryptor: AsyncTraitDeps + Clone + 'static {
     fn decrypt_event_impl(
         &self,
         raw: &Raw<AnySyncTimelineEvent>,
+        push_context: Option<&PushContext>,
     ) -> impl Future<Output = Result<TimelineEvent>> + SendOutsideWasm;
 }
 
 impl Decryptor for Room {
-    async fn decrypt_event_impl(&self, raw: &Raw<AnySyncTimelineEvent>) -> Result<TimelineEvent> {
-        let push_action_ctx = self.push_context().await?;
-        self.decrypt_event(raw.cast_ref(), push_action_ctx.as_ref()).await
+    async fn decrypt_event_impl(
+        &self,
+        raw: &Raw<AnySyncTimelineEvent>,
+        push_context: Option<&PushContext>,
+    ) -> Result<TimelineEvent> {
+        self.decrypt_event(raw.cast_ref(), push_context).await
     }
 }
 
 #[cfg(test)]
 impl Decryptor for (matrix_sdk_base::crypto::OlmMachine, ruma::OwnedRoomId) {
-    async fn decrypt_event_impl(&self, raw: &Raw<AnySyncTimelineEvent>) -> Result<TimelineEvent> {
+    async fn decrypt_event_impl(
+        &self,
+        raw: &Raw<AnySyncTimelineEvent>,
+        push_context: Option<&PushContext>,
+    ) -> Result<TimelineEvent> {
         let (olm_machine, room_id) = self;
         let decryption_settings =
             DecryptionSettings { sender_device_trust_requirement: TrustRequirement::Untrusted };
-        match olm_machine
+
+        let mut timeline_event = match olm_machine
             .try_decrypt_room_event(raw.cast_ref(), room_id, &decryption_settings)
             .await?
         {
-            RoomEventDecryptionResult::Decrypted(decrypted) => Ok(decrypted.into()),
+            RoomEventDecryptionResult::Decrypted(decrypted) => decrypted.into(),
             RoomEventDecryptionResult::UnableToDecrypt(utd_info) => {
-                Ok(TimelineEvent::new_utd_event(raw.clone(), utd_info))
+                TimelineEvent::new_utd_event(raw.clone(), utd_info)
             }
-        }
+        };
+
+        // Fill the push actions here, to mimic what `Room::decrypt_event` does.
+        timeline_event.push_actions = push_context.map(|ctx| ctx.for_event(timeline_event.raw()));
+
+        Ok(timeline_event)
     }
 }