Skip to main content

matrix_sdk/event_cache/
redecryptor.rs

1// Copyright 2025 The Matrix.org Foundation C.I.C.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! The Redecryptor (affectionately known as R2D2) is a layer and long-running
16//! background task which handles redecryption of events in case we couldn't
17//! decrypt them immediately.
18//!
19//! There are various reasons why a room key might not be available immediately
20//! when the event becomes available:
21//!     - The to-device message containing the room key just arrives late, i.e.
22//!       after the room event.
23//!     - The event is a historic event and we need to first download the room
24//!       key from the backup.
25//!     - The event is a historic event in a previously unjoined room, we need
26//!       to receive historic room keys as defined in [MSC3061].
27//!
28//! R2D2 listens to the [`OlmMachine`] for received room keys and new
29//! m.room_key.withheld events.
30//!
31//! If a new room key has been received, it attempts to find any UTDs in the
32//! [`EventCache`]. If R2D2 decrypts any UTDs from the event cache, it will
33//! replace the events in the cache and send out new [`RoomEventCacheUpdate`]s
34//! to any of its listeners.
35//!
36//! If a new withheld info has been received, it attempts to find any relevant
37//! events and updates the [`EncryptionInfo`] of an event.
38//!
39//! There's an additional gotcha: the [`OlmMachine`] might get recreated by
40//! calls to [`BaseClient::regenerate_olm()`]. When this happens, we will
41//! receive a `None` on the room keys stream and we need to re-listen to it.
42//!
43//! Another gotcha is that room keys might be received on another process if the
44//! [`Client`] is operating on a Apple iOS device. A separate process is used
45//! in this case to receive push notifications. In this case, the room key will
46//! be received and R2D2 won't get notified about it. To work around this,
47//! decryption requests can be explicitly sent to R2D2.
48//!
49//! The final gotcha is that a room key might be received just in between the
50//! time the event was initially tried to be decrypted and the time it took to
51//! persist it in the event cache. To handle this race condition, R2D2 listens
52//! to the event cache and attempts to decrypt any UTDs the event cache
53//! persists.
54//!
55//! In the graph below, the Timeline block is meant to be the `Timeline` from
56//! the `matrix-sdk-ui` crate, but it could be any other listener that
57//! subscribes to [`RedecryptorReport`] stream.
58//!
59//! ```markdown
60//! 
61//!      .----------------------.
62//!     |                        |
63//!     |      Beeb, boop!       |
64//!     |                        .
65//!      ----------------------._ \
66//!                               -;  _____
67//!                                 .`/L|__`.
68//!                                / =[_]O|` \
69//!                                |"+_____":|
70//!                              __:='|____`-:__
71//!                             ||[] ||====|| []||
72//!                             ||[] ||====|| []||
73//!                             |:== ||====|| ==:|
74//!                             ||[] ||====|| []||
75//!                             ||[] ||====|| []||
76//!                            _||_  ||====||  _||_
77//!                           (====) |:====:| (====)
78//!                            }--{  | |  | |  }--{
79//!                           (____) |_|  |_| (____)
80//!
81//!                              ┌─────────────┐
82//!                              │             │
83//!                  ┌───────────┤   Timeline  │◄────────────┐
84//!                  │           │             │             │
85//!                  │           └──────▲──────┘             │
86//!                  │                  │                    │
87//!                  │                  │                    │
88//!                  │                  │                    │
89//!              Decryption             │                Redecryptor
90//!                request              │                  report
91//!                  │        RoomEventCacheUpdates          │
92//!                  │                  │                    │
93//!                  │                  │                    │
94//!                  │      ┌───────────┴───────────┐        │
95//!                  │      │                       │        │
96//!                  └──────►         R2D2          │────────┘
97//!                         │                       │
98//!                         └──▲─────────────────▲──┘
99//!                            │                 │
100//!                            │                 │
101//!                            │                 │
102//!                         Received        Received room
103//!                          events          keys stream
104//!                            │                 │
105//!                            │                 │
106//!                            │                 │
107//!                    ┌───────┴──────┐  ┌───────┴──────┐
108//!                    │              │  │              │
109//!                    │  Event Cache │  │  OlmMachine  │
110//!                    │              │  │              │
111//!                    └──────────────┘  └──────────────┘
112//! ```
113//!
114//! [MSC3061]: https://github.com/matrix-org/matrix-spec/pull/1655#issuecomment-2213152255
115
116use std::{
117    borrow::Cow,
118    collections::{BTreeMap, BTreeSet},
119    pin::{Pin, pin},
120    sync::Weak,
121};
122
123use as_variant::as_variant;
124use futures_core::Stream;
125use futures_util::{StreamExt, future::try_join_all};
126#[cfg(feature = "e2e-encryption")]
127use matrix_sdk_base::deserialized_responses::UnableToDecryptReason;
128#[cfg(doc)]
129use matrix_sdk_base::{BaseClient, crypto::OlmMachine};
130use matrix_sdk_base::{
131    crypto::{
132        store::types::{RoomKeyInfo, RoomKeyWithheldInfo},
133        types::events::room::encrypted::EncryptedEvent,
134    },
135    deserialized_responses::{DecryptedRoomEvent, TimelineEvent, TimelineEventKind},
136    locks::Mutex,
137    task_monitor::BackgroundTaskHandle,
138    timer,
139};
140#[cfg(doc)]
141use matrix_sdk_common::deserialized_responses::EncryptionInfo;
142use ruma::{
143    OwnedEventId, OwnedRoomId, RoomId,
144    events::{AnySyncTimelineEvent, room::encrypted::OriginalSyncRoomEncryptedEvent},
145    push::Action,
146    serde::Raw,
147};
148use tokio::sync::{
149    broadcast::{self, Sender},
150    mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel},
151};
152use tokio_stream::wrappers::{
153    BroadcastStream, UnboundedReceiverStream, errors::BroadcastStreamRecvError,
154};
155use tracing::{info, instrument, trace, warn};
156
157#[cfg(doc)]
158use super::RoomEventCache;
159use super::{
160    EventCache, EventCacheError, EventCacheInner, EventsOrigin, RoomEventCacheGenericUpdate,
161    RoomEventCacheUpdate, TimelineVectorDiffs,
162    caches::{
163        EventLocation,
164        event_linked_chunk::EventLinkedChunk,
165        room::{LinkedChunkUpdateSender, RoomEventCacheLinkedChunkUpdate},
166    },
167};
168use crate::{Client, Result, Room, encryption::backups::BackupState, room::PushContext};
169
170type SessionId<'a> = &'a str;
171type OwnedSessionId = String;
172
173type EventIdAndUtd = (OwnedEventId, Raw<AnySyncTimelineEvent>);
174type EventIdAndEvent = (OwnedEventId, DecryptedRoomEvent);
175
176#[derive(Clone)]
177pub(super) struct ResolvedUtd {
178    pub event_id: OwnedEventId,
179    decrypted_event: DecryptedRoomEvent,
180    actions: Option<Vec<Action>>,
181}
182
183#[derive(Clone)]
184pub(super) enum MaybeResolvedEvent {
185    NotYet(ResolvedUtd),
186    Resolved(TimelineEvent),
187}
188
189impl MaybeResolvedEvent {
190    pub fn try_resolve_event(self, mut unresolved_event: TimelineEvent) -> Self {
191        match self {
192            Self::NotYet(resolved_utd) => {
193                // There is a race between the multiple sources of updates. It's
194                // possible that two sources trigger a decryption for the same
195                // event (for example, the room key stream and the event cache
196                // updates). It is then likely that the event has been already
197                // resolved. This race is fine, but we should avoid to replace
198                // an event that has already been resolved as it is a
199                // non-negligible operation.
200                //
201                // Note that a simple check like “event's kind is
202                // `UnableToDecrypt`” is not enough. The event can already be
203                // decrypted but its encryption info can change. So we must
204                // ensure they are also different.
205                if matches!(unresolved_event.kind, TimelineEventKind::UnableToDecrypt { .. })
206                    || unresolved_event.encryption_info()
207                        != Some(&resolved_utd.decrypted_event.encryption_info)
208                {
209                    unresolved_event.kind =
210                        TimelineEventKind::Decrypted(resolved_utd.decrypted_event);
211
212                    if let Some(actions) = resolved_utd.actions {
213                        unresolved_event.set_push_actions(actions);
214                    }
215
216                    // The unresolved event becomes resolved :-].
217                    Self::Resolved(unresolved_event)
218                } else {
219                    Self::NotYet(resolved_utd)
220                }
221            }
222
223            Self::Resolved(event) => Self::Resolved(event),
224        }
225    }
226
227    pub fn as_resolved(&self) -> Option<&TimelineEvent> {
228        if let Self::Resolved(event) = self { Some(event) } else { None }
229    }
230}
231
232/// Internal trait to resolve events on a `&[MaybeResolvedEvent]` with the help
233/// of `Cow` to avoid copying if no event is newly resolved.
234pub(super) trait TryResolveEvents {
235    fn try_resolve_events(
236        &self,
237        event_linked_chunk: &EventLinkedChunk,
238    ) -> Cow<'_, [MaybeResolvedEvent]>;
239}
240
241impl TryResolveEvents for [MaybeResolvedEvent] {
242    fn try_resolve_events(
243        &self,
244        event_linked_chunk: &EventLinkedChunk,
245    ) -> Cow<'_, [MaybeResolvedEvent]> {
246        let mut new_resolved_events = Cow::Borrowed(self);
247
248        for (nth, resolved_event) in self.iter().enumerate() {
249            match resolved_event {
250                MaybeResolvedEvent::NotYet(resolved_utd) => {
251                    // Event has not been resolved. Let's try to locate the
252                    // corresponding event with the provided `EventLinkedChunk`
253                    // and try to resolve it.
254
255                    if let Some((_location, event)) =
256                        event_linked_chunk.find_event(&resolved_utd.event_id)
257                    {
258                        let new_resolved_event = MaybeResolvedEvent::NotYet(resolved_utd.clone())
259                            .try_resolve_event(event);
260
261                        if matches!(new_resolved_event, MaybeResolvedEvent::Resolved(_)) {
262                            // Use `slice::get_unchecked_mut` to avoid a bounds
263                            // check.
264                            //
265                            // SAFETY: `self` and `new_resolved_events` have the
266                            // same size and represent the same data. Thus, the
267                            // index `nth` exists in `new_resolved_events`.
268                            unsafe {
269                                *new_resolved_events.to_mut().get_unchecked_mut(nth) =
270                                    new_resolved_event;
271                            }
272                        }
273                    }
274                }
275
276                MaybeResolvedEvent::Resolved(_event) => {
277                    // Event has already been resolved. Nothing to do.
278                }
279            }
280        }
281
282        new_resolved_events
283    }
284}
285
286/// The information sent across the channel to the long-running task requesting
287/// that the supplied set of sessions be retried.
288#[derive(Debug, Clone)]
289pub struct DecryptionRetryRequest {
290    /// The room ID of the room the events belong to.
291    pub room_id: OwnedRoomId,
292    /// Events that are not decrypted.
293    pub utd_session_ids: BTreeSet<OwnedSessionId>,
294    /// Events that are decrypted but might need to have their
295    /// [`EncryptionInfo`] refreshed.
296    pub refresh_info_session_ids: BTreeSet<OwnedSessionId>,
297}
298
299/// A report coming from the redecryptor.
300#[derive(Debug, Clone)]
301pub enum RedecryptorReport {
302    /// Events which we were able to decrypt.
303    ResolvedUtds {
304        /// The room ID of the room the events belong to.
305        room_id: OwnedRoomId,
306        /// The list of event IDs of the decrypted events.
307        events: BTreeSet<OwnedEventId>,
308    },
309    /// The redecryptor might have missed some room keys so it might not have
310    /// re-decrypted events that are now decryptable.
311    Lagging,
312    /// A room key backup has become available.
313    ///
314    /// This means that components might want to tell R2D2 about events they
315    /// care about to attempt a decryption.
316    BackupAvailable,
317}
318
319pub(super) struct RedecryptorChannels {
320    utd_reporter: Sender<RedecryptorReport>,
321    pub(super) decryption_request_sender: UnboundedSender<DecryptionRetryRequest>,
322    pub(super) decryption_request_receiver:
323        Mutex<Option<UnboundedReceiver<DecryptionRetryRequest>>>,
324}
325
326impl RedecryptorChannels {
327    pub(super) fn new() -> Self {
328        let (utd_reporter, _) = broadcast::channel(100);
329        let (decryption_request_sender, decryption_request_receiver) = unbounded_channel();
330
331        Self {
332            utd_reporter,
333            decryption_request_sender,
334            decryption_request_receiver: Mutex::new(Some(decryption_request_receiver)),
335        }
336    }
337}
338
339/// A function which can be used to filter and map [`TimelineEvent`]s into a
340/// tuple of event ID and raw [`AnySyncTimelineEvent`].
341///
342/// The tuple can be used to attempt to redecrypt events.
343fn filter_timeline_event_to_utd(
344    event: TimelineEvent,
345) -> Option<(OwnedEventId, Raw<AnySyncTimelineEvent>)> {
346    let event_id = event.event_id().map(ToOwned::to_owned);
347
348    // Only pick out events that are UTDs, get just the Raw event as this is
349    // what the OlmMachine needs.
350    let event = as_variant!(event.kind, TimelineEventKind::UnableToDecrypt { event, .. } => event);
351    // Zip the event ID and event together so we don't have to pick out the
352    // event ID again. We need the event ID to replace the event in the cache.
353    event_id.zip(event)
354}
355
356/// A function which can be used to filter an map [`TimelineEvent`]s into a
357/// tuple of event ID and [`DecryptedRoomEvent`].
358///
359/// The tuple can be used to attempt to update the encryption info of the
360/// decrypted event.
361fn filter_timeline_event_to_decrypted(
362    event: TimelineEvent,
363) -> Option<(OwnedEventId, DecryptedRoomEvent)> {
364    let event_id = event.event_id().map(ToOwned::to_owned);
365
366    let event = as_variant!(event.kind, TimelineEventKind::Decrypted(event) => event);
367    // Zip the event ID and event together so we don't have to pick out the
368    // event ID again. We need the event ID to replace the event in the cache.
369    event_id.zip(event)
370}
371
372impl EventCacheInner {
373    /// Retrieve a set of events that we weren't able to decrypt.
374    ///
375    /// # Arguments
376    ///
377    /// - `room_id` - The ID of the room where the events were sent to.
378    /// - `session_id` - The unique ID of the room key that was used to encrypt
379    ///   the event.
380    async fn all_encrypted_events(
381        &self,
382        room_id: &RoomId,
383        session_id: SessionId<'_>,
384    ) -> Result<Vec<EventIdAndUtd>, EventCacheError> {
385        let caches = self.all_caches_for_room(room_id).await?;
386
387        Ok(caches
388            .all_events_of_type(Some("m.room.encrypted"), Some(session_id))
389            .await?
390            .filter_map(filter_timeline_event_to_utd)
391            .collect())
392    }
393
394    /// Retrieve a set of events that we weren't able to decrypt from the memory
395    /// of the event cache.
396    async fn all_in_memory_encrypted_events(&self) -> BTreeMap<OwnedRoomId, Vec<EventIdAndUtd>> {
397        let mut utds = BTreeMap::new();
398
399        for (room_id, caches) in self.by_room.read().await.iter() {
400            let room_utds: Vec<_> = caches
401                .all_in_memory_events()
402                .await
403                .into_iter()
404                .flatten()
405                .filter_map(filter_timeline_event_to_utd)
406                .collect();
407
408            utds.insert(room_id.to_owned(), room_utds);
409        }
410
411        utds
412    }
413
414    async fn all_decrypted_events(
415        &self,
416        room_id: &RoomId,
417        session_id: SessionId<'_>,
418    ) -> Result<Vec<EventIdAndEvent>, EventCacheError> {
419        let caches = self.all_caches_for_room(room_id).await?;
420
421        Ok(caches
422            .all_events_of_type(None, Some(session_id))
423            .await?
424            .filter_map(filter_timeline_event_to_decrypted)
425            .collect())
426    }
427
428    async fn all_in_memory_decrypted_events(&self) -> BTreeMap<OwnedRoomId, Vec<EventIdAndEvent>> {
429        let mut decrypted_events = BTreeMap::new();
430
431        for (room_id, caches) in self.by_room.read().await.iter() {
432            let room_utds: Vec<_> = caches
433                .all_in_memory_events()
434                .await
435                .into_iter()
436                .flatten()
437                .filter_map(filter_timeline_event_to_decrypted)
438                .collect();
439
440            decrypted_events.insert(room_id.to_owned(), room_utds);
441        }
442
443        decrypted_events
444    }
445
446    /// Handle a chunk of events that we were previously unable to decrypt but
447    /// have now successfully decrypted.
448    ///
449    /// This function will replace the existing UTD events in memory and the
450    /// store and send out a [`RoomEventCacheUpdate`] for the newly decrypted
451    /// events.
452    ///
453    /// # Arguments
454    ///
455    /// - `room_id` - The ID of the room where the events were sent to.
456    /// - `events` - A chunk of events that were successfully decrypted.
457    #[instrument(skip_all, fields(room_id))]
458    async fn on_resolved_utds(
459        &self,
460        room_id: &RoomId,
461        resolved_utds: Vec<ResolvedUtd>,
462    ) -> Result<(), EventCacheError> {
463        if resolved_utds.is_empty() {
464            trace!("No events were redecrypted or updated, nothing to replace");
465            return Ok(());
466        }
467
468        timer!("Resolving UTDs");
469
470        let event_ids: BTreeSet<_> =
471            resolved_utds.iter().map(|resolved_utd| resolved_utd.event_id.clone()).collect();
472
473        let all_caches = self.all_caches_for_room(room_id).await?;
474        let mut maybe_resolved_events = Vec::with_capacity(resolved_utds.len());
475
476        // # Room cache, thread caches, and pinned-event cache
477        //
478        // For each resolved UTD, find the corresponding event (either in-store
479        // or in-memory of the room cache), build the resolved event, and
480        // replace it in the store. We use the room cache for that because it
481        // contains all events (there is an exception with the event-focused
482        // cache, see below).
483        //
484        // # Event-focused cache
485        //
486        // Events received by the sync or pagination are not forwarded to the
487        // event-focused cache: it handles its own set of events. All of them
488        // live in-memory, there are not put in the store by any cache. So this
489        // cache will miss all UTD resolutions. To address that, the
490        // event-focused cache resolves UTD on its own, without the general
491        // logic described above.
492        {
493            let room_cache = &all_caches.room;
494            let mut state = room_cache.state().write().await?;
495
496            let mut maybe_resolved_in_memory_events = Vec::new();
497
498            for resolved_utd in resolved_utds {
499                // Try to locate the event (either in-store or in-memory).
500                if let Some((location, event)) = state.find_event(&resolved_utd.event_id).await? {
501                    let maybe_resolved_event =
502                        MaybeResolvedEvent::NotYet(resolved_utd).try_resolve_event(event);
503
504                    // It is an in-memory event, let's keep it apart to replace
505                    // in-memory UTDs.
506                    if matches!(location, EventLocation::Memory(_)) {
507                        maybe_resolved_in_memory_events.push(maybe_resolved_event.clone());
508                    }
509
510                    // Even is known, let's keep it for later, even if
511                    // unresolved.
512                    maybe_resolved_events.push(maybe_resolved_event);
513                } else {
514                    // Event is unknown by the room cache, the thread caches,
515                    // nor the pinned-events cache. However, it might be known
516                    // by an event-focused cache! So let's keep it for later.
517                    maybe_resolved_events.push(MaybeResolvedEvent::NotYet(resolved_utd));
518                }
519            }
520
521            // Replace all resolved events in the store.
522            state
523                .save_events(
524                    maybe_resolved_events
525                        .iter()
526                        .filter_map(|resolved_event| resolved_event.as_resolved())
527                        .cloned(),
528                )
529                .await?;
530
531            // Saving events doesn't notify the linked chunk update observers,
532            // so index the resolved events explicitly.
533            #[cfg(feature = "experimental-search")]
534            self.linked_chunk_update_sender.send_to_search_index(
535                room_id.to_owned(),
536                maybe_resolved_events
537                    .iter()
538                    .filter_map(|resolved_event| resolved_event.as_resolved())
539                    .cloned()
540                    .collect(),
541            );
542
543            // Now, replace the in-memory events.
544            let timeline_event_diffs = state
545                .replace_in_memory_utds(&maybe_resolved_in_memory_events)
546                .await?
547                .unwrap_or_default();
548
549            if !timeline_event_diffs.is_empty() {
550                state.update_sender.send(
551                    RoomEventCacheUpdate::UpdateTimelineEvents(TimelineVectorDiffs {
552                        diffs: timeline_event_diffs,
553                        origin: EventsOrigin::Cache,
554                    }),
555                    Some(RoomEventCacheGenericUpdate { room_id: room_id.to_owned() }),
556                );
557            }
558        }
559
560        // Resolve in-memory UTDs on the thread caches.
561        {
562            // TODO: This ain't great for performance; there shouldn't be that
563            // many thread caches alive at the same time, but they could
564            // accumulate over time. Consider keeping track of which linked
565            // chunk contains which event ID, to avoid doing the linear searches
566            // here.
567
568            // Replaces UTDs in each thread, and maybe update the thread
569            // summary.
570            for (thread_id, thread_cache) in try_join_all(
571                all_caches.threads.read().await.iter().map(|(thread_id, thread_cache)| async {
572                    Result::<_, EventCacheError>::Ok(
573                        // If at least one event has been replaced, return the
574                        // `thread_id` and the `thread_cache` to update the
575                        // thread summary later.
576                        thread_cache
577                            .replace_in_memory_utds(&maybe_resolved_events)
578                            .await?
579                            .then(|| (thread_id.clone(), thread_cache.clone())),
580                    )
581                }),
582            )
583            .await?
584            .into_iter()
585            // Filter out results that are `None`, i.e. a thread where no UTD has been replaced.
586            .flatten()
587            {
588                let new_thread_summary = thread_cache.update_thread_summary().await?;
589
590                all_caches.room.update_thread_summary(&thread_id, new_thread_summary)?;
591            }
592        }
593
594        // Resolve in-memory UTDs on the pinned-events cache.
595        if let Some(pinned_events_cache) = all_caches.pinned_events.get() {
596            pinned_events_cache.replace_in_memory_utds(&maybe_resolved_events).await?;
597        }
598
599        // Resolve in-memory UTDs on the event-focused caches.
600        {
601            // TODO: This ain't great for performance; there shouldn't be that
602            // many event-focused caches alive at the same time, but they could
603            // accumulate over time. Consider keeping track of which linked
604            // chunk contains which event ID, to avoid doing the linear searches
605            // here.
606            try_join_all(all_caches.event_focused.read().await.values().map(
607                |event_focused_cache| {
608                    event_focused_cache.replace_in_memory_utds(&maybe_resolved_events)
609                },
610            ))
611            .await?;
612        }
613
614        let report =
615            RedecryptorReport::ResolvedUtds { room_id: room_id.to_owned(), events: event_ids };
616        let _ = self.redecryption_channels.utd_reporter.send(report);
617
618        Ok(())
619    }
620
621    /// Attempt to decrypt a single event.
622    async fn decrypt_event(
623        &self,
624        room_id: &RoomId,
625        room: Option<&Room>,
626        push_context: Option<&PushContext>,
627        event: &Raw<EncryptedEvent>,
628    ) -> Option<(DecryptedRoomEvent, Option<Vec<Action>>)> {
629        if let Some(room) = room {
630            match room
631                .decrypt_event(
632                    event.cast_ref_unchecked::<OriginalSyncRoomEncryptedEvent>(),
633                    push_context,
634                )
635                .await
636            {
637                Ok(maybe_decrypted) => {
638                    let actions = maybe_decrypted.push_actions().map(|a| a.to_vec());
639
640                    if let TimelineEventKind::Decrypted(decrypted) = maybe_decrypted.kind {
641                        Some((decrypted, actions))
642                    } else {
643                        warn!(
644                            "Failed to redecrypt an event despite receiving a room key or request to redecrypt"
645                        );
646                        None
647                    }
648                }
649                Err(e) => {
650                    warn!(
651                        "Failed to redecrypt an event despite receiving a room key or request to redecrypt {e:?}"
652                    );
653                    None
654                }
655            }
656        } else {
657            let client = self.client().ok()?;
658            let machine = client.olm_machine().await;
659            let machine = machine.as_ref()?;
660
661            match machine.decrypt_room_event(event, room_id, client.decryption_settings()).await {
662                Ok(decrypted) => Some((decrypted, None)),
663                Err(e) => {
664                    warn!(
665                        "Failed to redecrypt an event despite receiving a room key or a request to redecrypt {e:?}"
666                    );
667                    None
668                }
669            }
670        }
671    }
672
673    /// Attempt to redecrypt events after a room key with the given session ID
674    /// has been received.
675    #[instrument(skip_all, fields(room_id, session_id))]
676    async fn retry_decryption(
677        &self,
678        room_id: &RoomId,
679        session_id: SessionId<'_>,
680    ) -> Result<(), EventCacheError> {
681        // Get all the relevant UTDs.
682        let events = self.all_encrypted_events(room_id, session_id).await?;
683        self.retry_decryption_for_events(room_id, events).await
684    }
685
686    /// Attempt to redecrypt events that were persisted in the event cache.
687    #[instrument(skip_all, fields(updates.linked_chunk_id))]
688    async fn retry_decryption_for_event_cache_updates(
689        &self,
690        updates: RoomEventCacheLinkedChunkUpdate,
691    ) -> Result<(), EventCacheError> {
692        let room_id = updates.linked_chunk_id.room_id();
693        let events: Vec<_> = updates
694            .updates
695            .into_iter()
696            .flat_map(|updates| updates.into_items())
697            .filter_map(filter_timeline_event_to_utd)
698            .collect();
699
700        self.retry_decryption_for_events(room_id, events).await
701    }
702
703    async fn retry_decryption_for_in_memory_events(&self) {
704        let utds = self.all_in_memory_encrypted_events().await;
705
706        for (room_id, utds) in utds.into_iter() {
707            if let Err(e) = self.retry_decryption_for_events(&room_id, utds).await {
708                warn!(%room_id, "Failed to redecrypt in-memory events {e:?}");
709            }
710        }
711    }
712
713    /// Attempt to redecrypt a chunk of UTDs.
714    #[instrument(skip_all, fields(room_id, session_id))]
715    async fn retry_decryption_for_events(
716        &self,
717        room_id: &RoomId,
718        events: Vec<EventIdAndUtd>,
719    ) -> Result<(), EventCacheError> {
720        trace!("Retrying to decrypt");
721
722        if events.is_empty() {
723            trace!("No relevant events found.");
724            return Ok(());
725        }
726
727        let room = self.client().ok().and_then(|client| client.get_room(room_id));
728        let push_context =
729            if let Some(room) = &room { room.push_context().await.ok().flatten() } else { None };
730
731        // Let's attempt to decrypt them them.
732        let mut decrypted_events = Vec::with_capacity(events.len());
733
734        for (event_id, event) in events {
735            // If we managed to decrypt the event, and we should have to since
736            // we received the room key for this specific event, then replace
737            // the event.
738            if let Some((decrypted_event, actions)) = self
739                .decrypt_event(
740                    room_id,
741                    room.as_ref(),
742                    push_context.as_ref(),
743                    event.cast_ref_unchecked(),
744                )
745                .await
746            {
747                decrypted_events.push(ResolvedUtd { event_id, decrypted_event, actions });
748            }
749        }
750
751        if !decrypted_events.is_empty() && tracing::level_enabled!(tracing::Level::TRACE) {
752            let event_ids: BTreeSet<_> =
753                decrypted_events.iter().map(|resolved_utd| &resolved_utd.event_id).collect();
754
755            trace!(?event_ids, "Successfully redecrypted events");
756        }
757
758        // Replace the events and notify listeners that UTDs have been replaced
759        // with decrypted events.
760        self.on_resolved_utds(room_id, decrypted_events).await?;
761
762        Ok(())
763    }
764
765    /// Attempt to update the encryption info for the given list of events.
766    async fn update_encryption_info_for_events(
767        &self,
768        room: &Room,
769        events: Vec<EventIdAndEvent>,
770    ) -> Result<(), EventCacheError> {
771        // Let's attempt to update their encryption info.
772        let mut updated_events = Vec::with_capacity(events.len());
773
774        for (event_id, mut event) in events {
775            if let Some(session_id) = event.encryption_info.session_id() {
776                let new_encryption_info =
777                    room.get_encryption_info(session_id, &event.encryption_info.sender).await;
778
779                // Only create a replacement if the encryption info actually
780                // changed.
781                if let Some(new_encryption_info) = new_encryption_info
782                    && event.encryption_info != new_encryption_info
783                {
784                    event.encryption_info = new_encryption_info;
785                    updated_events.push(ResolvedUtd {
786                        event_id,
787                        decrypted_event: event,
788                        actions: None,
789                    });
790                }
791            }
792        }
793
794        if !updated_events.is_empty() && tracing::level_enabled!(tracing::Level::TRACE) {
795            let event_ids: BTreeSet<_> =
796                updated_events.iter().map(|resolved_utd| &resolved_utd.event_id).collect();
797
798            trace!(?event_ids, "Replacing the encryption info of some events");
799        }
800
801        self.on_resolved_utds(room.room_id(), updated_events).await
802    }
803
804    #[instrument(skip_all, fields(room_id, session_id))]
805    async fn update_encryption_info(
806        &self,
807        room_id: &RoomId,
808        session_id: SessionId<'_>,
809    ) -> Result<(), EventCacheError> {
810        trace!("Updating encryption info");
811
812        let Ok(client) = self.client() else {
813            return Ok(());
814        };
815
816        let Some(room) = client.get_room(room_id) else {
817            return Ok(());
818        };
819
820        // Get all the relevant events.
821        let events = self.all_decrypted_events(room_id, session_id).await?;
822
823        if events.is_empty() {
824            trace!("No relevant events found.");
825            return Ok(());
826        }
827
828        // Let's attempt to update their encryption info.
829        self.update_encryption_info_for_events(&room, events).await
830    }
831
832    async fn retry_update_encryption_info_for_in_memory_events(&self) {
833        let decrypted_events = self.all_in_memory_decrypted_events().await;
834
835        for (room_id, events) in decrypted_events.into_iter() {
836            let Some(room) = self.client().ok().and_then(|c| c.get_room(&room_id)) else {
837                continue;
838            };
839
840            if let Err(e) = self.update_encryption_info_for_events(&room, events).await {
841                warn!(
842                    %room_id,
843                    "Failed to replace the encryption info for in-memory events {e:?}"
844                );
845            }
846        }
847    }
848
849    /// Retry to decrypt and update the encryption info of all the events
850    /// contained in the memory part of the event cache.
851    ///
852    /// This list of events will map one-to-one to the events components
853    /// subscribed to the event cache are have received and are keeping cached.
854    ///
855    /// If components subscribed to the event cache are doing additional
856    /// caching, they'll need to listen to [`RedecryptorReport`]s and explicitly
857    /// request redecryption attempts using [`EventCache::request_decryption`].
858    async fn retry_in_memory_events(&self) {
859        self.retry_decryption_for_in_memory_events().await;
860        self.retry_update_encryption_info_for_in_memory_events().await;
861    }
862}
863
864impl EventCache {
865    /// Explicitly request the redecryption of a set of events.
866    ///
867    /// The redecryption logic in the event cache might sometimes miss that a
868    /// room key has become available and that a certain set of events has
869    /// become decryptable.
870    ///
871    /// This might happen because some room keys might arrive in a separate
872    /// process handling push notifications or if a room key arrives but the
873    /// process shuts down before we could have decrypted the events.
874    ///
875    /// For this reason it is useful to tell the event cache explicitly that
876    /// some events should be retried to be redecrypted.
877    ///
878    /// This method allows you to do so. The events that get decrypted, if any,
879    /// will be advertised over the usual event cache subscription mechanism
880    /// which can be accessed using the [`RoomEventCache::subscribe()`] method.
881    ///
882    /// # Examples
883    ///
884    /// ```no_run
885    /// # use matrix_sdk::{Client, event_cache::DecryptionRetryRequest};
886    /// # use url::Url;
887    /// # use ruma::owned_room_id;
888    /// # use std::collections::BTreeSet;
889    /// # async {
890    /// # let homeserver = Url::parse("http://localhost:8080")?;
891    /// # let client = Client::new(homeserver).await?;
892    /// let event_cache = client.event_cache().unwrap();
893    /// let room_id = owned_room_id!("!my_room:localhost");
894    ///
895    /// let request = DecryptionRetryRequest {
896    ///     room_id,
897    ///     utd_session_ids: BTreeSet::from(["session_id".into()]),
898    ///     refresh_info_session_ids: BTreeSet::new(),
899    /// };
900    ///
901    /// event_cache.request_decryption(request);
902    /// # anyhow::Ok(()) };
903    /// ```
904    pub fn request_decryption(&self, request: DecryptionRetryRequest) {
905        let _ =
906            self.inner.redecryption_channels.decryption_request_sender.send(request).inspect_err(
907                |_| warn!("Requesting a decryption while the redecryption task has been shut down"),
908            );
909    }
910
911    /// Subscribe to reports that the redecryptor generates.
912    ///
913    /// The redecryption logic in the event cache might sometimes miss that a
914    /// room key has become available and that a certain set of events has
915    /// become decryptable.
916    ///
917    /// This might happen because some room keys might arrive in a separate
918    /// process handling push notifications or if room keys arrive faster than
919    /// we can handle them.
920    ///
921    /// This stream can be used to get notified about such situations as well as
922    /// a general channel where the event cache reports which events got
923    /// successfully redecrypted.
924    ///
925    /// # Examples
926    ///
927    /// ```no_run
928    /// # use matrix_sdk::{Client, event_cache::RedecryptorReport};
929    /// # use url::Url;
930    /// # use tokio_stream::StreamExt;
931    /// # async {
932    /// # let homeserver = Url::parse("http://localhost:8080")?;
933    /// # let client = Client::new(homeserver).await?;
934    /// let event_cache = client.event_cache().unwrap();
935    ///
936    /// let mut stream = event_cache.subscribe_to_decryption_reports();
937    ///
938    /// while let Some(Ok(report)) = stream.next().await {
939    ///     match report {
940    ///         RedecryptorReport::Lagging => {
941    ///             // The event cache might have missed to redecrypt some events. We should tell
942    ///             // it which events we care about, i.e. which events we're displaying to the
943    ///             // user, and let it redecrypt things with an explicit request.
944    ///         }
945    ///         RedecryptorReport::BackupAvailable => {
946    ///             // A backup has become available. We can, just like in the Lagging case, tell
947    ///             // the event cache to attempt to redecrypt some events.
948    ///             //
949    ///             // This is only necessary with the BackupDownloadStrategy::OnDecryptionFailure
950    ///             // as the decryption attempt in this case will trigger the download of the
951    ///             // room key from the backup.
952    ///         }
953    ///         RedecryptorReport::ResolvedUtds { .. } => {
954    ///             // This may be interesting for statistical reasons or in case we'd like to
955    ///             // fetch and inspect these events in some manner.
956    ///         }
957    ///     }
958    /// }
959    /// # anyhow::Ok(()) };
960    /// ```
961    pub fn subscribe_to_decryption_reports(
962        &self,
963    ) -> impl Stream<Item = Result<RedecryptorReport, BroadcastStreamRecvError>> + use<> {
964        BroadcastStream::new(self.inner.redecryption_channels.utd_reporter.subscribe())
965    }
966}
967
968async fn send_report_and_retry_memory_events(
969    cache: &Weak<EventCacheInner>,
970    report: RedecryptorReport,
971) -> Result<(), ()> {
972    let Some(cache) = cache.upgrade() else {
973        return Err(());
974    };
975
976    cache.retry_in_memory_events().await;
977    let _ = cache.redecryption_channels.utd_reporter.send(report);
978
979    Ok(())
980}
981
982/// Struct holding on to the redecryption task.
983///
984/// This struct implements the bulk of the redecryption task. It listens to the
985/// various streams that should trigger redecryption attempts.
986///
987/// For more info see the [module level docs](self).
988pub(crate) struct Redecryptor {
989    _task: BackgroundTaskHandle,
990}
991
992impl Redecryptor {
993    /// Create a new [`Redecryptor`].
994    ///
995    /// This creates a task that listens to various streams and attempts to
996    /// redecrypt UTDs that can be found inside the [`EventCache`].
997    pub(super) fn new(
998        client: &Client,
999        cache: Weak<EventCacheInner>,
1000        receiver: UnboundedReceiver<DecryptionRetryRequest>,
1001        linked_chunk_update_sender: &LinkedChunkUpdateSender,
1002    ) -> Self {
1003        let linked_chunk_stream = BroadcastStream::new(linked_chunk_update_sender.subscribe());
1004        let backup_state_stream = client.encryption().backups().state_stream();
1005
1006        let task = client
1007            .task_monitor()
1008            .spawn_infinite_task("event_cache::redecryptor", async {
1009                let request_redecryption_stream = UnboundedReceiverStream::new(receiver);
1010
1011                Self::listen_for_room_keys_task(
1012                    cache,
1013                    request_redecryption_stream,
1014                    linked_chunk_stream,
1015                    backup_state_stream,
1016                )
1017                .await;
1018            })
1019            .abort_on_drop();
1020
1021        Self { _task: task }
1022    }
1023
1024    /// (Re)-subscribe to the room key stream from the [`OlmMachine`].
1025    ///
1026    /// This needs to happen any time this stream returns a `None` meaning that
1027    /// the sending part of the stream has been dropped.
1028    async fn subscribe_to_room_key_stream(
1029        cache: &Weak<EventCacheInner>,
1030    ) -> Option<(
1031        impl Stream<Item = Result<Vec<RoomKeyInfo>, BroadcastStreamRecvError>>,
1032        impl Stream<Item = Vec<RoomKeyWithheldInfo>>,
1033    )> {
1034        let event_cache = cache.upgrade()?;
1035        let client = event_cache.client().ok()?;
1036        let machine = client.olm_machine().await;
1037
1038        machine.as_ref().map(|m| {
1039            (m.store().room_keys_received_stream(), m.store().room_keys_withheld_received_stream())
1040        })
1041    }
1042
1043    async fn redecryption_loop(
1044        cache: &Weak<EventCacheInner>,
1045        decryption_request_stream: &mut Pin<&mut impl Stream<Item = DecryptionRetryRequest>>,
1046        events_stream: &mut Pin<
1047            &mut impl Stream<Item = Result<RoomEventCacheLinkedChunkUpdate, BroadcastStreamRecvError>>,
1048        >,
1049        backup_state_stream: &mut Pin<
1050            &mut impl Stream<Item = Result<BackupState, BroadcastStreamRecvError>>,
1051        >,
1052    ) -> bool {
1053        let Some((room_key_stream, withheld_stream)) =
1054            Self::subscribe_to_room_key_stream(cache).await
1055        else {
1056            return false;
1057        };
1058
1059        let mut room_key_stream = pin!(room_key_stream);
1060        let mut withheld_stream = pin!(withheld_stream);
1061
1062        loop {
1063            tokio::select! {
1064                // An explicit request, presumably from the timeline, has been
1065                // received to decrypt events that were encrypted with a certain
1066                // room key.
1067                Some(request) = decryption_request_stream.next() => {
1068                        let Some(cache) = cache.upgrade() else {
1069                            break false;
1070                        };
1071
1072                        trace!(?request, "Received a redecryption request");
1073
1074                        for session_id in request.utd_session_ids {
1075                            let _ = cache
1076                                .retry_decryption(&request.room_id, &session_id)
1077                                .await
1078                                .inspect_err(|e| warn!("Error redecrypting after an explicit request was received {e:?}"));
1079                        }
1080
1081                        for session_id in request.refresh_info_session_ids {
1082                            let _ = cache.update_encryption_info(&request.room_id, &session_id).await.inspect_err(|e|
1083                                warn!(
1084                                    room_id = %request.room_id,
1085                                    session_id = session_id,
1086                                    "Unable to update the encryption info {e:?}",
1087                            ));
1088                        }
1089                }
1090                // The room key stream from the OlmMachine. Needs to be
1091                // recreated every time we receive a `None` from the stream.
1092                room_keys = room_key_stream.next() => {
1093                    match room_keys {
1094                        Some(Ok(room_keys)) => {
1095                            // Alright, some room keys were received and
1096                            // persisted in our store, let's attempt to
1097                            // redecrypt events that were encrypted using these
1098                            // room keys.
1099                            let Some(cache) = cache.upgrade() else {
1100                                break false;
1101                            };
1102
1103                            trace!(?room_keys, "Received new room keys");
1104
1105                            for key in &room_keys {
1106                                let _ = cache
1107                                    .retry_decryption(&key.room_id, &key.session_id)
1108                                    .await
1109                                    .inspect_err(|e| warn!("Error redecrypting {e:?}"));
1110                            }
1111
1112                            for key in room_keys {
1113                                let _ = cache.update_encryption_info(&key.room_id, &key.session_id).await.inspect_err(|e|
1114                                    warn!(
1115                                        room_id = %key.room_id,
1116                                        session_id = key.session_id,
1117                                        "Unable to update the encryption info {e:?}",
1118                                ));
1119                            }
1120                        },
1121                        Some(Err(_)) => {
1122                            // We missed some room keys, we need to report this
1123                            // in case a listener has and idea which UTDs we
1124                            // should attempt to redecrypt.
1125                            //
1126                            // This would most likely be the timeline from the
1127                            // UI crate. The timeline might attempt to redecrypt
1128                            // all UTDs it is showing to the user.
1129                            warn!("The room key stream lagged, reporting the lag to our listeners");
1130
1131                            if send_report_and_retry_memory_events(cache, RedecryptorReport::Lagging).await.is_err() {
1132                                break false;
1133                            }
1134                        },
1135                        // The stream got closed, this could mean that our
1136                        // OlmMachine got regenerated, let's return true and try
1137                        // to recreate the stream.
1138                        None => {
1139                            break true;
1140                        }
1141                    }
1142                }
1143                withheld_info = withheld_stream.next() => {
1144                    match withheld_info {
1145                        Some(infos) => {
1146                            let Some(cache) = cache.upgrade() else {
1147                                break false;
1148                            };
1149
1150                            trace!(?infos, "Received new withheld infos");
1151
1152                            for RoomKeyWithheldInfo { room_id, session_id, .. } in &infos {
1153                                let _ = cache.update_encryption_info(room_id, session_id).await.inspect_err(|e|
1154                                    warn!(
1155                                        room_id = %room_id,
1156                                        session_id = session_id,
1157                                        "Unable to update the encryption info {e:?}",
1158                                ));
1159                            }
1160                        }
1161                        // The stream got closed, same as for the room key
1162                        // stream, we'll try to recreate the streams.
1163                        None => break true,
1164                    }
1165                }
1166                // Events that the event cache handled. If the event cache
1167                // received any UTDs, let's attempt to redecrypt them in case
1168                // the room key was received before the event cache was able to
1169                // return them using `get_utds()`.
1170                Some(event_updates) = events_stream.next() => {
1171                    match event_updates {
1172                        Ok(updates) => {
1173                            let Some(cache) = cache.upgrade() else {
1174                                break false;
1175                            };
1176
1177                            let linked_chunk_id = updates.linked_chunk_id.to_owned();
1178
1179                            let _ = cache.retry_decryption_for_event_cache_updates(updates).await.inspect_err(|e|
1180                                warn!(
1181                                    %linked_chunk_id,
1182                                    "Unable to handle UTDs from event cache updates {e:?}",
1183                                )
1184                            );
1185                        }
1186                        Err(_) => {
1187                            if send_report_and_retry_memory_events(cache, RedecryptorReport::Lagging).await.is_err() {
1188                                break false;
1189                            }
1190                        }
1191                    }
1192                }
1193                Some(backup_state_update) = backup_state_stream.next() => {
1194                    match backup_state_update {
1195                        Ok(state) => {
1196                            match state {
1197                                BackupState::Unknown |
1198                                BackupState::Creating |
1199                                BackupState::Enabling |
1200                                BackupState::Resuming |
1201                                BackupState::Downloading |
1202                                BackupState::Disabling => {
1203                                    // Those states aren't particularly
1204                                    // interesting to components listening to
1205                                    // R2D2 reports.
1206                                }
1207                                BackupState::Enabled => {
1208                                    // Alright, the backup got enabled, we might
1209                                    // or might not have downloaded the room
1210                                    // keys from the backup. In case they get
1211                                    // downloaded on-demand, let's try to
1212                                    // decrypt all the events we have cached
1213                                    // in-memory.
1214                                    if send_report_and_retry_memory_events(cache, RedecryptorReport::BackupAvailable).await.is_err() {
1215                                        break false;
1216                                    }
1217                                }
1218                            }
1219                        }
1220                        Err(_) => {
1221                            if send_report_and_retry_memory_events(cache, RedecryptorReport::Lagging).await.is_err() {
1222                                break false;
1223                            }
1224                        }
1225                    }
1226                }
1227                else => break false,
1228            }
1229        }
1230    }
1231
1232    async fn listen_for_room_keys_task(
1233        cache: Weak<EventCacheInner>,
1234        decryption_request_stream: UnboundedReceiverStream<DecryptionRetryRequest>,
1235        events_stream: BroadcastStream<RoomEventCacheLinkedChunkUpdate>,
1236        backup_state_stream: impl Stream<Item = Result<BackupState, BroadcastStreamRecvError>>,
1237    ) {
1238        // We pin the decryption request stream here since that one doesn't need
1239        // to be recreated and we don't want to miss messages coming from the
1240        // stream while recreating it unnecessarily.
1241        let mut decryption_request_stream = pin!(decryption_request_stream);
1242        let mut events_stream = pin!(events_stream);
1243        let mut backup_state_stream = pin!(backup_state_stream);
1244
1245        while Self::redecryption_loop(
1246            &cache,
1247            &mut decryption_request_stream,
1248            &mut events_stream,
1249            &mut backup_state_stream,
1250        )
1251        .await
1252        {
1253            info!("Regenerating the re-decryption streams");
1254
1255            // Report that the stream got recreated so listeners know about it,
1256            // at the same time retry to decrypt anything we have cached in
1257            // memory.
1258            if send_report_and_retry_memory_events(&cache, RedecryptorReport::Lagging)
1259                .await
1260                .is_err()
1261            {
1262                break;
1263            }
1264        }
1265
1266        info!("Shutting down the event cache redecryptor");
1267    }
1268}
1269
1270/// Helper method used by `RoomEventCacheState` and `ThreadEventCacheState` to
1271/// attempt to decrypt a bundled thread event we haven't tried to decrypt yet.
1272pub(super) async fn try_decrypt_in_place(event: &mut TimelineEvent, room: Option<&Room>) {
1273    if let TimelineEventKind::UnableToDecrypt { utd_info, .. } = &event.kind
1274        && let Some(room) = room
1275        && utd_info.reason == UnableToDecryptReason::Unknown
1276    {
1277        match room
1278            // Cast safety: a `TimelineEventKind::UnableToDecrypt` always holds the
1279            // `m.room.encrypted` event, and the `Unknown` reason above narrows this
1280            // further to the branch that has positively checked the `type` field in
1281            // `TimelineEvent::from_bundled_latest_event`.
1282            .decrypt_event(event.raw().cast_ref_unchecked::<OriginalSyncRoomEncryptedEvent>(), None)
1283            .await
1284        {
1285            Ok(decrypted) => *event = decrypted,
1286            Err(e) => warn!(
1287                "Failed to decrypt a bundled thread event we haven't tried to decrypt yet {e:?}"
1288            ),
1289        }
1290    }
1291}
1292
1293#[cfg(not(target_family = "wasm"))]
1294#[cfg(test)]
1295mod tests {
1296    use std::{assert_matches, collections::BTreeSet, time::Duration};
1297
1298    use eyeball_im::VectorDiff;
1299    use futures_util::future::join_all;
1300    use matrix_sdk_base::{
1301        crypto::types::events::room::encrypted::EncryptedToDeviceEvent,
1302        deserialized_responses::{TimelineEventKind, UnableToDecryptReason, VerificationState},
1303        event_cache::store::ControlledMemoryStore,
1304        locks::Mutex,
1305        sleep::sleep,
1306        store::StoreConfig,
1307    };
1308    use matrix_sdk_common::{cross_process_lock::CrossProcessLockConfig, executor::spawn};
1309    use matrix_sdk_test::{
1310        JoinedRoomBuilder, SyncResponseBuilder, async_test, event_factory::EventFactory,
1311    };
1312    use ruma::{
1313        EventId, RoomId, RoomVersionId, UserId, device_id, event_id,
1314        events::{
1315            AnySyncTimelineEvent, MessageLikeEventContent,
1316            relation::Thread,
1317            room::message::{
1318                AddMentions, ReplyMetadata, ReplyWithinThread, RoomMessageEventContent,
1319            },
1320        },
1321        room_id,
1322        serde::Raw,
1323        user_id,
1324    };
1325    use serde_json::json;
1326    use strass::assert_let;
1327    use tracing::{Instrument, info};
1328    use wiremock::{Request, Respond, ResponseTemplate};
1329
1330    use crate::{
1331        Client, assert_let_timeout,
1332        encryption::EncryptionSettings,
1333        event_cache::{
1334            DecryptionRetryRequest, RoomEventCacheGenericUpdate, RoomEventCacheUpdate,
1335            TimelineVectorDiffs,
1336        },
1337        test_utils::mocks::MatrixMockServer,
1338    };
1339
1340    async fn set_up_clients(
1341        room_id: &RoomId,
1342        alice_enables_cross_signing: bool,
1343        use_delayed_store: bool,
1344    ) -> (Client, Client, MatrixMockServer, Option<ControlledMemoryStore>) {
1345        let alice_span = tracing::info_span!("alice");
1346        let bob_span = tracing::info_span!("bob");
1347
1348        let alice_user_id = user_id!("@alice:localhost");
1349        let alice_device_id = device_id!("ALICEDEVICE");
1350        let bob_user_id = user_id!("@bob:localhost");
1351        let bob_device_id = device_id!("BOBDEVICE");
1352
1353        let matrix_mock_server = MatrixMockServer::new().await;
1354        matrix_mock_server.mock_crypto_endpoints_preset().await;
1355
1356        let encryption_settings = EncryptionSettings {
1357            auto_enable_cross_signing: alice_enables_cross_signing,
1358            ..Default::default()
1359        };
1360
1361        // Create some clients for Alice and Bob.
1362
1363        let alice = matrix_mock_server
1364            .client_builder_for_crypto_end_to_end(alice_user_id, alice_device_id)
1365            .on_builder(|builder| {
1366                builder
1367                    .with_enable_share_history_on_invite(true)
1368                    .with_encryption_settings(encryption_settings)
1369            })
1370            .build()
1371            .instrument(alice_span.clone())
1372            .await;
1373
1374        let encryption_settings =
1375            EncryptionSettings { auto_enable_cross_signing: true, ..Default::default() };
1376
1377        let (store_config, store) = if use_delayed_store {
1378            let store = ControlledMemoryStore::new();
1379            store.block_writes();
1380
1381            (
1382                StoreConfig::new(CrossProcessLockConfig::multi_process(
1383                    "delayed_store_event_cache_test",
1384                ))
1385                .event_cache_store(store.clone()),
1386                Some(store),
1387            )
1388        } else {
1389            (
1390                StoreConfig::new(CrossProcessLockConfig::multi_process(
1391                    "normal_store_event_cache_test",
1392                )),
1393                None,
1394            )
1395        };
1396
1397        let bob = matrix_mock_server
1398            .client_builder_for_crypto_end_to_end(bob_user_id, bob_device_id)
1399            .on_builder(|builder| {
1400                builder
1401                    .with_enable_share_history_on_invite(true)
1402                    .with_encryption_settings(encryption_settings)
1403                    .store_config(store_config)
1404            })
1405            .build()
1406            .instrument(bob_span.clone())
1407            .await;
1408
1409        // Ensure that Alice and Bob are aware of their devices and identities.
1410        matrix_mock_server.exchange_e2ee_identities(&alice, &bob).await;
1411
1412        let event_factory = EventFactory::new().room(room_id).sender(alice_user_id);
1413
1414        // Let us now create a room for them.
1415        let room_builder = JoinedRoomBuilder::new(room_id)
1416            .add_state_event(event_factory.create(alice_user_id, RoomVersionId::V1))
1417            .add_state_event(event_factory.room_encryption());
1418
1419        matrix_mock_server
1420            .mock_sync()
1421            .ok_and_run(&alice, |builder| {
1422                builder.add_joined_room(room_builder.clone());
1423            })
1424            .instrument(alice_span)
1425            .await;
1426
1427        matrix_mock_server
1428            .mock_sync()
1429            .ok_and_run(&bob, |builder| {
1430                builder.add_joined_room(room_builder);
1431            })
1432            .instrument(bob_span)
1433            .await;
1434
1435        // Mock the members endpoint, so that Alice can share room keys with
1436        // Bob.
1437        let member_factory = EventFactory::new().room(room_id);
1438        matrix_mock_server
1439            .mock_get_members()
1440            .ok(vec![
1441                member_factory.member(alice_user_id).into_raw(),
1442                member_factory.member(bob_user_id).into_raw(),
1443            ])
1444            .mount()
1445            .await;
1446
1447        (alice, bob, matrix_mock_server, store)
1448    }
1449
1450    #[async_test]
1451    async fn test_redecryptor() {
1452        let room_id = room_id!("!test:localhost");
1453
1454        let (alice, bob, matrix_mock_server, _) = set_up_clients(room_id, true, false).await;
1455
1456        let (event, room_key) = send_encrypted(
1457            &matrix_mock_server,
1458            &alice,
1459            room_id,
1460            event_id!("$some_id"),
1461            RoomMessageEventContent::text_plain("It's a secret to everybody"),
1462        )
1463        .await;
1464
1465        // Let's now see what Bob's event cache does.
1466
1467        let event_cache = bob.event_cache().unwrap();
1468        let (room_cache, _) = event_cache
1469            .room(room_id)
1470            .await
1471            .expect("We should be able to get to the event cache for a specific room");
1472
1473        let (_, mut subscriber) = room_cache.subscribe().await.unwrap();
1474        let mut generic_stream = event_cache.subscribe_to_room_generic_updates();
1475
1476        // We regenerate the Olm machine to check if the room key stream is
1477        // recreated to correctly.
1478        bob.inner
1479            .base_client
1480            .regenerate_olm(None)
1481            .await
1482            .expect("We should be able to regenerate the Olm machine");
1483
1484        // Let us forward the event to Bob.
1485        sync_room(
1486            &matrix_mock_server,
1487            &bob,
1488            JoinedRoomBuilder::new(room_id).add_timeline_event(event),
1489        )
1490        .await;
1491
1492        // Alright, Bob has received an update from the cache.
1493
1494        assert_let_timeout!(
1495            Ok(RoomEventCacheUpdate::UpdateTimelineEvents(TimelineVectorDiffs { diffs, .. })) =
1496                subscriber.recv()
1497        );
1498
1499        // There should be a single new event, and it should be a UTD as we did
1500        // not receive the room key yet.
1501        assert_eq!(diffs.len(), 1);
1502        assert_let!(VectorDiff::Append { values } = &diffs[0]);
1503        assert_eq!(values.len(), 1);
1504        assert_matches!(&values[0].kind, TimelineEventKind::UnableToDecrypt { .. });
1505
1506        assert_let_timeout!(
1507            Ok(RoomEventCacheGenericUpdate { room_id: expected_room_id }) = generic_stream.recv()
1508        );
1509        assert_eq!(expected_room_id, room_id);
1510        assert!(generic_stream.is_empty());
1511
1512        // Now we send the room key to Bob.
1513        sync_room_keys(&matrix_mock_server, &bob, &[room_key]).await;
1514
1515        // Bob should receive a new update from the cache.
1516        assert_let_timeout!(
1517            Duration::from_secs(1),
1518            Ok(RoomEventCacheUpdate::UpdateTimelineEvents(TimelineVectorDiffs { diffs, .. })) =
1519                subscriber.recv()
1520        );
1521
1522        // It should replace the UTD with a decrypted event.
1523        assert_eq!(diffs.len(), 1);
1524        assert_let!(VectorDiff::Set { index, value } = &diffs[0]);
1525        assert_eq!(*index, 0);
1526        assert_matches!(&value.kind, TimelineEventKind::Decrypted { .. });
1527
1528        assert_let_timeout!(
1529            Ok(RoomEventCacheGenericUpdate { room_id: expected_room_id }) = generic_stream.recv()
1530        );
1531        assert_eq!(expected_room_id, room_id);
1532        assert!(generic_stream.is_empty());
1533    }
1534
1535    #[cfg(feature = "experimental-search")]
1536    #[async_test]
1537    async fn test_redecrypted_event_is_indexed() {
1538        let room_id = room_id!("!test:localhost");
1539        let event_id = event_id!("$some_id");
1540
1541        let (alice, bob, matrix_mock_server, _) = set_up_clients(room_id, true, false).await;
1542
1543        let (event, room_key) = send_encrypted(
1544            &matrix_mock_server,
1545            &alice,
1546            room_id,
1547            event_id,
1548            RoomMessageEventContent::text_plain("It's a secret to everybody"),
1549        )
1550        .await;
1551
1552        let event_cache = bob.event_cache().unwrap();
1553        let (room_cache, _) = event_cache.room(room_id).await.unwrap();
1554        let (_, mut subscriber) = room_cache.subscribe().await.unwrap();
1555
1556        // Bob receives the event as a UTD, which can't be indexed.
1557        sync_room(
1558            &matrix_mock_server,
1559            &bob,
1560            JoinedRoomBuilder::new(room_id).add_timeline_event(event),
1561        )
1562        .await;
1563        assert_let_timeout!(Ok(RoomEventCacheUpdate::UpdateTimelineEvents(_)) = subscriber.recv());
1564
1565        let room = bob.get_room(room_id).unwrap();
1566        assert!(room.search("secret", 5, None).await.unwrap().events.is_empty());
1567
1568        // Once the room key arrives, the event is redecrypted…
1569        sync_room_keys(&matrix_mock_server, &bob, &[room_key]).await;
1570        assert_let_timeout!(
1571            Duration::from_secs(1),
1572            Ok(RoomEventCacheUpdate::UpdateTimelineEvents(_)) = subscriber.recv()
1573        );
1574
1575        // … and indexed.
1576        let mut results = Vec::new();
1577        for _ in 0..50 {
1578            results = room.search("secret", 5, None).await.unwrap().events;
1579            if !results.is_empty() {
1580                break;
1581            }
1582            sleep(Duration::from_millis(20)).await;
1583        }
1584        assert_eq!(results.len(), 1, "the redecrypted event should be indexed");
1585        assert_eq!(results[0].1, event_id);
1586    }
1587
1588    #[async_test]
1589    async fn test_redecryptor_updating_encryption_info() {
1590        let bob_span = tracing::info_span!("bob");
1591
1592        let room_id = room_id!("!test:localhost");
1593
1594        let (alice, bob, matrix_mock_server, _) = set_up_clients(room_id, false, false).await;
1595
1596        let (event, room_key) = send_encrypted(
1597            &matrix_mock_server,
1598            &alice,
1599            room_id,
1600            event_id!("$some_id"),
1601            RoomMessageEventContent::text_plain("It's a secret to everybody"),
1602        )
1603        .await;
1604
1605        // Let's now see what Bob's event cache does.
1606
1607        let event_cache = bob.event_cache().unwrap();
1608        let (room_cache, _) = event_cache
1609            .room(room_id)
1610            .instrument(bob_span.clone())
1611            .await
1612            .expect("We should be able to get to the event cache for a specific room");
1613
1614        let (_, mut subscriber) = room_cache.subscribe().await.unwrap();
1615        let mut generic_stream = event_cache.subscribe_to_room_generic_updates();
1616
1617        // Let us forward the event to Bob.
1618        sync_room(
1619            &matrix_mock_server,
1620            &bob,
1621            JoinedRoomBuilder::new(room_id).add_timeline_event(event),
1622        )
1623        .instrument(bob_span.clone())
1624        .await;
1625
1626        // Alright, Bob has received an update from the cache.
1627
1628        assert_let_timeout!(
1629            Ok(RoomEventCacheUpdate::UpdateTimelineEvents(TimelineVectorDiffs { diffs, .. })) =
1630                subscriber.recv()
1631        );
1632
1633        // There should be a single new event, and it should be a UTD as we did
1634        // not receive the room key yet.
1635        assert_eq!(diffs.len(), 1);
1636        assert_let!(VectorDiff::Append { values } = &diffs[0]);
1637        assert_eq!(values.len(), 1);
1638        assert_matches!(&values[0].kind, TimelineEventKind::UnableToDecrypt { .. });
1639
1640        assert_let_timeout!(
1641            Ok(RoomEventCacheGenericUpdate { room_id: expected_room_id }) = generic_stream.recv()
1642        );
1643        assert_eq!(expected_room_id, room_id);
1644        assert!(generic_stream.is_empty());
1645
1646        // Now we send the room key to Bob.
1647        sync_room_keys(&matrix_mock_server, &bob, &[room_key]).instrument(bob_span.clone()).await;
1648
1649        // Bob should receive a new update from the cache.
1650        assert_let_timeout!(
1651            Duration::from_secs(1),
1652            Ok(RoomEventCacheUpdate::UpdateTimelineEvents(TimelineVectorDiffs { diffs, .. })) =
1653                subscriber.recv()
1654        );
1655
1656        // It should replace the UTD with a decrypted event.
1657        assert_eq!(diffs.len(), 1);
1658        assert_let!(VectorDiff::Set { index: 0, value } = &diffs[0]);
1659        assert_matches!(&value.kind, TimelineEventKind::Decrypted { .. });
1660
1661        let encryption_info = value.encryption_info().unwrap();
1662        assert_matches!(&encryption_info.verification_state, VerificationState::Unverified(_));
1663
1664        assert_let_timeout!(
1665            Ok(RoomEventCacheGenericUpdate { room_id: expected_room_id }) = generic_stream.recv()
1666        );
1667        assert_eq!(expected_room_id, room_id);
1668        assert!(generic_stream.is_empty());
1669
1670        let session_id = encryption_info.session_id().unwrap().to_owned();
1671        let alice_user_id = alice.user_id().unwrap();
1672
1673        // Alice now creates the identity.
1674        alice
1675            .encryption()
1676            .bootstrap_cross_signing(None)
1677            .await
1678            .expect("Alice should be able to create the cross-signing keys");
1679
1680        bob.update_tracked_users_for_testing([alice_user_id]).instrument(bob_span.clone()).await;
1681        matrix_mock_server
1682            .mock_sync()
1683            .ok_and_run(&bob, |builder| {
1684                builder.add_change_device(alice_user_id);
1685            })
1686            .instrument(bob_span.clone())
1687            .await;
1688
1689        bob.event_cache().unwrap().request_decryption(DecryptionRetryRequest {
1690            room_id: room_id.into(),
1691            utd_session_ids: BTreeSet::new(),
1692            refresh_info_session_ids: BTreeSet::from([session_id]),
1693        });
1694
1695        // Bob should again receive a new update from the cache, this time
1696        // updating the encryption info.
1697        assert_let_timeout!(
1698            Duration::from_secs(1),
1699            Ok(RoomEventCacheUpdate::UpdateTimelineEvents(TimelineVectorDiffs { diffs, .. })) =
1700                subscriber.recv()
1701        );
1702
1703        assert_eq!(diffs.len(), 1);
1704        assert_let!(VectorDiff::Set { index: 0, value } = &diffs[0]);
1705        assert_matches!(&value.kind, TimelineEventKind::Decrypted { .. });
1706        let encryption_info = value.encryption_info().unwrap();
1707
1708        assert_matches!(
1709            &encryption_info.verification_state,
1710            VerificationState::Unverified(_),
1711            "The event should now know about the identity but still be unverified"
1712        );
1713
1714        assert_let_timeout!(
1715            Ok(RoomEventCacheGenericUpdate { room_id: expected_room_id }) = generic_stream.recv()
1716        );
1717        assert_eq!(expected_room_id, room_id);
1718        assert!(generic_stream.is_empty());
1719    }
1720
1721    #[async_test]
1722    async fn test_event_is_redecrypted_even_if_key_arrives_while_event_processing() {
1723        let room_id = room_id!("!test:localhost");
1724
1725        let (alice, bob, matrix_mock_server, delayed_store) =
1726            set_up_clients(room_id, true, true).await;
1727
1728        let delayed_store = delayed_store.unwrap();
1729
1730        let (event, room_key) = send_encrypted(
1731            &matrix_mock_server,
1732            &alice,
1733            room_id,
1734            event_id!("$some_id"),
1735            RoomMessageEventContent::text_plain("It's a secret to everybody"),
1736        )
1737        .await;
1738
1739        let event_cache = bob.event_cache().unwrap();
1740
1741        // Let's now see what Bob's event cache does.
1742        let (room_cache, _) = event_cache
1743            .room(room_id)
1744            .await
1745            .expect("We should be able to get to the event cache for a specific room");
1746
1747        let (_, mut subscriber) = room_cache.subscribe().await.unwrap();
1748        let mut generic_stream = event_cache.subscribe_to_room_generic_updates();
1749
1750        // Custom type to mock the `/sync` endpoints.
1751        struct Responder<'a> {
1752            sync_response_builder: Mutex<(usize, SyncResponseBuilder)>,
1753            room_id: &'a RoomId,
1754            event: Raw<AnySyncTimelineEvent>,
1755            room_key: Raw<EncryptedToDeviceEvent>,
1756        }
1757
1758        impl<'a> Responder<'a> {
1759            fn new(
1760                room_id: &'a RoomId,
1761                event: Raw<AnySyncTimelineEvent>,
1762                room_key: Raw<EncryptedToDeviceEvent>,
1763            ) -> Self {
1764                Self {
1765                    sync_response_builder: Mutex::new((0, SyncResponseBuilder::new())),
1766                    room_id,
1767                    event,
1768                    room_key,
1769                }
1770            }
1771        }
1772
1773        impl<'a> Respond for Responder<'a> {
1774            fn respond(&self, _request: &Request) -> ResponseTemplate {
1775                let (nth_sync, sync_response_builder) = &mut *self.sync_response_builder.lock();
1776
1777                match nth_sync {
1778                    // First request for `sync_room`.
1779                    0 => {
1780                        *nth_sync = 1;
1781                        ResponseTemplate::new(200).set_body_json(
1782                            sync_response_builder
1783                                .add_joined_room(
1784                                    JoinedRoomBuilder::new(self.room_id)
1785                                        .add_timeline_event(self.event.clone()),
1786                                )
1787                                .build_json_sync_response(),
1788                        )
1789                    }
1790
1791                    // Second request for `sync_room_keys`.
1792                    1 => {
1793                        *nth_sync = 2;
1794                        ResponseTemplate::new(200).set_body_json(
1795                            sync_response_builder
1796                                .add_to_device_event(self.room_key.deserialize_as().unwrap())
1797                                .build_json_sync_response(),
1798                        )
1799                    }
1800
1801                    _ => panic!("`/sync` is called once too much"),
1802                }
1803            }
1804        }
1805
1806        let mock = matrix_mock_server
1807            .mock_sync()
1808            .respond_with(Responder::new(room_id, event, room_key))
1809            .mount_as_scoped()
1810            .await;
1811
1812        let sync_room = spawn({
1813            let bob = bob.clone();
1814
1815            // Let us forward the event to Bob, immediately. It will stop
1816            // before persisting everything in the store.
1817            async move {
1818                let _ = bob.sync_once(Default::default()).await.unwrap();
1819            }
1820        });
1821
1822        let sync_room_keys = spawn({
1823            let bob = bob.clone();
1824
1825            // Now we send the room key to Bob, let's say after a delay of
1826            // 1sec, so that the Event Cache had time to start handling the
1827            // event and being blocked when persisting it.
1828            async move {
1829                // Wait on `sync_room` to hang.
1830                sleep(Duration::from_secs(1)).await;
1831
1832                let _ = bob.sync_once(Default::default()).await.unwrap();
1833
1834                // Allow the redecryptor some time to retry any UTDs.
1835                sleep(Duration::from_millis(100)).await;
1836            }
1837        });
1838
1839        let delay = spawn({
1840            // Unblock `sync_room` after 2 secs, so that `sync_room_keys` has
1841            // finished before `sync_room` finishes.
1842            async move {
1843                sleep(Duration::from_secs(2)).await;
1844                info!("Stopping the delay");
1845                delayed_store.unblock_writes();
1846                delayed_store.wait_for_write_to_finish().await;
1847            }
1848        });
1849
1850        assert!(
1851            join_all([sync_room, sync_room_keys, delay]).await.into_iter().all(|res| res.is_ok()),
1852            "all tasks should succeed"
1853        );
1854
1855        drop(mock);
1856
1857        // The first decryption attempt has failed because the first sync (the
1858        // one with the event) did not contain the room key. The decryptor has
1859        // later received the room key.
1860
1861        // Alright, Bob has received an update from the cache.
1862        assert_let_timeout!(
1863            Ok(RoomEventCacheUpdate::UpdateTimelineEvents(TimelineVectorDiffs { diffs, .. })) =
1864                subscriber.recv()
1865        );
1866
1867        // There should be a single new event, and it should be a UTD as we did
1868        // not receive the room key yet.
1869        assert_eq!(diffs.len(), 1);
1870        assert_let!(VectorDiff::Append { values } = &diffs[0]);
1871        assert_eq!(values.len(), 1);
1872        assert_matches!(&values[0].kind, TimelineEventKind::UnableToDecrypt { .. });
1873
1874        // And the companion generic update.
1875        assert_let_timeout!(
1876            Ok(RoomEventCacheGenericUpdate { room_id: expected_room_id }) = generic_stream.recv()
1877        );
1878        assert_eq!(expected_room_id, room_id);
1879
1880        // Bob should receive a new update from the cache.
1881        assert_let_timeout!(
1882            Duration::from_secs(1),
1883            Ok(RoomEventCacheUpdate::UpdateTimelineEvents(TimelineVectorDiffs { diffs, .. })) =
1884                subscriber.recv()
1885        );
1886
1887        // It should replace the UTD with a decrypted event.
1888        assert_eq!(diffs.len(), 1);
1889        assert_let!(VectorDiff::Set { index, value } = &diffs[0]);
1890        assert_eq!(*index, 0);
1891        assert_matches!(&value.kind, TimelineEventKind::Decrypted { .. });
1892
1893        // And the companion generic update.
1894        assert_let_timeout!(
1895            Ok(RoomEventCacheGenericUpdate { room_id: expected_room_id }) = generic_stream.recv()
1896        );
1897        assert_eq!(expected_room_id, room_id);
1898        assert!(generic_stream.is_empty());
1899    }
1900
1901    /// Helper method to send an encrypted message to a room, and capture both
1902    /// the sent event and the to-device message sharing its room key.
1903    ///
1904    /// The message must share a new room key: it has to be the first message in
1905    /// the room, or the first one after the room key was discarded.
1906    async fn send_encrypted(
1907        matrix_mock_server: &MatrixMockServer,
1908        client: &Client,
1909        room_id: &RoomId,
1910        event_id: &EventId,
1911        content: impl MessageLikeEventContent,
1912    ) -> (Raw<AnySyncTimelineEvent>, Raw<EncryptedToDeviceEvent>) {
1913        let user_id = client.user_id().unwrap();
1914        let room = client.get_room(room_id).unwrap();
1915
1916        let (event_receiver, mock) =
1917            matrix_mock_server.mock_room_send().ok_with_capture(event_id, user_id);
1918        // Keep the guards bound until the end of the method, or the mocks are
1919        // dropped immediately.
1920        let _send_guard = mock.mock_once().mount_as_scoped().await;
1921        let (_to_device_guard, room_key) =
1922            matrix_mock_server.mock_capture_put_to_device(user_id).await;
1923
1924        room.send(content).await.unwrap();
1925
1926        (event_receiver.await.unwrap(), room_key.await)
1927    }
1928
1929    /// Bundle `child` as the latest thread event of `root`.
1930    fn with_bundled_thread(
1931        root: &Raw<AnySyncTimelineEvent>,
1932        child: &Raw<AnySyncTimelineEvent>,
1933    ) -> Raw<AnySyncTimelineEvent> {
1934        let mut root: serde_json::Value = root.deserialize_as().unwrap();
1935        root["unsigned"] = json!({
1936            "m.relations": {
1937                "m.thread": {
1938                    "latest_event": child,
1939                    "count": 1,
1940                    "current_user_participated": false,
1941                },
1942            },
1943        });
1944        Raw::new(&root).unwrap().cast_unchecked()
1945    }
1946
1947    /// Helper method to sync room keys to a client, in a single sync response.
1948    async fn sync_room_keys(
1949        matrix_mock_server: &MatrixMockServer,
1950        client: &Client,
1951        room_keys: &[Raw<EncryptedToDeviceEvent>],
1952    ) {
1953        matrix_mock_server
1954            .mock_sync()
1955            .ok_and_run(client, |builder| {
1956                for room_key in room_keys {
1957                    builder.add_to_device_event(room_key.deserialize_as().unwrap());
1958                }
1959            })
1960            .await;
1961        // Allow the redecryptor some time to retry any UTDs.
1962        sleep(Duration::from_millis(100)).await;
1963    }
1964
1965    /// Helper method to sync a joined room update to a client.
1966    async fn sync_room(
1967        matrix_mock_server: &MatrixMockServer,
1968        client: &Client,
1969        room: JoinedRoomBuilder,
1970    ) {
1971        matrix_mock_server
1972            .mock_sync()
1973            .ok_and_run(client, |builder| {
1974                builder.add_joined_room(room);
1975            })
1976            .await;
1977        // Allow the redecryptor some time to retry any UTDs.
1978        sleep(Duration::from_millis(200)).await;
1979    }
1980
1981    /// Get the cached timeline event kind for a given event ID in a room, if it
1982    /// exists. Returns `None` if the event is not in the cache.
1983    async fn get_cached_timeline_event_kind(
1984        client: &Client,
1985        room_id: &RoomId,
1986        event_id: &EventId,
1987    ) -> Option<TimelineEventKind> {
1988        client
1989            .get_room(room_id)
1990            .unwrap()
1991            .load_or_fetch_event(event_id, None)
1992            .await
1993            .map(|event| event.kind)
1994            // An error here means the event is not in the cache, so we return `None`.
1995            .ok()
1996    }
1997
1998    /// Helper method to create a thread reply event contents.
1999    fn make_thread_reply_event_content(
2000        root_id: &EventId,
2001        sender: &UserId,
2002    ) -> RoomMessageEventContent {
2003        RoomMessageEventContent::text_plain("In the thread").make_for_thread(
2004            ReplyMetadata::new(
2005                root_id,
2006                sender,
2007                Some(&Thread::plain(root_id.to_owned(), root_id.to_owned())),
2008            ),
2009            ReplyWithinThread::No,
2010            AddMentions::No,
2011        )
2012    }
2013
2014    /// Case 1 (normal): the key arrives first, then the root and the reply as
2015    /// separate events.
2016    #[async_test]
2017    async fn test_thread_reply_in_timeline_after_key() {
2018        let room_id = room_id!("!test:localhost");
2019        let root_id = event_id!("$root");
2020        let reply_id = event_id!("$reply");
2021
2022        // Alice and Bob are in an encrypted room ...
2023        let (alice, bob, server, _) = set_up_clients(room_id, true, false).await;
2024        let alice_id = alice.user_id().unwrap();
2025
2026        // ... in which Alice sends a message.
2027        let (root, root_key) = send_encrypted(
2028            &server,
2029            &alice,
2030            room_id,
2031            root_id,
2032            RoomMessageEventContent::text_plain("Root"),
2033        )
2034        .await;
2035
2036        // Alice discards the room key ...
2037        alice.get_room(room_id).unwrap().discard_room_key().await.unwrap();
2038
2039        // ... then sends a reply in a new thread using a newly generated room
2040        // key.
2041        let (reply, reply_key) = send_encrypted(
2042            &server,
2043            &alice,
2044            room_id,
2045            reply_id,
2046            make_thread_reply_event_content(root_id, alice_id),
2047        )
2048        .await;
2049
2050        // Bob receives both room keys first, then the root and reply events.
2051        sync_room_keys(&server, &bob, &[root_key, reply_key]).await;
2052        sync_room(&server, &bob, JoinedRoomBuilder::new(room_id).add_timeline_bulk([root, reply]))
2053            .await;
2054
2055        // Bob should now have both events in his cache, and they should be
2056        // decrypted.
2057        assert_matches!(
2058            get_cached_timeline_event_kind(&bob, room_id, root_id).await,
2059            Some(TimelineEventKind::Decrypted { .. })
2060        );
2061        assert_matches!(
2062            get_cached_timeline_event_kind(&bob, room_id, reply_id).await,
2063            Some(TimelineEventKind::Decrypted { .. })
2064        );
2065    }
2066
2067    /// Case 2: The key arrives first, then the root and bundled reply arrive in
2068    /// a single sync.
2069    #[async_test]
2070    async fn test_bundled_thread_reply_after_key_decryptable_root() {
2071        let room_id = room_id!("!test:localhost");
2072        let root_id = event_id!("$root");
2073        let reply_id = event_id!("$reply");
2074
2075        // Alice and Bob are in an encrypted room ...
2076        let (alice, bob, server, _) = set_up_clients(room_id, true, false).await;
2077        let alice_id = alice.user_id().unwrap();
2078
2079        // ... in which Alice sends a message.
2080        let (root, root_key) = send_encrypted(
2081            &server,
2082            &alice,
2083            room_id,
2084            root_id,
2085            RoomMessageEventContent::text_plain("Root"),
2086        )
2087        .await;
2088
2089        // Alice discards the room key ...
2090        alice.get_room(room_id).unwrap().discard_room_key().await.unwrap();
2091
2092        // ... then sends a reply in a new thread using a newly generated room
2093        // key.
2094        let (reply, reply_key) = send_encrypted(
2095            &server,
2096            &alice,
2097            room_id,
2098            reply_id,
2099            make_thread_reply_event_content(root_id, alice_id),
2100        )
2101        .await;
2102
2103        // Bob receives both room keys first, then the reply bundled as the
2104        // latest thread event of the root.
2105        sync_room_keys(&server, &bob, &[root_key, reply_key]).await;
2106        sync_room(
2107            &server,
2108            &bob,
2109            // Servers only bundle aggregations into limited timelines, per https://spec.matrix.org/latest/client-server-api/#aggregations-of-child-events
2110            JoinedRoomBuilder::new(room_id)
2111                .add_timeline_event(with_bundled_thread(&root, &reply))
2112                .set_timeline_limited()
2113                .set_timeline_prev_batch("prev_batch"),
2114        )
2115        .await;
2116
2117        // Bob should have both events in his cache, and they should be
2118        // decrypted.
2119        assert_matches!(
2120            get_cached_timeline_event_kind(&bob, room_id, root_id).await,
2121            Some(TimelineEventKind::Decrypted { .. })
2122        );
2123        assert_matches!(
2124            get_cached_timeline_event_kind(&bob, room_id, reply_id).await,
2125            Some(TimelineEventKind::Decrypted { .. })
2126        );
2127    }
2128
2129    /// Case 3: The root key arrives first, then the root and bundled reply
2130    /// arrive in a single sync, then the reply key arrives.
2131    #[async_test]
2132    async fn test_bundled_thread_reply_key_arrives_later() {
2133        let room_id = room_id!("!test:localhost");
2134        let root_id = event_id!("$root");
2135        let reply_id = event_id!("$reply");
2136
2137        // Alice and Bob are in an encrypted room ...
2138        let (alice, bob, server, _) = set_up_clients(room_id, true, false).await;
2139        let alice_id = alice.user_id().unwrap();
2140
2141        // ... in which Alice sends a message.
2142        let (root, root_key) = send_encrypted(
2143            &server,
2144            &alice,
2145            room_id,
2146            root_id,
2147            RoomMessageEventContent::text_plain("Root"),
2148        )
2149        .await;
2150
2151        // Alice discards the room key ...
2152        alice.get_room(room_id).unwrap().discard_room_key().await.unwrap();
2153
2154        // ... then sends a reply in a new thread using a newly generated room
2155        // key.
2156        let (reply, reply_key) = send_encrypted(
2157            &server,
2158            &alice,
2159            room_id,
2160            reply_id,
2161            make_thread_reply_event_content(root_id, alice_id),
2162        )
2163        .await;
2164
2165        // Bob receives the root key first, then the root and reply events
2166        // bundled together, but not the reply key.
2167        sync_room_keys(&server, &bob, &[root_key]).await;
2168        sync_room(
2169            &server,
2170            &bob,
2171            // Servers only bundle aggregations into limited timelines, per https://spec.matrix.org/latest/client-server-api/#aggregations-of-child-events
2172            JoinedRoomBuilder::new(room_id)
2173                .add_timeline_event(with_bundled_thread(&root, &reply))
2174                .set_timeline_limited()
2175                .set_timeline_prev_batch("prev_batch"),
2176        )
2177        .await;
2178
2179        // Bob should have both events in his cache, but should only be able to
2180        // decrypt the root.
2181        assert_matches!(
2182            get_cached_timeline_event_kind(&bob, room_id, root_id).await,
2183            Some(TimelineEventKind::Decrypted { .. })
2184        );
2185        // Bob should also know why the decryption failed, rather than a generic
2186        // `UnableToDecryptReason::Unknown`.
2187        assert_let!(
2188            Some(TimelineEventKind::UnableToDecrypt { utd_info, .. }) =
2189                get_cached_timeline_event_kind(&bob, room_id, reply_id).await
2190        );
2191        assert_matches!(
2192            utd_info.reason,
2193            UnableToDecryptReason::MissingMegolmSession { withheld_code: None }
2194        );
2195
2196        // Now Bob receives the reply key ...
2197        sync_room_keys(&server, &bob, &[reply_key]).await;
2198
2199        // ... and should therefore be able to decrypt the reply as well.
2200        assert_matches!(
2201            get_cached_timeline_event_kind(&bob, room_id, reply_id).await,
2202            Some(TimelineEventKind::Decrypted { .. })
2203        );
2204    }
2205
2206    /// Case 4: The reply key arrives first, then the root and bundled reply
2207    /// arrive in a single sync, then the root key arrives.
2208    #[async_test]
2209    async fn test_bundled_thread_reply_root_key_arrives_later() {
2210        let room_id = room_id!("!test:localhost");
2211        let root_id = event_id!("$root");
2212        let reply_id = event_id!("$reply");
2213
2214        // Alice and Bob are in an encrypted room ...
2215        let (alice, bob, server, _) = set_up_clients(room_id, true, false).await;
2216        let alice_id = alice.user_id().unwrap();
2217
2218        // ... in which Alice sends a message.
2219        let (root, root_key) = send_encrypted(
2220            &server,
2221            &alice,
2222            room_id,
2223            root_id,
2224            RoomMessageEventContent::text_plain("Root"),
2225        )
2226        .await;
2227
2228        // Alice discards the room key ...
2229        alice.get_room(room_id).unwrap().discard_room_key().await.unwrap();
2230
2231        // ... then sends a reply in a new thread using a newly generated room
2232        // key.
2233        let (reply, reply_key) = send_encrypted(
2234            &server,
2235            &alice,
2236            room_id,
2237            reply_id,
2238            make_thread_reply_event_content(root_id, alice_id),
2239        )
2240        .await;
2241
2242        // Bob receives the reply key first, then the root and reply events
2243        // bundled together, but not the root key.
2244        sync_room_keys(&server, &bob, &[reply_key]).await;
2245        sync_room(
2246            &server,
2247            &bob,
2248            // Servers only bundle aggregations into limited timelines, per https://spec.matrix.org/latest/client-server-api/#aggregations-of-child-events
2249            JoinedRoomBuilder::new(room_id)
2250                .add_timeline_event(with_bundled_thread(&root, &reply))
2251                .set_timeline_limited()
2252                .set_timeline_prev_batch("prev_batch"),
2253        )
2254        .await;
2255
2256        // Bob should have both events in his cache. The root stays a UTD, since
2257        // its room key is still missing, but the reply must be
2258        // decrypted: it was bundled inside a UTD, so nothing decrypted
2259        // it on the way in, but Bob does have its room key, and the
2260        // redecryptor is asked to retry bundled UTDs as soon as they're
2261        // saved.
2262        assert_matches!(
2263            get_cached_timeline_event_kind(&bob, room_id, root_id).await,
2264            Some(TimelineEventKind::UnableToDecrypt { .. })
2265        );
2266        assert_matches!(
2267            get_cached_timeline_event_kind(&bob, room_id, reply_id).await,
2268            Some(TimelineEventKind::Decrypted { .. })
2269        );
2270
2271        // Now Bob receives the root key ...
2272        sync_room_keys(&server, &bob, &[root_key]).await;
2273
2274        // ... and should therefore be able to decrypt both the root and the
2275        // reply.
2276        assert_matches!(
2277            get_cached_timeline_event_kind(&bob, room_id, root_id).await,
2278            Some(TimelineEventKind::Decrypted { .. })
2279        );
2280        assert_matches!(
2281            get_cached_timeline_event_kind(&bob, room_id, reply_id).await,
2282            Some(TimelineEventKind::Decrypted { .. })
2283        );
2284    }
2285}