fix(crypto): When a session updates, re-fetch encryption info for non-UTDs

andybalaam · andybalaam · commit 414f03d59110 · 2025-02-14T11:59:31.000Z
Fixes element-hq/element-meta#2697
diff --git a/crates/matrix-sdk-ui/src/timeline/controller/mod.rs b/crates/matrix-sdk-ui/src/timeline/controller/mod.rs
@@ -1062,7 +1062,7 @@ impl<P: RoomDataProvider> TimelineController<P> {
         decryptor: impl Decryptor,
         session_ids: Option<BTreeSet<String>>,
     ) {
-        let state = self.state.clone().write_owned().await;
+        let mut state = self.state.clone().write_owned().await;
 
         // We should retry an event if its session is included in the list, or the list
         // is None.
@@ -1074,11 +1074,25 @@ impl<P: RoomDataProvider> TimelineController<P> {
             }
         };
 
-        // Find which messages need retrying
-        let retry_indices = event_indices_to_retry_decryption(&state, &should_retry);
+        // Find which events need retrying
+        let (retry_decryption_indices, retry_info_indices) =
+            event_indices_to_retry_decryption(&state, &should_retry);
+
+        // Retry fetching encryption info for events that are already decrypted
+        let room_data_provider = self.room_data_provider.clone();
+        matrix_sdk::executor::spawn(async move {
+            state.retry_event_encryption_info(retry_info_indices, &room_data_provider).await;
+        });
 
-        // Retry decrypting
-        self.retry_event_decryption_by_index(state, retry_indices, should_retry, decryptor).await;
+        // Retry decrypting UTDs
+        let state = self.state.clone().write_owned().await;
+        self.retry_event_decryption_by_index(
+            state,
+            retry_decryption_indices,
+            should_retry,
+            decryptor,
+        )
+        .await;
     }
 
     /// Retry decryption of the supplied events, which are expected to be UTDs.
@@ -1438,24 +1452,40 @@ impl<P: RoomDataProvider> TimelineController<P> {
     }
 }
 
