matrix_sdk/encryption/

tasks.rs

// Copyright 2023-2024 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{collections::BTreeMap, sync::Arc, time::Duration};

use futures_core::Stream;
use futures_util::{StreamExt, pin_mut};
use matrix_sdk_base::crypto::store::types::{RoomKeyBundleInfo, RoomPendingKeyBundleDetails};
#[cfg(feature = "experimental-encrypted-state-events")]
use matrix_sdk_base::crypto::types::events::room::encrypted::{
    EncryptedEvent, RoomEventEncryptionScheme,
};
use matrix_sdk_common::failures_cache::FailuresCache;
#[cfg(not(feature = "experimental-encrypted-state-events"))]
use ruma::events::room::encrypted::{EncryptedEventScheme, OriginalSyncRoomEncryptedEvent};
#[cfg(feature = "experimental-encrypted-state-events")]
use ruma::serde::JsonCastable;
use ruma::{OwnedEventId, OwnedRoomId, serde::Raw};
use tokio::sync::{Mutex, mpsc};
use tracing::{debug, info, instrument, trace, warn};

use crate::{
    Client, Room,
    client::WeakClient,
    encryption::backups::UploadState,
    executor::{JoinHandle, spawn},
    room::shared_room_history,
};

/// A cache of room keys we already downloaded.
type DownloadCache = FailuresCache<RoomKeyInfo>;

#[derive(Default)]
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<()>>,
}

pub(crate) struct BackupUploadingTask {
    sender: mpsc::UnboundedSender<()>,
    #[allow(dead_code)]
    join_handle: JoinHandle<()>,
}

impl Drop for BackupUploadingTask {
    fn drop(&mut self) {
        #[cfg(not(target_family = "wasm"))]
        self.join_handle.abort();
    }
}

impl BackupUploadingTask {
    pub(crate) fn new(client: WeakClient) -> Self {
        let (sender, receiver) = mpsc::unbounded_channel();

        let join_handle = spawn(async move {
            Self::listen(client, receiver).await;
        });

        Self { sender, join_handle }
    }

    pub(crate) fn trigger_upload(&self) {
        let _ = self.sender.send(());
    }

    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;

                if let Err(e) = client.encryption().backups().backup_room_keys().await {
                    upload_progress.set(UploadState::Error);
                    warn!("Error backing up room keys {e:?}");
                    // Note: it's expected we're not `continue`ing here, because
                    // *every* single state update
                    // is propagated to the caller.
                }

                upload_progress.set(UploadState::Idle);
            } else {
                trace!("Client got dropped, shutting down the task");
                break;
            }
        }
    }
}

/// Information about a request for a backup download for an undecryptable
/// event.
#[derive(Debug)]
struct RoomKeyDownloadRequest {
    /// The room in which the event was sent.
    room_id: OwnedRoomId,

    /// The ID of the event we could not decrypt.
    event_id: OwnedEventId,

    /// The event we could not decrypt.
    #[cfg(not(feature = "experimental-encrypted-state-events"))]
    event: Raw<OriginalSyncRoomEncryptedEvent>,

    /// The event we could not decrypt.
    #[cfg(feature = "experimental-encrypted-state-events")]
    event: Raw<EncryptedEvent>,

    /// The unique ID of the room key that the event was encrypted with.
    megolm_session_id: String,
}

impl RoomKeyDownloadRequest {
    pub fn to_room_key_info(&self) -> RoomKeyInfo {
        (self.room_id.clone(), self.megolm_session_id.clone())
    }
}

pub type RoomKeyInfo = (OwnedRoomId, String);

pub(crate) struct BackupDownloadTask {
    sender: mpsc::UnboundedSender<RoomKeyDownloadRequest>,
    #[allow(dead_code)]
    join_handle: JoinHandle<()>,
}

impl Drop for BackupDownloadTask {
    fn drop(&mut self) {
        #[cfg(not(target_family = "wasm"))]
        self.join_handle.abort();
    }
}

impl BackupDownloadTask {
    #[cfg(not(test))]
    const DOWNLOAD_DELAY_MILLIS: u64 = 100;

    pub(crate) fn new(client: WeakClient) -> Self {
        let (sender, receiver) = mpsc::unbounded_channel();

        let join_handle = spawn(async move {
            Self::listen(client, receiver).await;
        });

        Self { sender, join_handle }
    }

