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