diff --git a/crates/matrix-sdk/src/event_cache/mod.rs b/crates/matrix-sdk/src/event_cache/mod.rs index 50c67c1b7..16d0dba55 100644 --- a/crates/matrix-sdk/src/event_cache/mod.rs +++ b/crates/matrix-sdk/src/event_cache/mod.rs @@ -280,7 +280,7 @@ impl EventCache { .take() .expect("We should have initialized the channel an subscribing should happen only once"); - redecryptor::Redecryptor::new(Arc::downgrade(&self.inner), receiver) + redecryptor::Redecryptor::new(Arc::downgrade(&self.inner), receiver, &self.inner.linked_chunk_update_sender) }; diff --git a/crates/matrix-sdk/src/event_cache/redecryptor.rs b/crates/matrix-sdk/src/event_cache/redecryptor.rs index 2c21f40a8..48d4707b6 100644 --- a/crates/matrix-sdk/src/event_cache/redecryptor.rs +++ b/crates/matrix-sdk/src/event_cache/redecryptor.rs @@ -130,7 +130,7 @@ use ruma::{ serde::Raw, }; use tokio::sync::{ - broadcast, + broadcast::{self, Sender}, mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel}, }; use tokio_stream::wrappers::{ @@ -141,7 +141,7 @@ use tracing::{info, instrument, trace, warn}; #[cfg(doc)] use super::RoomEventCache; use super::{EventCache, EventCacheError, EventCacheInner, EventsOrigin, RoomEventCacheUpdate}; -use crate::{Room, room::PushContext}; +use crate::{Room, event_cache::RoomEventCacheLinkedChunkUpdate, room::PushContext}; type SessionId<'a> = &'a str; type OwnedSessionId = String; @@ -179,7 +179,7 @@ pub enum RedecryptorReport { } pub(super) struct RedecryptorChannels { - utd_reporter: broadcast::Sender, + utd_reporter: Sender, pub(super) decryption_request_sender: UnboundedSender, pub(super) decryption_request_receiver: Mutex>>, @@ -391,10 +391,36 @@ impl EventCache { room_id: &RoomId, session_id: SessionId<'_>, ) -> Result<(), EventCacheError> { - trace!("Retrying to decrypt"); - // Get all the relevant UTDs. let events = self.get_utds(room_id, session_id).await?; + self.retry_decryption_for_events(room_id, events).await + } + + /// Attempt to redecrypt events that were persisted in the event cache. + #[instrument(skip_all, fields(updates.linked_chunk_id))] + async fn retry_decryption_for_event_cache_updates( + &self, + updates: RoomEventCacheLinkedChunkUpdate, + ) -> Result<(), EventCacheError> { + let room_id = updates.linked_chunk_id.room_id(); + let events: Vec<_> = updates + .updates + .into_iter() + .flat_map(|updates| updates.into_items()) + .filter_map(filter_timeline_event_to_utd) + .collect(); + + self.retry_decryption_for_events(room_id, events).await + } + + /// Attempt to redecrypt a chunk of UTDs. + #[instrument(skip_all, fields(room_id, session_id))] + async fn retry_decryption_for_events( + &self, + room_id: &RoomId, + events: Vec, + ) -> Result<(), EventCacheError> { + trace!("Retrying to decrypt"); if events.is_empty() { trace!("No relevant events found."); @@ -604,11 +630,19 @@ impl Redecryptor { pub(super) fn new( cache: Weak, receiver: UnboundedReceiver, + linked_chunk_update_sender: &Sender, ) -> Self { + let linked_chunk_stream = BroadcastStream::new(linked_chunk_update_sender.subscribe()); + let task = spawn(async { let request_redecryption_stream = UnboundedReceiverStream::new(receiver); - Self::listen_for_room_keys_task(cache, request_redecryption_stream).await; + Self::listen_for_room_keys_task( + cache, + request_redecryption_stream, + linked_chunk_stream, + ) + .await; }) .abort_on_drop(); @@ -642,6 +676,9 @@ impl Redecryptor { async fn redecryption_loop( cache: &Weak, decryption_request_stream: &mut Pin<&mut impl Stream>, + events_stream: &mut Pin< + &mut impl Stream>, + >, ) -> bool { let Some((room_key_stream, withheld_stream)) = Self::subscribe_to_room_key_stream(cache).await @@ -753,6 +790,35 @@ impl Redecryptor { None => break true } } + // Events that the event cache handled. If the event cache received any UTDs, let's + // attempt to redecrypt them in case the room key was received before the event + // cache was able to return them using `get_utds()`. + Some(event_updates) = events_stream.next() => { + match event_updates { + Ok(updates) => { + let Some(cache) = Self::upgrade_event_cache(cache) else { + break false; + }; + + let linked_chunk_id = updates.linked_chunk_id.to_owned(); + + let _ = cache.retry_decryption_for_event_cache_updates(updates).await.inspect_err(|e| + warn!( + %linked_chunk_id, + "Unable to handle UTDs from event cache updates {e:?}", + ) + ); + } + Err(_) => { + let Some(cache) = Self::upgrade_event_cache(cache) else { + break false; + }; + + let message = RedecryptorReport::Lagging; + let _ = cache.inner.redecryption_channels.utd_reporter.send(message); + } + } + } else => break false, } } @@ -761,13 +827,17 @@ impl Redecryptor { async fn listen_for_room_keys_task( cache: Weak, decryption_request_stream: UnboundedReceiverStream, + events_stream: BroadcastStream, ) { // We pin the decryption request stream here since that one doesn't need to be // recreated and we don't want to miss messages coming from the stream // while recreating it unnecessarily. pin_mut!(decryption_request_stream); + pin_mut!(events_stream); - while Self::redecryption_loop(&cache, &mut decryption_request_stream).await { + while Self::redecryption_loop(&cache, &mut decryption_request_stream, &mut events_stream) + .await + { info!("Regenerating the re-decryption streams"); let Some(cache) = Self::upgrade_event_cache(&cache) else {