crypto: Add new method Store::merge_received_group_session

Add a method which can be used to merge a received `InboundGroupSession` into
whatever we find in the store.

crates/matrix-sdk-crypto/src/olm/group_sessions/inbound.rs

@@ -340,6 +340,36 @@ impl InboundGroupSession {
         Self::try_from(exported_session)
     }
 
+    /// Create a new [`InboundGroupSession`] which is a copy of this one, except
+    /// that its Megolm ratchet is replaced with a copy of that from another
+    /// [`InboundGroupSession`].
+    ///
+    /// This can be useful, for example, when we receive a new copy of the room
+    /// key, but at an earlier ratchet index.
+    ///
+    /// # Panics
+    ///
+    /// If the two sessions are for different room IDs, or have different
+    /// session IDs, this function will panic. It is up to the caller to ensure
+    /// that it only attempts to merge related sessions.
+    pub(crate) fn with_ratchet(mut self, other: &InboundGroupSession) -> Self {
+        if self.session_id != other.session_id {
+            panic!(
+                "Attempt to merge Megolm sessions with different session IDs: {} vs {}",
+                self.session_id, other.session_id
+            );
+        }
+        if self.room_id != other.room_id {
+            panic!(
+                "Attempt to merge Megolm sessions with different room IDs: {} vs {}",
+                self.room_id, other.room_id,
+            );
+        }
+        self.inner = other.inner.clone();
+        self.first_known_index = other.first_known_index;
+        self
+    }
+
     /// Convert the [`InboundGroupSession`] into a
     /// [`PickledInboundGroupSession`] which can be serialized.
     pub async fn pickle(&self) -> PickledInboundGroupSession {

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

@@ -657,6 +657,94 @@ impl Store {
         })
     }
 
