use our device keys from storage

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
     }
 

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<OwnedDeviceKeyId, Raw<ruma::encryption::OneTimeKey>>,
+        our_device_keys: DeviceKeys,
     ) -> Result<Session, SessionCreationError> {
         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<InboundCreationResult, SessionCreationError> {
         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.

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());
 

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<PrivateCrossSigningIdentity>,
-        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

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<DeviceKeys> = 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!(

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

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

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;