crypto: Merge inbound Megolm sessions [#5865]

When we receive two copies of the same inbound Megolm session from two sources, merge them together intelligently.

Fixes: #5108, #4698

Author: richvdh

crates/matrix-sdk-crypto/CHANGELOG.md

@@ -8,6 +8,8 @@ All notable changes to this project will be documented in this file.
 
 ### Features
 
+- When we receive an inbound Megolm session from two different sources, merge the two copies together to get the best of both.
+  ([#5865](https://github.com/matrix-org/matrix-rust-sdk/pull/5865)
 - When constructing a key bundle for history sharing, if we had received a key bundle ourselves, in which one or more sessions was marked as "history not shared", pass that on to the new user.
   ([#5820](https://github.com/matrix-org/matrix-rust-sdk/pull/5820)
 - Expose new method `CryptoStore::get_withheld_sessions_by_room_id`.

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

@@ -39,7 +39,7 @@ use ruma::{
     TransactionId, UserId,
 };
 use tracing::{debug, field::debug, info, instrument, trace, warn, Span};
-use vodozemac::{megolm::SessionOrdering, Curve25519PublicKey};
+use vodozemac::Curve25519PublicKey;
 
 use super::{GossipRequest, GossippedSecret, RequestEvent, RequestInfo, SecretInfo, WaitQueue};
 use crate::{
@@ -964,6 +964,7 @@ impl GossipMachine {
         Ok(name)
     }
 
+    #[tracing::instrument(skip_all)]
     async fn accept_forwarded_room_key(
         &self,
         info: &GossipRequest,
@@ -972,33 +973,11 @@ impl GossipMachine {
     ) -> Result<Option<InboundGroupSession>, CryptoStoreError> {
         match InboundGroupSession::try_from(event) {
             Ok(session) => {
-                if self.inner.store.compare_group_session(&session).await?
-                    == SessionOrdering::Better
-                {
+                let new_session = self.inner.store.merge_received_group_session(session).await?;
+                if new_session.is_some() {
                     self.mark_as_done(info).await?;
-
-                    info!(
-                        ?sender_key,
-                        claimed_sender_key = ?session.sender_key(),
-                        room_id = ?session.room_id(),
-                        session_id = session.session_id(),
-                        algorithm = ?session.algorithm(),
-                        "Received a forwarded room key",
-                    );
-
-                    Ok(Some(session))
-                } else {
-                    info!(
-                        ?sender_key,
-                        claimed_sender_key = ?session.sender_key(),
-                        room_id = ?session.room_id(),
-                        session_id = session.session_id(),
-                        algorithm = ?session.algorithm(),
-                        "Received a forwarded room key but we already have a better version of it",
-                    );
-
-                    Ok(None)
                 }
+                Ok(new_session)
             }
             Err(e) => {
                 warn!(?sender_key, "Couldn't create a group session from a received room key");

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

@@ -62,10 +62,7 @@ use tracing::{
     field::{debug, display},
     info, instrument, trace, warn, Span,
 };
-use vodozemac::{
-    megolm::{DecryptionError, SessionOrdering},
-    Curve25519PublicKey, Ed25519Signature,
-};
+use vodozemac::{megolm::DecryptionError, Curve25519PublicKey, Ed25519Signature};
 
 #[cfg(feature = "experimental-send-custom-to-device")]
 use crate::session_manager::split_devices_for_share_strategy;
@@ -916,25 +913,9 @@ impl OlmMachine {
                 let sender_data =
                     SenderDataFinder::find_using_event(self.store(), sender_key, event, &session)
                         .await?;
-
                 session.sender_data = sender_data;
 
-                match self.store().compare_group_session(&session).await? {
-                    SessionOrdering::Better => {
-                        info!("Received a new megolm room key");
-
-                        Ok(Some(session))
-                    }
-                    comparison_result => {
-                        warn!(
-                            ?comparison_result,
-                            "Received a megolm room key that we already have a better version \
-                             of, discarding"
-                        );
-
-                        Ok(None)
-                    }
-                }
+                Ok(self.store().merge_received_group_session(session).await?)
             }
             Err(e) => {
                 Span::current().record("session_id", &content.session_id);

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 {
@@ -488,7 +518,34 @@ impl InboundGroupSession {
 
     /// Check if the [`InboundGroupSession`] is better than the given other
     /// [`InboundGroupSession`]
+    #[deprecated(
+        note = "Sessions cannot be compared on a linear scale. Consider calling `compare_ratchet`, as well as comparing the `sender_data`."
+    )]
     pub async fn compare(&self, other: &InboundGroupSession) -> SessionOrdering {
+        match self.compare_ratchet(other).await {
+            SessionOrdering::Equal => {
+                match self.sender_data.compare_trust_level(&other.sender_data) {
+                    Ordering::Less => SessionOrdering::Worse,
+                    Ordering::Equal => SessionOrdering::Equal,
+                    Ordering::Greater => SessionOrdering::Better,
+                }
+            }
+            result => result,
+        }
+    }
+
+    /// Check if the [`InboundGroupSession`]'s ratchet index is better than that
+    /// of the given other [`InboundGroupSession`].
+    ///
+    /// If the two sessions are not connected (i.e., they are from different
+    /// senders, or if advancing the ratchets to the same index does not
+    /// give the same ratchet value), returns [`SessionOrdering::Unconnected`].
+    ///
+    /// Otherwise, returns [`SessionOrdering::Equal`],
+    /// [`SessionOrdering::Better`], or [`SessionOrdering::Worse`] respectively
+    /// depending on whether this session's first known index is equal to,
+    /// lower than, or higher than, that of `other`.
+    pub async fn compare_ratchet(&self, other: &InboundGroupSession) -> SessionOrdering {
         // If this is the same object the ordering is the same, we can't compare because
         // we would deadlock while trying to acquire the same lock twice.
         if Arc::ptr_eq(&self.inner, &other.inner) {
@@ -501,17 +558,7 @@ impl InboundGroupSession {
             SessionOrdering::Unconnected
         } else {
             let mut other_inner = other.inner.lock().await;
-
-            match self.inner.lock().await.compare(&mut other_inner) {
-                SessionOrdering::Equal => {
-                    match self.sender_data.compare_trust_level(&other.sender_data) {
-                        Ordering::Less => SessionOrdering::Worse,
-                        Ordering::Equal => SessionOrdering::Equal,
-                        Ordering::Greater => SessionOrdering::Better,
-                    }
-                }
-                result => result,
-            }
+            self.inner.lock().await.compare(&mut other_inner)
         }
     }
 
@@ -1057,6 +1104,7 @@ mod tests {
     }
 
     #[async_test]
+    #[allow(deprecated)]
     async fn test_session_comparison() {
         let alice = Account::with_device_id(alice_id(), alice_device_id());
         let room_id = room_id!("!test:localhost");
@@ -1067,18 +1115,24 @@ mod tests {
         let mut copy = InboundGroupSession::from_pickle(inbound.pickle().await).unwrap();
 
         assert_eq!(inbound.compare(&worse).await, SessionOrdering::Better);
+        assert_eq!(inbound.compare_ratchet(&worse).await, SessionOrdering::Better);
         assert_eq!(worse.compare(&inbound).await, SessionOrdering::Worse);
+        assert_eq!(worse.compare_ratchet(&inbound).await, SessionOrdering::Worse);
         assert_eq!(inbound.compare(&inbound).await, SessionOrdering::Equal);
+        assert_eq!(inbound.compare_ratchet(&inbound).await, SessionOrdering::Equal);
         assert_eq!(inbound.compare(&copy).await, SessionOrdering::Equal);
+        assert_eq!(inbound.compare_ratchet(&copy).await, SessionOrdering::Equal);
 
         copy.creator_info.curve25519_key =
             Curve25519PublicKey::from_base64("XbmrPa1kMwmdtNYng1B2gsfoo8UtF+NklzsTZiaVKyY")
                 .unwrap();
 
         assert_eq!(inbound.compare(&copy).await, SessionOrdering::Unconnected);
+        assert_eq!(inbound.compare_ratchet(&copy).await, SessionOrdering::Unconnected);
     }
 
     #[async_test]
+    #[allow(deprecated)]
     async fn test_session_comparison_sender_data() {
         let alice = Account::with_device_id(alice_id(), alice_device_id());
         let room_id = room_id!("!test:localhost");