Merge pull request #5545 from matrix-org/kaylendog/msc3414/json-castable

feat(sdk): Support room key downloading using `JsonCastable<EncryptedEvent>`

crates/matrix-sdk-crypto/Cargo.toml

@@ -22,7 +22,8 @@ experimental-send-custom-to-device = []
 # Enable experimental support for encrypting state events; see
 # https://github.com/matrix-org/matrix-rust-sdk/issues/5397.
 experimental-encrypted-state-events = [
-    "matrix-sdk-common/experimental-encrypted-state-events"
+    "matrix-sdk-common/experimental-encrypted-state-events",
+    "ruma/unstable-msc3414"
 ]
 
 js = ["ruma/js", "vodozemac/js", "matrix-sdk-common/js"]

crates/matrix-sdk-crypto/src/types/events/room/encrypted.rs

@@ -70,6 +70,12 @@ impl JsonCastable<EncryptedEvent>
 {
 }
 
+#[cfg(feature = "experimental-encrypted-state-events")]
+impl JsonCastable<EncryptedEvent>
+    for ruma::events::room::encrypted::unstable_state::OriginalSyncStateRoomEncryptedEvent
+{
+}
+
 /// An m.room.encrypted to-device event.
 pub type EncryptedToDeviceEvent = ToDeviceEvent<ToDeviceEncryptedEventContent>;
 
@@ -230,6 +236,9 @@ impl EventType for RoomEncryptedEventContent {
 
 impl JsonCastable<ruma::events::AnyMessageLikeEventContent> for RoomEncryptedEventContent {}
 
+#[cfg(feature = "experimental-encrypted-state-events")]
+impl JsonCastable<ruma::events::AnyStateEventContent> for RoomEncryptedEventContent {}
+
 /// An enum for per encryption algorithm event contents.
 #[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
 #[serde(try_from = "Helper")]

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

@@ -24,12 +24,16 @@ use std::collections::{BTreeMap, BTreeSet};
 
 use futures_core::Stream;
 use futures_util::StreamExt;
+#[cfg(feature = "experimental-encrypted-state-events")]
+use matrix_sdk_base::crypto::types::events::room::encrypted::EncryptedEvent;
 use matrix_sdk_base::crypto::{
     backups::MegolmV1BackupKey,
     store::types::BackupDecryptionKey,
     types::{requests::KeysBackupRequest, RoomKeyBackupInfo},
     OlmMachine, RoomKeyImportResult,
 };
+#[cfg(feature = "experimental-encrypted-state-events")]
+use ruma::serde::JsonCastable;
 use ruma::{
     api::client::{
         backup::{
@@ -1002,6 +1006,7 @@ impl Backups {
 
     /// Send a notification to the task responsible for key backup downloads
     /// that it should attempt to download the keys for the given event.
+    #[cfg(not(feature = "experimental-encrypted-state-events"))]
     pub(crate) fn maybe_download_room_key(
         &self,
         room_id: OwnedRoomId,
@@ -1013,6 +1018,20 @@ impl Backups {
         }
     }
 
+    /// Send a notification to the task responsible for key backup downloads
+    /// that it should attempt to download the keys for the given event.
+    #[cfg(feature = "experimental-encrypted-state-events")]
+    pub(crate) fn maybe_download_room_key<T: JsonCastable<EncryptedEvent>>(
+        &self,
+        room_id: OwnedRoomId,
+        event: Raw<T>,
+    ) {
+        let tasks = self.client.inner.e2ee.tasks.lock();
+        if let Some(task) = tasks.download_room_keys.as_ref() {
+            task.trigger_download_for_utd_event(room_id, event);
+        }
+    }
+
     /// Send a notification to the task which is responsible for uploading room
     /// keys to the backup that it might have new room keys to back up.
     pub(crate) fn maybe_trigger_backup(&self) {

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

@@ -16,15 +16,19 @@ use std::{collections::BTreeMap, sync::Arc, time::Duration};
 
 use futures_core::Stream;
 use futures_util::{pin_mut, StreamExt};
+#[cfg(feature = "experimental-encrypted-state-events")]
+use matrix_sdk_base::crypto::types::events::room::encrypted::{
+    EncryptedEvent, RoomEventEncryptionScheme,
+};
 use matrix_sdk_base::{
     crypto::store::types::RoomKeyBundleInfo, InviteAcceptanceDetails, RoomState,
 };
 use matrix_sdk_common::failures_cache::FailuresCache;
-use ruma::{
-    events::room::encrypted::{EncryptedEventScheme, OriginalSyncRoomEncryptedEvent},
-    serde::Raw,
-    OwnedEventId, OwnedRoomId,
-};
+#[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::{serde::Raw, OwnedEventId, OwnedRoomId};
 use tokio::sync::{mpsc, Mutex};
 use tracing::{debug, info, instrument, trace, warn};
 
@@ -109,8 +113,13 @@ struct RoomKeyDownloadRequest {
     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,
 }
@@ -155,6 +164,7 @@ impl BackupDownloadTask {
     /// 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,
@@ -172,6 +182,30 @@ impl BackupDownloadTask {
         }
     }
 
+    /// 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
@@ -356,7 +390,13 @@ impl BackupDownloadTaskListenerState {
         // (though if the store is returning errors, probably something else is
         // going to go wrong very soon).
         if machine
-            .is_room_key_available(download_request.event.cast_ref(), &download_request.room_id)
+            .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)
         {
@@ -477,6 +517,8 @@ mod test {
     use matrix_sdk_test::{
         async_test, event_factory::EventFactory, InvitedRoomBuilder, JoinedRoomBuilder,
     };
+    #[cfg(not(feature = "experimental-encrypted-state-events"))]
+    use ruma::events::room::encrypted::OriginalSyncRoomEncryptedEvent;
     use ruma::{event_id, room_id, user_id};
     use serde_json::json;
     use vodozemac::Curve25519PublicKey;
@@ -515,9 +557,13 @@ mod test {
             }
         });
 
+        #[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(),