Embed device keys in Olm-encrypted messages

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;

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,

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,16 +311,15 @@ 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
             .unwrap()
             .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"));
     }

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