Crypto: Fixed a bug where room keys would be rotated unecessarily in the presence of blacklisted/withheld devices

crates/matrix-sdk-crypto/CHANGELOG.md

@@ -16,6 +16,10 @@ All notable changes to this project will be documented in this file.
   (as per [MSC4147](https://github.com/matrix-org/matrix-spec-proposals/pull/4147)).
   ([#4922](https://github.com/matrix-org/matrix-rust-sdk/pull/4922))
 
+- Fix bug which caused room keys to be unnecessarily rotated on every send in the
+  presence of blacklisted/withheld devices in the room.
+  ([#4954](https://github.com/matrix-org/matrix-rust-sdk/pull/4954))
+
 ## [0.11.0] - 2025-04-11
 
 ### Features

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

@@ -160,7 +160,7 @@ pub struct OutboundGroupSession {
     pub(crate) shared_with_set:
         Arc<StdRwLock<BTreeMap<OwnedUserId, BTreeMap<OwnedDeviceId, ShareInfo>>>>,
     #[allow(clippy::type_complexity)]
-    to_share_with_set:
+    pub(crate) to_share_with_set:
         Arc<StdRwLock<BTreeMap<OwnedTransactionId, (Arc<ToDeviceRequest>, ShareInfoSet)>>>,
 }
 

crates/matrix-sdk-crypto/src/session_manager/group_sessions/share_strategy.rs

@@ -27,6 +27,7 @@ use tracing::{debug, instrument, trace};
 use super::OutboundGroupSession;
 use crate::{
     error::{OlmResult, SessionRecipientCollectionError},
+    olm::ShareInfo,
     store::Store,
     DeviceData, EncryptionSettings, LocalTrust, OlmError, OwnUserIdentityData, UserIdentityData,
 };
@@ -426,18 +427,55 @@ fn is_session_overshared_for_user(
     let recipient_device_ids: BTreeSet<&DeviceId> =
         recipient_devices.iter().map(|d| d.device_id()).collect();
 
+    let mut shared: Vec<&DeviceId> = Vec::new();
+
+    // This duplicates a conservative subset of the logic in
+    // `OutboundGroupSession::is_shared_with`, because we
+    // don't have corresponding DeviceData at hand
+    fn is_actually_shared(info: &ShareInfo) -> bool {
+        match info {
+            ShareInfo::Shared(_) => true,
+            ShareInfo::Withheld(_) => false,
+        }
+    }
+
+    // Collect the devices that have definitely received the session already
     let guard = outbound_session.shared_with_set.read();
+    if let Some(for_user) = guard.get(user_id) {
+        shared.extend(for_user.iter().filter_map(|(d, info)| {
+            if is_actually_shared(info) {
+                Some(AsRef::<DeviceId>::as_ref(d))
+            } else {
+                None
+            }
+        }));
+    }
 
-    let Some(shared) = guard.get(user_id) else {
+    // To be conservative, also collect the devices that would still receive the
+    // session from a pending to-device request if we don't rotate beforehand
+    let guard = outbound_session.to_share_with_set.read();
+    for (_txid, share_infos) in guard.values() {
+        if let Some(for_user) = share_infos.get(user_id) {
+            shared.extend(for_user.iter().filter_map(|(d, info)| {
+                if is_actually_shared(info) {
+                    Some(AsRef::<DeviceId>::as_ref(d))
+                } else {
+                    None
+                }
+            }));
+        }
+    }
+
+    if shared.is_empty() {
         return false;
-    };
+    }
 
-    // Devices that received this session
-    let shared: BTreeSet<&DeviceId> = shared.keys().map(|d| d.as_ref()).collect();
+    let shared: BTreeSet<&DeviceId> = shared.into_iter().collect();
 
     // The set difference between
     //
-    // 1. Devices that had previously received the session, and
+    // 1. Devices that had previously received (or are queued to receive) the
+    //    session, and
     // 2. Devices that would now receive the session
     //
     // Represents newly deleted or blacklisted devices. If this
@@ -2095,9 +2133,55 @@ mod tests {
         let machine = test_machine().await;
         import_known_users_to_test_machine(&machine).await;
 
-        let fake_room_id = room_id!("!roomid:localhost");
         let encryption_settings = all_devices_strategy_settings();
+        let group_session = create_test_outbound_group_session(&machine, &encryption_settings);
+        let sender_key = machine.identity_keys().curve25519;
+
+        group_session
+            .mark_shared_with(
+                KeyDistributionTestData::dan_id(),
+                KeyDistributionTestData::dan_signed_device_id(),
+                sender_key,
+            )
+            .await;
+        group_session
+            .mark_shared_with(
+                KeyDistributionTestData::dan_id(),
+                KeyDistributionTestData::dan_unsigned_device_id(),
+                sender_key,
+            )
+            .await;
+
+        // Try to share again after dan has removed one of his devices
+        let keys_query = KeyDistributionTestData::dan_keys_query_response_device_loggedout();
+        machine.mark_request_as_sent(&TransactionId::new(), &keys_query).await.unwrap();
+
+        // share again
+        let share_result = collect_session_recipients(
+            machine.store(),
+            vec![KeyDistributionTestData::dan_id()].into_iter(),
+            &encryption_settings,
+            &group_session,
+        )
+        .await
+        .unwrap();
+
+        assert!(share_result.should_rotate);
+    }
+
+    /// Test that the session is not rotated if a devices is removed
+    /// but was already withheld from receiving the session.
+    #[async_test]
+    async fn test_should_not_rotate_if_keys_were_withheld() {
+        let machine = test_machine().await;
+        import_known_users_to_test_machine(&machine).await;
+
+        let encryption_settings = all_devices_strategy_settings();
+        let group_session = create_test_outbound_group_session(&machine, &encryption_settings);
+        let fake_room_id = group_session.room_id();
 
+        // Because we don't have Olm sessions initialized, this will contain
+        // withheld requests for both of Dan's devices
         let requests = machine
             .share_room_key(
                 fake_room_id,
@@ -2115,13 +2199,11 @@ mod tests {
                 .await
                 .unwrap();
         }
+
         // Try to share again after dan has removed one of his devices
         let keys_query = KeyDistributionTestData::dan_keys_query_response_device_loggedout();
-        let txn_id = TransactionId::new();
-        machine.mark_request_as_sent(&txn_id, &keys_query).await.unwrap();
+        machine.mark_request_as_sent(&TransactionId::new(), &keys_query).await.unwrap();
 
-        let group_session =
-            machine.store().get_outbound_group_session(fake_room_id).await.unwrap().unwrap();
         // share again
         let share_result = collect_session_recipients(
             machine.store(),
@@ -2132,7 +2214,7 @@ mod tests {
         .await
         .unwrap();
 
-        assert!(share_result.should_rotate);
+        assert!(!share_result.should_rotate);
     }
 
     /// Common setup for tests which require a verified user to have unsigned