    /// Trigger a backup download for the keys for the given event.
    ///
    /// Does nothing unless the event is encrypted using `m.megolm.v1.aes-sha2`.
    /// Otherwise, tells the listener task to set off a task to do a backup
    /// download, unless there is one already running.
    #[cfg(not(feature = "experimental-encrypted-state-events"))]
    pub(crate) fn trigger_download_for_utd_event(
        &self,
        room_id: OwnedRoomId,
        event: Raw<OriginalSyncRoomEncryptedEvent>,
    ) {
        if let Ok(deserialized_event) = event.deserialize()
            && let EncryptedEventScheme::MegolmV1AesSha2(c) = deserialized_event.content.scheme
        {
            let _ = self.sender.send(RoomKeyDownloadRequest {
                room_id,
                event_id: deserialized_event.event_id,
                event,
                megolm_session_id: c.session_id,
            });
        }
    }

    /// Trigger a backup download for the keys for the given event.
    ///
    /// Does nothing unless the event is encrypted using `m.megolm.v1.aes-sha2`.
    /// Otherwise, tells the listener task to set off a task to do a backup
    /// download, unless there is one already running.
    #[cfg(feature = "experimental-encrypted-state-events")]
    pub(crate) fn trigger_download_for_utd_event<T: JsonCastable<EncryptedEvent>>(
        &self,
        room_id: OwnedRoomId,
        event: Raw<T>,
    ) {
        if let Ok(deserialized_event) = event.deserialize_as::<EncryptedEvent>() {
            if let RoomEventEncryptionScheme::MegolmV1AesSha2(c) = deserialized_event.content.scheme
            {
                let _ = self.sender.send(RoomKeyDownloadRequest {
                    room_id,
                    event_id: deserialized_event.event_id,
                    event: event.cast(),
                    megolm_session_id: c.session_id,
                });
            }
        }
    }

    /// Listen for incoming [`RoomKeyDownloadRequest`]s and process them.
    ///
    /// This will keep running until either the request channel is closed, or
    /// all other references to `Client` are dropped.
    ///
    /// # Arguments
    ///
    /// * `receiver` - The source of incoming [`RoomKeyDownloadRequest`]s.
    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 {
            let mut state_guard = state.lock().await;

            if state_guard.client.strong_count() == 0 {
                trace!("Client got dropped, shutting down the task");
                break;
            }

            // Check that we don't already have a task to process this event, and fire one
            // off else if not.
            let event_id = &room_key_download_request.event_id;
            if !state_guard.active_tasks.contains_key(event_id) {
                let event_id = event_id.to_owned();
                let task =
                    spawn(Self::handle_download_request(state.clone(), room_key_download_request));
                state_guard.active_tasks.insert(event_id, task);
            }
        }
    }

    /// Handle a request to download a room key for a given event.
    ///
    /// Sleeps for a while to see if the key turns up; then checks if we still
    /// want to do a download, and does the download if so.
    async fn handle_download_request(
        state: Arc<Mutex<BackupDownloadTaskListenerState>>,
        download_request: RoomKeyDownloadRequest,
    ) {
        // Wait a bit, perhaps the room key will arrive in the meantime.
        #[cfg(not(test))]
        crate::sleep::sleep(Duration::from_millis(Self::DOWNLOAD_DELAY_MILLIS)).await;

        // Now take the lock, and check that we still want to do a download. If we do,
        // keep hold of a strong reference to the `Client`.
        let client = {
            let mut state = state.lock().await;

            let Some(client) = state.client.get() else {
                // The client was dropped while we were sleeping. We should just bail out;
                // the main BackupDownloadTask loop will bail out too.
                return;
            };

            // Check that we still want to do a download.
            if !state.should_download(&client, &download_request).await {
                // We decided against doing a download. Mark the job done for this event before
                // dropping the lock.
                state.active_tasks.remove(&download_request.event_id);
                return;
            }

            // Before we drop the lock, indicate to other tasks that may be considering this
            // room key, that we're going to go ahead and do a download.
            state.downloaded_room_keys.insert(download_request.to_room_key_info());

            client
        };

        // Do the download without holding the lock.
        let result = client
            .encryption()
            .backups()
            .download_room_key(&download_request.room_id, &download_request.megolm_session_id)
            .await;

        // Then take the lock again to update the state.
        {
            let mut state = state.lock().await;
            let room_key_info = download_request.to_room_key_info();

            match result {
                Ok(true) => {
                    // We successfully downloaded the room key. We can clear any record of previous
                    // backoffs from the failures cache, because we won't be needing them again.
                    state.failures_cache.remove(std::iter::once(&room_key_info))
                }
                Ok(false) => {
                    // We did not find a valid backup decryption key or backup version, we did not
                    // even attempt to download the room key.
                    state.downloaded_room_keys.remove(std::iter::once(&room_key_info));
                }
                Err(_) => {
                    // We were unable to download the room key. Update the failure cache so that we
                    // back off from more requests, and also remove the entry from the list of
                    // room keys that we are downloading.
                    state.downloaded_room_keys.remove(std::iter::once(&room_key_info));
                    state.failures_cache.insert(room_key_info);
                }
            }

            state.active_tasks.remove(&download_request.event_id);
        }
    }
}