+/// Decide which events should be retried, either for re-decryption, or, if they
+/// are already decrypted, for re-checking their encryption info.
+///
+/// Returns a tuple `(retry_decryption_indices, retry_info_indices)` where
+/// `retry_decryption_indices` is a list of UTDs to try decrypting, and
+/// retry_info_indices is a list of already-decrypted events whose encryption
+/// info we can re-fetch.
 fn event_indices_to_retry_decryption(
     state: &tokio::sync::OwnedRwLockWriteGuard<TimelineState>,
     should_retry: impl Fn(&str) -> bool,
-) -> Vec<usize> {
-    let retry_indices: Vec<_> = state
-        .items
-        .iter()
-        .enumerate()
-        .filter_map(|(idx, item)| match item.as_event()?.content().as_unable_to_decrypt()? {
-            EncryptedMessage::MegolmV1AesSha2 { session_id, .. } if should_retry(session_id) => {
-                Some(idx)
+) -> (Vec<usize>, Vec<usize>) {
+    let mut retry_decryption_indices = Vec::new();
+    let mut retry_info_indices = Vec::new();
+
+    for (idx, item) in state.items.iter().enumerate() {
+        if let Some(event) = item.as_event() {
+            if let Some(remote_event) = event.as_remote() {
+                if let Some(session_id) = &remote_event.session_id {
+                    if should_retry(session_id) {
+                        match event.content() {
+                            TimelineItemContent::UnableToDecrypt(_) => {
+                                retry_decryption_indices.push(idx);
+                            }
+                            _ => {
+                                retry_info_indices.push(idx);
+                            }
+                        }
+                    }
+                }
             }
-            EncryptedMessage::MegolmV1AesSha2 { .. }
-            | EncryptedMessage::OlmV1Curve25519AesSha2 { .. }
-            | EncryptedMessage::Unknown => None,
-        })
-        .collect();
-    retry_indices
+        }
+    }
+
+    (retry_decryption_indices, retry_info_indices)
 }
 
 impl TimelineController {
diff --git a/crates/matrix-sdk-ui/src/timeline/controller/state.rs b/crates/matrix-sdk-ui/src/timeline/controller/state.rs
@@ -46,7 +46,10 @@ use super::{
     DateDividerMode, TimelineFocusKind, TimelineMetadata, TimelineSettings,
     TimelineStateTransaction,
 };
-use crate::unable_to_decrypt_hook::UtdHookManager;
+use crate::{
+    timeline::{event_item::EventTimelineItemKind, TimelineItemKind},
+    unable_to_decrypt_hook::UtdHookManager,
+};
 
 #[derive(Debug)]
 pub(in crate::timeline) struct TimelineState {
@@ -225,6 +228,45 @@ impl TimelineState {
         txn.commit();
     }
 
+    /// Try to fetch [`EncryptionInfo`] for the events with the supplied
+    /// indices, and update them where we succeed.
+    pub(super) async fn retry_event_encryption_info<P: RoomDataProvider>(
+        &mut self,
+        retry_indices: Vec<usize>,
+        room_data_provider: &P,
+    ) {
+        let mut new_info = Vec::new();
+        for idx in retry_indices {
+            let item = self.items.get(idx);
+            if let Some(item) = item {
+                if let Some(event) = item.as_event() {
+                    let sender = event.sender.clone();
+                    if let Some(remote) = event.as_remote() {
+                        if let Some(session_id) = &remote.session_id {
+                            let mut new_remote = remote.clone();
+
+                            new_remote.encryption_info =
+                                room_data_provider.get_encryption_info(session_id, &sender).await;
+
+                            let new_event =
+                                event.with_kind(EventTimelineItemKind::Remote(new_remote));
+
+                            let new_item = item.with_kind(TimelineItemKind::Event(new_event));
+
+                            new_info.push((idx, new_item));
+                        }
+                    }
+                }
+            }
+        }
+
+        let mut txn = self.transaction();
+        for (idx, item) in new_info {
+            txn.replace(idx, item);
+        }
+        txn.commit();
+    }
+
     #[cfg(test)]
     pub(super) fn handle_read_receipts(
         &mut self,
diff --git a/crates/matrix-sdk-ui/src/timeline/tests/encryption.rs b/crates/matrix-sdk-ui/src/timeline/tests/encryption.rs
@@ -28,6 +28,9 @@ use eyeball_im::VectorDiff;
 use matrix_sdk::{
     assert_next_matches_with_timeout,
     crypto::{decrypt_room_key_export, types::events::UtdCause, OlmMachine},
+    deserialized_responses::{
+        AlgorithmInfo, DecryptedRoomEvent, EncryptionInfo, VerificationLevel, VerificationState,
+    },
     test_utils::test_client_builder,
 };
 use matrix_sdk_base::deserialized_responses::{TimelineEvent, UnableToDecryptReason};
@@ -38,18 +41,19 @@ use ruma::{
         EncryptedEventScheme, MegolmV1AesSha2ContentInit, Relation, Replacement,
         RoomEncryptedEventContent,
     },
-    room_id,
+    owned_device_id, room_id,
     serde::Raw,
     user_id,
 };
 use serde_json::{json, value::to_raw_value};
-use stream_assert::assert_next_matches;
+use stream_assert::{assert_next_matches, assert_pending};
 use tokio::time::sleep;
 
 use super::TestTimeline;
 use crate::{
     timeline::{
-        tests::TestTimelineBuilder, EncryptedMessage, TimelineDetails, TimelineItemContent,
+        tests::{TestRoomDataProvider, TestTimelineBuilder},
+        EncryptedMessage, TimelineDetails, TimelineItemContent,
     },
     unable_to_decrypt_hook::{UnableToDecryptHook, UnableToDecryptInfo, UtdHookManager},
 };
@@ -550,6 +554,94 @@ async fn test_retry_message_decryption_highlighted() {
     assert!(event.is_highlighted());
 }
 
+#[async_test]
+async fn test_retry_fetching_encryption_info() {
+    const SESSION_ID: &str = "C25PoE+4MlNidQD0YU5ibZqHawV0zZ/up7R8vYJBYTY";
+    let sender = user_id!("@sender:s.co");
+    let room_id = room_id!("!room:s.co");
+
+    // Given when I ask the room for new encryption info for any session, it will
+    // say "verified"
+    let verified_encryption_info = make_encryption_info(SESSION_ID, VerificationState::Verified);
+    let provider = TestRoomDataProvider::default().with_encryption_info(verified_encryption_info);
+    let timeline = TestTimelineBuilder::new().provider(provider).build();
+    let f = &timeline.factory;
+    let mut stream = timeline.subscribe().await;
+
+    // But right now the timeline contains 2 events whose info says "unverified"
+    // One is linked to SESSION_ID, the other is linked to some other session.
+    let timeline_event_this_session: TimelineEvent = DecryptedRoomEvent {
+        event: f.text_msg("foo").sender(sender).room(room_id).into_raw(),
+        encryption_info: make_encryption_info(
+            SESSION_ID,
+            VerificationState::Unverified(VerificationLevel::UnsignedDevice),
+        ),
+        unsigned_encryption_info: None,
+    }
+    .into();
+    let timeline_event_other_session: TimelineEvent = DecryptedRoomEvent {
+        event: f.text_msg("foo").sender(sender).room(room_id).into_raw(),
+        encryption_info: make_encryption_info(
+            "other_session_id",
+            VerificationState::Unverified(VerificationLevel::UnsignedDevice),
+        ),
+        unsigned_encryption_info: None,
+    }
+    .into();
+    timeline.handle_live_event(timeline_event_this_session).await;
+    timeline.handle_live_event(timeline_event_other_session).await;
+
+    // Sanity: the events come through as unverified
+    assert_eq!(timeline.controller.items().await.len(), 3);
+    let item_1 = assert_next_matches!(stream, VectorDiff::PushBack { value } => value);
+    let fetched_encryption_info_1 =
+        item_1.as_event().unwrap().as_remote().unwrap().encryption_info.as_ref().unwrap();
+    assert_matches!(fetched_encryption_info_1.verification_state, VerificationState::Unverified(_));
+    // (Plus a date divider is emitted - not sure why it's in the middle...)
+    let date_divider = assert_next_matches!(stream, VectorDiff::PushFront { value } => value);
+    assert!(date_divider.is_date_divider());
+    let item_2 = assert_next_matches!(stream, VectorDiff::PushBack { value } => value);
+    let fetched_encryption_info_2 =
+        item_2.as_event().unwrap().as_remote().unwrap().encryption_info.as_ref().unwrap();
+    assert_matches!(fetched_encryption_info_2.verification_state, VerificationState::Unverified(_));
+
+    // When we retry the session with ID SESSION_ID
+    let own_user_id = user_id!("@me:s.co");
+    let olm_machine = OlmMachine::new(own_user_id, "SomeDeviceId".into()).await;
+    timeline
+        .controller
+        .retry_event_decryption_test(
+            room_id,
+            olm_machine,
+            Some(iter::once(SESSION_ID.to_owned()).collect()),
+        )
+        .await;
+
+    // Then the event in that session has been updated to be verified
+    assert_eq!(timeline.controller.items().await.len(), 3);
+    let item = assert_next_matches!(stream, VectorDiff::Set { index: 1, value } => value);
+    let event = item.as_event().unwrap();
+    let fetched_encryption_info = event.as_remote().unwrap().encryption_info.as_ref().unwrap();
+    assert_matches!(fetched_encryption_info.verification_state, VerificationState::Verified);
+
+    // But the other one is unchanged because it was for a different session - no
+    // other updates are waiting
+    assert_pending!(stream);
+}
+
+fn make_encryption_info(session_id: &str, verification_state: VerificationState) -> EncryptionInfo {
+    EncryptionInfo {
+        sender: BOB.to_owned(),
+        sender_device: Some(owned_device_id!("BOBDEVICE")),
+        algorithm_info: AlgorithmInfo::MegolmV1AesSha2 {
+            curve25519_key: Default::default(),
+            sender_claimed_keys: Default::default(),
+        },
+        verification_state,
+        session_id: session_id.to_owned(),
+    }
+}
+
 #[async_test]
 async fn test_utd_cause_for_nonmember_event_is_found() {
     // Given a timline
diff --git a/crates/matrix-sdk-ui/src/timeline/tests/mod.rs b/crates/matrix-sdk-ui/src/timeline/tests/mod.rs
@@ -28,7 +28,7 @@ use imbl::vector;
 use indexmap::IndexMap;
 use matrix_sdk::{
     config::RequestConfig,
-    deserialized_responses::TimelineEvent,
+    deserialized_responses::{EncryptionInfo, TimelineEvent},
     event_cache::paginator::{PaginableRoom, PaginatorError},
     room::{EventWithContextResponse, Messages, MessagesOptions},
     send_queue::RoomSendQueueUpdate,
@@ -267,6 +267,9 @@ struct TestRoomDataProvider {
 
     /// Events redacted with that room data providier.
     pub redacted: Arc<RwLock<Vec<OwnedEventId>>>,
+
+    /// Returned from get_encryption_info method
+    pub encryption_info: Option<EncryptionInfo>,
 }
 
 impl TestRoomDataProvider {
@@ -278,6 +281,14 @@ impl TestRoomDataProvider {
         self.fully_read_marker = Some(event_id);
         self
     }
+
+    pub(crate) fn with_encryption_info(
+        mut self,
+        encryption_info: EncryptionInfo,
+    ) -> TestRoomDataProvider {
+        self.encryption_info = Some(encryption_info);
+        self
+    }
 }
 
 impl PaginableRoom for TestRoomDataProvider {
@@ -419,7 +430,7 @@ impl RoomDataProvider for TestRoomDataProvider {
         &self,
         _session_id: &str,
         _sender: &UserId,
-    ) -> Option<matrix_sdk::deserialized_responses::EncryptionInfo> {
-        None
+    ) -> Option<EncryptionInfo> {
+        self.encryption_info.clone()
     }
 }
