refactor(crypto): Properly encapsulate internal OutboundGroupSession state

Previously, the `share_strategy` was breaking the abstraction provided
by `OutboundGroupSession` by accessing its internal fields in an
inconsistent and adhoc way. Now all fields are private and a proper
abstraction was added to access the required state in a consistent API.

Author: procr1337

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

@@ -645,7 +645,7 @@ impl GossipMachine {
         // at. For this, we need an outbound session because this
         // information is recorded there.
         } else if let Some(outbound) = outbound_session {
-            match outbound.is_shared_with(&device.inner) {
+            match outbound.sharing_view().get_share_state(&device.inner) {
                 ShareState::Shared { message_index, olm_wedging_index: _ } => {
                     Ok(Some(message_index))
                 }

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

@@ -16,9 +16,10 @@ use std::{
     cmp::max,
     collections::{BTreeMap, BTreeSet},
     fmt,
+    ops::Bound,
     sync::{
         atomic::{AtomicBool, AtomicU64, Ordering},
-        Arc,
+        Arc, RwLockReadGuard,
     },
     time::Duration,
 };
@@ -70,7 +71,7 @@ const ONE_WEEK: Duration = Duration::from_secs(60 * 60 * 24 * 7);
 const ROTATION_PERIOD: Duration = ONE_WEEK;
 const ROTATION_MESSAGES: u64 = 100;
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
 /// Information about whether a session was shared with a device.
 pub(crate) enum ShareState {
     /// The session was not shared with the device.
@@ -157,11 +158,8 @@ pub struct OutboundGroupSession {
     shared: Arc<AtomicBool>,
     invalidated: Arc<AtomicBool>,
     settings: Arc<EncryptionSettings>,
-    pub(crate) shared_with_set:
-        Arc<StdRwLock<BTreeMap<OwnedUserId, BTreeMap<OwnedDeviceId, ShareInfo>>>>,
-    #[allow(clippy::type_complexity)]
-    pub(crate) to_share_with_set:
-        Arc<StdRwLock<BTreeMap<OwnedTransactionId, (Arc<ToDeviceRequest>, ShareInfoSet)>>>,
+    shared_with_set: Arc<StdRwLock<ShareInfoSet>>,
+    to_share_with_set: Arc<StdRwLock<ToShareMap>>,
 }
 
 /// A a map of userid/device it to a `ShareInfo`.
@@ -170,6 +168,8 @@ pub struct OutboundGroupSession {
 /// room key.
 pub type ShareInfoSet = BTreeMap<OwnedUserId, BTreeMap<OwnedDeviceId, ShareInfo>>;
 
+type ToShareMap = BTreeMap<OwnedTransactionId, (Arc<ToDeviceRequest>, ShareInfoSet)>;
+
 /// Struct holding info about the share state of a outbound group session.
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub enum ShareInfo {
@@ -206,6 +206,90 @@ pub struct SharedWith {
     pub olm_wedging_index: SequenceNumber,
 }
 
+/// A read-only view into the device sharing state of an
+/// [`OutboundGroupSession`].
+pub(crate) struct SharingView<'a> {
+    shared_with_set: RwLockReadGuard<'a, ShareInfoSet>,
+    to_share_with_set: RwLockReadGuard<'a, ToShareMap>,
+}
+
+impl SharingView<'_> {
+    /// Has the session been shared with the given user/device pair (or if not,
+    /// is there such a request pending).
+    pub(crate) fn get_share_state(&self, device: &DeviceData) -> ShareState {
+        self.iter_shares(Some(device.user_id()), Some(device.device_id()))
+            .map(|(_, _, info)| match info {
+                ShareInfo::Shared(info) => {
+                    if device.curve25519_key() == Some(info.sender_key) {
+                        ShareState::Shared {
+                            message_index: info.message_index,
+                            olm_wedging_index: info.olm_wedging_index,
+                        }
+                    } else {
+                        ShareState::SharedButChangedSenderKey
+                    }
+                }
+                ShareInfo::Withheld(_) => ShareState::NotShared,
+            })
+            // Return the most "definitive" ShareState found (in case there
+            // are multiple entries for the same device).
+            .max()
+            .unwrap_or(ShareState::NotShared)
+    }
+
+    /// Has the session been withheld for the given user/device pair (or if not,
+    /// is there such a request pending).
+    pub(crate) fn is_withheld_to(&self, device: &DeviceData, code: &WithheldCode) -> bool {
+        self.iter_shares(Some(device.user_id()), Some(device.device_id()))
+            .any(|(_, _, info)| matches!(info, ShareInfo::Withheld(c) if c == code))
+    }
+
+    /// Enumerate all sent or pending sharing requests for the given device (or
+    /// for all devices if not specified).  This can yield the same device
+    /// multiple times.
+    pub(crate) fn iter_shares<'b, 'c>(
+        &self,
+        user_id: Option<&'b UserId>,
+        device_id: Option<&'c DeviceId>,
+    ) -> impl Iterator<Item = (&UserId, &DeviceId, &ShareInfo)> + use<'_, 'b, 'c> {
+        fn iter_share_info_set<'a, 'b, 'c>(
+            set: &'a ShareInfoSet,
+            user_ids: (Bound<&'b UserId>, Bound<&'b UserId>),
+            device_ids: (Bound<&'c DeviceId>, Bound<&'c DeviceId>),
+        ) -> impl Iterator<Item = (&'a UserId, &'a DeviceId, &'a ShareInfo)> + use<'a, 'b, 'c>
+        {
+            set.range::<UserId, _>(user_ids).flat_map(move |(uid, d)| {
+                d.range::<DeviceId, _>(device_ids)
+                    .map(|(id, info)| (uid.as_ref(), id.as_ref(), info))
+            })
+        }
+
+        let user_ids = user_id
+            .map(|u| (Bound::Included(u), Bound::Included(u)))
+            .unwrap_or((Bound::Unbounded, Bound::Unbounded));
+        let device_ids = device_id
+            .map(|d| (Bound::Included(d), Bound::Included(d)))
+            .unwrap_or((Bound::Unbounded, Bound::Unbounded));
+
+        let already_shared = iter_share_info_set(&self.shared_with_set, user_ids, device_ids);
+        let pending = self
+            .to_share_with_set
+            .values()
+            .flat_map(move |(_, set)| iter_share_info_set(set, user_ids, device_ids));
+        already_shared.chain(pending)
+    }
+
+    /// Enumerate all users that have received the session, or have pending
+    /// requests to receive it.  This can yield the same user multiple times,
+    /// so you may want to `collect()` the result into a `BTreeSet`.
+    pub(crate) fn shared_with_users(&self) -> impl Iterator<Item = &UserId> {
+        self.iter_shares(None, None).filter_map(|(u, _, info)| match info {
+            ShareInfo::Shared(_) => Some(u),
+            ShareInfo::Withheld(_) => None,
+        })
+    }
+}
+
 impl OutboundGroupSession {
     pub(super) fn session_config(
         algorithm: &EventEncryptionAlgorithm,
@@ -541,78 +625,16 @@ impl OutboundGroupSession {
         )
     }
 
-    /// Has or will the session be shared with the given user/device pair.
-    pub(crate) fn is_shared_with(&self, device: &DeviceData) -> ShareState {
-        // Check if we shared the session.
-        let shared_state = self.shared_with_set.read().get(device.user_id()).and_then(|d| {
-            d.get(device.device_id()).map(|s| match s {
-                ShareInfo::Shared(s) => {
-                    if device.curve25519_key() == Some(s.sender_key) {
-                        ShareState::Shared {
-                            message_index: s.message_index,
-                            olm_wedging_index: s.olm_wedging_index,
-                        }
-                    } else {
-                        ShareState::SharedButChangedSenderKey
-                    }
-                }
-                ShareInfo::Withheld(_) => ShareState::NotShared,
-            })
-        });
-
-        if let Some(state) = shared_state {
-            state
-        } else {
-            // If we haven't shared the session, check if we're going to share
-            // the session.
-
-            // Find the first request that contains the given user id and
-            // device ID.
-            let shared = self.to_share_with_set.read().values().find_map(|(_, share_info)| {
-                let d = share_info.get(device.user_id())?;
-                let info = d.get(device.device_id())?;
-                Some(match info {
-                    ShareInfo::Shared(info) => {
-                        if device.curve25519_key() == Some(info.sender_key) {
-                            ShareState::Shared {
-                                message_index: info.message_index,
-                                olm_wedging_index: info.olm_wedging_index,
-                            }
-                        } else {
-                            ShareState::SharedButChangedSenderKey
-                        }
-                    }
-                    ShareInfo::Withheld(_) => ShareState::NotShared,
-                })
-            });
-
-            shared.unwrap_or(ShareState::NotShared)
+    /// Create a read-only view into the device sharing state of this session.
+    /// This view includes pending requests, so it is not guaranteed that the
+    /// represented state has been fully propagated yet.
+    pub(crate) fn sharing_view(&self) -> SharingView<'_> {
+        SharingView {
+            shared_with_set: self.shared_with_set.read(),
+            to_share_with_set: self.to_share_with_set.read(),
         }
     }
 
-    pub(crate) fn is_withheld_to(&self, device: &DeviceData, code: &WithheldCode) -> bool {
-        self.shared_with_set
-            .read()
-            .get(device.user_id())
-            .and_then(|d| {
-                let info = d.get(device.device_id())?;
-                Some(matches!(info, ShareInfo::Withheld(c) if c == code))
-            })
-            .unwrap_or_else(|| {
-                // If we haven't yet withheld, check if we're going to withheld
-                // the session.
-
-                // Find the first request that contains the given user id and
-                // device ID.
-                self.to_share_with_set.read().values().any(|(_, share_info)| {
-                    share_info
-                        .get(device.user_id())
-                        .and_then(|d| d.get(device.device_id()))
-                        .is_some_and(|info| matches!(info, ShareInfo::Withheld(c) if c == code))
-                })
-            })
-    }
-
     /// Mark the session as shared with the given user/device pair, starting
     /// from some message index.
     #[cfg(test)]
@@ -782,7 +804,7 @@ mod tests {
         uint, EventEncryptionAlgorithm,
     };
 
-    use super::{EncryptionSettings, ROTATION_MESSAGES, ROTATION_PERIOD};
+    use super::{EncryptionSettings, ShareState, ROTATION_MESSAGES, ROTATION_PERIOD};
     use crate::CollectStrategy;
 
     #[test]
@@ -811,6 +833,24 @@ mod tests {
         assert_eq!(settings.rotation_period_msgs, 500);
     }
 
+    /// Ensure that the `ShareState` PartialOrd instance orders according to
+    /// specificity of the value.
+    #[test]
+    fn test_share_state_ordering() {
+        let values = [
+            ShareState::NotShared,
+            ShareState::SharedButChangedSenderKey,
+            ShareState::Shared { message_index: 1, olm_wedging_index: Default::default() },
+        ];
+        // Make sure our test case of possible variants is exhaustive
+        match values[0] {
+            ShareState::NotShared
+            | ShareState::SharedButChangedSenderKey
+            | ShareState::Shared { .. } => {}
+        }
+        assert!(values.is_sorted());
+    }
+
     #[cfg(any(target_os = "linux", target_os = "macos", target_arch = "wasm32"))]
     mod expiration {
         use std::{sync::atomic::Ordering, time::Duration};

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

@@ -122,7 +122,7 @@ impl GroupSessionCache {
 
     /// Returns whether any session is withheld with the given device and code.
     fn has_session_withheld_to(&self, device: &DeviceData, code: &WithheldCode) -> bool {
-        self.sessions.read().values().any(|s| s.is_withheld_to(device, code))
+        self.sessions.read().values().any(|s| s.sharing_view().is_withheld_to(device, code))
     }
 
     fn remove_from_being_shared(&self, id: &TransactionId) -> Option<OutboundGroupSession> {
@@ -500,7 +500,7 @@ impl GroupSessionManager {
         if code == &WithheldCode::NoOlm {
             device.was_withheld_code_sent() || self.sessions.has_session_withheld_to(device, code)
         } else {
-            group_session.is_withheld_to(device, code)
+            group_session.sharing_view().is_withheld_to(device, code)
         }
     }
 
@@ -696,7 +696,7 @@ impl GroupSessionManager {
         let devices: Vec<_> = devices
             .into_iter()
             .flat_map(|(_, d)| {
-                d.into_iter().filter(|d| match outbound.is_shared_with(d) {
+                d.into_iter().filter(|d| match outbound.sharing_view().get_share_state(d) {
                     ShareState::NotShared => true,
                     ShareState::Shared { message_index: _, olm_wedging_index } => {
                         // If the recipient device's Olm wedging index is higher

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

@@ -15,7 +15,6 @@
 use std::{
     collections::{BTreeMap, BTreeSet, HashMap},
     default::Default,
-    ops::Deref,
 };
 
 use itertools::{Either, Itertools};
@@ -216,10 +215,8 @@ pub(crate) async fn collect_recipients_for_share_strategy(
     // users that should get the session but is in the set of users that
     // received the session.
     if let Some(outbound) = outbound {
-        let users_shared_with: BTreeSet<OwnedUserId> =
-            outbound.shared_with_set.read().keys().cloned().collect();
-        let users_shared_with: BTreeSet<&UserId> =
-            users_shared_with.iter().map(Deref::deref).collect();
+        let view = outbound.sharing_view();
+        let users_shared_with = view.shared_with_users().collect::<BTreeSet<_>>();
         let left_users = users_shared_with.difference(&users).collect::<BTreeSet<_>>();
         if !left_users.is_empty() {
             trace!(?left_users, "Some users have left the chat: session must be rotated");
@@ -427,61 +424,20 @@ 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))
+    let view = outbound_session.sharing_view();
+    let newly_deleted_or_blacklisted: BTreeSet<&DeviceId> = view
+        .iter_shares(Some(user_id), None)
+        .filter_map(|(_user_id, device_id, info)| {
+            // If a devices who we've shared the session with before is not in the
+            // list of devices that should receive the session, we need to rotate.
+            // We also collect all of those device IDs to log them out.
+            if matches!(info, ShareInfo::Shared(_)) && !recipient_device_ids.contains(device_id) {
+                Some(device_id)
             } else {
                 None
             }
-        }));
-    }
-
-    // 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;
-    }
-
-    let shared: BTreeSet<&DeviceId> = shared.into_iter().collect();
-
-    // The set difference between
-    //
-    // 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
-    // set is non-empty, we must rotate.
-    let newly_deleted_or_blacklisted =
-        shared.difference(&recipient_device_ids).collect::<BTreeSet<_>>();
+        })
+        .collect();
 
     let should_rotate = !newly_deleted_or_blacklisted.is_empty();
     if should_rotate {