feat(cache): Let the redecryptor listen to room key withheld updates

Author: poljar

crates/matrix-sdk/src/event_cache/redecryptor.rs

@@ -15,10 +15,16 @@
 //! The Redecryptor (Rd) is a layer and long-running background task which
 //! handles redecryption of events in case we couldn't decrypt them imediatelly.
 //!
-//! Rd listens to the OlmMachine for received room keys. If a new room key has
-//! been received it attempts to find any UTDs in the [`EventCache`]. If Rd
-//! decrypts any UTDs from the event cache it will replace the events in the
-//! cache and send out new [`RoomEventCacheUpdates`] to any of its listeners.
+//! Rd listens to the OlmMachine for received room keys and new
+//! m.room_key.withheld events.
+//!
+//! If a new room key has been received it attempts to find any UTDs in the
+//! [`EventCache`]. If Rd decrypts any UTDs from the event cache it will replace
+//! the events in the cache and send out new [`RoomEventCacheUpdates`] to any of
+//! its listeners.
+//!
+//! If a new withheld info has been received it attempts to find any relevant
+//! events and updates the [`EncryptionInfo`] of an event.
 //!
 //! There's an additional gotcha, the [`OlmMachine`] might get recreated by
 //! calls to [`BaseClient::regenerate_olm()`]. When this happens we will receive
@@ -70,12 +76,15 @@ use futures_util::{StreamExt, pin_mut};
 #[cfg(doc)]
 use matrix_sdk_base::{BaseClient, crypto::OlmMachine};
 use matrix_sdk_base::{
-    crypto::{store::types::RoomKeyInfo, types::events::room::encrypted::EncryptedEvent},
+    crypto::{
+        store::types::{RoomKeyInfo, RoomKeyWithheldInfo},
+        types::events::room::encrypted::EncryptedEvent,
+    },
     deserialized_responses::{DecryptedRoomEvent, TimelineEvent, TimelineEventKind},
     locks::Mutex,
 };
 use matrix_sdk_common::executor::spawn;
