feat(sdk): Add a tasks that listens for historic room keys if they arrive out of order

Historic room key bundles are uploaded as an encrypted file to the media
repo and the key to decrypt the file is sent as a to-device message to
the recipient device.

In the nominal case, the invite and this to-device message should arrive
at the same time and accepting the invite would download and import the
bundle.

If the to-device message arrives after the invite has already been
accepted we would never download and import the bundle.

To mitigate this problem, this patch introduces a task that listens for
bundles that arrive. If the bundle is for a room that we have joined we
will consider importing the bundle.

Author: poljar

crates/matrix-sdk/CHANGELOG.md

@@ -14,6 +14,9 @@ All notable changes to this project will be documented in this file.
 
 ### Features
 
+- Add support to accept historic room key bundles that arrive out of order, i.e.
+  the bundle arrives after the invite has already been accepted.
+  ([#5322](https://github.com/matrix-org/matrix-rust-sdk/pull/5322))
 - `Client::add_event_handler`: Set `Option<EncryptionInfo>` in `EventHandlerData` for to-device messages.
   If the to-device message was encrypted, the `EncryptionInfo` will be set. If it is `None` the message was sent in clear.
   ([#5099](https://github.com/matrix-org/matrix-rust-sdk/pull/5099))

crates/matrix-sdk/src/authentication/matrix/mod.rs

@@ -802,7 +802,7 @@ impl MatrixAuth {
                 _ => None,
             };
 
-            self.client.encryption().spawn_initialization_task(auth_data);
+            self.client.encryption().spawn_initialization_task(auth_data).await;
         }
 
         Ok(())

crates/matrix-sdk/src/authentication/oauth/mod.rs

@@ -818,7 +818,7 @@ impl OAuth {
         }
 
         #[cfg(feature = "e2e-encryption")]
-        self.client.encryption().spawn_initialization_task(None);
+        self.client.encryption().spawn_initialization_task(None).await;
 
         Ok(())
     }
@@ -1019,7 +1019,7 @@ impl OAuth {
             self.enable_cross_process_lock().await.map_err(OAuthError::from)?;
 
             #[cfg(feature = "e2e-encryption")]
-            self.client.encryption().spawn_initialization_task(None);
+            self.client.encryption().spawn_initialization_task(None).await;
         }
 
         Ok(())

crates/matrix-sdk/src/authentication/oauth/qrcode/login.rs

@@ -253,7 +253,7 @@ impl<'a> IntoFuture for LoginWithQrCode<'a> {
             // ourselves see us as verified and the recovery/backup states will
             // be known. If we did receive all the secrets in the secrets
             // bundle, then backups will be enabled after this step as well.
-            self.client.encryption().spawn_initialization_task(None);
+            self.client.encryption().spawn_initialization_task(None).await;
             self.client.encryption().wait_for_e2ee_initialization_tasks().await;
 
             trace!("successfully logged in and enabled E2EE.");

crates/matrix-sdk/src/client/mod.rs

@@ -415,7 +415,7 @@ impl ClientInner {
         let client = Arc::new(client);
 
         #[cfg(feature = "e2e-encryption")]
-        client.e2ee.initialize_room_key_tasks(&client);
+        client.e2ee.initialize_tasks(&client);
 
         let _ = client
             .event_cache

crates/matrix-sdk/src/encryption/mod.rs

@@ -62,6 +62,7 @@ use ruma::{
 #[cfg(feature = "experimental-send-custom-to-device")]
 use ruma::{events::AnyToDeviceEventContent, serde::Raw, to_device::DeviceIdOrAllDevices};
 use serde::Deserialize;
+use tasks::BundleReceiverTask;
 use tokio::sync::{Mutex, RwLockReadGuard};
 use tokio_stream::wrappers::errors::BroadcastStreamRecvError;
 use tracing::{debug, error, instrument, trace, warn};
@@ -134,7 +135,7 @@ impl EncryptionData {
         }
     }
 
-    pub fn initialize_room_key_tasks(&self, client: &Arc<ClientInner>) {
+    pub fn initialize_tasks(&self, client: &Arc<ClientInner>) {
         let weak_client = WeakClient::from_inner(client);
 
         let mut tasks = self.tasks.lock();
@@ -1685,10 +1686,13 @@ impl Encryption {
     ///   there is a proposal (MSC3967) to remove this requirement, which would
     ///   allow for the initial upload of cross-signing keys without
     ///   authentication, rendering this parameter obsolete.
-    pub(crate) fn spawn_initialization_task(&self, auth_data: Option<AuthData>) {
+    pub(crate) async fn spawn_initialization_task(&self, auth_data: Option<AuthData>) {
+        let bundle_receiver_task = BundleReceiverTask::new(&self.client).await;
+
         let mut tasks = self.client.inner.e2ee.tasks.lock();
 
         let this = self.clone();
+
         tasks.setup_e2ee = Some(spawn(async move {
             // Update the current state first, so we don't have to wait for the result of
             // network requests
@@ -1707,6 +1711,11 @@ impl Encryption {
                 error!("Couldn't setup and resume recovery {e:?}");
             }
         }));
+
+        // It's fine to be async here as we're only getting the lock protecting the
+        // `OlmMachine`. Since the lock shouldn't be that contested right after logging
+        // in we won't delay the login or restoration of the Client.
+        tasks.receive_historic_room_key_bundles = Some(bundle_receiver_task);
     }
 
     /// Waits for end-to-end encryption initialization tasks to finish, if any

crates/matrix-sdk/src/encryption/tasks.rs

@@ -14,23 +14,24 @@
 
 use std::{collections::BTreeMap, sync::Arc, time::Duration};
 
+use futures_core::Stream;
+use futures_util::{pin_mut, StreamExt};
+use matrix_sdk_base::{crypto::store::types::RoomKeyBundleInfo, RoomState};
 use matrix_sdk_common::failures_cache::FailuresCache;
 use ruma::{
     events::room::encrypted::{EncryptedEventScheme, OriginalSyncRoomEncryptedEvent},
     serde::Raw,
     OwnedEventId, OwnedRoomId,
 };
-use tokio::sync::{
-    mpsc::{self, UnboundedReceiver},
-    Mutex,
-};
-use tracing::{debug, trace, warn};
+use tokio::sync::{mpsc, Mutex};
+use tracing::{debug, info, instrument, trace, warn};
 
 use crate::{
     client::WeakClient,
     encryption::backups::UploadState,
     executor::{spawn, JoinHandle},
-    Client,
+    room::shared_room_history,
+    Client, Room,
 };
 
 /// A cache of room keys we already downloaded.
@@ -41,6 +42,7 @@ pub(crate) struct ClientTasks {
     pub(crate) upload_room_keys: Option<BackupUploadingTask>,
     pub(crate) download_room_keys: Option<BackupDownloadTask>,
     pub(crate) update_recovery_state_after_backup: Option<JoinHandle<()>>,
+    pub(crate) receive_historic_room_key_bundles: Option<BundleReceiverTask>,
     pub(crate) setup_e2ee: Option<JoinHandle<()>>,
 }
 
@@ -72,7 +74,7 @@ impl BackupUploadingTask {
         let _ = self.sender.send(());
     }
 
-    pub(crate) async fn listen(client: WeakClient, mut receiver: UnboundedReceiver<()>) {
+    pub(crate) async fn listen(client: WeakClient, mut receiver: mpsc::UnboundedReceiver<()>) {
         while receiver.recv().await.is_some() {
             if let Some(client) = client.get() {
                 let upload_progress = &client.inner.e2ee.backup_state.upload_progress;
@@ -176,7 +178,10 @@ impl BackupDownloadTask {
     /// # Arguments
     ///
     /// * `receiver` - The source of incoming [`RoomKeyDownloadRequest`]s.
-    async fn listen(client: WeakClient, mut receiver: UnboundedReceiver<RoomKeyDownloadRequest>) {
+    async fn listen(
+        client: WeakClient,
+        mut receiver: mpsc::UnboundedReceiver<RoomKeyDownloadRequest>,
+    ) {
         let state = Arc::new(Mutex::new(BackupDownloadTaskListenerState::new(client)));
 
         while let Some(room_key_download_request) = receiver.recv().await {
@@ -385,6 +390,68 @@ impl BackupDownloadTaskListenerState {
     }
 }
 
+pub(crate) struct BundleReceiverTask {
+    _handle: JoinHandle<()>,
+}
+
+impl BundleReceiverTask {
+    pub async fn new(client: &Client) -> Self {
+        let stream = client.encryption().historic_room_key_stream().await.expect("E2EE tasks should only be initialized once we have logged in and have access to an OlmMachine");
+        let weak_client = WeakClient::from_client(client);
+        let handle = spawn(Self::listen_task(weak_client, stream));
+
+        Self { _handle: handle }
+    }
+
+    async fn listen_task(client: WeakClient, stream: impl Stream<Item = RoomKeyBundleInfo>) {
+        pin_mut!(stream);
+
+        // TODO: Listening to this stream is not enough for iOS due to the NSE killing
+        // our OlmMachine and thus also this stream. We need to add an event handler
+        // that will listen for the bundle event. To be able to add an event handler,
+        // we'll have to implement the bundle event in Ruma.
+        while let Some(bundle_info) = stream.next().await {
+            let Some(client) = client.get() else {
+                // The client was dropped while we were waiting on the stream. Let's end the
+                // loop, since this means that the application has shut down.
+                break;
+            };
+
+            let Some(room) = client.get_room(&bundle_info.room_id) else {
+                warn!(room_id = %bundle_info.room_id, "Received a historic room key bundle for an unknown room");
+                continue;
+            };
+
+            Self::handle_bundle(&room, &bundle_info).await;
+        }
+    }
+
+    #[instrument(skip(room), fields(room_id = %room.room_id()))]
+    async fn handle_bundle(room: &Room, bundle_info: &RoomKeyBundleInfo) {
+        if Self::should_accept_bundle(room, bundle_info) {
+            info!("Accepting an out of order key bundle.");
+
+            if let Err(e) =
+                shared_room_history::maybe_accept_key_bundle(room, &bundle_info.sender).await
+            {
+                warn!("Couldn't accept a late room key bundle {e:?}");
+            }
+        } else {
+            info!("Refusing to accept a historic room key bundle.");
+        }
+    }
+
+    fn should_accept_bundle(room: &Room, _bundle_info: &RoomKeyBundleInfo) -> bool {
+        // TODO: Check that the person that invited us to this room is the same as the
+        // sender, of the bundle. Otherwise don't ignore the bundle.
+        // TODO: Check that we joined the room "recently". (How do you do this if you
+        // accept the invite on another client? I guess we remember when the transition
+        // from Invited to Joined happened, but can't the server force us into a joined
+        // state if we do this?
+        room.state() == RoomState::Joined
+    }
+}
+
 #[cfg(all(test, not(target_family = "wasm")))]
 mod test {
     use matrix_sdk_test::async_test;