From c3f4a738a4803d97dac83261e935997794876c59 Mon Sep 17 00:00:00 2001 From: Hubert Chathi Date: Mon, 10 Jun 2024 13:56:00 -0400 Subject: [PATCH] use our device keys from storage --- .../src/identities/device.rs | 2 +- crates/matrix-sdk-crypto/src/olm/account.rs | 66 ++++++++++++------- crates/matrix-sdk-crypto/src/olm/mod.rs | 37 +++++------ crates/matrix-sdk-crypto/src/olm/session.rs | 33 ++++------ .../src/session_manager/sessions.rs | 27 +++++++- crates/matrix-sdk-crypto/src/store/caches.rs | 6 +- .../src/store/integration_tests.rs | 3 +- .../src/store/memorystore.rs | 6 +- crates/matrix-sdk-indexeddb/Cargo.toml | 1 - .../src/crypto_store/mod.rs | 37 ++++++----- crates/matrix-sdk-sqlite/src/crypto_store.rs | 28 ++++++-- 11 files changed, 150 insertions(+), 96 deletions(-) diff --git a/crates/matrix-sdk-crypto/src/identities/device.rs b/crates/matrix-sdk-crypto/src/identities/device.rs index 79cd50eea52..bc159b63a3d 100644 --- a/crates/matrix-sdk-crypto/src/identities/device.rs +++ b/crates/matrix-sdk-crypto/src/identities/device.rs @@ -871,7 +871,7 @@ impl ReadOnlyDevice { } } - pub(crate) fn as_device_keys(&self) -> &DeviceKeys { + pub fn as_device_keys(&self) -> &DeviceKeys { &self.inner } diff --git a/crates/matrix-sdk-crypto/src/olm/account.rs b/crates/matrix-sdk-crypto/src/olm/account.rs index 9220e90f6ab..7fff50fa24e 100644 --- a/crates/matrix-sdk-crypto/src/olm/account.rs +++ b/crates/matrix-sdk-crypto/src/olm/account.rs @@ -902,26 +902,24 @@ impl Account { /// created and shared with us. /// /// * `fallback_used` - Was the one-time key a fallback key. - pub async fn create_outbound_session_helper( + pub fn create_outbound_session_helper( &self, config: SessionConfig, identity_key: Curve25519PublicKey, one_time_key: Curve25519PublicKey, fallback_used: bool, + our_device_keys: DeviceKeys, ) -> Session { let session = self.inner.create_outbound_session(config, identity_key, one_time_key); let now = SecondsSinceUnixEpoch::now(); let session_id = session.session_id(); - let device_keys = self.device_keys(); - // FIXME: sign with cross signing identity - Session { inner: Arc::new(Mutex::new(session)), session_id: session_id.into(), sender_key: identity_key, - device_keys, + device_keys: our_device_keys, created_using_fallback_key: fallback_used, creation_time: now, last_use_time: now, @@ -977,10 +975,11 @@ impl Account { /// * `key_map` - A map from the algorithm and device ID to the one-time key /// that the other account created and shared with us. #[allow(clippy::result_large_err)] - pub async fn create_outbound_session( + pub fn create_outbound_session( &self, device: &ReadOnlyDevice, key_map: &BTreeMap>, + our_device_keys: DeviceKeys, ) -> Result { let pre_key_bundle = Self::find_pre_key_bundle(device, key_map)?; @@ -1005,9 +1004,13 @@ impl Account { let one_time_key = key.key(); let config = device.olm_session_config(); - Ok(self - .create_outbound_session_helper(config, identity_key, one_time_key, is_fallback) - .await) + Ok(self.create_outbound_session_helper( + config, + identity_key, + one_time_key, + is_fallback, + our_device_keys, + )) } } } @@ -1022,9 +1025,10 @@ impl Account { /// /// * `message` - A pre-key Olm message that was sent to us by the other /// account. - pub async fn create_inbound_session( + pub fn create_inbound_session( &mut self, their_identity_key: Curve25519PublicKey, + our_device_keys: DeviceKeys, message: &PreKeyMessage, ) -> Result { Span::current().record("session_id", debug(message.session_id())); @@ -1036,14 +1040,11 @@ impl Account { debug!(session=?result.session, "Decrypted an Olm message from a new Olm session"); - let device_keys = self.device_keys(); - // FIXME: sign with cross signing identity - let session = Session { inner: Arc::new(Mutex::new(result.session)), session_id: session_id.into(), sender_key: their_identity_key, - device_keys, + device_keys: our_device_keys, created_using_fallback_key: false, creation_time: now, last_use_time: now, @@ -1067,7 +1068,8 @@ impl Account { let one_time_map = other.signed_one_time_keys(); let device = ReadOnlyDevice::from_account(other); - let mut our_session = self.create_outbound_session(&device, &one_time_map).await.unwrap(); + let mut our_session = + self.create_outbound_session(&device, &one_time_map, self.device_keys()).unwrap(); other.mark_keys_as_published(); @@ -1100,8 +1102,11 @@ impl Account { let our_device = ReadOnlyDevice::from_account(self); let other_session = other - .create_inbound_session(our_device.curve25519_key().unwrap(), &prekey) - .await + .create_inbound_session( + our_device.curve25519_key().unwrap(), + self.device_keys(), + &prekey, + ) .unwrap(); (our_session, other_session.session) @@ -1294,13 +1299,28 @@ impl Account { } // We didn't find a matching session; try to create a new session. - let result = match self.create_inbound_session(sender_key, prekey_message).await { - Ok(r) => r, - Err(e) => { - warn!("Failed to create a new Olm session from a pre-key message: {e:?}"); - return Err(OlmError::SessionWedged(sender.to_owned(), sender_key)); - } + // try to get our own stored device keys + let device_keys = store + .get_device(&self.user_id, &self.device_id) + .await + .unwrap_or(None) + .map(|read_only_device| read_only_device.as_device_keys().clone()); + // if we don't have it stored, fall back to generating a fresh + // device keys from our own Account + let device_keys = match device_keys { + Some(device_keys) => device_keys, + None => self.device_keys(), }; + let result = + match self.create_inbound_session(sender_key, device_keys, prekey_message) { + Ok(r) => r, + Err(e) => { + warn!( + "Failed to create a new Olm session from a pre-key message: {e:?}" + ); + return Err(OlmError::SessionWedged(sender.to_owned(), sender_key)); + } + }; // We need to add the new session to the session cache, otherwise // we might try to create the same session again. diff --git a/crates/matrix-sdk-crypto/src/olm/mod.rs b/crates/matrix-sdk-crypto/src/olm/mod.rs index 4c1e2e671e0..bbb9061665d 100644 --- a/crates/matrix-sdk-crypto/src/olm/mod.rs +++ b/crates/matrix-sdk-crypto/src/olm/mod.rs @@ -81,21 +81,20 @@ pub(crate) mod tests { device_id!("BOBDEVICE") } - pub(crate) async fn get_account_and_session_test_helper() -> (Account, Session) { + pub(crate) fn get_account_and_session_test_helper() -> (Account, Session) { let alice = Account::with_device_id(alice_id(), alice_device_id()); let mut bob = Account::with_device_id(bob_id(), bob_device_id()); bob.generate_one_time_keys(1); let one_time_key = *bob.one_time_keys().values().next().unwrap(); let sender_key = bob.identity_keys().curve25519; - let session = alice - .create_outbound_session_helper( - SessionConfig::default(), - sender_key, - one_time_key, - false, - ) - .await; + let session = alice.create_outbound_session_helper( + SessionConfig::default(), + sender_key, + one_time_key, + false, + alice.device_keys(), + ); (alice, session) } @@ -141,14 +140,13 @@ pub(crate) mod tests { let one_time_key = *one_time_keys.values().next().unwrap(); - let mut bob_session = bob - .create_outbound_session_helper( - SessionConfig::default(), - alice_keys.curve25519, - one_time_key, - false, - ) - .await; + let mut bob_session = bob.create_outbound_session_helper( + SessionConfig::default(), + alice_keys.curve25519, + one_time_key, + false, + bob.device_keys(), + ); let plaintext = "Hello world"; @@ -160,8 +158,9 @@ pub(crate) mod tests { }; let bob_keys = bob.identity_keys(); - let result = - alice.create_inbound_session(bob_keys.curve25519, &prekey_message).await.unwrap(); + let result = alice + .create_inbound_session(bob_keys.curve25519, alice.device_keys(), &prekey_message) + .unwrap(); assert_eq!(bob_session.session_id(), result.session.session_id()); diff --git a/crates/matrix-sdk-crypto/src/olm/session.rs b/crates/matrix-sdk-crypto/src/olm/session.rs index b2c6eecfd64..abbdeec8ea1 100644 --- a/crates/matrix-sdk-crypto/src/olm/session.rs +++ b/crates/matrix-sdk-crypto/src/olm/session.rs @@ -28,7 +28,6 @@ use vodozemac::{ use crate::types::events::room::encrypted::OlmV2Curve25519AesSha2Content; use crate::{ error::{EventError, OlmResult}, - olm::PrivateCrossSigningIdentity, types::{ events::room::encrypted::{OlmV1Curve25519AesSha2Content, ToDeviceEncryptedEventContent}, DeviceKeys, EventEncryptionAlgorithm, @@ -234,19 +233,11 @@ impl Session { /// /// * `pickle_mode` - The mode that was used to pickle the session, either /// an unencrypted mode or an encrypted using passphrase. - pub async fn from_pickle( - mut device_keys: DeviceKeys, - identity: Option, - pickle: PickledSession, - ) -> Self { + pub fn from_pickle(device_keys: DeviceKeys, pickle: PickledSession) -> Self { // FIXME: assert that device_keys has curve25519 and ed25519 keys let session: vodozemac::olm::Session = pickle.pickle.into(); let session_id = session.session_id(); - if let Some(identity) = identity { - let _ = identity.sign_device_keys(&mut device_keys).await; - } - Session { inner: Arc::new(Mutex::new(session)), session_id: session_id.into(), @@ -309,14 +300,13 @@ mod tests { bob.generate_one_time_keys(1); let one_time_key = *bob.one_time_keys().values().next().unwrap(); let sender_key = bob.identity_keys().curve25519; - let mut alice_session = alice - .create_outbound_session_helper( - SessionConfig::default(), - sender_key, - one_time_key, - false, - ) - .await; + let mut alice_session = alice.create_outbound_session_helper( + SessionConfig::default(), + sender_key, + one_time_key, + false, + alice.device_keys(), + ); let alice_device = ReadOnlyDevice::from_account(&alice); // let bob_device = ReadOnlyDevice::from_account(&bob); @@ -341,8 +331,11 @@ mod tests { }; let bob_session_result = bob - .create_inbound_session(alice_device.curve25519_key().unwrap(), &prekey) - .await + .create_inbound_session( + alice_device.curve25519_key().unwrap(), + bob.device_keys(), + &prekey, + ) .unwrap(); // check that the payload has the device keys diff --git a/crates/matrix-sdk-crypto/src/session_manager/sessions.rs b/crates/matrix-sdk-crypto/src/session_manager/sessions.rs index b44ffc8d4c3..dc1cbba46bd 100644 --- a/crates/matrix-sdk-crypto/src/session_manager/sessions.rs +++ b/crates/matrix-sdk-crypto/src/session_manager/sessions.rs @@ -36,7 +36,7 @@ use crate::{ gossiping::GossipMachine, requests::{OutgoingRequest, ToDeviceRequest}, store::{Changes, Result as StoreResult, Store}, - types::{events::EventType, EventEncryptionAlgorithm}, + types::{events::EventType, DeviceKeys, EventEncryptionAlgorithm}, ReadOnlyDevice, }; @@ -514,6 +514,8 @@ impl SessionManager { let mut new_sessions: BTreeMap<&UserId, BTreeMap<&DeviceId, SessionInfo>> = BTreeMap::new(); let mut store_transaction = self.store.transaction().await; + let mut our_device_keys: Option = None; + for (user_id, user_devices) in &response.one_time_keys { for (device_id, key_map) in user_devices { let device = match self.store.get_readonly_device(user_id, device_id).await { @@ -537,7 +539,28 @@ impl SessionManager { }; let account = store_transaction.account().await?; - let session = match account.create_outbound_session(&device, key_map).await { + let device_keys = match our_device_keys { + Some(ref device_keys) => device_keys.clone(), + None => { + let device_keys = self + .store + .get_device(&account.user_id, &account.device_id) + .await + .unwrap_or(None) + .map(|read_only_device| read_only_device.as_device_keys().clone()); + // if we don't have it stored, fall back to generating a fresh + // device keys from our own Account + let device_keys = match device_keys { + Some(device_keys) => device_keys, + None => account.device_keys(), + }; + + our_device_keys = Some(device_keys.clone()); + + device_keys + } + }; + let session = match account.create_outbound_session(&device, key_map, device_keys) { Ok(s) => s, Err(e) => { warn!( diff --git a/crates/matrix-sdk-crypto/src/store/caches.rs b/crates/matrix-sdk-crypto/src/store/caches.rs index 55eb3081211..e1a85805ba8 100644 --- a/crates/matrix-sdk-crypto/src/store/caches.rs +++ b/crates/matrix-sdk-crypto/src/store/caches.rs @@ -397,7 +397,7 @@ mod tests { #[async_test] async fn test_session_store() { - let (_, session) = get_account_and_session_test_helper().await; + let (_, session) = get_account_and_session_test_helper(); let store = SessionStore::new(); @@ -414,7 +414,7 @@ mod tests { #[async_test] async fn test_session_store_bulk_storing() { - let (_, session) = get_account_and_session_test_helper().await; + let (_, session) = get_account_and_session_test_helper(); let store = SessionStore::new(); store.set_for_sender(&session.sender_key.to_base64(), vec![session.clone()]); @@ -429,7 +429,7 @@ mod tests { #[async_test] async fn test_group_session_store() { - let (account, _) = get_account_and_session_test_helper().await; + let (account, _) = get_account_and_session_test_helper(); let room_id = room_id!("!test:localhost"); let curve_key = "Nn0L2hkcCMFKqynTjyGsJbth7QrVmX3lbrksMkrGOAw"; diff --git a/crates/matrix-sdk-crypto/src/store/integration_tests.rs b/crates/matrix-sdk-crypto/src/store/integration_tests.rs index b657a78e107..fb297f88ef6 100644 --- a/crates/matrix-sdk-crypto/src/store/integration_tests.rs +++ b/crates/matrix-sdk-crypto/src/store/integration_tests.rs @@ -85,7 +85,8 @@ macro_rules! cryptostore_integration_tests { sender_key, one_time_key, false, - ).await; + alice.device_keys(), + ); (alice, session) } diff --git a/crates/matrix-sdk-crypto/src/store/memorystore.rs b/crates/matrix-sdk-crypto/src/store/memorystore.rs index d56ed9d8452..cace29e1220 100644 --- a/crates/matrix-sdk-crypto/src/store/memorystore.rs +++ b/crates/matrix-sdk-crypto/src/store/memorystore.rs @@ -626,7 +626,7 @@ mod tests { #[async_test] async fn test_session_store() { - let (account, session) = get_account_and_session_test_helper().await; + let (account, session) = get_account_and_session_test_helper(); let store = MemoryStore::new(); assert!(store.load_account().await.unwrap().is_none()); @@ -644,7 +644,7 @@ mod tests { #[async_test] async fn test_inbound_group_session_store() { - let (account, _) = get_account_and_session_test_helper().await; + let (account, _) = get_account_and_session_test_helper(); let room_id = room_id!("!test:localhost"); let curve_key = "Nn0L2hkcCMFKqynTjyGsJbth7QrVmX3lbrksMkrGOAw"; @@ -833,7 +833,7 @@ mod tests { #[async_test] async fn test_outbound_group_session_store() { // Given an outbound session - let (account, _) = get_account_and_session_test_helper().await; + let (account, _) = get_account_and_session_test_helper(); let room_id = room_id!("!test:localhost"); let (outbound, _) = account.create_group_session_pair_with_defaults(room_id).await; diff --git a/crates/matrix-sdk-indexeddb/Cargo.toml b/crates/matrix-sdk-indexeddb/Cargo.toml index 2faeadaaa22..8063b55aca4 100644 --- a/crates/matrix-sdk-indexeddb/Cargo.toml +++ b/crates/matrix-sdk-indexeddb/Cargo.toml @@ -23,7 +23,6 @@ testing = ["matrix-sdk-crypto?/testing"] anyhow = { workspace = true } async-trait = { workspace = true } base64 = { workspace = true } -futures-util = { workspace = true } gloo-utils = { version = "0.2.0", features = ["serde"] } indexed_db_futures = "0.4.1" js-sys = { version = "0.3.58" } diff --git a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs index 204b09085ab..171846846d6 100644 --- a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs +++ b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs @@ -18,7 +18,6 @@ use std::{ }; use async_trait::async_trait; -use futures_util::stream::{FuturesUnordered, StreamExt}; use gloo_utils::format::JsValueSerdeExt; use hkdf::Hkdf; use indexed_db_futures::prelude::*; @@ -31,7 +30,7 @@ use matrix_sdk_crypto::{ caches::SessionStore, BackupKeys, Changes, CryptoStore, CryptoStoreError, PendingChanges, RoomKeyCounts, RoomSettings, }, - types::events::room_key_withheld::RoomKeyWithheldEvent, + types::{events::room_key_withheld::RoomKeyWithheldEvent, DeviceKeys}, vodozemac::base64_encode, Account, GossipRequest, GossippedSecret, ReadOnlyDevice, ReadOnlyUserIdentities, SecretInfo, TrackedUser, @@ -861,20 +860,26 @@ impl_crypto_store! { } async fn get_sessions(&self, sender_key: &str) -> Result>>>> { + let account_info = self.get_static_account().ok_or(CryptoStoreError::AccountUnset)?; if self.session_cache.get(sender_key).is_none() { - let account = self.load_account() - .await - .or(Err(CryptoStoreError::AccountUnset))? - .ok_or(CryptoStoreError::AccountUnset)?; - let identity = self.load_identity() + // try to get our own stored device keys + let device_keys = self.get_device(&account_info.user_id, &account_info.device_id) .await - .unwrap_or(None); - let mut device_keys = account.device_keys(); - // FIXME: we could avoid the FuturesUnordered and the .clone for - // device_keys and identity if we can sign the device_keys here, - // but the function for signing isn't visible to this crate. (It - // would also be more efficient since we will only need to do the - // signing once, rather than for each session) + .unwrap_or(None) + .map(|read_only_device| read_only_device.as_device_keys().clone()); + + // if we don't have it stored, fall back to generating a fresh + // device keys from our own Account + let device_keys = match device_keys { + Some(device_keys) => device_keys, + None => { + let account = self.load_account() + .await + .or(Err(CryptoStoreError::AccountUnset))? + .ok_or(CryptoStoreError::AccountUnset)?; + account.device_keys() + }, + }; let range = self.serializer.encode_to_range(keys::SESSION, sender_key)?; let sessions: Vec = self @@ -887,12 +892,10 @@ impl_crypto_store! { .filter_map(|f| self.serializer.deserialize_value(f).ok().map(|p| { Session::from_pickle( device_keys.clone(), - identity.clone(), p, ) })) - .collect::>() - .collect::>().await; + .collect::>(); self.session_cache.set_for_sender(sender_key, sessions); } diff --git a/crates/matrix-sdk-sqlite/src/crypto_store.rs b/crates/matrix-sdk-sqlite/src/crypto_store.rs index 29842dbc596..1acb6babc53 100644 --- a/crates/matrix-sdk-sqlite/src/crypto_store.rs +++ b/crates/matrix-sdk-sqlite/src/crypto_store.rs @@ -908,6 +908,27 @@ impl CryptoStore for SqliteCryptoStore { let account_info = self.get_static_account().ok_or(Error::AccountUnset)?; if self.session_cache.get(sender_key).is_none() { + // try to get our own stored device keys + let device_keys = self + .get_device(&account_info.user_id, &account_info.device_id) + .await + .unwrap_or(None) + .map(|read_only_device| read_only_device.as_device_keys().clone()); + + // if we don't have it stored, fall back to generating a fresh + // device keys from our own Account + let device_keys = match device_keys { + Some(device_keys) => device_keys, + None => { + let account = self + .load_account() + .await + .or(Err(Error::AccountUnset))? + .ok_or(Error::AccountUnset)?; + account.device_keys() + } + }; + let sessions = self .acquire() .await? @@ -916,12 +937,7 @@ impl CryptoStore for SqliteCryptoStore { .into_iter() .map(|bytes| { let pickle = self.deserialize_value(&bytes)?; - Ok(Session::from_pickle( - account_info.user_id.clone(), - account_info.device_id.clone(), - account_info.identity_keys.clone(), - pickle, - )) + Ok(Session::from_pickle(device_keys.clone(), pickle)) }) .collect::>()?;