-use ruma::{OwnedEventId, OwnedRoomId, RoomId, events::AnySyncTimelineEvent, serde::Raw};
+use ruma::{OwnedEventId, OwnedRoomId, RoomId, serde::Raw};
 use tokio::{
     sync::{
         broadcast,
@@ -151,36 +160,21 @@ impl EventCache {
     /// * `room_id` - The ID of the room where the events were sent to.
     /// * `session_id` - The unique ID of the room key that was used to encrypt
     ///   the event.
-    async fn get_utds(
+    async fn get_utds<B>(
         &self,
         room_id: &RoomId,
-        session_id: SessionId<'_>,
-    ) -> Result<Vec<(OwnedEventId, Raw<AnySyncTimelineEvent>)>, EventCacheError> {
-        let filter_non_utds = |event: TimelineEvent| {
-            let event_id = event.event_id();
-
-            // We only care about events fort his particular room key, identified by the
-            // session ID.
-            if event.kind.session_id() == Some(session_id) {
-                // Only pick out events that are UTDs, get just the Raw event as this is what
-                // the OlmMachine needs.
-                let event = as_variant!(event.kind, TimelineEventKind::UnableToDecrypt { event, .. } => event);
-                // Zip the event ID and event together so we don't have to pick out the event ID
-                // again. We need the event ID to replace the event in the cache.
-                event_id.zip(event)
-            } else {
-                None
-            }
-        };
-
+        _session_id: SessionId<'_>,
+        filter: impl Fn(TimelineEvent) -> Option<B>,
+    ) -> Result<Vec<B>, EventCacheError> {
         // Load the relevant events from the event cache store and attempt to redecrypt
         // things.
         //
         // TODO: We can't load **all** events all the time.
+        // TODO: Use the session ID to filter things.
         let store = self.inner.store.lock().await?;
         let events = store.get_room_events(&room_id).await?;
 
-        Ok(events.into_iter().filter_map(filter_non_utds).collect())
+        Ok(events.into_iter().filter_map(filter).collect())
     }
 
     /// Handle a chunk of events that we were previously unable to decrypt but
@@ -272,8 +266,25 @@ impl EventCache {
     ) -> Result<(), EventCacheError> {
         trace!("Retrying to decrypt");
 
+        let filter_non_utds = |event: TimelineEvent| {
+            let event_id = event.event_id();
+
+            // We only care about events fort his particular room key, identified by the
+            // session ID.
+            if event.kind.session_id() == Some(session_id) {
+                // Only pick out events that are UTDs, get just the Raw event as this is what
+                // the OlmMachine needs.
+                let event = as_variant!(event.kind, TimelineEventKind::UnableToDecrypt { event, .. } => event);
+                // Zip the event ID and event together so we don't have to pick out the event ID
+                // again. We need the event ID to replace the event in the cache.
+                event_id.zip(event)
+            } else {
+                None
+            }
+        };
+
         // Get all the relevant UTDs.
-        let events = self.get_utds(room_id, session_id).await?;
+        let events = self.get_utds(room_id, session_id, filter_non_utds).await?;
 
         // Let's attempt to decrypt them them.
         let mut decrypted_events = Vec::with_capacity(events.len());
@@ -293,6 +304,54 @@ impl EventCache {
         Ok(())
     }
 
+    async fn update_encryption_info(
+        &self,
+        room_id: &RoomId,
+        session_id: SessionId<'_>,
+    ) -> Result<(), EventCacheError> {
+        let filter_non_utds = |event: TimelineEvent| {
+            let event_id = event.event_id();
+
+            // We only care about events fort his particular room key, identified by the
+            // session ID.
+            if event.kind.session_id() == Some(session_id) {
+                let event = as_variant!(event.kind, TimelineEventKind::Decrypted(event) => event);
+                // Zip the event ID and event together so we don't have to pick out the event ID
+                // again. We need the event ID to replace the event in the cache.
+                event_id.zip(event)
+            } else {
+                None
+            }
+        };
+
+        let client = self.inner.client().ok().unwrap();
+        let room = client.get_room(room_id).unwrap();
+
+        // Get all the relevant events.
+        let events = self.get_utds(room_id, session_id, filter_non_utds).await?;
+
+        // Let's attempt to update their encryption info.
+        let mut updated_events = Vec::with_capacity(events.len());
+
+        for (event_id, mut event) in events {
+            let new_encryption_info =
+                room.get_encryption_info(session_id, &event.encryption_info.sender).await;
+
+            // Only create a replacement if the encryption info actually changed.
+            if let Some(new_encryption_info) = new_encryption_info {
+                if event.encryption_info != new_encryption_info {
+                    event.encryption_info = new_encryption_info;
+
+                    updated_events.push((event_id, event));
+                }
+            }
+        }
+
+        self.on_resolved_utds(room_id, updated_events).await?;
+
+        Ok(())
+    }
+
     /// Explicitly request the redecryption of a set of events.
     ///
     /// TODO: Explain when and why this might be useful.
@@ -347,12 +406,17 @@ impl Redecryptor {
     /// the sending part of the stream has been dropped.
     async fn subscribe_to_room_key_stream(
         cache: &Weak<EventCacheInner>,
-    ) -> Option<impl Stream<Item = Result<Vec<RoomKeyInfo>, BroadcastStreamRecvError>>> {
+    ) -> Option<(
+        impl Stream<Item = Result<Vec<RoomKeyInfo>, BroadcastStreamRecvError>>,
+        impl Stream<Item = Vec<RoomKeyWithheldInfo>>,
+    )> {
         let event_cache = cache.upgrade()?;
         let client = event_cache.client().ok()?;
         let machine = client.olm_machine().await;
 
-        machine.as_ref().map(|m| m.store().room_keys_received_stream())
+        machine.as_ref().map(|m| {
+            (m.store().room_keys_received_stream(), m.store().room_keys_withheld_received_stream())
+        })
     }
 
     fn upgrade_event_cache(cache: &Weak<EventCacheInner>) -> Option<EventCache> {
@@ -363,11 +427,14 @@ impl Redecryptor {
         cache: &Weak<EventCacheInner>,
         decryption_request_stream: &mut Pin<&mut impl Stream<Item = DecryptionRetryRequest>>,
     ) -> bool {
-        let Some(room_key_stream) = Self::subscribe_to_room_key_stream(cache).await else {
+        let Some((room_key_stream, withheld_stream)) =
+            Self::subscribe_to_room_key_stream(cache).await
+        else {
             return false;
         };
 
         pin_mut!(room_key_stream);
+        pin_mut!(withheld_stream);
 
         loop {
             tokio::select! {
@@ -385,7 +452,14 @@ impl Redecryptor {
                                 .inspect_err(|e| warn!("Error redecrypting {e:?}"));
                         }
 
-                        // TODO: Deal with encryption info updating as well.
+                        for session_id in request.refresh_info_session_ids {
+                            let _ = cache.update_encryption_info(&request.room_id, &session_id).await.inspect_err(|e|
+                                warn!(
+                                    room_id = %request.room_id,
+                                    session_id = session_id,
+                                    "Unable to update the encryption info {e:?}",
+                            ));
+                        }
                 }
                 // The room key stream from the OlmMachine. Needs to be recreated every time we
                 // receive a `None` from the stream.
@@ -399,14 +473,21 @@ impl Redecryptor {
                                 break false;
                             };
 
-                            for key in room_keys {
+                            for key in &room_keys {
                                 let _ = cache
                                     .retry_decryption(&key.room_id, &key.session_id)
                                     .await
                                     .inspect_err(|e| warn!("Error redecrypting {e:?}"));
                             }
 
-                            // TODO: Deal with encryption info updating as well.
+                            for key in room_keys {
+                                let _ = cache.update_encryption_info(&key.room_id, &key.session_id).await.inspect_err(|e|
+                                    warn!(
+                                        room_id = %key.room_id,
+                                        session_id = key.session_id,
+                                        "Unable to update the encryption info {e:?}",
+                                ));
+                            }
                         },
                         Some(Err(_)) => {
                             // We missed some room keys, we need to report this in case a listener
@@ -428,6 +509,27 @@ impl Redecryptor {
                         }
                     }
                 }
+                withheld_info = withheld_stream.next() => {
+                    match withheld_info {
+                        Some(infos) => {
+                            let Some(cache) = Self::upgrade_event_cache(cache) else {
+                                break false;
+                            };
+
+                            for RoomKeyWithheldInfo { room_id, session_id, .. } in &infos {
+                                let _ = cache.update_encryption_info(room_id, session_id).await.inspect_err(|e|
+                                    warn!(
+                                        room_id = %room_id,
+                                        session_id = session_id,
+                                        "Unable to update the encryption info {e:?}",
+                                ));
+                            }
+                        }
+                        // The stream got closed, same as for the room key stream, we'll try to
+                        // recreate the streams.
+                        None => break true
+                    }
+                }
                 else => break false,
             }
         }