From efb72063ac45a7b0039b91897dba32db482986d8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Damir=20Jeli=C4=87?= Date: Tue, 3 Oct 2023 13:16:38 +0200 Subject: [PATCH] Add a backup specific method to import room keys --- bindings/matrix-sdk-crypto-ffi/src/machine.rs | 1 + crates/matrix-sdk-crypto/CHANGELOG.md | 5 + crates/matrix-sdk-crypto/src/backups/mod.rs | 99 ++++++++++++++- .../src/file_encryption/key_export.rs | 13 +- crates/matrix-sdk-crypto/src/machine.rs | 79 ++---------- crates/matrix-sdk-crypto/src/store/mod.rs | 117 +++++++++++++++++- .../src/timeline/tests/encryption.rs | 8 +- .../src/timeline/tests/read_receipts.rs | 2 +- crates/matrix-sdk/src/encryption/mod.rs | 2 +- 9 files changed, 244 insertions(+), 82 deletions(-) diff --git a/bindings/matrix-sdk-crypto-ffi/src/machine.rs b/bindings/matrix-sdk-crypto-ffi/src/machine.rs index 685144034..b598ff769 100644 --- a/bindings/matrix-sdk-crypto-ffi/src/machine.rs +++ b/bindings/matrix-sdk-crypto-ffi/src/machine.rs @@ -1464,6 +1464,7 @@ impl OlmMachine { progress_listener.on_progress(progress as i32, total as i32) }; + #[allow(deprecated)] let result = self.runtime.block_on(self.inner.import_room_keys(keys, from_backup, listener))?; diff --git a/crates/matrix-sdk-crypto/CHANGELOG.md b/crates/matrix-sdk-crypto/CHANGELOG.md index 23e6d8487..314d000b3 100644 --- a/crates/matrix-sdk-crypto/CHANGELOG.md +++ b/crates/matrix-sdk-crypto/CHANGELOG.md @@ -1,5 +1,10 @@ # unreleased +- Add two new methods to import room keys, + `OlmMachine::store()::import_exported_room_keys()` for file exports and + `OlmMachine::backup_machine()::import_backed_up_room_keys()` for backups. The + `OlmMachine::import_room_keys()` method is now deprecated. + - Add support for secret storage. - Add initial support for MSC3814 - dehydrated devices. diff --git a/crates/matrix-sdk-crypto/src/backups/mod.rs b/crates/matrix-sdk-crypto/src/backups/mod.rs index 47440c1c6..7920e34bc 100644 --- a/crates/matrix-sdk-crypto/src/backups/mod.rs +++ b/crates/matrix-sdk-crypto/src/backups/mod.rs @@ -36,10 +36,10 @@ use tokio::sync::RwLock; use tracing::{debug, info, instrument, trace, warn}; use crate::{ - olm::{InboundGroupSession, SignedJsonObject}, + olm::{BackedUpRoomKey, ExportedRoomKey, InboundGroupSession, SignedJsonObject}, store::{BackupDecryptionKey, BackupKeys, Changes, RoomKeyCounts, Store}, types::{MegolmV1AuthData, RoomKeyBackupInfo, Signatures}, - CryptoStoreError, Device, KeysBackupRequest, OutgoingRequest, + CryptoStoreError, Device, KeysBackupRequest, OutgoingRequest, RoomKeyImportResult, }; mod keys; @@ -555,15 +555,73 @@ impl BackupMachine { (backup, session_record) } + + /// Import the given room keys into our store. + /// + /// # Arguments + /// + /// * `room_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 [`RoomKeyImportResult`] containing information about room keys + /// which were imported. + pub async fn import_backed_up_room_keys( + &self, + room_keys: BTreeMap>, + progress_listener: impl Fn(usize, usize), + ) -> Result { + let mut decrypted_room_keys = vec![]; + + for (room_id, room_keys) in room_keys { + for (session_id, room_key) in room_keys { + let room_key = ExportedRoomKey { + algorithm: room_key.algorithm, + room_id: room_id.to_owned(), + sender_key: room_key.sender_key, + session_id, + session_key: room_key.session_key, + sender_claimed_keys: room_key.sender_claimed_keys, + forwarding_curve25519_key_chain: room_key.forwarding_curve25519_key_chain, + }; + + decrypted_room_keys.push(room_key); + } + } + + self.store.import_room_keys(decrypted_room_keys, true, progress_listener).await + } } #[cfg(test)] mod tests { + use std::collections::BTreeMap; + + use assert_matches2::assert_let; use matrix_sdk_test::async_test; use ruma::{device_id, room_id, user_id, CanonicalJsonValue, DeviceId, RoomId, UserId}; use serde_json::json; - use crate::{store::BackupDecryptionKey, types::RoomKeyBackupInfo, OlmError, OlmMachine}; + use crate::{ + olm::BackedUpRoomKey, store::BackupDecryptionKey, types::RoomKeyBackupInfo, OlmError, + OlmMachine, + }; + + fn room_key() -> BackedUpRoomKey { + let json = json!({ + "algorithm": "m.megolm.v1.aes-sha2", + "sender_key": "DeHIg4gwhClxzFYcmNntPNF9YtsdZbmMy8+3kzCMXHA", + "session_key": "AQAAAABvWMNZjKFtebYIePKieQguozuoLgzeY6wKcyJjLJcJtQgy1dPqTBD12U+XrYLrRHn\ + lKmxoozlhFqJl456+9hlHCL+yq+6ScFuBHtJepnY1l2bdLb4T0JMDkNsNErkiLiLnD6yp3J\ + DSjIhkdHxmup/huygrmroq6/L5TaThEoqvW4DPIuO14btKudsS34FF82pwjKS4p6Mlch+0e\ + fHAblQV", + "sender_claimed_keys":{}, + "forwarding_curve25519_key_chain":[] + }); + + serde_json::from_value(json) + .expect("We should be able to deserialize our backed up room key") + } fn alice_id() -> &'static UserId { user_id!("@alice:example.org") @@ -717,4 +775,39 @@ mod tests { Ok(()) } + + #[async_test] + async fn import_backed_up_room_keys() { + let machine = OlmMachine::new(alice_id(), alice_device_id()).await; + let backup_machine = machine.backup_machine(); + + let room_id = room_id!("!DovneieKSTkdHKpIXy:morpheus.localhost"); + let session_id = "gM8i47Xhu0q52xLfgUXzanCMpLinoyVyH7R58cBuVBU"; + let room_key = room_key(); + + let room_keys: BTreeMap<_, BTreeMap<_, _>> = BTreeMap::from([( + room_id.to_owned(), + BTreeMap::from([(session_id.to_owned(), room_key)]), + )]); + + let session = + machine.store().get_inbound_group_session(room_id, &session_id).await.unwrap(); + + assert!(session.is_none(), "Initially we should not have the session in the store"); + + backup_machine + .import_backed_up_room_keys(room_keys, |_, _| {}) + .await + .expect("We should be able to import a room key"); + + let session = + machine.store().get_inbound_group_session(room_id, &session_id).await.unwrap(); + + assert_let!(Some(session) = session); + assert!( + session.backed_up(), + "If a session was imported from a backup, it should be considered to be backed up" + ); + assert!(session.has_been_imported()); + } } diff --git a/crates/matrix-sdk-crypto/src/file_encryption/key_export.rs b/crates/matrix-sdk-crypto/src/file_encryption/key_export.rs index af3e0b77c..c5cd9b1aa 100644 --- a/crates/matrix-sdk-crypto/src/file_encryption/key_export.rs +++ b/crates/matrix-sdk-crypto/src/file_encryption/key_export.rs @@ -305,7 +305,7 @@ mod tests { } assert_eq!( - machine.import_room_keys(decrypted, false, |_, _| {}).await.unwrap(), + machine.store().import_exported_room_keys(decrypted, |_, _| {}).await.unwrap(), RoomKeyImportResult::new(0, 1, BTreeMap::new()) ); } @@ -332,17 +332,20 @@ mod tests { )]), ); - assert_eq!(machine.import_room_keys(export, false, |_, _| {}).await?, keys); + assert_eq!(machine.store().import_exported_room_keys(export, |_, _| {}).await?, keys); let export = vec![session.export_at_index(10).await]; assert_eq!( - machine.import_room_keys(export, false, |_, _| {}).await?, + machine.store().import_exported_room_keys(export, |_, _| {}).await?, RoomKeyImportResult::new(0, 1, BTreeMap::new()) ); let better_export = vec![session.export().await]; - assert_eq!(machine.import_room_keys(better_export, false, |_, _| {}).await?, keys); + assert_eq!( + machine.store().import_exported_room_keys(better_export, |_, _| {}).await?, + keys + ); let another_session = machine.create_inbound_session(room_id).await?; let export = vec![another_session.export_at_index(10).await]; @@ -359,7 +362,7 @@ mod tests { )]), ); - assert_eq!(machine.import_room_keys(export, false, |_, _| {}).await?, keys); + assert_eq!(machine.store().import_exported_room_keys(export, |_, _| {}).await?, keys); Ok(()) } diff --git a/crates/matrix-sdk-crypto/src/machine.rs b/crates/matrix-sdk-crypto/src/machine.rs index 9052b1204..74f8bb137 100644 --- a/crates/matrix-sdk-crypto/src/machine.rs +++ b/crates/matrix-sdk-crypto/src/machine.rs @@ -13,7 +13,7 @@ // limitations under the License. use std::{ - collections::{BTreeMap, BTreeSet, HashSet}, + collections::{BTreeMap, HashSet}, sync::{Arc, RwLock as StdRwLock}, time::Duration, }; @@ -1741,77 +1741,24 @@ impl OlmMachine { /// machine.import_room_keys(exported_keys, false, |_, _| {}).await.unwrap(); /// # }; /// ``` + #[deprecated( + since = "0.7.0", + note = "Use the OlmMachine::store::import_exported_room_keys method instead" + )] pub async fn import_room_keys( &self, exported_keys: Vec, #[allow(unused_variables)] from_backup: bool, progress_listener: impl Fn(usize, usize), ) -> StoreResult { - let mut sessions = Vec::new(); - - async fn new_session_better( - session: &InboundGroupSession, - old_session: Option, - ) -> bool { - if let Some(old_session) = &old_session { - session.compare(old_session).await == SessionOrdering::Better - } else { - true - } - } - - let total_count = exported_keys.len(); - let mut keys = BTreeMap::new(); - - for (i, key) in exported_keys.into_iter().enumerate() { - match InboundGroupSession::from_export(&key) { - Ok(session) => { - let old_session = self - .inner - .store - .get_inbound_group_session(session.room_id(), session.session_id()) - .await?; - - // Only import the session if we didn't have this session or - // if it's a better version of the same session. - if new_session_better(&session, old_session).await { - #[cfg(feature = "backups_v1")] - if from_backup { - session.mark_as_backed_up(); - } - - keys.entry(session.room_id().to_owned()) - .or_insert_with(BTreeMap::new) - .entry(session.sender_key().to_base64()) - .or_insert_with(BTreeSet::new) - .insert(session.session_id().to_owned()); - - sessions.push(session); - } - } - Err(e) => { - warn!( - 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." - ); - } - } - - progress_listener(i, total_count); - } - - let imported_count = sessions.len(); - - let changes = Changes { inbound_group_sessions: sessions, ..Default::default() }; - - self.store().save_changes(changes).await?; - - info!(total_count, imported_count, room_keys = ?keys, "Successfully imported room keys"); - - Ok(RoomKeyImportResult::new(imported_count, total_count, keys)) + self.store() + .import_room_keys( + exported_keys, + #[cfg(feature = "backups_v1")] + from_backup, + progress_listener, + ) + .await } /// Export the keys that match the given predicate. diff --git a/crates/matrix-sdk-crypto/src/store/mod.rs b/crates/matrix-sdk-crypto/src/store/mod.rs index 7add8ce9f..4ffad47f3 100644 --- a/crates/matrix-sdk-crypto/src/store/mod.rs +++ b/crates/matrix-sdk-crypto/src/store/mod.rs @@ -66,12 +66,12 @@ use crate::{ user::UserIdentities, Device, ReadOnlyDevice, ReadOnlyUserIdentities, UserDevices, }, olm::{ - Account, InboundGroupSession, OlmMessageHash, OutboundGroupSession, + Account, ExportedRoomKey, InboundGroupSession, OlmMessageHash, OutboundGroupSession, PrivateCrossSigningIdentity, Session, StaticAccountData, }, types::{events::room_key_withheld::RoomKeyWithheldEvent, EventEncryptionAlgorithm}, verification::VerificationMachine, - CrossSigningStatus, ReadOnlyOwnUserIdentity, + CrossSigningStatus, ReadOnlyOwnUserIdentity, RoomKeyImportResult, }; pub mod caches; @@ -1489,6 +1489,119 @@ impl Store { pub fn secrets_stream(&self) -> impl Stream { self.inner.store.secrets_stream() } + + pub(crate) async fn import_room_keys( + &self, + exported_keys: Vec, + #[cfg(feature = "backups_v1")] from_backup: bool, + progress_listener: impl Fn(usize, usize), + ) -> Result { + let mut sessions = Vec::new(); + + async fn new_session_better( + session: &InboundGroupSession, + old_session: Option, + ) -> bool { + if let Some(old_session) = &old_session { + session.compare(old_session).await == SessionOrdering::Better + } else { + true + } + } + + let total_count = exported_keys.len(); + let mut keys = BTreeMap::new(); + + for (i, key) in exported_keys.into_iter().enumerate() { + match InboundGroupSession::from_export(&key) { + Ok(session) => { + let old_session = self + .inner + .store + .get_inbound_group_session(session.room_id(), session.session_id()) + .await?; + + // Only import the session if we didn't have this session or + // if it's a better version of the same session. + if new_session_better(&session, old_session).await { + #[cfg(feature = "backups_v1")] + if from_backup { + session.mark_as_backed_up(); + } + + keys.entry(session.room_id().to_owned()) + .or_insert_with(BTreeMap::new) + .entry(session.sender_key().to_base64()) + .or_insert_with(BTreeSet::new) + .insert(session.session_id().to_owned()); + + sessions.push(session); + } + } + Err(e) => { + warn!( + 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." + ); + } + } + + progress_listener(i, total_count); + } + + let imported_count = sessions.len(); + + let changes = Changes { inbound_group_sessions: sessions, ..Default::default() }; + + self.save_changes(changes).await?; + + info!(total_count, imported_count, room_keys = ?keys, "Successfully imported room keys"); + + 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.import_room_keys(exported_keys, false, |_, _| {}).await.unwrap(); + /// # }; + /// ``` + pub async fn import_exported_room_keys( + &self, + exported_keys: Vec, + progress_listener: impl Fn(usize, usize), + ) -> Result { + self.import_room_keys( + exported_keys, + #[cfg(feature = "backups_v1")] + false, + progress_listener, + ) + .await + } } impl Deref for Store { diff --git a/crates/matrix-sdk-ui/src/timeline/tests/encryption.rs b/crates/matrix-sdk-ui/src/timeline/tests/encryption.rs index a284e90b0..be2290e69 100644 --- a/crates/matrix-sdk-ui/src/timeline/tests/encryption.rs +++ b/crates/matrix-sdk-ui/src/timeline/tests/encryption.rs @@ -96,7 +96,7 @@ async fn retry_message_decryption() { let exported_keys = decrypt_room_key_export(Cursor::new(SESSION_KEY), "1234").unwrap(); let olm_machine = OlmMachine::new(own_user_id, "SomeDeviceId".into()).await; - olm_machine.import_room_keys(exported_keys, false, |_, _| {}).await.unwrap(); + olm_machine.store().import_exported_room_keys(exported_keys, |_, _| {}).await.unwrap(); timeline .inner @@ -201,7 +201,7 @@ async fn retry_edit_decryption() { let own_user_id = user_id!("@example:morheus.localhost"); let olm_machine = OlmMachine::new(own_user_id, "SomeDeviceId".into()).await; - olm_machine.import_room_keys(keys, false, |_, _| {}).await.unwrap(); + olm_machine.store().import_exported_room_keys(keys, |_, _| {}).await.unwrap(); timeline .inner @@ -304,7 +304,7 @@ async fn retry_edit_and_more() { let olm_machine = OlmMachine::new(user_id!("@jptest:matrix.org"), DEVICE_ID.into()).await; let keys = decrypt_room_key_export(Cursor::new(SESSION_KEY), "testing").unwrap(); - olm_machine.import_room_keys(keys, false, |_, _| {}).await.unwrap(); + olm_machine.store().import_exported_room_keys(keys, |_, _| {}).await.unwrap(); timeline .inner @@ -389,7 +389,7 @@ async fn retry_message_decryption_highlighted() { let exported_keys = decrypt_room_key_export(Cursor::new(SESSION_KEY), "1234").unwrap(); let olm_machine = OlmMachine::new(own_user_id, "SomeDeviceId".into()).await; - olm_machine.import_room_keys(exported_keys, false, |_, _| {}).await.unwrap(); + olm_machine.store().import_exported_room_keys(exported_keys, |_, _| {}).await.unwrap(); timeline .inner diff --git a/crates/matrix-sdk-ui/src/timeline/tests/read_receipts.rs b/crates/matrix-sdk-ui/src/timeline/tests/read_receipts.rs index 74afcfe59..ec99f7360 100644 --- a/crates/matrix-sdk-ui/src/timeline/tests/read_receipts.rs +++ b/crates/matrix-sdk-ui/src/timeline/tests/read_receipts.rs @@ -438,7 +438,7 @@ async fn read_receipts_updates_on_message_decryption() { let exported_keys = decrypt_room_key_export(Cursor::new(SESSION_KEY), "1234").unwrap(); let olm_machine = OlmMachine::new(own_user_id, "SomeDeviceId".into()).await; - olm_machine.import_room_keys(exported_keys, false, |_, _| {}).await.unwrap(); + olm_machine.store().import_exported_room_keys(exported_keys, |_, _| {}).await.unwrap(); timeline .inner diff --git a/crates/matrix-sdk/src/encryption/mod.rs b/crates/matrix-sdk/src/encryption/mod.rs index af5f57f21..6201914c4 100644 --- a/crates/matrix-sdk/src/encryption/mod.rs +++ b/crates/matrix-sdk/src/encryption/mod.rs @@ -1052,7 +1052,7 @@ impl Encryption { let task = tokio::task::spawn_blocking(decrypt); let import = task.await.expect("Task join error")?; - Ok(olm.import_room_keys(import, false, |_, _| {}).await?) + Ok(olm.store().import_exported_room_keys(import, |_, _| {}).await?) } /// Enables the crypto-store cross-process lock.