From 00584ffc8b78109e65553dafcb68ae80dcc051f2 Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Wed, 5 Jun 2024 20:51:20 -0400
Subject: [PATCH 01/19] initial work to embed device keys in Olm-encrypted
 messages

---
 crates/matrix-sdk-crypto/src/olm/account.rs   |  41 +++----
 crates/matrix-sdk-crypto/src/olm/mod.rs       |  33 +++---
 crates/matrix-sdk-crypto/src/olm/session.rs   | 105 ++++++++++++++----
 .../src/session_manager/sessions.rs           |   2 +-
 crates/matrix-sdk-crypto/src/store/caches.rs  |   6 +-
 .../src/store/integration_tests.rs            |   2 +-
 .../src/store/memorystore.rs                  |   6 +-
 crates/matrix-sdk-indexeddb/Cargo.toml        |   1 +
 .../src/crypto_store/mod.rs                   |  25 ++++-
 9 files changed, 153 insertions(+), 68 deletions(-)

diff --git a/crates/matrix-sdk-crypto/src/olm/account.rs b/crates/matrix-sdk-crypto/src/olm/account.rs
index 81c58c0e071..9220e90f6ab 100644
--- a/crates/matrix-sdk-crypto/src/olm/account.rs
+++ b/crates/matrix-sdk-crypto/src/olm/account.rs
@@ -902,7 +902,7 @@ impl Account {
     /// created and shared with us.
     ///
     /// * `fallback_used` - Was the one-time key a fallback key.
-    pub fn create_outbound_session_helper(
+    pub async fn create_outbound_session_helper(
         &self,
         config: SessionConfig,
         identity_key: Curve25519PublicKey,
@@ -914,13 +914,14 @@ impl Account {
         let now = SecondsSinceUnixEpoch::now();
         let session_id = session.session_id();
 
+        let device_keys = self.device_keys();
+        // FIXME: sign with cross signing identity
+
         Session {
-            user_id: self.static_data.user_id.clone(),
-            device_id: self.static_data.device_id.clone(),
-            our_identity_keys: self.static_data.identity_keys.clone(),
             inner: Arc::new(Mutex::new(session)),
             session_id: session_id.into(),
             sender_key: identity_key,
+            device_keys,
             created_using_fallback_key: fallback_used,
             creation_time: now,
             last_use_time: now,
@@ -976,7 +977,7 @@ 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 fn create_outbound_session(
+    pub async fn create_outbound_session(
         &self,
         device: &ReadOnlyDevice,
         key_map: &BTreeMap<OwnedDeviceKeyId, Raw<ruma::encryption::OneTimeKey>>,
@@ -1004,12 +1005,9 @@ 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,
-                ))
+                Ok(self
+                    .create_outbound_session_helper(config, identity_key, one_time_key, is_fallback)
+                    .await)
             }
         }
     }
@@ -1024,7 +1022,7 @@ impl Account {
     ///
     /// * `message` - A pre-key Olm message that was sent to us by the other
     /// account.
-    pub fn create_inbound_session(
+    pub async fn create_inbound_session(
         &mut self,
         their_identity_key: Curve25519PublicKey,
         message: &PreKeyMessage,
@@ -1038,13 +1036,14 @@ 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 {
-            user_id: self.static_data.user_id.clone(),
-            device_id: self.static_data.device_id.clone(),
-            our_identity_keys: self.static_data.identity_keys.clone(),
             inner: Arc::new(Mutex::new(result.session)),
             session_id: session_id.into(),
             sender_key: their_identity_key,
+            device_keys,
             created_using_fallback_key: false,
             creation_time: now,
             last_use_time: now,
@@ -1068,7 +1067,7 @@ 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).unwrap();
+        let mut our_session = self.create_outbound_session(&device, &one_time_map).await.unwrap();
 
         other.mark_keys_as_published();
 
@@ -1100,8 +1099,10 @@ impl Account {
         };
 
         let our_device = ReadOnlyDevice::from_account(self);
-        let other_session =
-            other.create_inbound_session(our_device.curve25519_key().unwrap(), &prekey).unwrap();
+        let other_session = other
+            .create_inbound_session(our_device.curve25519_key().unwrap(), &prekey)
+            .await
+            .unwrap();
 
         (our_session, other_session.session)
     }
@@ -1286,14 +1287,14 @@ impl Account {
                         );
 
                         return Err(OlmError::SessionWedged(
-                            session.user_id.to_owned(),
+                            session.device_keys.user_id.to_owned(),
                             session.sender_key(),
                         ));
                     }
                 }
 
                 // We didn't find a matching session; try to create a new session.
-                let result = match self.create_inbound_session(sender_key, prekey_message) {
+                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:?}");
diff --git a/crates/matrix-sdk-crypto/src/olm/mod.rs b/crates/matrix-sdk-crypto/src/olm/mod.rs
index 4e09600d3dd..4c1e2e671e0 100644
--- a/crates/matrix-sdk-crypto/src/olm/mod.rs
+++ b/crates/matrix-sdk-crypto/src/olm/mod.rs
@@ -81,19 +81,21 @@ pub(crate) mod tests {
         device_id!("BOBDEVICE")
     }
 
-    pub(crate) fn get_account_and_session_test_helper() -> (Account, Session) {
+    pub(crate) async 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,
-        );
+        let session = alice
+            .create_outbound_session_helper(
+                SessionConfig::default(),
+                sender_key,
+                one_time_key,
+                false,
+            )
+            .await;
 
         (alice, session)
     }
@@ -139,12 +141,14 @@ 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,
-        );
+        let mut bob_session = bob
+            .create_outbound_session_helper(
+                SessionConfig::default(),
+                alice_keys.curve25519,
+                one_time_key,
+                false,
+            )
+            .await;
 
         let plaintext = "Hello world";
 
@@ -156,7 +160,8 @@ pub(crate) mod tests {
         };
 
         let bob_keys = bob.identity_keys();
-        let result = alice.create_inbound_session(bob_keys.curve25519, &prekey_message).unwrap();
+        let result =
+            alice.create_inbound_session(bob_keys.curve25519, &prekey_message).await.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 d6fa10a1c53..b2c6eecfd64 100644
--- a/crates/matrix-sdk-crypto/src/olm/session.rs
+++ b/crates/matrix-sdk-crypto/src/olm/session.rs
@@ -14,7 +14,7 @@
 
 use std::{fmt, sync::Arc};
 
-use ruma::{serde::Raw, OwnedDeviceId, OwnedUserId, SecondsSinceUnixEpoch};
+use ruma::{serde::Raw, SecondsSinceUnixEpoch};
 use serde::{Deserialize, Serialize};
 use serde_json::json;
 use tokio::sync::Mutex;
@@ -24,14 +24,14 @@ use vodozemac::{
     Curve25519PublicKey,
 };
 
-use super::IdentityKeys;
 #[cfg(feature = "experimental-algorithms")]
 use crate::types::events::room::encrypted::OlmV2Curve25519AesSha2Content;
 use crate::{
     error::{EventError, OlmResult},
+    olm::PrivateCrossSigningIdentity,
     types::{
         events::room::encrypted::{OlmV1Curve25519AesSha2Content, ToDeviceEncryptedEventContent},
-        EventEncryptionAlgorithm,
+        DeviceKeys, EventEncryptionAlgorithm,
     },
     ReadOnlyDevice,
 };
