refactor(crypto): Make the room key importing logic more generic

Author: poljar

crates/matrix-sdk-crypto/src/olm/group_sessions/mod.rs

@@ -35,8 +35,8 @@ use crate::types::events::forwarded_room_key::ForwardedMegolmV2AesSha2Content;
 use crate::types::{
     deserialize_curve_key, deserialize_curve_key_vec,
     events::forwarded_room_key::{ForwardedMegolmV1AesSha2Content, ForwardedRoomKeyContent},
-    serialize_curve_key, serialize_curve_key_vec, EventEncryptionAlgorithm, SigningKey,
-    SigningKeys,
+    serialize_curve_key, serialize_curve_key_vec, EventEncryptionAlgorithm, RoomKeyExport,
+    SigningKey, SigningKeys,
 };
 
 /// An error type for the creation of group sessions.
@@ -141,6 +141,20 @@ impl ExportedRoomKey {
     }
 }
 
+impl RoomKeyExport for &ExportedRoomKey {
+    fn room_id(&self) -> &ruma::RoomId {
+        &self.room_id
+    }
+
+    fn session_id(&self) -> &str {
+        &self.session_id
+    }
+
+    fn sender_key(&self) -> Curve25519PublicKey {
+        self.sender_key
+    }
+}
+
 /// A backed up version of an [`InboundGroupSession`].
 ///
 /// This can be used to back up the [`InboundGroupSession`] to the server using

crates/matrix-sdk-crypto/src/store/mod.rs

@@ -74,7 +74,8 @@ use crate::{
     },
     types::{
         events::room_key_withheld::RoomKeyWithheldEvent, BackupSecrets, CrossSigningSecrets,
-        EventEncryptionAlgorithm, MegolmBackupV1Curve25519AesSha2Secrets, SecretsBundle,
+        EventEncryptionAlgorithm, MegolmBackupV1Curve25519AesSha2Secrets, RoomKeyExport,
+        SecretsBundle,
     },
     verification::VerificationMachine,
     CrossSigningStatus, OwnUserIdentityData, RoomKeyImportResult,
