feat: Redecryptor start to send out redecryptor reports

poljar · poljar · commit e8c882f916f6 · 2025-10-10T12:17:26.000+02:00
diff --git a/crates/matrix-sdk/src/event_cache/mod.rs b/crates/matrix-sdk/src/event_cache/mod.rs
@@ -69,16 +69,18 @@ use tracing::{Instrument as _, Span, debug, error, info, info_span, instrument,
 use crate::{
     Client,
     client::WeakClient,
-    event_cache::redecryptor::Redecryptor,
     send_queue::{LocalEchoContent, RoomSendQueueUpdate, SendQueueUpdate},
 };
 
 mod deduplicator;
 mod pagination;
+#[cfg(feature = "e2e-encryption")]
 mod redecryptor;
 mod room;
 
 pub use pagination::{RoomPagination, RoomPaginationStatus};
+#[cfg(feature = "e2e-encryption")]
+pub use redecryptor::{DecryptionRetryRequest, RedecryptorReport};
 pub use room::{RoomEventCache, RoomEventCacheSubscriber, ThreadEventCacheUpdate};
 
 /// An error observed in the [`EventCache`].
@@ -152,7 +154,8 @@ pub struct EventCacheDropHandles {
     auto_shrink_linked_chunk_task: JoinHandle<()>,
 
     /// The task used to automatically redecrypt UTDs.
-    redecryptor: Redecryptor,
+    #[cfg(feature = "e2e-encryption")]
+    _redecryptor: redecryptor::Redecryptor,
 }
 
 impl fmt::Debug for EventCacheDropHandles {
@@ -205,6 +208,9 @@ impl EventCache {
             linked_chunk_update_sender.clone(),
         )));
 
+        #[cfg(feature = "e2e-encryption")]
+        let redecryption_channels = redecryptor::RedecryptorChannels::new();
+
         Self {
             inner: Arc::new(EventCacheInner {
                 client,
@@ -218,6 +224,8 @@ impl EventCache {
                 _thread_subscriber_task: thread_subscriber_task,
                 #[cfg(feature = "experimental-search")]
                 _search_indexing_task: search_indexing_task,
+                #[cfg(feature = "e2e-encryption")]
+                redecryption_channels,
                 thread_subscriber_receiver,
             }),
         }
@@ -262,13 +270,25 @@ impl EventCache {
                 auto_shrink_receiver,
             ));
 
-            let redecryptor = Redecryptor::new(Arc::downgrade(&self.inner));
+            #[cfg(feature = "e2e-encryption")]
+            let redecryptor = {
+                let receiver = self
+                    .inner
+                    .redecryption_channels
+                    .decryption_request_receiver
+                    .lock()
+                    .take()
+                    .expect("We should have initialized the channel an subscribing should happen only once");
+
+                redecryptor::Redecryptor::new(Arc::downgrade(&self.inner), receiver)
+            };
 
             Arc::new(EventCacheDropHandles {
                 listen_updates_task,
                 ignore_user_list_update_task,
                 auto_shrink_linked_chunk_task,
-                redecryptor,
+                #[cfg(feature = "e2e-encryption")]
+                _redecryptor: redecryptor,
             })
         });
 
@@ -848,6 +868,9 @@ struct EventCacheInner {
     /// This is helpful for tests to coordinate that a new thread subscription
     /// has been sent or not.
     thread_subscriber_receiver: Receiver<()>,
+
+    #[cfg(feature = "e2e-encryption")]
+    redecryption_channels: redecryptor::RedecryptorChannels,
 }
 
 type AutoShrinkChannelPayload = OwnedRoomId;
