From 2fc4aacdd0026d3e2de4ed93f314b563e6386337 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Damir=20Jeli=C4=87?= <poljar@termina.org.uk>
Date: Wed, 25 Sep 2024 15:26:55 +0200
Subject: [PATCH] feat: Prefer room keys with better SenderData when comparing
 duplicate room keys

---
 .../src/olm/group_sessions/inbound.rs         | 44 ++++++++++++++++---
 1 file changed, 39 insertions(+), 5 deletions(-)

diff --git a/crates/matrix-sdk-crypto/src/olm/group_sessions/inbound.rs b/crates/matrix-sdk-crypto/src/olm/group_sessions/inbound.rs
index f43e013e4..c42d7ba88 100644
--- a/crates/matrix-sdk-crypto/src/olm/group_sessions/inbound.rs
+++ b/crates/matrix-sdk-crypto/src/olm/group_sessions/inbound.rs
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::{
+    cmp::Ordering,
     fmt,
     ops::Deref,
     sync::{
@@ -375,8 +376,8 @@ impl InboundGroupSession {
         self.imported
     }
 
-    /// Check if the `InboundGroupSession` is better than the given other
-    /// `InboundGroupSession`
+    /// Check if the [`InboundGroupSession`] is better than the given other
+    /// [`InboundGroupSession`]
     pub async fn compare(&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.
@@ -389,8 +390,18 @@ impl InboundGroupSession {
         {
             SessionOrdering::Unconnected
         } else {
-            let mut other = other.inner.lock().await;
-            self.inner.lock().await.compare(&mut other)
+            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,
+            }
         }
     }
 
@@ -648,7 +659,7 @@ mod tests {
     };
 
     use crate::{
-        olm::{InboundGroupSession, SenderData},
+        olm::{InboundGroupSession, KnownSenderData, SenderData},
         types::EventEncryptionAlgorithm,
         Account,
     };
@@ -868,6 +879,29 @@ mod tests {
         assert_eq!(inbound.compare(&copy).await, SessionOrdering::Unconnected);
     }
 
+    #[async_test]
+    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");
+
+        let (_, mut inbound) = alice.create_group_session_pair_with_defaults(room_id).await;
+
+        let sender_data = SenderData::SenderVerified(KnownSenderData {
+            user_id: alice.user_id().into(),
+            device_id: Some(alice.device_id().into()),
+            master_key: alice.identity_keys().ed25519.into(),
+        });
+
+        let mut better = InboundGroupSession::from_pickle(inbound.pickle().await).unwrap();
+        better.sender_data = sender_data.clone();
+
+        assert_eq!(inbound.compare(&better).await, SessionOrdering::Worse);
+        assert_eq!(better.compare(&inbound).await, SessionOrdering::Better);
+
+        inbound.sender_data = sender_data;
+        assert_eq!(better.compare(&inbound).await, SessionOrdering::Equal);
+    }
+
     fn create_session_key() -> SessionKey {
         SessionKey::from_base64(
             "\