+    /// Given an `InboundGroupSession` which we have just received, see if we
+    /// have a matching session already in the store, and determine how to
+    /// handle it.
+    ///
+    /// If the store already has everything we can gather from the new session,
+    /// returns `None`. Otherwise, returns a merged session which should be
+    /// persisted to the store.
+    pub(crate) async fn merge_received_group_session(
+        &self,
+        session: InboundGroupSession,
+    ) -> Result<Option<InboundGroupSession>> {
+        let old_session = self
+            .inner
+            .store
+            .get_inbound_group_session(session.room_id(), session.session_id())
+            .await?;
+
+        // If there is no old session, just use the new session.
+        let Some(old_session) = old_session else {
+            info!("Received a new megolm room key");
+            return Ok(Some(session));
+        };
+
+        let index_comparison = session.compare_ratchet(&old_session).await;
+        let trust_level_comparison =
+            session.sender_data.compare_trust_level(&old_session.sender_data);
+
+        let result = match (index_comparison, trust_level_comparison) {
+            (SessionOrdering::Unconnected, _) => {
+                // If this happens, it means that we have two sessions purporting to have the
+                // same session id, but where the ratchets do not match up.
+                // In other words, someone is playing silly buggers.
+                warn!("Received a group session with an ratchet that does not connect to the one in the store, discarding");
+                None
+            }
+
+            (SessionOrdering::Better, std::cmp::Ordering::Greater)
+            | (SessionOrdering::Better, std::cmp::Ordering::Equal)
+            | (SessionOrdering::Equal, std::cmp::Ordering::Greater) => {
+                // The new session is unambiguously better than what we have in the store.
+                info!(
+                    ?index_comparison,
+                    ?trust_level_comparison,
+                    "Received a megolm room key that we have a worse version of, merging"
+                );
+                Some(session)
+            }
+
+            (SessionOrdering::Worse, std::cmp::Ordering::Less)
+            | (SessionOrdering::Worse, std::cmp::Ordering::Equal)
+            | (SessionOrdering::Equal, std::cmp::Ordering::Less) => {
+                // The new session is unambiguously worse than the one we have in the store.
+                warn!(
+                    ?index_comparison,
+                    ?trust_level_comparison,
+                    "Received a megolm room key that we already have a better version \
+                     of, discarding"
+                );
+                None
+            }
+
+            (SessionOrdering::Equal, std::cmp::Ordering::Equal) => {
+                // The new session is the same as what we have.
+                info!("Received a megolm room key that we already have, discarding");
+                None
+            }
+
+            (SessionOrdering::Better, std::cmp::Ordering::Less) => {
+                // We need to take the ratchet from the new session, and the
+                // sender data from the old session.
+                info!("Upgrading a previously-received megolm session with new ratchet");
+                let result = old_session.with_ratchet(&session);
+                // We'll need to back it up again.
+                result.reset_backup_state();
+                Some(result)
+            }
+
+            (SessionOrdering::Worse, std::cmp::Ordering::Greater) => {
+                // We need to take the ratchet from the old session, and the
+                // sender data from the new session.
+                info!("Upgrading a previously-received megolm session with new sender data");
+                Some(session.with_ratchet(&old_session))
+            }
+        };
+
+        Ok(result)
+    }
+
     #[cfg(test)]
     /// Testing helper to allow to save only a set of devices
     pub(crate) async fn save_device_data(&self, devices: &[DeviceData]) -> Result<()> {
@@ -1799,9 +1887,9 @@ impl matrix_sdk_common::cross_process_lock::TryLock for LockableCryptoStore {
 
 #[cfg(test)]
 mod tests {
-    use std::pin::pin;
+    use std::{collections::BTreeMap, pin::pin};
 
-    use assert_matches2::assert_matches;
+    use assert_matches2::{assert_let, assert_matches};
     use futures_util::StreamExt;
     use insta::{_macro_support::Content, assert_json_snapshot, internals::ContentPath};
     use matrix_sdk_test::async_test;
@@ -1813,7 +1901,7 @@ mod tests {
         user_id, RoomId,
     };
     use serde_json::json;
-    use vodozemac::megolm::SessionKey;
+    use vodozemac::{megolm::SessionKey, Ed25519Keypair};
 
     use crate::{
         machine::test_helpers::get_machine_pair,
@@ -1826,9 +1914,144 @@ mod tests {
             },
             EventEncryptionAlgorithm,
         },
-        OlmMachine,
+        Account, OlmMachine,
     };
 
+    #[async_test]
+    async fn test_merge_received_group_session() {
+        let alice_account = Account::with_device_id(user_id!("@a:s.co"), device_id!("ABC"));
+        let bob = OlmMachine::new(user_id!("@b:s.co"), device_id!("DEF")).await;
+
+        let room_id = room_id!("!test:localhost");
+
+        let megolm_signing_key = Ed25519Keypair::new();
+        let inbound = make_inbound_group_session(&alice_account, &megolm_signing_key, room_id);
+
+        // Bob already knows about the session, at index 5, with the device keys.
+        let mut inbound_at_index_5 =
+            InboundGroupSession::from_export(&inbound.export_at_index(5).await).unwrap();
+        inbound_at_index_5.sender_data = inbound.sender_data.clone();
+        bob.store().save_inbound_group_sessions(&[inbound_at_index_5.clone()]).await.unwrap();
+
+        // No changes if we get a disconnected session.
+        let disconnected = make_inbound_group_session(&alice_account, &megolm_signing_key, room_id);
+        assert_eq!(bob.store().merge_received_group_session(disconnected).await.unwrap(), None);
+
+        // No changes needed when we receive a worse copy of the session
+        let mut worse =
+            InboundGroupSession::from_export(&inbound.export_at_index(10).await).unwrap();
+        worse.sender_data = inbound.sender_data.clone();
+        assert_eq!(bob.store().merge_received_group_session(worse).await.unwrap(), None);
+
+        // Nor when we receive an exact copy of what we already have
+        let mut copy = InboundGroupSession::from_pickle(inbound_at_index_5.pickle().await).unwrap();
+        copy.sender_data = inbound.sender_data.clone();
+        assert_eq!(bob.store().merge_received_group_session(copy).await.unwrap(), None);
+
+        // But when we receive a better copy of the session, we should get it back
+        let mut better =
+            InboundGroupSession::from_export(&inbound.export_at_index(0).await).unwrap();
+        better.sender_data = inbound.sender_data.clone();
+        assert_let!(Some(update) = bob.store().merge_received_group_session(better).await.unwrap());
+        assert_eq!(update.first_known_index(), 0);
+
+        // A worse copy of the ratchet, but better trust data
+        {
+            let mut worse_ratchet_better_trust =
+                InboundGroupSession::from_export(&inbound.export_at_index(10).await).unwrap();
+            let updated_sender_data = SenderData::sender_verified(
+                alice_account.user_id(),
+                alice_account.device_id(),
+                Ed25519Keypair::new().public_key(),
+            );
+            worse_ratchet_better_trust.sender_data = updated_sender_data.clone();
+            assert_let!(
+                Some(update) = bob
+                    .store()
+                    .merge_received_group_session(worse_ratchet_better_trust)
+                    .await
+                    .unwrap()
+            );
+            assert_eq!(update.sender_data, updated_sender_data);
+            assert_eq!(update.first_known_index(), 5);
+            assert_eq!(
+                update.export_at_index(0).await.session_key.to_bytes(),
+                inbound.export_at_index(5).await.session_key.to_bytes()
+            );
+        }
+
+        // A better copy of the ratchet, but worse trust data
+        {
+            let mut better_ratchet_worse_trust =
+                InboundGroupSession::from_export(&inbound.export_at_index(0).await).unwrap();
+            let updated_sender_data = SenderData::unknown();
+            better_ratchet_worse_trust.sender_data = updated_sender_data.clone();
+            assert_let!(
+                Some(update) = bob
+                    .store()
+                    .merge_received_group_session(better_ratchet_worse_trust)
+                    .await
+                    .unwrap()
+            );
+            assert_eq!(update.sender_data, inbound.sender_data);
+            assert_eq!(update.first_known_index(), 0);
+            assert_eq!(
+                update.export_at_index(0).await.session_key.to_bytes(),
+                inbound.export_at_index(0).await.session_key.to_bytes()
+            );
+        }
+    }
+
+    /// Create an [`InboundGroupSession`] for the given room, using the given
+    /// Ed25519 key as the signing key/session ID.
+    fn make_inbound_group_session(
+        sender_account: &Account,
+        signing_key: &Ed25519Keypair,
+        room_id: &RoomId,
+    ) -> InboundGroupSession {
+        InboundGroupSession::new(
+            sender_account.identity_keys.curve25519,
+            sender_account.identity_keys.ed25519,
+            room_id,
+            &make_session_key(signing_key),
+            SenderData::device_info(crate::types::DeviceKeys::new(
+                sender_account.user_id().to_owned(),
+                sender_account.device_id().to_owned(),
+                vec![],
+                BTreeMap::new(),
+                crate::types::Signatures::new(),
+            )),
+            EventEncryptionAlgorithm::MegolmV1AesSha2,
+            Some(ruma::events::room::history_visibility::HistoryVisibility::Shared),
+            true,
+        )
+        .unwrap()
+    }
+
+    /// Make a Megolm [`SessionKey`] using the given Ed25519 key as a signing
+    /// key/session ID.
+    fn make_session_key(signing_key: &Ed25519Keypair) -> SessionKey {
+        use rand::Rng;
+
+        // `SessionKey::new` is not public, so the easiest way to construct a Megolm
+        // session using a known Ed25519 key is to build a byte array in the export
+        // format.
+
+        let mut session_key_bytes = vec![0u8; 229];
+        // 0: version
+        session_key_bytes[0] = 2;
+        // 1..5: index
+        // 5..133: ratchet key
+        rand::thread_rng().fill(&mut session_key_bytes[5..133]);
+        // 133..165: public ed25519 key
+        session_key_bytes[133..165].copy_from_slice(signing_key.public_key().as_bytes());
+        // 165..229: signature
+        let sig = signing_key.sign(&session_key_bytes[0..165]);
+        session_key_bytes[165..229].copy_from_slice(&sig.to_bytes());
+
+        SessionKey::from_bytes(&session_key_bytes).unwrap()
+    }
+
     #[async_test]
     async fn test_import_room_keys_notifies_stream() {
         use futures_util::FutureExt;