diff --git a/crates/matrix-sdk/src/event_cache/redecryptor.rs b/crates/matrix-sdk/src/event_cache/redecryptor.rs
@@ -72,25 +72,76 @@ use matrix_sdk_base::{BaseClient, crypto::OlmMachine};
 use matrix_sdk_base::{
     crypto::{store::types::RoomKeyInfo, 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 tokio::{sync::mpsc::UnboundedSender, task::JoinHandle};
-use tokio_stream::wrappers::{UnboundedReceiverStream, errors::BroadcastStreamRecvError};
+use tokio::{
+    sync::{
+        broadcast,
+        mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel},
+    },
+    task::JoinHandle,
+};
+use tokio_stream::wrappers::{
+    BroadcastStream, UnboundedReceiverStream, errors::BroadcastStreamRecvError,
+};
 use tracing::{info, instrument, trace, warn};
 
 use crate::event_cache::{
     EventCache, EventCacheError, EventCacheInner, EventsOrigin, RoomEventCacheUpdate,
 };
 
+type SessionId<'a> = &'a str;
+type OwnedSessionId = String;
+
 /// The information sent across the channel to the long-running task requesting
 /// that the supplied set of sessions be retried.
+#[derive(Debug, Clone)]
 pub struct DecryptionRetryRequest {
-    room_id: OwnedRoomId,
-    session_ids: BTreeSet<String>,
+    /// The room ID of the room the events belong to.
+    pub room_id: OwnedRoomId,
+    /// Events that are not decrypted.
+    pub utd_session_ids: BTreeSet<OwnedSessionId>,
+    /// Events that are decrypted but might need to have their
+    /// [`EncryptionInfo`] refreshed.
+    pub refresh_info_session_ids: BTreeSet<OwnedSessionId>,
 }
 
-type SessionId<'a> = &'a str;
+/// A report coming from the redecryptor.
+#[derive(Debug, Clone)]
+pub enum RedecryptorReport {
+    /// Events which we were able to decrypt.
+    ResolvedUtds {
+        /// The room ID of the room the events belong to.
+        room_id: OwnedRoomId,
+        /// The list of event IDs of the decrypted events.
+        events: BTreeSet<OwnedEventId>,
+    },
+    /// The redecryptor might have missed some room keys so it might not have
+    /// re-decrypted events that are now decryptable.
+    Lagging,
+}
+
+pub(super) struct RedecryptorChannels {
+    utd_reporter: broadcast::Sender<RedecryptorReport>,
+    pub(super) decryption_request_sender: UnboundedSender<DecryptionRetryRequest>,
+    pub(super) decryption_request_receiver:
+        Mutex<Option<UnboundedReceiver<DecryptionRetryRequest>>>,
+}
+
+impl RedecryptorChannels {
+    pub(super) fn new() -> Self {
+        let (utd_reporter, _) = broadcast::channel(100);
+        let (decryption_request_sender, decryption_request_receiver) = unbounded_channel();
+
+        Self {
+            utd_reporter,
+            decryption_request_sender,
+            decryption_request_receiver: Mutex::new(Some(decryption_request_receiver)),
+        }
+    }
+}
 
 impl EventCache {
     /// Retrieve a set of events that we weren't able to decrypt.
@@ -154,7 +205,7 @@ impl EventCache {
         let (room_cache, _) = self.for_room(room_id).await?;
         let mut state = room_cache.inner.state.write().await;
 
-        let event_ids: Vec<_> = events.iter().map(|(event_id, _)| event_id).collect();
+        let event_ids: BTreeSet<_> = events.iter().cloned().map(|(event_id, _)| event_id).collect();
 
         trace!(?event_ids, "Replacing successfully re-decrypted events");
 
@@ -177,6 +228,14 @@ impl EventCache {
             origin: EventsOrigin::Cache,
         });
 
+        // We report that we resolved some UTDs, this is mainly for listeners that don't
+        // care about the actual events, just about the fact that UTDs got
+        // resolved. Not sure if we'll have more such listeners but the UTD hook
+        // is one such thing.
+        let report =
+            RedecryptorReport::ResolvedUtds { room_id: room_id.to_owned(), events: event_ids };
+        let _ = self.inner.redecryption_channels.utd_reporter.send(report.into());
+
         Ok(())
     }
 
@@ -233,10 +292,28 @@ impl EventCache {
 
         Ok(())
     }
+
+    /// Explicitly request the redecryption of a set of events.
+    ///
+    /// TODO: Explain when and why this might be useful.
+    pub fn request_decryption(&self, request: DecryptionRetryRequest) {
+        let _ =
+            self.inner.redecryption_channels.decryption_request_sender.send(request).inspect_err(
+                |_| warn!("Requesting a decryption while the redecryption task has been shut down"),
+            );
+    }
+
+    /// Subscribe to reports that the redecryptor generates.
+    ///
+    /// TODO: Explain when the redecryptor might send such reports.
+    pub fn subscrube_to_decryption_reports(
+        &self,
+    ) -> impl Stream<Item = Result<RedecryptorReport, BroadcastStreamRecvError>> {
+        BroadcastStream::new(self.inner.redecryption_channels.utd_reporter.subscribe())
+    }
 }
 
 pub(crate) struct Redecryptor {
-    request_decryption_sender: UnboundedSender<DecryptionRetryRequest>,
     task: JoinHandle<()>,
 }
 