/// The state for an active [`BackupDownloadTask`].
struct BackupDownloadTaskListenerState {
    /// Reference to the `Client`, which will be used to fire off the download
    /// requests.
    client: WeakClient,

    /// A record of backup download attempts that have recently failed.
    failures_cache: FailuresCache<RoomKeyInfo>,

    /// Map from event ID to download task
    active_tasks: BTreeMap<OwnedEventId, JoinHandle<()>>,

    /// A list of room keys that we have already downloaded, or are about to
    /// download.
    ///
    /// The idea here is that once we've (successfully) downloaded a room key
    /// from the backup, there's not much point trying again even if we get
    /// another UTD event that uses the same room key.
    downloaded_room_keys: DownloadCache,
}

impl BackupDownloadTaskListenerState {
    /// Prepare a new `BackupDownloadTaskListenerState`.
    ///
    /// # Arguments
    ///
    /// * `client` - A reference to the `Client`, which is used to fire off the
    ///   backup download request.
    pub fn new(client: WeakClient) -> Self {
        Self {
            client,
            failures_cache: FailuresCache::with_settings(Duration::from_secs(60 * 60 * 24), 60),
            active_tasks: Default::default(),
            downloaded_room_keys: DownloadCache::with_settings(
                Duration::from_secs(60 * 60 * 24),
                60,
            ),
        }
    }

    /// Check if we should set off a download for the given request.
    ///
    /// Checks if:
    ///  * we already have the key,
    ///  * we have already downloaded this room key, or are about to do so, or
    ///  * we've backed off from trying to download this room key.
    ///
    /// If any of the above are true, returns `false`. Otherwise, returns
    /// `true`.
    pub async fn should_download(
        &self,
        client: &Client,
        download_request: &RoomKeyDownloadRequest,
    ) -> bool {
        // Check that the Client has an OlmMachine
        let machine_guard = client.olm_machine().await;
        let Some(machine) = machine_guard.as_ref() else {
            return false;
        };

        // If backups aren't enabled, there's no point in trying to download a room key.
        if !client.encryption().backups().are_enabled().await {
            debug!(
                ?download_request,
                "Not performing backup download because backups are not enabled"
            );

            return false;
        }

        // Check if the keys for this message have arrived in the meantime.
        // If we get a StoreError doing the lookup, we assume the keys haven't arrived
        // (though if the store is returning errors, probably something else is
        // going to go wrong very soon).
        if machine
            .is_room_key_available(
                #[cfg(not(feature = "experimental-encrypted-state-events"))]
                download_request.event.cast_ref(),
                #[cfg(feature = "experimental-encrypted-state-events")]
                &download_request.event,
                &download_request.room_id,
            )
            .await
            .unwrap_or(false)
        {
            debug!(
                ?download_request,
                "Not performing backup download because key became available while we were sleeping"
            );
            return false;
        }

        // Check if we already downloaded this room key, or another task is in the
        // process of doing so.
        let room_key_info = download_request.to_room_key_info();
        if self.downloaded_room_keys.contains(&room_key_info) {
            debug!(
                ?download_request,
                "Not performing backup download because this room key has already been downloaded recently"
            );
            return false;
        }

        // Check if we're backing off from attempts to download this room key
        if self.failures_cache.contains(&room_key_info) {
            debug!(
                ?download_request,
                "Not performing backup download because this room key failed to download recently"
            );
            return false;
        }

        debug!(?download_request, "Performing backup download");
        true
    }
}