@@ -1832,6 +1833,54 @@ impl Store {
         from_backup_version: Option<&str>,
         progress_listener: impl Fn(usize, usize),
     ) -> Result<RoomKeyImportResult> {
+        let exported_keys: Vec<&ExportedRoomKey> = exported_keys.iter().collect();
+        self.import_sessions_impl(exported_keys, from_backup_version, progress_listener).await
+    }
+
+    /// Import the given room keys into our store.
+    ///
+    /// # Arguments
+    ///
+    /// * `exported_keys` - A list of previously exported keys that should be
+    ///   imported into our store. If we already have a better version of a key
+    ///   the key will *not* be imported.
+    ///
+    /// Returns a tuple of numbers that represent the number of sessions that
+    /// were imported and the total number of sessions that were found in the
+    /// key export.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use std::io::Cursor;
+    /// # use matrix_sdk_crypto::{OlmMachine, decrypt_room_key_export};
+    /// # use ruma::{device_id, user_id};
+    /// # let alice = user_id!("@alice:example.org");
+    /// # async {
+    /// # let machine = OlmMachine::new(&alice, device_id!("DEVICEID")).await;
+    /// # let export = Cursor::new("".to_owned());
+    /// let exported_keys = decrypt_room_key_export(export, "1234").unwrap();
+    /// machine.store().import_exported_room_keys(exported_keys, |_, _| {}).await.unwrap();
+    /// # };
+    /// ```
+    pub async fn import_exported_room_keys(
+        &self,
+        exported_keys: Vec<ExportedRoomKey>,
+        progress_listener: impl Fn(usize, usize),
+    ) -> Result<RoomKeyImportResult> {
+        self.import_room_keys(exported_keys, None, progress_listener).await
+    }
+
+    async fn import_sessions_impl<T>(
+        &self,
+        room_keys: Vec<T>,
+        from_backup_version: Option<&str>,
+        progress_listener: impl Fn(usize, usize),
+    ) -> Result<RoomKeyImportResult>
+    where
+        T: TryInto<InboundGroupSession> + RoomKeyExport + Copy,
+        T::Error: Debug,
+    {
         let mut sessions = Vec::new();
 
         async fn new_session_better(
@@ -1845,11 +1894,11 @@ impl Store {
             }
         }
 
-        let total_count = exported_keys.len();
+        let total_count = room_keys.len();
         let mut keys = BTreeMap::new();
 
-        for (i, key) in exported_keys.into_iter().enumerate() {
-            match InboundGroupSession::from_export(&key) {
+        for (i, key) in room_keys.into_iter().enumerate() {
+            match key.try_into() {
                 Ok(session) => {
                     let old_session = self
                         .inner
@@ -1875,9 +1924,9 @@ impl Store {
                 }
                 Err(e) => {
                     warn!(
-                        sender_key= key.sender_key.to_base64(),
-                        room_id = ?key.room_id,
-                        session_id = key.session_id,
+                        sender_key = key.sender_key().to_base64(),
+                        room_id = ?key.room_id(),
+                        session_id = key.session_id(),
                         error = ?e,
                         "Couldn't import a room key from a file export."
                     );
@@ -1896,40 +1945,6 @@ impl Store {
         Ok(RoomKeyImportResult::new(imported_count, total_count, keys))
     }
 
-    /// Import the given room keys into our store.
-    ///
-    /// # Arguments
-    ///
-    /// * `exported_keys` - A list of previously exported keys that should be
-    ///   imported into our store. If we already have a better version of a key
-    ///   the key will *not* be imported.
-    ///
-    /// Returns a tuple of numbers that represent the number of sessions that
-    /// were imported and the total number of sessions that were found in the
-    /// key export.
-    ///
-    /// # Examples
-    ///
-    /// ```no_run
-    /// # use std::io::Cursor;
-    /// # use matrix_sdk_crypto::{OlmMachine, decrypt_room_key_export};
-    /// # use ruma::{device_id, user_id};
-    /// # let alice = user_id!("@alice:example.org");
-    /// # async {
-    /// # let machine = OlmMachine::new(&alice, device_id!("DEVICEID")).await;
-    /// # let export = Cursor::new("".to_owned());
-    /// let exported_keys = decrypt_room_key_export(export, "1234").unwrap();
-    /// machine.store().import_exported_room_keys(exported_keys, |_, _| {}).await.unwrap();
-    /// # };
-    /// ```
-    pub async fn import_exported_room_keys(
-        &self,
-        exported_keys: Vec<ExportedRoomKey>,
-        progress_listener: impl Fn(usize, usize),
-    ) -> Result<RoomKeyImportResult> {
-        self.import_room_keys(exported_keys, None, progress_listener).await
-    }
-
     pub(crate) fn crypto_store(&self) -> Arc<CryptoStoreWrapper> {
         self.inner.store.clone()
     }

crates/matrix-sdk-crypto/src/types/mod.rs

@@ -38,7 +38,7 @@ use matrix_sdk_common::deserialized_responses::PrivOwnedStr;
 use ruma::{
     events::AnyToDeviceEvent,
     serde::{Raw, StringEnum},
-    DeviceKeyAlgorithm, DeviceKeyId, OwnedDeviceKeyId, OwnedUserId, UserId,
+    DeviceKeyAlgorithm, DeviceKeyId, OwnedDeviceKeyId, OwnedUserId, RoomId, UserId,
 };
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use vodozemac::{Curve25519PublicKey, Ed25519PublicKey, Ed25519Signature, KeyError};
@@ -660,3 +660,15 @@ impl ProcessedToDeviceEvent {
         }
     }
 }
+
+/// Trait to express the various room key export formats we have in a unified
+/// manner.
+pub trait RoomKeyExport {
+    /// The ID of the room where the exported room key was used.
+    fn room_id(&self) -> &RoomId;
+    /// The unique ID of the exported room key.
+    fn session_id(&self) -> &str;
+    /// The [Curve25519PublicKey] long-term identity key of the sender of this
+    /// room key.
+    fn sender_key(&self) -> Curve25519PublicKey;
+}