@@ -251,23 +328,17 @@ impl Redecryptor {
     ///
     /// This creates a task that listens to various streams and attempts to
     /// redecrypt UTDs that can be found inside the [`EventCache`].
-    pub(super) fn new(cache: Weak<EventCacheInner>) -> Self {
-        let (request_decryption_sender, receiver) = tokio::sync::mpsc::unbounded_channel();
-
+    pub(super) fn new(
+        cache: Weak<EventCacheInner>,
+        receiver: UnboundedReceiver<DecryptionRetryRequest>,
+    ) -> Self {
         let task = spawn(async {
             let request_redecryption_stream = UnboundedReceiverStream::new(receiver);
 
             Self::listen_for_room_keys_task(cache, request_redecryption_stream).await;
         });
 
-        Self { task, request_decryption_sender }
-    }
-
-    #[allow(dead_code)]
-    pub(super) fn request_decryption(&self, request: DecryptionRetryRequest) {
-        let _ = self.request_decryption_sender.send(request).inspect_err(|_| {
-            warn!("Requesting a decryption while the redecryption task has been shut down")
-        });
+        Self { task }
     }
 
     /// (Re)-subscribe to the room key stream from the [`OlmMachine`].
@@ -307,18 +378,23 @@ impl Redecryptor {
                             break false;
                         };
 
-                        for session_id in request.session_ids {
+                        for session_id in request.utd_session_ids {
                             let _ = cache
                                 .retry_decryption(&request.room_id, &session_id)
                                 .await
                                 .inspect_err(|e| warn!("Error redecrypting {e:?}"));
                         }
+
+                        // TODO: Deal with encryption info updating as well.
                 }
                 // The room key stream from the OlmMachine. Needs to be recreated every time we
                 // receive a `None` from the stream.
                 room_keys = room_key_stream.next() => {
                     match room_keys {
                         Some(Ok(room_keys)) => {
+                            // Alright, some room keys were received and persisted in our store,
+                            // let's attempt to redecrypt events that were encrypted using these
+                            // room keys.
                             let Some(cache) = Self::upgrade_event_cache(cache) else {
                                 break false;
                             };
@@ -329,9 +405,21 @@ impl Redecryptor {
                                     .await
                                     .inspect_err(|e| warn!("Error redecrypting {e:?}"));
                             }
+
+                            // TODO: Deal with encryption info updating as well.
                         },
                         Some(Err(_)) => {
-                            todo!("Handle lagging here, how?")
+                            // We missed some room keys, we need to report this in case a listener
+                            // has and idea which UTDs we should attempt to redecrypt.
+                            //
+                            // This would most likely be the timeline. The timeline might attempt
+                            // to redecrypt all UTDs it is showing to the user.
+                            let Some(cache) = Self::upgrade_event_cache(cache) else {
+                                break false;
+                            };
+
+                            let message = RedecryptorReport::Lagging;
+                            let _ = cache.inner.redecryption_channels.utd_reporter.send(message);
                         },
                         // The stream got closed, this could mean that our OlmMachine got
                         // regenerated, let's return true and try to recreate the stream.
@@ -355,7 +443,16 @@ impl Redecryptor {
         pin_mut!(decryption_request_stream);
 
         while Self::redecryption_loop(&cache, &mut decryption_request_stream).await {
-            info!("Regenerating the re-decryption streams")
+            info!("Regenerating the re-decryption streams");
+
+            let Some(cache) = Self::upgrade_event_cache(&cache) else {
+                break;
+            };
+
+            // Report that the stream got recreated so listeners can attempt to redecrypt
+            // any UTDs they might be seeing.
+            let message = RedecryptorReport::Lagging;
+            let _ = cache.inner.redecryption_channels.utd_reporter.send(message);
         }
 
         info!("Shutting down the event cache redecryptor");