@@ -40,18 +40,14 @@ use crate::{
 /// `Account`s
 #[derive(Clone)]
 pub struct Session {
-    /// The `UserId` associated with this session
-    pub user_id: OwnedUserId,
-    /// The specific `DeviceId` associated with this session
-    pub device_id: OwnedDeviceId,
-    /// The `IdentityKeys` associated with this session
-    pub our_identity_keys: Arc<IdentityKeys>,
     /// The OlmSession
     pub inner: Arc<Mutex<InnerSession>>,
     /// Our sessionId
     pub session_id: Arc<str>,
     /// The Key of the sender
     pub sender_key: Curve25519PublicKey,
+    /// The signed device keys
+    pub device_keys: DeviceKeys,
     /// Has this been created using the fallback key
     pub created_using_fallback_key: bool,
     /// When the session was created
@@ -156,11 +152,12 @@ impl Session {
                 recipient_device.ed25519_key().ok_or(EventError::MissingSigningKey)?;
 
             let payload = json!({
-                "sender": &self.user_id,
-                "sender_device": &self.device_id,
+                "sender": &self.device_keys.user_id,
+                "sender_device": &self.device_keys.device_id,
                 "keys": {
-                    "ed25519": self.our_identity_keys.ed25519.to_base64(),
+                    "ed25519": self.device_keys.ed25519_key().unwrap().to_base64(),
                 },
+                "device_keys": self.device_keys,
                 "recipient": recipient_device.user_id(),
                 "recipient_keys": {
                     "ed25519": recipient_signing_key.to_base64(),
@@ -178,7 +175,7 @@ impl Session {
             EventEncryptionAlgorithm::OlmV1Curve25519AesSha2 => OlmV1Curve25519AesSha2Content {
                 ciphertext,
                 recipient_key: self.sender_key,
-                sender_key: self.our_identity_keys.curve25519,
+                sender_key: self.device_keys.curve25519_key().unwrap(),
                 message_id,
             }
             .into(),
@@ -237,23 +234,25 @@ impl Session {
     ///
     /// * `pickle_mode` - The mode that was used to pickle the session, either
     /// an unencrypted mode or an encrypted using passphrase.
-    pub fn from_pickle(
-        user_id: OwnedUserId,
-        device_id: OwnedDeviceId,
-        our_identity_keys: Arc<IdentityKeys>,
+    pub async fn from_pickle(
+        mut device_keys: DeviceKeys,
+        identity: Option<PrivateCrossSigningIdentity>,
         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 {
-            user_id,
-            device_id,
-            our_identity_keys,
             inner: Arc::new(Mutex::new(session)),
             session_id: session_id.into(),
             created_using_fallback_key: pickle.created_using_fallback_key,
             sender_key: pickle.sender_key,
+            device_keys,
             creation_time: pickle.creation_time,
             last_use_time: pickle.last_use_time,
         }
@@ -285,3 +284,69 @@ pub struct PickledSession {
     /// The Unix timestamp when the session was last used.
     pub last_use_time: SecondsSinceUnixEpoch,
 }
+
+#[cfg(test)]
+mod tests {
+    use matrix_sdk_test::async_test;
+
+    use ruma::{device_id, user_id};
+    use serde_json::{self, Value};
+    use vodozemac::olm::{OlmMessage, SessionConfig};
+
+    use crate::{
+        identities::ReadOnlyDevice, olm::Account,
+        types::events::room::encrypted::ToDeviceEncryptedEventContent,
+    };
+
+    #[async_test]
+    async fn test_encryption_and_decryption() {
+        use ruma::events::dummy::ToDeviceDummyEventContent;
+
+        let alice =
+            Account::with_device_id(user_id!("@alice:localhost"), device_id!("ALICEDEVICE"));
+        let mut bob = Account::with_device_id(user_id!("@bob:localhost"), device_id!("BOBDEVICE"));
+
+        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 alice_device = ReadOnlyDevice::from_account(&alice);
+        // let bob_device = ReadOnlyDevice::from_account(&bob);
+
+        let message = alice_session
+            .encrypt(&alice_device, "m.dummy", ToDeviceDummyEventContent::new(), None)
+            .await
+            .unwrap()
+            .deserialize()
+            .unwrap();
+
+        let content = if let ToDeviceEncryptedEventContent::OlmV1Curve25519AesSha2(c) = message {
+            c
+        } else {
+            panic!("Invalid encrypted event algorithm {}", message.algorithm());
+        };
+
+        let prekey = if let OlmMessage::PreKey(m) = content.ciphertext {
+            m
+        } else {
+            panic!("Wrong Olm message type");
+        };
+
+        let bob_session_result = bob
+            .create_inbound_session(alice_device.curve25519_key().unwrap(), &prekey)
+            .await
+            .unwrap();
+
+        // check that the payload has the device keys
+        let plaintext: Value = serde_json::from_str(&bob_session_result.plaintext).unwrap();
+        assert_eq!(plaintext["device_keys"]["user_id"].as_str(), Some("@alice:localhost"));
+    }
+}
diff --git a/crates/matrix-sdk-crypto/src/session_manager/sessions.rs b/crates/matrix-sdk-crypto/src/session_manager/sessions.rs
index e81f32ded10..b44ffc8d4c3 100644
--- a/crates/matrix-sdk-crypto/src/session_manager/sessions.rs
+++ b/crates/matrix-sdk-crypto/src/session_manager/sessions.rs
@@ -537,7 +537,7 @@ impl SessionManager {
                 };
 
                 let account = store_transaction.account().await?;
-                let session = match account.create_outbound_session(&device, key_map) {
+                let session = match account.create_outbound_session(&device, key_map).await {
                     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 e1a85805ba8..55eb3081211 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();
+        let (_, session) = get_account_and_session_test_helper().await;
 
         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();
+        let (_, session) = get_account_and_session_test_helper().await;
 
         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();
+        let (account, _) = get_account_and_session_test_helper().await;
         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 0f457b48062..b657a78e107 100644
--- a/crates/matrix-sdk-crypto/src/store/integration_tests.rs
+++ b/crates/matrix-sdk-crypto/src/store/integration_tests.rs
@@ -85,7 +85,7 @@ macro_rules! cryptostore_integration_tests {
                     sender_key,
                     one_time_key,
                     false,
-                );
+                ).await;
 
                 (alice, session)
             }
diff --git a/crates/matrix-sdk-crypto/src/store/memorystore.rs b/crates/matrix-sdk-crypto/src/store/memorystore.rs
index cace29e1220..d56ed9d8452 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();
+        let (account, session) = get_account_and_session_test_helper().await;
         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();
+        let (account, _) = get_account_and_session_test_helper().await;
         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();
+        let (account, _) = get_account_and_session_test_helper().await;
         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 8063b55aca4..2faeadaaa22 100644
--- a/crates/matrix-sdk-indexeddb/Cargo.toml
+++ b/crates/matrix-sdk-indexeddb/Cargo.toml
@@ -23,6 +23,7 @@ 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 8c2f371954d..204b09085ab 100644
--- a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
+++ b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
@@ -18,6 +18,7 @@ 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::*;
@@ -860,9 +861,21 @@ impl_crypto_store! {
     }
 
     async fn get_sessions(&self, sender_key: &str) -> Result<Option<Arc<Mutex<Vec<Session>>>>> {
-        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()
+                .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)
+
             let range = self.serializer.encode_to_range(keys::SESSION, sender_key)?;
             let sessions: Vec<Session> = self
                 .inner
@@ -873,13 +886,13 @@ impl_crypto_store! {
                 .iter()
                 .filter_map(|f| self.serializer.deserialize_value(f).ok().map(|p| {
                     Session::from_pickle(
-                        account_info.user_id.clone(),
-                        account_info.device_id.clone(),
-                        account_info.identity_keys.clone(),
+                        device_keys.clone(),
+                        identity.clone(),
                         p,
                     )
                 }))
-                .collect::<Vec<Session>>();
+                .collect::<FuturesUnordered<_>>()
+                .collect::<Vec<Session>>().await;
 
             self.session_cache.set_for_sender(sender_key, sessions);
         }

From c3f4a738a4803d97dac83261e935997794876c59 Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Mon, 10 Jun 2024 13:56:00 -0400
Subject: [PATCH 02/19] 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<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.
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<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
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<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!(
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<Option<Arc<Mutex<Vec<Session>>>>> {
+        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<Session> = 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::<FuturesUnordered<_>>()
-                .collect::<Vec<Session>>().await;
+                .collect::<Vec<Session>>();
 
             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::<Result<_>>()?;
 

From 304bcf324594d60964d9d206e4eb14527a45dff1 Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Mon, 10 Jun 2024 19:26:14 -0400
Subject: [PATCH 03/19] clear caches when our device key is updated

---
 .../src/crypto_store/mod.rs                   | 29 ++++++++--
 crates/matrix-sdk-sqlite/src/crypto_store.rs  | 58 ++++++++++++++++++-
 2 files changed, 78 insertions(+), 9 deletions(-)

diff --git a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
index 171846846d6..66a103b5483 100644
--- a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
+++ b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
@@ -447,7 +447,10 @@ impl IndexeddbCryptoStore {
 
     /// Process all the changes and do all encryption/serialization before the
     /// actual transaction.
-    async fn prepare_for_transaction(&self, changes: &Changes) -> Result<PendingIndexeddbChanges> {
+    async fn prepare_for_transaction(
+        &self,
+        changes: &Changes,
+    ) -> Result<(PendingIndexeddbChanges, bool)> {
         let mut indexeddb_changes = PendingIndexeddbChanges::new();
 
         let private_identity_pickle =
@@ -534,7 +537,17 @@ impl IndexeddbCryptoStore {
 
         let mut device_store = indexeddb_changes.get(keys::DEVICES);
 
+        let account_info = self.get_static_account();
+        let mut clear_caches = false;
         for device in device_changes.new.iter().chain(&device_changes.changed) {
+            // if our own device key changes, we need to clear the
+            // session cache because the sessions contain a copy of our
+            // device key
+            if account_info.clone().is_some_and(|info| {
+                info.user_id == device.user_id() && info.device_id == device.device_id()
+            }) {
+                clear_caches = true;
+            }
             let key =
                 self.serializer.encode_key(keys::DEVICES, (device.user_id(), device.device_id()));
             let device = self.serializer.serialize_value(&device)?;
@@ -617,7 +630,7 @@ impl IndexeddbCryptoStore {
             }
         }
 
-        Ok(indexeddb_changes)
+        Ok((indexeddb_changes, clear_caches))
     }
 }
 
@@ -697,7 +710,7 @@ impl_crypto_store! {
         // TODO: #2000 should make this lock go away, or change its shape.
         let _guard = self.save_changes_lock.lock().await;
 
-        let indexeddb_changes = self.prepare_for_transaction(&changes).await?;
+        let (indexeddb_changes, clear_caches) = self.prepare_for_transaction(&changes).await?;
 
         let stores = indexeddb_changes.touched_stores();
 
@@ -713,9 +726,13 @@ impl_crypto_store! {
 
         tx.await.into_result()?;
 
-        // all good, let's update our caches:indexeddb
-        for session in changes.sessions {
-            self.session_cache.add(session).await;
+        if clear_caches {
+            self.clear_caches().await;
+        } else {
+            // all good, let's update our caches:indexeddb
+            for session in changes.sessions {
+                self.session_cache.add(session).await;
+            }
         }
 
         Ok(())
diff --git a/crates/matrix-sdk-sqlite/src/crypto_store.rs b/crates/matrix-sdk-sqlite/src/crypto_store.rs
index 1acb6babc53..ebda1abd575 100644
--- a/crates/matrix-sdk-sqlite/src/crypto_store.rs
+++ b/crates/matrix-sdk-sqlite/src/crypto_store.rs
@@ -780,9 +780,11 @@ impl CryptoStore for SqliteCryptoStore {
         }
 
         let this = self.clone();
-        self.acquire()
+        let clear_caches = self
+            .acquire()
             .await?
             .with_transaction(move |txn| {
+                let mut clear_caches = false;
                 if let Some(pickled_private_identity) = &pickled_private_identity {
                     let serialized_private_identity =
                         this.serialize_value(pickled_private_identity)?;
@@ -804,7 +806,16 @@ impl CryptoStore for SqliteCryptoStore {
                     txn.set_kv("backup_version_v1", &serialized_backup_version)?;
                 }
 
+                let account_info = this.get_static_account();
                 for device in changes.devices.new.iter().chain(&changes.devices.changed) {
+                    // if our own device key changes, we need to clear the
+                    // session cache because the sessions contain a copy of our
+                    // device key
+                    if account_info.clone().is_some_and(|info| {
+                        info.user_id == device.user_id() && info.device_id == device.device_id()
+                    }) {
+                        clear_caches = true;
+                    }
                     let user_id = this.encode_key("device", device.user_id().as_bytes());
                     let device_id = this.encode_key("device", device.device_id().as_bytes());
                     let data = this.serialize_value(&device)?;
@@ -875,10 +886,14 @@ impl CryptoStore for SqliteCryptoStore {
                     txn.set_secret(&secret_name, &value)?;
                 }
 
-                Ok::<_, Error>(())
+                Ok::<_, Error>(clear_caches)
             })
             .await?;
 
+        if clear_caches {
+            self.clear_caches().await;
+        }
+
         Ok(())
     }
 
@@ -1347,7 +1362,8 @@ mod tests {
 mod encrypted_tests {
     use matrix_sdk_crypto::{
         cryptostore_integration_tests, cryptostore_integration_tests_time,
-        store::{Changes, CryptoStore as _, PendingChanges},
+        store::{Changes, CryptoStore as _, DeviceChanges, PendingChanges},
+        ReadOnlyDevice,
     };
     use matrix_sdk_test::async_test;
     use once_cell::sync::Lazy;
@@ -1393,6 +1409,42 @@ mod encrypted_tests {
         );
     }
 
+    #[async_test]
+    async fn cache_cleared_after_device_update() {
+        let store = get_store("cache_cleared_after_device_update", None).await;
+        // Given we created a session and saved it in the store
+        let (account, session) = cryptostore_integration_tests::get_account_and_session().await;
+        let sender_key = session.sender_key.to_base64();
+
+        store
+            .save_pending_changes(PendingChanges { account: Some(account.deep_clone()) })
+            .await
+            .expect("Can't save account");
+
+        let changes = Changes { sessions: vec![session.clone()], ..Default::default() };
+        store.save_changes(changes).await.unwrap();
+
+        store.session_cache.get(&sender_key).expect("We should have a session");
+
+        // When we save a new version of our device keys
+        store
+            .save_changes(Changes {
+                devices: DeviceChanges {
+                    new: vec![ReadOnlyDevice::from_account(&account)],
+                    ..Default::default()
+                },
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        // Then the session is no longer in the cache
+        assert!(
+            store.session_cache.get(&sender_key).is_none(),
+            "Session should not be in the cache!"
+        );
+    }
+
     cryptostore_integration_tests!();
     cryptostore_integration_tests_time!();
 }

From 2e0d21bb5de49e2611aa8a612fff091f73c9031b Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Tue, 11 Jun 2024 14:54:10 -0400
Subject: [PATCH 04/19] cleanups

---
 crates/matrix-sdk-crypto/src/identities/device.rs    |  1 +
 crates/matrix-sdk-crypto/src/olm/account.rs          |  6 +++---
 crates/matrix-sdk-crypto/src/olm/session.rs          |  3 +--
 .../src/session_manager/sessions.rs                  |  7 +++++--
 crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs  | 12 ++++++------
 crates/matrix-sdk-sqlite/src/crypto_store.rs         | 10 +++++-----
 6 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/crates/matrix-sdk-crypto/src/identities/device.rs b/crates/matrix-sdk-crypto/src/identities/device.rs
index bc159b63a3d..225c8c8bf49 100644
--- a/crates/matrix-sdk-crypto/src/identities/device.rs
+++ b/crates/matrix-sdk-crypto/src/identities/device.rs
@@ -871,6 +871,7 @@ impl ReadOnlyDevice {
         }
     }
 
+    /// Return the device keys
     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 7fff50fa24e..9b5360b4521 100644
--- a/crates/matrix-sdk-crypto/src/olm/account.rs
+++ b/crates/matrix-sdk-crypto/src/olm/account.rs
@@ -1299,14 +1299,14 @@ impl Account {
                 }
 
                 // We didn't find a matching session; try to create a new session.
-                // try to get our own stored device keys
+                // 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
+                // If we don't have our device keys 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(),
diff --git a/crates/matrix-sdk-crypto/src/olm/session.rs b/crates/matrix-sdk-crypto/src/olm/session.rs
index abbdeec8ea1..d8717ed3516 100644
--- a/crates/matrix-sdk-crypto/src/olm/session.rs
+++ b/crates/matrix-sdk-crypto/src/olm/session.rs
@@ -309,7 +309,6 @@ mod tests {
         );
 
         let alice_device = ReadOnlyDevice::from_account(&alice);
-        // let bob_device = ReadOnlyDevice::from_account(&bob);
 
         let message = alice_session
             .encrypt(&alice_device, "m.dummy", ToDeviceDummyEventContent::new(), None)
@@ -338,7 +337,7 @@ mod tests {
             )
             .unwrap();
 
-        // check that the payload has the device keys
+        // Check that the encrypted payload has the device keys.
         let plaintext: Value = serde_json::from_str(&bob_session_result.plaintext).unwrap();
         assert_eq!(plaintext["device_keys"]["user_id"].as_str(), Some("@alice:localhost"));
     }
diff --git a/crates/matrix-sdk-crypto/src/session_manager/sessions.rs b/crates/matrix-sdk-crypto/src/session_manager/sessions.rs
index dc1cbba46bd..0b36c16a122 100644
--- a/crates/matrix-sdk-crypto/src/session_manager/sessions.rs
+++ b/crates/matrix-sdk-crypto/src/session_manager/sessions.rs
@@ -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;
 
+        // Cache of our own device keys, so we only need to fetch it once, if
+        // we need it at all.
         let mut our_device_keys: Option<DeviceKeys> = None;
 
         for (user_id, user_devices) in &response.one_time_keys {
@@ -542,14 +544,15 @@ impl SessionManager {
                 let device_keys = match our_device_keys {
                     Some(ref device_keys) => device_keys.clone(),
                     None => {
+                        // Try to get our own stored device keys.
                         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
+                        // 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(),
diff --git a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
index 66a103b5483..73506de13c5 100644
--- a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
+++ b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
@@ -540,9 +540,9 @@ impl IndexeddbCryptoStore {
         let account_info = self.get_static_account();
         let mut clear_caches = false;
         for device in device_changes.new.iter().chain(&device_changes.changed) {
-            // if our own device key changes, we need to clear the
-            // session cache because the sessions contain a copy of our
-            // device key
+            // If our own device key changes, we need to clear the session
+            // cache because the sessions contain a copy of our device key, and
+            // we want the sessions to use the new version.
             if account_info.clone().is_some_and(|info| {
                 info.user_id == device.user_id() && info.device_id == device.device_id()
             }) {
@@ -879,14 +879,14 @@ impl_crypto_store! {
     async fn get_sessions(&self, sender_key: &str) -> Result<Option<Arc<Mutex<Vec<Session>>>>> {
         let account_info = self.get_static_account().ok_or(CryptoStoreError::AccountUnset)?;
         if self.session_cache.get(sender_key).is_none() {
-            // try to get our own stored device keys
+            // 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
+            // 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 => {
diff --git a/crates/matrix-sdk-sqlite/src/crypto_store.rs b/crates/matrix-sdk-sqlite/src/crypto_store.rs
index ebda1abd575..0096f264cd2 100644
--- a/crates/matrix-sdk-sqlite/src/crypto_store.rs
+++ b/crates/matrix-sdk-sqlite/src/crypto_store.rs
@@ -808,9 +808,9 @@ impl CryptoStore for SqliteCryptoStore {
 
                 let account_info = this.get_static_account();
                 for device in changes.devices.new.iter().chain(&changes.devices.changed) {
-                    // if our own device key changes, we need to clear the
+                    // If our own device key changes, we need to clear the
                     // session cache because the sessions contain a copy of our
-                    // device key
+                    // device key.
                     if account_info.clone().is_some_and(|info| {
                         info.user_id == device.user_id() && info.device_id == device.device_id()
                     }) {
@@ -923,15 +923,15 @@ 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
+            // 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
+            // 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 => {

From 4d0b45a70796dedc9f877060b2d5c2a79864717a Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Thu, 13 Jun 2024 16:15:48 -0400
Subject: [PATCH 05/19] fix CI issues and address code review comments

---
 bindings/matrix-sdk-crypto-ffi/src/lib.rs     | 42 ++++++++++++++++---
 crates/matrix-sdk-crypto/src/olm/account.rs   |  3 ++
 crates/matrix-sdk-crypto/src/olm/session.rs   | 13 +++---
 .../src/crypto_store/mod.rs                   | 12 ++++--
 4 files changed, 55 insertions(+), 15 deletions(-)

diff --git a/bindings/matrix-sdk-crypto-ffi/src/lib.rs b/bindings/matrix-sdk-crypto-ffi/src/lib.rs
index dc48c07f503..f0020f3e88a 100644
--- a/bindings/matrix-sdk-crypto-ffi/src/lib.rs
+++ b/bindings/matrix-sdk-crypto-ffi/src/lib.rs
@@ -16,7 +16,11 @@ mod responses;
 mod users;
 mod verification;
 
-use std::{collections::HashMap, sync::Arc, time::Duration};
+use std::{
+    collections::{BTreeMap, HashMap},
+    sync::Arc,
+    time::Duration,
+};
 
 use anyhow::Context as _;
 pub use backup_recovery_key::{
@@ -33,7 +37,9 @@ use matrix_sdk_common::deserialized_responses::ShieldState as RustShieldState;
 use matrix_sdk_crypto::{
     olm::{IdentityKeys, InboundGroupSession, Session},
     store::{Changes, CryptoStore, PendingChanges, RoomSettings as RustRoomSettings},
-    types::{EventEncryptionAlgorithm as RustEventEncryptionAlgorithm, SigningKey},
+    types::{
+        DeviceKey, DeviceKeys, EventEncryptionAlgorithm as RustEventEncryptionAlgorithm, SigningKey,
+    },
     EncryptionSettings as RustEncryptionSettings,
 };
 use matrix_sdk_sqlite::SqliteCryptoStore;
@@ -43,8 +49,8 @@ pub use responses::{
 };
 use ruma::{
     events::room::history_visibility::HistoryVisibility as RustHistoryVisibility,
-    DeviceKeyAlgorithm, MilliSecondsSinceUnixEpoch, OwnedDeviceId, OwnedUserId, RoomId,
-    SecondsSinceUnixEpoch, UserId,
+    DeviceKeyAlgorithm, DeviceKeyId, MilliSecondsSinceUnixEpoch, OwnedDeviceId, OwnedUserId,
+    RoomId, SecondsSinceUnixEpoch, UserId,
 };
 use serde::{Deserialize, Serialize};
 use tokio::runtime::Runtime;
@@ -332,6 +338,10 @@ async fn save_changes(
     processed_steps += 1;
     listener(processed_steps, total_steps);
 
+    // The Sessions were created with incorrect device keys, so clear the cache
+    // so that they'll get recreated with correct ones.
+    store.clear_caches().await;
+
     Ok(())
 }
 
@@ -419,6 +429,27 @@ fn collect_sessions(
 ) -> anyhow::Result<(Vec<Session>, Vec<InboundGroupSession>)> {
     let mut sessions = Vec::new();
 
+    // Create a DeviceKeys struct with enough information to get a working
+    // Session, but we will won't actually use the Sessions (and we'll clear
+    // the session cache after migration) so we don't need to worry about
+    // signatures.
+    let device_keys = DeviceKeys::new(
+        user_id.clone(),
+        device_id.clone(),
+        Default::default(),
+        BTreeMap::from([
+            (
+                DeviceKeyId::from_parts(DeviceKeyAlgorithm::Ed25519, &device_id),
+                DeviceKey::Ed25519(identity_keys.ed25519),
+            ),
+            (
+                DeviceKeyId::from_parts(DeviceKeyAlgorithm::Curve25519, &device_id),
+                DeviceKey::Curve25519(identity_keys.curve25519),
+            ),
+        ]),
+        Default::default(),
+    );
+
     for session_pickle in session_pickles {
         let pickle =
             vodozemac::olm::Session::from_libolm_pickle(&session_pickle.pickle, pickle_key)?
@@ -439,8 +470,7 @@ fn collect_sessions(
             last_use_time,
         };
 
-        let session =
-            Session::from_pickle(user_id.clone(), device_id.clone(), identity_keys.clone(), pickle);
+        let session = Session::from_pickle(device_keys.clone(), pickle);
 
         sessions.push(session);
         processed_steps += 1;
diff --git a/crates/matrix-sdk-crypto/src/olm/account.rs b/crates/matrix-sdk-crypto/src/olm/account.rs
index 9b5360b4521..0aa00e3ce07 100644
--- a/crates/matrix-sdk-crypto/src/olm/account.rs
+++ b/crates/matrix-sdk-crypto/src/olm/account.rs
@@ -974,6 +974,9 @@ 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.
+    ///
+    /// * `our_device_keys` - Our own `DeviceKeys`, including cross-signing
+    ///   signatures if applicable, for embedding in encrypted messages.
     #[allow(clippy::result_large_err)]
     pub fn create_outbound_session(
         &self,
diff --git a/crates/matrix-sdk-crypto/src/olm/session.rs b/crates/matrix-sdk-crypto/src/olm/session.rs
index d8717ed3516..a4d023ae708 100644
--- a/crates/matrix-sdk-crypto/src/olm/session.rs
+++ b/crates/matrix-sdk-crypto/src/olm/session.rs
@@ -181,7 +181,7 @@ impl Session {
             #[cfg(feature = "experimental-algorithms")]
             EventEncryptionAlgorithm::OlmV2Curve25519AesSha2 => OlmV2Curve25519AesSha2Content {
                 ciphertext,
-                sender_key: self.our_identity_keys.curve25519,
+                sender_key: self.device_keys.curve25519_key().unwrap(),
                 message_id,
             }
             .into(),
@@ -279,7 +279,6 @@ pub struct PickledSession {
 #[cfg(test)]
 mod tests {
     use matrix_sdk_test::async_test;
-
     use ruma::{device_id, user_id};
     use serde_json::{self, Value};
     use vodozemac::olm::{OlmMessage, SessionConfig};
@@ -293,10 +292,12 @@ mod tests {
     async fn test_encryption_and_decryption() {
         use ruma::events::dummy::ToDeviceDummyEventContent;
 
+        // Given users Alice and Bob
         let alice =
             Account::with_device_id(user_id!("@alice:localhost"), device_id!("ALICEDEVICE"));
         let mut bob = Account::with_device_id(user_id!("@bob:localhost"), device_id!("BOBDEVICE"));
 
+        // When Alice creates an Olm session with Bob
         bob.generate_one_time_keys(1);
         let one_time_key = *bob.one_time_keys().values().next().unwrap();
         let sender_key = bob.identity_keys().curve25519;
@@ -310,6 +311,7 @@ mod tests {
 
         let alice_device = ReadOnlyDevice::from_account(&alice);
 
+        // and encrypts a message
         let message = alice_session
             .encrypt(&alice_device, "m.dummy", ToDeviceDummyEventContent::new(), None)
             .await
@@ -317,9 +319,7 @@ mod tests {
             .deserialize()
             .unwrap();
 
-        let content = if let ToDeviceEncryptedEventContent::OlmV1Curve25519AesSha2(c) = message {
-            c
-        } else {
+        if let ToDeviceEncryptedEventContent::OlmV1Curve25519AesSha2(content) = else {
             panic!("Invalid encrypted event algorithm {}", message.algorithm());
         };
 
@@ -329,6 +329,7 @@ mod tests {
             panic!("Wrong Olm message type");
         };
 
+        // Then Bob should be able to create a session from the message and decrypt it.
         let bob_session_result = bob
             .create_inbound_session(
                 alice_device.curve25519_key().unwrap(),
@@ -337,7 +338,7 @@ mod tests {
             )
             .unwrap();
 
-        // Check that the encrypted payload has the device keys.
+        // Also ensure that the encrypted payload has the device keys.
         let plaintext: Value = serde_json::from_str(&bob_session_result.plaintext).unwrap();
         assert_eq!(plaintext["device_keys"]["user_id"].as_str(), Some("@alice:localhost"));
     }
diff --git a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
index 73506de13c5..0e916612f83 100644
--- a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
+++ b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
@@ -30,7 +30,7 @@ use matrix_sdk_crypto::{
         caches::SessionStore, BackupKeys, Changes, CryptoStore, CryptoStoreError, PendingChanges,
         RoomKeyCounts, RoomSettings,
     },
-    types::{events::room_key_withheld::RoomKeyWithheldEvent, DeviceKeys},
+    types::events::room_key_withheld::RoomKeyWithheldEvent,
     vodozemac::base64_encode,
     Account, GossipRequest, GossippedSecret, ReadOnlyDevice, ReadOnlyUserIdentities, SecretInfo,
     TrackedUser,
@@ -447,6 +447,10 @@ impl IndexeddbCryptoStore {
 
     /// Process all the changes and do all encryption/serialization before the
     /// actual transaction.
+    ///
+    /// Returns a tuple where the first item is a `PendingIndexeddbChanges`
+    /// struct, and the second item is a boolean indicating whether the session
+    /// cache should be cleared.
     async fn prepare_for_transaction(
         &self,
         changes: &Changes,
@@ -543,7 +547,7 @@ impl IndexeddbCryptoStore {
             // If our own device key changes, we need to clear the session
             // cache because the sessions contain a copy of our device key, and
             // we want the sessions to use the new version.
-            if account_info.clone().is_some_and(|info| {
+            if account_info.as_ref().is_some_and(|info| {
                 info.user_id == device.user_id() && info.device_id == device.device_id()
             }) {
                 clear_caches = true;
@@ -729,7 +733,9 @@ impl_crypto_store! {
         if clear_caches {
             self.clear_caches().await;
         } else {
-            // all good, let's update our caches:indexeddb
+            // All good, let's update our caches:indexeddb.
+            // We only do this if clear_caches is false, because the sessions may
+            // have been created using old device_keys.
             for session in changes.sessions {
                 self.session_cache.add(session).await;
             }

From dce9bafd237ee52102f4ebff0a54ce3769f51279 Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Thu, 13 Jun 2024 16:31:05 -0400
Subject: [PATCH 06/19] fix let-else statement

---
 crates/matrix-sdk-crypto/src/olm/session.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/matrix-sdk-crypto/src/olm/session.rs b/crates/matrix-sdk-crypto/src/olm/session.rs
index a4d023ae708..69de5aca011 100644
--- a/crates/matrix-sdk-crypto/src/olm/session.rs
+++ b/crates/matrix-sdk-crypto/src/olm/session.rs
@@ -319,7 +319,7 @@ mod tests {
             .deserialize()
             .unwrap();
 
-        if let ToDeviceEncryptedEventContent::OlmV1Curve25519AesSha2(content) = else {
+        let ToDeviceEncryptedEventContent::OlmV1Curve25519AesSha2(content) = message else {
             panic!("Invalid encrypted event algorithm {}", message.algorithm());
         };
 

From 657204edd1c6c984c7999692690afc5ed483ef4f Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Fri, 14 Jun 2024 17:15:29 -0400
Subject: [PATCH 07/19] always use stored device, and fix test with
 experimental algorithms

---
 crates/matrix-sdk-crypto/src/olm/session.rs   |  5 +++
 .../src/store/memorystore.rs                  | 12 +++++++
 crates/matrix-sdk-crypto/src/store/traits.rs  | 10 ++++++
 .../src/crypto_store/mod.rs                   | 30 +++++++----------
 crates/matrix-sdk-sqlite/src/crypto_store.rs  | 33 +++++++------------
 5 files changed, 49 insertions(+), 41 deletions(-)

diff --git a/crates/matrix-sdk-crypto/src/olm/session.rs b/crates/matrix-sdk-crypto/src/olm/session.rs
index 69de5aca011..4feb2290ae9 100644
--- a/crates/matrix-sdk-crypto/src/olm/session.rs
+++ b/crates/matrix-sdk-crypto/src/olm/session.rs
@@ -319,6 +319,11 @@ mod tests {
             .deserialize()
             .unwrap();
 
+        #[cfg(feature = "experimental-algorithms")]
+        let ToDeviceEncryptedEventContent::OlmV2Curve25519AesSha2(content) = message else {
+            panic!("Invalid encrypted event algorithm {}", message.algorithm());
+        };
+        #[cfg(not(feature = "experimental-algorithms"))]
         let ToDeviceEncryptedEventContent::OlmV1Curve25519AesSha2(content) = message else {
             panic!("Invalid encrypted event algorithm {}", message.algorithm());
         };
diff --git a/crates/matrix-sdk-crypto/src/store/memorystore.rs b/crates/matrix-sdk-crypto/src/store/memorystore.rs
index cace29e1220..845dcd43a91 100644
--- a/crates/matrix-sdk-crypto/src/store/memorystore.rs
+++ b/crates/matrix-sdk-crypto/src/store/memorystore.rs
@@ -474,6 +474,14 @@ impl CryptoStore for MemoryStore {
         Ok(self.devices.user_devices(user_id))
     }
 
+    async fn get_own_device(&self) -> Result<ReadOnlyDevice> {
+        let account = self.load_account()
+            .await?
+            .unwrap();
+
+        Ok(self.devices.get(&account.user_id, &account.device_id).unwrap())
+    }
+
     async fn get_user_identity(&self, user_id: &UserId) -> Result<Option<ReadOnlyUserIdentities>> {
         Ok(self.identities.read().unwrap().get(user_id).cloned())
     }
@@ -1265,6 +1273,10 @@ mod integration_tests {
             self.0.get_user_devices(user_id).await
         }
 
+        async fn get_own_device(&self) -> Result<ReadOnlyDevice, Self::Error> {
+            self.0.get_own_device().await
+        }
+
         async fn get_user_identity(
             &self,
             user_id: &UserId,
diff --git a/crates/matrix-sdk-crypto/src/store/traits.rs b/crates/matrix-sdk-crypto/src/store/traits.rs
index 442b30c2543..2ce396c4b89 100644
--- a/crates/matrix-sdk-crypto/src/store/traits.rs
+++ b/crates/matrix-sdk-crypto/src/store/traits.rs
@@ -203,6 +203,12 @@ pub trait CryptoStore: AsyncTraitDeps {
         user_id: &UserId,
     ) -> Result<HashMap<OwnedDeviceId, ReadOnlyDevice>, Self::Error>;
 
+    /// Get the device for the current client.
+    ///
+    /// Since our own device is set when the store is created, this will always
+    /// return a device (unless there is an error).
+    async fn get_own_device(&self) -> Result<ReadOnlyDevice, Self::Error>;
+
     /// Get the user identity that is attached to the given user id.
     ///
     /// # Arguments
@@ -450,6 +456,10 @@ impl<T: CryptoStore> CryptoStore for EraseCryptoStoreError<T> {
         self.0.get_user_devices(user_id).await.map_err(Into::into)
     }
 
+    async fn get_own_device(&self) -> Result<ReadOnlyDevice> {
+        self.0.get_own_device().await.map_err(Into::into)
+    }
+
     async fn get_user_identity(&self, user_id: &UserId) -> Result<Option<ReadOnlyUserIdentities>> {
         self.0.get_user_identity(user_id).await.map_err(Into::into)
     }
diff --git a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
index 0e916612f83..7c533c5d170 100644
--- a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
+++ b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
@@ -883,26 +883,11 @@ impl_crypto_store! {
     }
 
     async fn get_sessions(&self, sender_key: &str) -> Result<Option<Arc<Mutex<Vec<Session>>>>> {
-        let account_info = self.get_static_account().ok_or(CryptoStoreError::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(CryptoStoreError::AccountUnset))?
-                        .ok_or(CryptoStoreError::AccountUnset)?;
-                    account.device_keys()
-                },
-            };
+            let device_keys = self.get_own_device()
+                .await?
+                .as_device_keys()
+                .clone();
 
             let range = self.serializer.encode_to_range(keys::SESSION, sender_key)?;
             let sessions: Vec<Session> = self
@@ -1133,6 +1118,13 @@ impl_crypto_store! {
             .collect::<HashMap<_, _>>())
     }
 
+    async fn get_own_device(&self) -> Result<ReadOnlyDevice> {
+        let account_info = self.get_static_account().ok_or(CryptoStoreError::AccountUnset)?;
+        Ok(self.get_device(&account_info.user_id, &account_info.device_id)
+           .await?
+           .unwrap())
+    }
+
     async fn get_user_identity(&self, user_id: &UserId) -> Result<Option<ReadOnlyUserIdentities>> {
         Ok(self
             .inner
diff --git a/crates/matrix-sdk-sqlite/src/crypto_store.rs b/crates/matrix-sdk-sqlite/src/crypto_store.rs
index 0096f264cd2..4a9085d2235 100644
--- a/crates/matrix-sdk-sqlite/src/crypto_store.rs
+++ b/crates/matrix-sdk-sqlite/src/crypto_store.rs
@@ -920,29 +920,11 @@ impl CryptoStore for SqliteCryptoStore {
     }
 
     async fn get_sessions(&self, sender_key: &str) -> Result<Option<Arc<Mutex<Vec<Session>>>>> {
-        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 device_keys = self.get_own_device()
+                .await?
+                .as_device_keys()
+                .clone();
 
             let sessions = self
                 .acquire()
@@ -1141,6 +1123,13 @@ impl CryptoStore for SqliteCryptoStore {
             .collect()
     }
 
+    async fn get_own_device(&self) -> Result<ReadOnlyDevice> {
+        let account_info = self.get_static_account().ok_or(Error::AccountUnset)?;
+        Ok(self.get_device(&account_info.user_id, &account_info.device_id)
+           .await?
+           .unwrap())
+    }
+
     async fn get_user_identity(&self, user_id: &UserId) -> Result<Option<ReadOnlyUserIdentities>> {
         let user_id = self.encode_key("identity", user_id.as_bytes());
         Ok(self

From 552e0bc053bc51e4fe478fef587d5d1ee3ccc48e Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Fri, 14 Jun 2024 17:28:44 -0400
Subject: [PATCH 08/19] cargo fmt

---
 crates/matrix-sdk-crypto/src/olm/session.rs       | 6 ++++--
 crates/matrix-sdk-crypto/src/store/memorystore.rs | 4 +---
 crates/matrix-sdk-sqlite/src/crypto_store.rs      | 9 ++-------
 3 files changed, 7 insertions(+), 12 deletions(-)

diff --git a/crates/matrix-sdk-crypto/src/olm/session.rs b/crates/matrix-sdk-crypto/src/olm/session.rs
index 4feb2290ae9..adc9fa9acdd 100644
--- a/crates/matrix-sdk-crypto/src/olm/session.rs
+++ b/crates/matrix-sdk-crypto/src/olm/session.rs
@@ -320,11 +320,13 @@ mod tests {
             .unwrap();
 
         #[cfg(feature = "experimental-algorithms")]
-        let ToDeviceEncryptedEventContent::OlmV2Curve25519AesSha2(content) = message else {
+        let ToDeviceEncryptedEventContent::OlmV2Curve25519AesSha2(content) = message
+        else {
             panic!("Invalid encrypted event algorithm {}", message.algorithm());
         };
         #[cfg(not(feature = "experimental-algorithms"))]
-        let ToDeviceEncryptedEventContent::OlmV1Curve25519AesSha2(content) = message else {
+        let ToDeviceEncryptedEventContent::OlmV1Curve25519AesSha2(content) = message
+        else {
             panic!("Invalid encrypted event algorithm {}", message.algorithm());
         };
 
diff --git a/crates/matrix-sdk-crypto/src/store/memorystore.rs b/crates/matrix-sdk-crypto/src/store/memorystore.rs
index 845dcd43a91..0b465dcbae8 100644
--- a/crates/matrix-sdk-crypto/src/store/memorystore.rs
+++ b/crates/matrix-sdk-crypto/src/store/memorystore.rs
@@ -475,9 +475,7 @@ impl CryptoStore for MemoryStore {
     }
 
     async fn get_own_device(&self) -> Result<ReadOnlyDevice> {
-        let account = self.load_account()
-            .await?
-            .unwrap();
+        let account = self.load_account().await?.unwrap();
 
         Ok(self.devices.get(&account.user_id, &account.device_id).unwrap())
     }
diff --git a/crates/matrix-sdk-sqlite/src/crypto_store.rs b/crates/matrix-sdk-sqlite/src/crypto_store.rs
index 4a9085d2235..525b314dad6 100644
--- a/crates/matrix-sdk-sqlite/src/crypto_store.rs
+++ b/crates/matrix-sdk-sqlite/src/crypto_store.rs
@@ -921,10 +921,7 @@ impl CryptoStore for SqliteCryptoStore {
 
     async fn get_sessions(&self, sender_key: &str) -> Result<Option<Arc<Mutex<Vec<Session>>>>> {
         if self.session_cache.get(sender_key).is_none() {
-            let device_keys = self.get_own_device()
-                .await?
-                .as_device_keys()
-                .clone();
+            let device_keys = self.get_own_device().await?.as_device_keys().clone();
 
             let sessions = self
                 .acquire()
@@ -1125,9 +1122,7 @@ impl CryptoStore for SqliteCryptoStore {
 
     async fn get_own_device(&self) -> Result<ReadOnlyDevice> {
         let account_info = self.get_static_account().ok_or(Error::AccountUnset)?;
-        Ok(self.get_device(&account_info.user_id, &account_info.device_id)
-           .await?
-           .unwrap())
+        Ok(self.get_device(&account_info.user_id, &account_info.device_id).await?.unwrap())
     }
 
     async fn get_user_identity(&self, user_id: &UserId) -> Result<Option<ReadOnlyUserIdentities>> {

From 7fcdd6f1c9b06caae091fd81e0544aa38fc2cfb9 Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Fri, 14 Jun 2024 17:39:36 -0400
Subject: [PATCH 09/19] fix more tests

---
 .../matrix-sdk-crypto/src/store/integration_tests.rs   | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/crates/matrix-sdk-crypto/src/store/integration_tests.rs b/crates/matrix-sdk-crypto/src/store/integration_tests.rs
index fb297f88ef6..d058e8f3c13 100644
--- a/crates/matrix-sdk-crypto/src/store/integration_tests.rs
+++ b/crates/matrix-sdk-crypto/src/store/integration_tests.rs
@@ -217,6 +217,16 @@ macro_rules! cryptostore_integration_tests {
                         })
                         .await
                         .expect("Can't save account");
+                    store
+                        .save_changes(Changes {
+                            devices: DeviceChanges {
+                                new: vec![ReadOnlyDevice::from_account(&account)],
+                                ..Default::default()
+                            },
+                            ..Default::default()
+                        })
+                        .await
+                        .unwrap();
 
                     let changes = Changes { sessions: vec![session.clone()], ..Default::default() };
                     store.save_changes(changes).await.unwrap();

From 7b2a05f14c07a31928e438a47f64adbc092201e3 Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Fri, 14 Jun 2024 19:23:33 -0400
Subject: [PATCH 10/19] add some more comments

---
 crates/matrix-sdk-crypto/src/olm/account.rs | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/crates/matrix-sdk-crypto/src/olm/account.rs b/crates/matrix-sdk-crypto/src/olm/account.rs
index 0aa00e3ce07..304a369c767 100644
--- a/crates/matrix-sdk-crypto/src/olm/account.rs
+++ b/crates/matrix-sdk-crypto/src/olm/account.rs
@@ -902,6 +902,9 @@ impl Account {
     /// created and shared with us.
     ///
     /// * `fallback_used` - Was the one-time key a fallback key.
+    ///
+    /// * `our_device_keys` - Our own `DeviceKeys`, including cross-signing
+    ///   signatures if applicable, for embedding in encrypted messages.
     pub fn create_outbound_session_helper(
         &self,
         config: SessionConfig,
@@ -1026,6 +1029,9 @@ impl Account {
     /// # Arguments
     /// * `their_identity_key` - The other account's identity/curve25519 key.
     ///
+    /// * `our_device_keys` - Our own `DeviceKeys`, including cross-signing
+    ///   signatures if applicable, for embedding in encrypted messages.
+    ///
     /// * `message` - A pre-key Olm message that was sent to us by the other
     /// account.
     pub fn create_inbound_session(

From 392eacc0616802e46f3a58f11aef7aa8d8a967a2 Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Fri, 14 Jun 2024 19:23:40 -0400
Subject: [PATCH 11/19] use the right account

---
 crates/matrix-sdk-crypto/src/olm/account.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/matrix-sdk-crypto/src/olm/account.rs b/crates/matrix-sdk-crypto/src/olm/account.rs
index 304a369c767..f65e10df08d 100644
--- a/crates/matrix-sdk-crypto/src/olm/account.rs
+++ b/crates/matrix-sdk-crypto/src/olm/account.rs
@@ -1113,7 +1113,7 @@ impl Account {
         let other_session = other
             .create_inbound_session(
                 our_device.curve25519_key().unwrap(),
-                self.device_keys(),
+                other.device_keys(),
                 &prekey,
             )
             .unwrap();

From 0b746549bc0f4dff9eb3ee57335423356431319c Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Mon, 17 Jun 2024 16:37:39 -0400
Subject: [PATCH 12/19] simplify logic

---
 .../src/session_manager/sessions.rs           | 47 ++++++++-----------
 1 file changed, 19 insertions(+), 28 deletions(-)

diff --git a/crates/matrix-sdk-crypto/src/session_manager/sessions.rs b/crates/matrix-sdk-crypto/src/session_manager/sessions.rs
index 0b36c16a122..5f513a05001 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, DeviceKeys, EventEncryptionAlgorithm},
+    types::{events::EventType, EventEncryptionAlgorithm},
     ReadOnlyDevice,
 };
 
@@ -514,10 +514,6 @@ impl SessionManager {
         let mut new_sessions: BTreeMap<&UserId, BTreeMap<&DeviceId, SessionInfo>> = BTreeMap::new();
         let mut store_transaction = self.store.transaction().await;
 
-        // Cache of our own device keys, so we only need to fetch it once, if
-        // we need it at all.
-        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 {
@@ -541,28 +537,12 @@ impl SessionManager {
                 };
 
                 let account = store_transaction.account().await?;
-                let device_keys = match our_device_keys {
-                    Some(ref device_keys) => device_keys.clone(),
-                    None => {
-                        // Try to get our own stored device keys.
-                        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 device_keys = self
+                    .store
+                    .get_own_device()
+                    .await?
+                    .as_device_keys()
+                    .clone();
                 let session = match account.create_outbound_session(&device, key_map, device_keys) {
                     Ok(s) => s,
                     Err(e) => {
@@ -657,7 +637,7 @@ mod tests {
         identities::{IdentityManager, ReadOnlyDevice},
         olm::{Account, PrivateCrossSigningIdentity},
         session_manager::GroupSessionCache,
-        store::{CryptoStoreWrapper, MemoryStore, PendingChanges, Store},
+        store::{Changes, CryptoStoreWrapper, DeviceChanges, MemoryStore, PendingChanges, Store},
         verification::VerificationMachine,
     };
 
@@ -714,7 +694,18 @@ mod tests {
         );
 
         let store = Store::new(account.static_data().clone(), identity, store, verification);
+        let device = ReadOnlyDevice::from_account(&account);
         store.save_pending_changes(PendingChanges { account: Some(account) }).await.unwrap();
+        store
+            .save_changes(Changes {
+                devices: DeviceChanges {
+                    new: vec![device],
+                    ..Default::default()
+                },
+                ..Default::default()
+            })
+            .await
+            .unwrap();
 
         let session_cache = GroupSessionCache::new(store.clone());
         let identity_manager = IdentityManager::new(store.clone());

From 91c2d9ab29099271ba935e3dcd871410022b6090 Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Mon, 17 Jun 2024 16:45:51 -0400
Subject: [PATCH 13/19] rust fmt

---
 .../src/session_manager/sessions.rs                  | 12 ++----------
 1 file changed, 2 insertions(+), 10 deletions(-)

diff --git a/crates/matrix-sdk-crypto/src/session_manager/sessions.rs b/crates/matrix-sdk-crypto/src/session_manager/sessions.rs
index 5f513a05001..c367954d316 100644
--- a/crates/matrix-sdk-crypto/src/session_manager/sessions.rs
+++ b/crates/matrix-sdk-crypto/src/session_manager/sessions.rs
@@ -537,12 +537,7 @@ impl SessionManager {
                 };
 
                 let account = store_transaction.account().await?;
-                let device_keys = self
-                    .store
-                    .get_own_device()
-                    .await?
-                    .as_device_keys()
-                    .clone();
+                let device_keys = self.store.get_own_device().await?.as_device_keys().clone();
                 let session = match account.create_outbound_session(&device, key_map, device_keys) {
                     Ok(s) => s,
                     Err(e) => {
@@ -698,10 +693,7 @@ mod tests {
         store.save_pending_changes(PendingChanges { account: Some(account) }).await.unwrap();
         store
             .save_changes(Changes {
-                devices: DeviceChanges {
-                    new: vec![device],
-                    ..Default::default()
-                },
+                devices: DeviceChanges { new: vec![device], ..Default::default() },
                 ..Default::default()
             })
             .await

From 50b9ae7ff0a2825f6a25ff2d32f2438fa7bc63f4 Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Wed, 19 Jun 2024 20:02:06 -0400
Subject: [PATCH 14/19] address some review comments

---
 crates/matrix-sdk-crypto/Cargo.toml           |  1 +
 crates/matrix-sdk-crypto/src/error.rs         | 13 ++++
 crates/matrix-sdk-crypto/src/machine.rs       |  7 +++
 crates/matrix-sdk-crypto/src/olm/account.rs   | 20 ++----
 crates/matrix-sdk-crypto/src/olm/session.rs   | 63 ++++++++++---------
 .../src/crypto_store/mod.rs                   |  3 +-
 crates/matrix-sdk-sqlite/src/crypto_store.rs  |  3 +-
 7 files changed, 62 insertions(+), 48 deletions(-)

diff --git a/crates/matrix-sdk-crypto/Cargo.toml b/crates/matrix-sdk-crypto/Cargo.toml
index e67f338cc23..1049301724f 100644
--- a/crates/matrix-sdk-crypto/Cargo.toml
+++ b/crates/matrix-sdk-crypto/Cargo.toml
@@ -30,6 +30,7 @@ testing = ["dep:http"]
 [dependencies]
 aes = "0.8.1"
 as_variant = { workspace = true }
+assert_matches2 = { workspace = true }
 async-trait = { workspace = true }
 bs58 = { version = "0.5.0" }
 byteorder = { workspace = true }
diff --git a/crates/matrix-sdk-crypto/src/error.rs b/crates/matrix-sdk-crypto/src/error.rs
index 3e4e0bcf836..90f5d86b01f 100644
--- a/crates/matrix-sdk-crypto/src/error.rs
+++ b/crates/matrix-sdk-crypto/src/error.rs
@@ -165,6 +165,19 @@ pub enum EventError {
     MismatchedRoom(OwnedRoomId, Option<OwnedRoomId>),
 }
 
+/// Error type describing different errors that can happen when we create an
+/// Olm session from a pickle.
+#[derive(Error, Debug)]
+pub enum SessionPickleError {
+    /// The device keys are missing the signing key
+    #[error("the device keys are missing the signing key")]
+    MissingSigningKey,
+
+    /// The device keys are missing the identity key
+    #[error("the device keys are missing the identity key")]
+    MissingIdentityKey,
+}
+
 /// Error type describing different errors that happen when we check or create
 /// signatures for a Matrix JSON object.
 #[derive(Error, Debug)]
diff --git a/crates/matrix-sdk-crypto/src/machine.rs b/crates/matrix-sdk-crypto/src/machine.rs
index 245f571e011..e9c7d8884f9 100644
--- a/crates/matrix-sdk-crypto/src/machine.rs
+++ b/crates/matrix-sdk-crypto/src/machine.rs
@@ -170,7 +170,14 @@ impl OlmMachine {
         let static_account = account.static_data().clone();
 
         let store = Arc::new(CryptoStoreWrapper::new(self.user_id(), MemoryStore::new()));
+        let device = ReadOnlyDevice::from_account(&account);
         store.save_pending_changes(PendingChanges { account: Some(account) }).await?;
+        store
+            .save_changes(Changes {
+                devices: DeviceChanges { new: vec![device], ..Default::default() },
+                ..Default::default()
+            })
+            .await?;
 
         Ok(Self::new_helper(
             device_id,
diff --git a/crates/matrix-sdk-crypto/src/olm/account.rs b/crates/matrix-sdk-crypto/src/olm/account.rs
index f65e10df08d..70096d7f341 100644
--- a/crates/matrix-sdk-crypto/src/olm/account.rs
+++ b/crates/matrix-sdk-crypto/src/olm/account.rs
@@ -922,7 +922,7 @@ impl Account {
             inner: Arc::new(Mutex::new(session)),
             session_id: session_id.into(),
             sender_key: identity_key,
-            device_keys: our_device_keys,
+            our_device_keys,
             created_using_fallback_key: fallback_used,
             creation_time: now,
             last_use_time: now,
@@ -1053,7 +1053,7 @@ impl Account {
             inner: Arc::new(Mutex::new(result.session)),
             session_id: session_id.into(),
             sender_key: their_identity_key,
-            device_keys: our_device_keys,
+            our_device_keys,
             created_using_fallback_key: false,
             creation_time: now,
             last_use_time: now,
@@ -1301,25 +1301,13 @@ impl Account {
                         );
 
                         return Err(OlmError::SessionWedged(
-                            session.device_keys.user_id.to_owned(),
+                            session.our_device_keys.user_id.to_owned(),
                             session.sender_key(),
                         ));
                     }
                 }
 
-                // We didn't find a matching session; try to create a new session.
-                // 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 our device keys 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 device_keys = store.get_own_device().await?.as_device_keys().clone();
                 let result =
                     match self.create_inbound_session(sender_key, device_keys, prekey_message) {
                         Ok(r) => r,
diff --git a/crates/matrix-sdk-crypto/src/olm/session.rs b/crates/matrix-sdk-crypto/src/olm/session.rs
index adc9fa9acdd..379f59ea7bb 100644
--- a/crates/matrix-sdk-crypto/src/olm/session.rs
+++ b/crates/matrix-sdk-crypto/src/olm/session.rs
@@ -27,7 +27,7 @@ use vodozemac::{
 #[cfg(feature = "experimental-algorithms")]
 use crate::types::events::room::encrypted::OlmV2Curve25519AesSha2Content;
 use crate::{
-    error::{EventError, OlmResult},
+    error::{EventError, OlmResult, SessionPickleError},
     types::{
         events::room::encrypted::{OlmV1Curve25519AesSha2Content, ToDeviceEncryptedEventContent},
         DeviceKeys, EventEncryptionAlgorithm,
@@ -45,8 +45,8 @@ pub struct Session {
     pub session_id: Arc<str>,
     /// The Key of the sender
     pub sender_key: Curve25519PublicKey,
-    /// The signed device keys
-    pub device_keys: DeviceKeys,
+    /// Our own signed device keys
+    pub our_device_keys: DeviceKeys,
     /// Has this been created using the fallback key
     pub created_using_fallback_key: bool,
     /// When the session was created
@@ -151,12 +151,12 @@ impl Session {
                 recipient_device.ed25519_key().ok_or(EventError::MissingSigningKey)?;
 
             let payload = json!({
-                "sender": &self.device_keys.user_id,
-                "sender_device": &self.device_keys.device_id,
+                "sender": &self.our_device_keys.user_id,
+                "sender_device": &self.our_device_keys.device_id,
                 "keys": {
-                    "ed25519": self.device_keys.ed25519_key().unwrap().to_base64(),
+                    "ed25519": self.our_device_keys.ed25519_key().expect("Device doesn't have ed25519 key").to_base64(),
                 },
-                "device_keys": self.device_keys,
+                "device_keys": self.our_device_keys,
                 "recipient": recipient_device.user_id(),
                 "recipient_keys": {
                     "ed25519": recipient_signing_key.to_base64(),
@@ -174,14 +174,20 @@ impl Session {
             EventEncryptionAlgorithm::OlmV1Curve25519AesSha2 => OlmV1Curve25519AesSha2Content {
                 ciphertext,
                 recipient_key: self.sender_key,
-                sender_key: self.device_keys.curve25519_key().unwrap(),
+                sender_key: self
+                    .our_device_keys
+                    .curve25519_key()
+                    .expect("Device doesn't have curve25519 key"),
                 message_id,
             }
             .into(),
             #[cfg(feature = "experimental-algorithms")]
             EventEncryptionAlgorithm::OlmV2Curve25519AesSha2 => OlmV2Curve25519AesSha2Content {
                 ciphertext,
-                sender_key: self.device_keys.curve25519_key().unwrap(),
+                sender_key: self
+                    .device_keys
+                    .curve25519_key()
+                    .expect("Device doesn't have curve25519 key"),
                 message_id,
             }
             .into(),
@@ -223,30 +229,32 @@ impl Session {
     ///
     /// # Arguments
     ///
-    /// * `user_id` - Our own user id that the session belongs to.
-    ///
-    /// * `device_id` - Our own device ID that the session belongs to.
-    ///
-    /// * `our_identity_keys` - An clone of the Arc to our own identity keys.
+    /// * `our_device_keys` - Our own signed device keys.
     ///
     /// * `pickle` - The pickled version of the `Session`.
-    ///
-    /// * `pickle_mode` - The mode that was used to pickle the session, either
-    /// an unencrypted mode or an encrypted using passphrase.
-    pub fn from_pickle(device_keys: DeviceKeys, pickle: PickledSession) -> Self {
-        // FIXME: assert that device_keys has curve25519 and ed25519 keys
+    pub fn from_pickle(
+        our_device_keys: DeviceKeys,
+        pickle: PickledSession,
+    ) -> Result<Self, SessionPickleError> {
+        if our_device_keys.curve25519_key().is_none() {
+            return Err(SessionPickleError::MissingIdentityKey);
+        }
+        if our_device_keys.ed25519_key().is_none() {
+            return Err(SessionPickleError::MissingSigningKey);
+        }
+
         let session: vodozemac::olm::Session = pickle.pickle.into();
         let session_id = session.session_id();
 
-        Session {
+        Ok(Session {
             inner: Arc::new(Mutex::new(session)),
             session_id: session_id.into(),
             created_using_fallback_key: pickle.created_using_fallback_key,
             sender_key: pickle.sender_key,
-            device_keys,
+            our_device_keys,
             creation_time: pickle.creation_time,
             last_use_time: pickle.last_use_time,
-        }
+        })
     }
 }
 
@@ -278,6 +286,7 @@ pub struct PickledSession {
 
 #[cfg(test)]
 mod tests {
+    use assert_matches2::assert_let;
     use matrix_sdk_test::async_test;
     use ruma::{device_id, user_id};
     use serde_json::{self, Value};
@@ -320,15 +329,9 @@ mod tests {
             .unwrap();
 
         #[cfg(feature = "experimental-algorithms")]
-        let ToDeviceEncryptedEventContent::OlmV2Curve25519AesSha2(content) = message
-        else {
-            panic!("Invalid encrypted event algorithm {}", message.algorithm());
-        };
+        assert_let!(ToDeviceEncryptedEventContent::OlmV2Curve25519AesSha2(content) = message);
         #[cfg(not(feature = "experimental-algorithms"))]
-        let ToDeviceEncryptedEventContent::OlmV1Curve25519AesSha2(content) = message
-        else {
-            panic!("Invalid encrypted event algorithm {}", message.algorithm());
-        };
+        assert_let!(ToDeviceEncryptedEventContent::OlmV1Curve25519AesSha2(content) = message);
 
         let prekey = if let OlmMessage::PreKey(m) = content.ciphertext {
             m
diff --git a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
index 7c533c5d170..243a7997711 100644
--- a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
+++ b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
@@ -902,8 +902,9 @@ impl_crypto_store! {
                         device_keys.clone(),
                         p,
                     )
+                        .map_err(|_| IndexeddbCryptoStoreError::CryptoStoreError(CryptoStoreError::AccountUnset))
                 }))
-                .collect::<Vec<Session>>();
+                .collect::<Result<Vec<Session>>>()?;
 
             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 525b314dad6..426446d5cc2 100644
--- a/crates/matrix-sdk-sqlite/src/crypto_store.rs
+++ b/crates/matrix-sdk-sqlite/src/crypto_store.rs
@@ -931,7 +931,8 @@ impl CryptoStore for SqliteCryptoStore {
                 .into_iter()
                 .map(|bytes| {
                     let pickle = self.deserialize_value(&bytes)?;
-                    Ok(Session::from_pickle(device_keys.clone(), pickle))
+                    Session::from_pickle(device_keys.clone(), pickle)
+                        .map_err(|_| Error::AccountUnset)
                 })
                 .collect::<Result<_>>()?;
 

From 6dfe86496bd62e530889ddd4206b2352dc7310a8 Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Wed, 19 Jun 2024 21:03:21 -0400
Subject: [PATCH 15/19] update binding

---
 bindings/matrix-sdk-crypto-ffi/src/lib.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/bindings/matrix-sdk-crypto-ffi/src/lib.rs b/bindings/matrix-sdk-crypto-ffi/src/lib.rs
index f0020f3e88a..947d5f1351e 100644
--- a/bindings/matrix-sdk-crypto-ffi/src/lib.rs
+++ b/bindings/matrix-sdk-crypto-ffi/src/lib.rs
@@ -470,7 +470,7 @@ fn collect_sessions(
             last_use_time,
         };
 
-        let session = Session::from_pickle(device_keys.clone(), pickle);
+        let session = Session::from_pickle(device_keys.clone(), pickle)?;
 
         sessions.push(session);
         processed_steps += 1;

From 447b212494fa700a928d142bf0f33d9cc4544677 Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Wed, 19 Jun 2024 21:12:45 -0400
Subject: [PATCH 16/19] update field name

---
 crates/matrix-sdk-crypto/src/olm/session.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/matrix-sdk-crypto/src/olm/session.rs b/crates/matrix-sdk-crypto/src/olm/session.rs
index 379f59ea7bb..b8a8c401215 100644
--- a/crates/matrix-sdk-crypto/src/olm/session.rs
+++ b/crates/matrix-sdk-crypto/src/olm/session.rs
@@ -185,7 +185,7 @@ impl Session {
             EventEncryptionAlgorithm::OlmV2Curve25519AesSha2 => OlmV2Curve25519AesSha2Content {
                 ciphertext,
                 sender_key: self
-                    .device_keys
+                    .our_device_keys
                     .curve25519_key()
                     .expect("Device doesn't have curve25519 key"),
                 message_id,

From 2d660419e5280ccc791206ce59d1b4e79d34cb9b Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Tue, 25 Jun 2024 17:33:33 -0400
Subject: [PATCH 17/19] fetch device keys and sessions in one query

---
 crates/matrix-sdk-sqlite/src/crypto_store.rs | 40 ++++++++++++++++----
 1 file changed, 33 insertions(+), 7 deletions(-)

diff --git a/crates/matrix-sdk-sqlite/src/crypto_store.rs b/crates/matrix-sdk-sqlite/src/crypto_store.rs
index 426446d5cc2..5f1156e8e28 100644
--- a/crates/matrix-sdk-sqlite/src/crypto_store.rs
+++ b/crates/matrix-sdk-sqlite/src/crypto_store.rs
@@ -447,10 +447,27 @@ impl SqliteConnectionExt for rusqlite::Connection {
 
 #[async_trait]
 trait SqliteObjectCryptoStoreExt: SqliteObjectExt {
-    async fn get_sessions_for_sender_key(&self, sender_key: Key) -> Result<Vec<Vec<u8>>> {
+    async fn get_sessions_for_sender_key(
+        &self,
+        sender_key: Key,
+        user_id: Key,
+        device_id: Key,
+    ) -> Result<(Vec<u8>, Vec<Vec<u8>>)> {
         Ok(self
-            .prepare("SELECT data FROM session WHERE sender_key = ?", |mut stmt| {
-                stmt.query((sender_key,))?.mapped(|row| row.get(0)).collect()
+            .prepare("SELECT 1, data FROM session WHERE sender_key = ? UNION SELECT 0, data FROM device WHERE user_id = ? AND device_id = ?", |mut stmt| {
+                let result = stmt.query((sender_key, user_id, device_id))?
+                    .mapped(|row| Ok((row.get(0)?, row.get(1)?)))
+                    .collect::<rusqlite::Result<Vec<(u8, Vec<u8>)>>>()?;
+                let mut sessions: Vec<Vec<u8>> = Vec::new();
+                let mut device: Option<Vec<u8>> = None;
+                for (rowtype, data) in result {
+                    if rowtype == 1 {
+                        sessions.push(data);
+                    } else {
+                        device = Some(data);
+                    }
+                }
+                Ok((device.expect("Our own device key is not stored"), sessions))
             })
             .await?)
     }
@@ -921,13 +938,22 @@ impl CryptoStore for SqliteCryptoStore {
 
     async fn get_sessions(&self, sender_key: &str) -> Result<Option<Arc<Mutex<Vec<Session>>>>> {
         if self.session_cache.get(sender_key).is_none() {
-            let device_keys = self.get_own_device().await?.as_device_keys().clone();
+            let account_info = self.get_static_account().ok_or(Error::AccountUnset)?;
 
-            let sessions = self
+            let (device_keys_data, sessions_data) = self
                 .acquire()
                 .await?
-                .get_sessions_for_sender_key(self.encode_key("session", sender_key.as_bytes()))
-                .await?
+                .get_sessions_for_sender_key(
+                    self.encode_key("session", sender_key.as_bytes()),
+                    self.encode_key("device", account_info.user_id.as_bytes()),
+                    self.encode_key("device", account_info.device_id.as_bytes()),
+                )
+                .await?;
+
+            let read_only_device = self.deserialize_value::<ReadOnlyDevice>(&device_keys_data)?;
+            let device_keys = read_only_device.as_device_keys();
+
+            let sessions = sessions_data
                 .into_iter()
                 .map(|bytes| {
                     let pickle = self.deserialize_value(&bytes)?;

From 0b1d2212f811b2c51567cdd02cab902ff766e142 Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Wed, 26 Jun 2024 11:26:52 -0400
Subject: [PATCH 18/19] Revert "fetch device keys and sessions in one query"

This reverts commit 2d660419e5280ccc791206ce59d1b4e79d34cb9b.

Since it doesn't make enough of a performance difference.
---
 crates/matrix-sdk-sqlite/src/crypto_store.rs | 40 ++++----------------
 1 file changed, 7 insertions(+), 33 deletions(-)

diff --git a/crates/matrix-sdk-sqlite/src/crypto_store.rs b/crates/matrix-sdk-sqlite/src/crypto_store.rs
index 5f1156e8e28..426446d5cc2 100644
--- a/crates/matrix-sdk-sqlite/src/crypto_store.rs
+++ b/crates/matrix-sdk-sqlite/src/crypto_store.rs
@@ -447,27 +447,10 @@ impl SqliteConnectionExt for rusqlite::Connection {
 
 #[async_trait]
 trait SqliteObjectCryptoStoreExt: SqliteObjectExt {
-    async fn get_sessions_for_sender_key(
-        &self,
-        sender_key: Key,
-        user_id: Key,
-        device_id: Key,
-    ) -> Result<(Vec<u8>, Vec<Vec<u8>>)> {
+    async fn get_sessions_for_sender_key(&self, sender_key: Key) -> Result<Vec<Vec<u8>>> {
         Ok(self
-            .prepare("SELECT 1, data FROM session WHERE sender_key = ? UNION SELECT 0, data FROM device WHERE user_id = ? AND device_id = ?", |mut stmt| {
-                let result = stmt.query((sender_key, user_id, device_id))?
-                    .mapped(|row| Ok((row.get(0)?, row.get(1)?)))
-                    .collect::<rusqlite::Result<Vec<(u8, Vec<u8>)>>>()?;
-                let mut sessions: Vec<Vec<u8>> = Vec::new();
-                let mut device: Option<Vec<u8>> = None;
-                for (rowtype, data) in result {
-                    if rowtype == 1 {
-                        sessions.push(data);
-                    } else {
-                        device = Some(data);
-                    }
-                }
-                Ok((device.expect("Our own device key is not stored"), sessions))
+            .prepare("SELECT data FROM session WHERE sender_key = ?", |mut stmt| {
+                stmt.query((sender_key,))?.mapped(|row| row.get(0)).collect()
             })
             .await?)
     }
@@ -938,22 +921,13 @@ impl CryptoStore for SqliteCryptoStore {
 
     async fn get_sessions(&self, sender_key: &str) -> Result<Option<Arc<Mutex<Vec<Session>>>>> {
         if self.session_cache.get(sender_key).is_none() {
-            let account_info = self.get_static_account().ok_or(Error::AccountUnset)?;
+            let device_keys = self.get_own_device().await?.as_device_keys().clone();
 
-            let (device_keys_data, sessions_data) = self
+            let sessions = self
                 .acquire()
                 .await?
-                .get_sessions_for_sender_key(
-                    self.encode_key("session", sender_key.as_bytes()),
-                    self.encode_key("device", account_info.user_id.as_bytes()),
-                    self.encode_key("device", account_info.device_id.as_bytes()),
-                )
-                .await?;
-
-            let read_only_device = self.deserialize_value::<ReadOnlyDevice>(&device_keys_data)?;
-            let device_keys = read_only_device.as_device_keys();
-
-            let sessions = sessions_data
+                .get_sessions_for_sender_key(self.encode_key("session", sender_key.as_bytes()))
+                .await?
                 .into_iter()
                 .map(|bytes| {
                     let pickle = self.deserialize_value(&bytes)?;

From b42acd04d3f5ce3dbaba07be0415bf298cfebd9c Mon Sep 17 00:00:00 2001
From: Hubert Chathi <hubertc@matrix.org>
Date: Wed, 26 Jun 2024 12:08:28 -0400
Subject: [PATCH 19/19] rename error struct and add indexeddb test

---
 crates/matrix-sdk-crypto/src/error.rs         |  2 +-
 crates/matrix-sdk-crypto/src/olm/session.rs   |  8 ++--
 .../src/crypto_store/mod.rs                   | 44 ++++++++++++++++++-
 3 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/crates/matrix-sdk-crypto/src/error.rs b/crates/matrix-sdk-crypto/src/error.rs
index 90f5d86b01f..cedb9712da5 100644
--- a/crates/matrix-sdk-crypto/src/error.rs
+++ b/crates/matrix-sdk-crypto/src/error.rs
@@ -168,7 +168,7 @@ pub enum EventError {
 /// Error type describing different errors that can happen when we create an
 /// Olm session from a pickle.
 #[derive(Error, Debug)]
-pub enum SessionPickleError {
+pub enum SessionUnpickleError {
     /// The device keys are missing the signing key
     #[error("the device keys are missing the signing key")]
     MissingSigningKey,
diff --git a/crates/matrix-sdk-crypto/src/olm/session.rs b/crates/matrix-sdk-crypto/src/olm/session.rs
index b8a8c401215..34281c6d2d1 100644
--- a/crates/matrix-sdk-crypto/src/olm/session.rs
+++ b/crates/matrix-sdk-crypto/src/olm/session.rs
@@ -27,7 +27,7 @@ use vodozemac::{
 #[cfg(feature = "experimental-algorithms")]
 use crate::types::events::room::encrypted::OlmV2Curve25519AesSha2Content;
 use crate::{
-    error::{EventError, OlmResult, SessionPickleError},
+    error::{EventError, OlmResult, SessionUnpickleError},
     types::{
         events::room::encrypted::{OlmV1Curve25519AesSha2Content, ToDeviceEncryptedEventContent},
         DeviceKeys, EventEncryptionAlgorithm,
@@ -235,12 +235,12 @@ impl Session {
     pub fn from_pickle(
         our_device_keys: DeviceKeys,
         pickle: PickledSession,
-    ) -> Result<Self, SessionPickleError> {
+    ) -> Result<Self, SessionUnpickleError> {
         if our_device_keys.curve25519_key().is_none() {
-            return Err(SessionPickleError::MissingIdentityKey);
+            return Err(SessionUnpickleError::MissingIdentityKey);
         }
         if our_device_keys.ed25519_key().is_none() {
-            return Err(SessionPickleError::MissingSigningKey);
+            return Err(SessionUnpickleError::MissingSigningKey);
         }
 
         let session: vodozemac::olm::Session = pickle.pickle.into();
diff --git a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
index 243a7997711..f14c01930fb 100644
--- a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
+++ b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
@@ -1724,7 +1724,12 @@ mod wasm_unit_tests {
 
 #[cfg(all(test, target_arch = "wasm32"))]
 mod tests {
-    use matrix_sdk_crypto::cryptostore_integration_tests;
+    use matrix_sdk_crypto::{
+        cryptostore_integration_tests,
+        store::{Changes, CryptoStore as _, DeviceChanges, PendingChanges},
+        ReadOnlyDevice,
+    };
+    use matrix_sdk_test::async_test;
 
     use super::IndexeddbCryptoStore;
 
@@ -1740,6 +1745,43 @@ mod tests {
                 .expect("Can't create store without passphrase"),
         }
     }
+
+    #[async_test]
+    async fn cache_cleared_after_device_update() {
+        let store = get_store("cache_cleared_after_device_update", None).await;
+        // Given we created a session and saved it in the store
+        let (account, session) = cryptostore_integration_tests::get_account_and_session().await;
+        let sender_key = session.sender_key.to_base64();
+
+        store
+            .save_pending_changes(PendingChanges { account: Some(account.deep_clone()) })
+            .await
+            .expect("Can't save account");
+
+        let changes = Changes { sessions: vec![session.clone()], ..Default::default() };
+        store.save_changes(changes).await.unwrap();
+
+        store.session_cache.get(&sender_key).expect("We should have a session");
+
+        // When we save a new version of our device keys
+        store
+            .save_changes(Changes {
+                devices: DeviceChanges {
+                    new: vec![ReadOnlyDevice::from_account(&account)],
+                    ..Default::default()
+                },
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        // Then the session is no longer in the cache
+        assert!(
+            store.session_cache.get(&sender_key).is_none(),
+            "Session should not be in the cache!"
+        );
+    }
+
     cryptostore_integration_tests!();
 }