pub(crate) struct BundleReceiverTask {
    _startup_handle: JoinHandle<()>,
    _listen_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);
        Self {
            _listen_handle: spawn(Self::listen_task(weak_client.clone(), stream)),
            _startup_handle: spawn(Self::startup_task(weak_client)),
        }
    }

    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;
        }
    }

    /// Retrieves a list of all rooms pending key bundles, then cross-references
    /// this with bundles held in the crypto store. If all conditions outlined
    /// in [`shared_room_history::maybe_accept_key_bundle`], then the bundle
    /// will be imported.
    #[tracing::instrument(skip_all)]
    async fn startup_task(client: WeakClient) {
        tracing::debug!("Checking for unimported stored room key bundles...");

        let Some(client) = client.get() else {
            // The client was dropped before the worker future was first polled.
            return;
        };

        let olm_machine = client.olm_machine().await;
        let Some(olm_machine) = olm_machine.as_ref() else {
            // The Olm machine was not initialized by the time this task is ready
            // to perform its work. This is likely a bug, as this worker is only
            // expected to be spawned once the client is fully ready and the
            // Olm machine is available.
            tracing::warn!("Skipping startup bundle checks because the Olm machine is unavailable");
            return;
        };

        let room_details = match olm_machine.store().get_all_rooms_pending_key_bundles().await {
            Ok(room_details) => room_details,
            Err(e) => {
                tracing::warn!("Error while fetching rooms pending key bundles: {e:?}");
                return;
            }
        };

        // Partition the room details into two categories: those that should be
        // processed and those that should be removed.
        let (valid, invalid): (Vec<_>, Vec<_>) = room_details.iter().partition(|details| {
            shared_room_history::should_process_room_pending_key_bundle_details(details)
        });

        tracing::debug!(
            "Found {} valid and {} invalid rooms that are still pending key bundles",
            valid.len(),
            invalid.len(),
        );

        // Iterate over the details that are valid for processing. For each valid
        // details, check if we have the corresponding key bundle data in the
        // store. If the data exists, attempt to re-import the bundle.
        for RoomPendingKeyBundleDetails { room_id, inviter, .. } in valid {
            let Some(room) = client.get_room(room_id) else {
                // Skip processing if the room is not cached in the state store.
                tracing::trace!(?room_id, "Room not available in state store, skipping...");
                continue;
            };
            let bundle =
                match olm_machine.store().get_received_room_key_bundle_data(room_id, inviter).await
                {
                    Ok(Some(bundle)) => bundle,
                    Ok(None) => {
                        // If the bundle data is not available, skip processing. The listener task
                        // will handle this case when the bundle arrives.
                        tracing::trace!(?room_id, "No bundle available, skipping...");
                        continue;
                    }
                    Err(err) => {
                        tracing::warn!(
                            ?room_id,
                            "Failed to fetch received room key bundle data: {err:?}"
                        );
                        continue;
                    }
                };
            Self::handle_bundle(&room, &(&bundle).into()).await;
        }

        // For each invalid details, clear the pending key bundle information from the
        // respective room to avoid re-checking it in the future.
        for RoomPendingKeyBundleDetails { room_id, .. } in &invalid {
            tracing::trace!(?room_id, "Clearing pending flag for room");
            if let Err(e) = olm_machine.store().clear_room_pending_key_bundle(room_id).await {
                tracing::warn!("Error clearing room pending key bundle: {e:?}");
            }
        }
    }

    /// We have received a key bundle for a given room: check if we recently
    /// accepted an invite from the sender of the bundle, and if so, join
    /// the room.
    ///
    /// Note that there is a potential race with the room-join logic
    /// ([`Client::finish_join_room`]), which, having joined a room after an
    /// invite, checks if we already have a bundle, and accepts it if so.
    /// However, here we have already recorded the receipt of a bundle, and are
    /// now checking if we have accepted an invite, whereas in
    /// `finish_join_room`, we first record the acceptance of an invite
    /// and then check for a bundle.
    ///
    /// In theory then, it is possible for both threads to decide to process the
    /// bundle at once; however this is (a) unlikely and (b) harmless. On the
    /// other hand, it is **not** possible for a race to mean that we end up
    /// processing the bundle zero times: we can be sure that at least one
    /// thread will process it.
    #[instrument(skip(room), fields(room_id = %room.room_id()))]
    async fn handle_bundle(room: &Room, bundle_info: &RoomKeyBundleInfo) {
        if shared_room_history::should_accept_key_bundle(room, bundle_info).await {
            info!(room_id = %room.room_id(), "Accepting a late 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.");
        }
    }

    #[cfg(any(feature = "testing", test))]
    pub(crate) fn abort(&self) {
        self._startup_handle.abort();
        self._listen_handle.abort();
    }
}

