Embed device keys in Olm-encrypted messages (#3517)

This patch implements MSC4147[1].

Signed-off-by: Hubert Chathi <hubertc@matrix.org>

[1]: matrix-org/matrix-spec-proposals#4147

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,
+    },
     CollectStrategy, 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/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 }

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 SessionUnpickleError {
+    /// 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)]

crates/matrix-sdk-crypto/src/identities/device.rs

@@ -871,7 +871,8 @@ impl ReadOnlyDevice {
         }
     }
 
-    pub(crate) fn as_device_keys(&self) -> &DeviceKeys {
+    /// Return the device keys
+    pub fn as_device_keys(&self) -> &DeviceKeys {
         &self.inner
     }
 

crates/matrix-sdk-crypto/src/machine.rs

@@ -173,7 +173,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,

crates/matrix-sdk-crypto/src/olm/account.rs

@@ -905,25 +905,27 @@ impl Account {
     ///   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,
         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();
 
         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,
+            our_device_keys,
             created_using_fallback_key: fallback_used,
             creation_time: now,
             last_use_time: now,
@@ -978,11 +980,15 @@ 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,
         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)?;
 
@@ -1012,6 +1018,7 @@ impl Account {
                     identity_key,
                     one_time_key,
                     is_fallback,
+                    our_device_keys,
                 ))
             }
         }
@@ -1026,11 +1033,15 @@ impl Account {
     ///
     /// * `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(
         &mut self,
         their_identity_key: Curve25519PublicKey,
+        our_device_keys: DeviceKeys,
         message: &PreKeyMessage,
     ) -> Result<InboundCreationResult, SessionCreationError> {
         Span::current().record("session_id", debug(message.session_id()));
@@ -1043,12 +1054,10 @@ impl Account {
         debug!(session=?result.session, "Decrypted an Olm message from a new Olm session");
 
         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,
+            our_device_keys,
             created_using_fallback_key: false,
             creation_time: now,
             last_use_time: now,
@@ -1072,7 +1081,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).unwrap();
+        let mut our_session =
+            self.create_outbound_session(&device, &one_time_map, self.device_keys()).unwrap();
 
         other.mark_keys_as_published();
 
@@ -1104,8 +1114,13 @@ 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(),
+                other.device_keys(),
+                &prekey,
+            )
+            .unwrap();
 
         (our_session, other_session.session)
     }
@@ -1290,20 +1305,23 @@ impl Account {
                         );
 
                         return Err(OlmError::SessionWedged(
-                            session.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.
-                let result = match self.create_inbound_session(sender_key, 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));
-                    }
-                };
+                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,
+                        Err(e) => {
+                            warn!(
+                                "Failed to create a new Olm session from a pre-key message: {e:?}"
+                            );
+                            return Err(OlmError::SessionWedged(sender.to_owned(), sender_key));
+                        }
+                    };
 
                 // We need to add the new session to the session cache, otherwise
                 // we might try to create the same session again.

crates/matrix-sdk-crypto/src/olm/mod.rs

@@ -93,6 +93,7 @@ pub(crate) mod tests {
             sender_key,
             one_time_key,
             false,
+            alice.device_keys(),
         );
 
         (alice, session)
@@ -144,6 +145,7 @@ pub(crate) mod tests {
             alice_keys.curve25519,
             one_time_key,
             false,
+            bob.device_keys(),
         );
 
         let plaintext = "Hello world";
@@ -156,7 +158,9 @@ 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, alice.device_keys(), &prekey_message)
+            .unwrap();
 
         assert_eq!(bob_session.session_id(), result.session.session_id());