#[cfg(all(test, not(target_family = "wasm")))]
mod test {
    use matrix_sdk_test::async_test;
    #[cfg(not(feature = "experimental-encrypted-state-events"))]
    use ruma::events::room::encrypted::OriginalSyncRoomEncryptedEvent;
    use ruma::{event_id, room_id};
    use serde_json::json;
    use wiremock::MockServer;

    use super::*;
    use crate::test_utils::logged_in_client;

    // Test that, if backups are not enabled, we don't incorrectly mark a room key
    // as downloaded.
    #[async_test]
    async fn test_disabled_backup_does_not_mark_room_key_as_downloaded() {
        let room_id = room_id!("!DovneieKSTkdHKpIXy:morpheus.localhost");
        let event_id = event_id!("$JbFHtZpEJiH8uaajZjPLz0QUZc1xtBR9rPGBOjF6WFM");
        let session_id = "session_id";

        let server = MockServer::start().await;
        let client = logged_in_client(Some(server.uri())).await;
        let weak_client = WeakClient::from_client(&client);

        let event_content = json!({
            "event_id": event_id,
            "origin_server_ts": 1698579035927u64,
            "sender": "@example2:morpheus.localhost",
            "type": "m.room.encrypted",
            "content": {
                "algorithm": "m.megolm.v1.aes-sha2",
                "ciphertext": "AwgAEpABhetEzzZzyYrxtEVUtlJnZtJcURBlQUQJ9irVeklCTs06LwgTMQj61PMUS4Vy\
                               YOX+PD67+hhU40/8olOww+Ud0m2afjMjC3wFX+4fFfSkoWPVHEmRVucfcdSF1RSB4EmK\
                               PIP4eo1X6x8kCIMewBvxl2sI9j4VNvDvAN7M3zkLJfFLOFHbBviI4FN7hSFHFeM739Zg\
                               iwxEs3hIkUXEiAfrobzaMEM/zY7SDrTdyffZndgJo7CZOVhoV6vuaOhmAy4X2t4UnbuV\
                               JGJjKfV57NAhp8W+9oT7ugwO",
                "device_id": "KIUVQQSDTM",
                "sender_key": "LvryVyoCjdONdBCi2vvoSbI34yTOx7YrCFACUEKoXnc",
                "session_id": "64H7XKokIx0ASkYDHZKlT5zd/Zccz/cQspPNdvnNULA"
            }
        });

        #[cfg(not(feature = "experimental-encrypted-state-events"))]
        let event: Raw<OriginalSyncRoomEncryptedEvent> =
            serde_json::from_value(event_content).expect("");

        #[cfg(feature = "experimental-encrypted-state-events")]
        let event: Raw<EncryptedEvent> = serde_json::from_value(event_content).expect("");

        let state = Arc::new(Mutex::new(BackupDownloadTaskListenerState::new(weak_client)));
        let download_request = RoomKeyDownloadRequest {
            room_id: room_id.into(),
            megolm_session_id: session_id.to_owned(),
            event,
            event_id: event_id.into(),
        };

        assert!(
            !client.encryption().backups().are_enabled().await,
            "Backups should not be enabled."
        );

        BackupDownloadTask::handle_download_request(state.clone(), download_request).await;

        {
            let state = state.lock().await;
            assert!(
                !state.downloaded_room_keys.contains(&(room_id.to_owned(), session_id.to_owned())),
                "Backups are not enabled, we should not mark any room keys as downloaded."
            )
        }
    }
}