crypto: prep work for sharing room key history bundles

crates/matrix-sdk-crypto/src/identities/device.rs

@@ -21,7 +21,7 @@ use std::{
     },
 };
 
-use matrix_sdk_common::{deserialized_responses::WithheldCode, locks::RwLock};
+use matrix_sdk_common::locks::RwLock;
 use ruma::{
     api::client::keys::upload_signatures::v3::Request as SignatureUploadRequest,
     events::{key::verification::VerificationMethod, AnyToDeviceEventContent},
@@ -64,9 +64,9 @@ pub enum MaybeEncryptedRoomKey {
         share_info: ShareInfo,
         message: Raw<AnyToDeviceEventContent>,
     },
-    Withheld {
-        code: WithheldCode,
-    },
+    /// We could not encrypt a message to this device because there is no active
+    /// Olm session.
+    MissingSession,
 }
 
 /// A read-only version of a `Device`.
@@ -841,9 +841,7 @@ impl DeviceData {
                 message: encrypted.cast(),
             }),
 
-            Err(OlmError::MissingSession) => {
-                Ok(MaybeEncryptedRoomKey::Withheld { code: WithheldCode::NoOlm })
-            }
+            Err(OlmError::MissingSession) => Ok(MaybeEncryptedRoomKey::MissingSession),
             Err(e) => Err(e),
         }
     }

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

@@ -26,7 +26,7 @@ use matrix_sdk_common::{
     deserialized_responses::WithheldCode, executor::spawn, locks::RwLock as StdRwLock,
 };
 use ruma::{
-    events::{AnyMessageLikeEventContent, ToDeviceEventType},
+    events::{AnyMessageLikeEventContent, AnyToDeviceEventContent, ToDeviceEventType},
     serde::Raw,
     to_device::DeviceIdOrAllDevices,
     OwnedDeviceId, OwnedRoomId, OwnedTransactionId, OwnedUserId, RoomId, TransactionId, UserId,
@@ -254,29 +254,24 @@ impl GroupSessionManager {
         }
     }
 
-    /// Encrypt the given content for the given devices and create a to-device
+    /// Encrypt the given content for the given devices and create to-device
     /// requests that sends the encrypted content to them.
     async fn encrypt_session_for(
         store: Arc<CryptoStoreWrapper>,
         group_session: OutboundGroupSession,
         devices: Vec<DeviceData>,
     ) -> OlmResult<(
-        OwnedTransactionId,
-        ToDeviceRequest,
+        EncryptForDevicesResult,
         BTreeMap<OwnedUserId, BTreeMap<OwnedDeviceId, ShareInfo>>,
-        Vec<Session>,
-        Vec<(DeviceData, WithheldCode)>,
     )> {
         // Use a named type instead of a tuple with rather long type name
         pub struct DeviceResult {
             device: DeviceData,
             maybe_encrypted_room_key: MaybeEncryptedRoomKey,
         }
 
-        let mut messages = BTreeMap::new();
-        let mut changed_sessions = Vec::new();
+        let mut result_builder = EncryptForDevicesResultBuilder::default();
         let mut share_infos = BTreeMap::new();
-        let mut withheld_devices = Vec::new();
 
         // XXX is there a way to do this that doesn't involve cloning the
         // `Arc<CryptoStoreWrapper>` for each device?
@@ -300,35 +295,22 @@ impl GroupSessionManager {
 
             match result.maybe_encrypted_room_key {
                 MaybeEncryptedRoomKey::Encrypted { used_session, share_info, message } => {
-                    changed_sessions.push(used_session);
+                    result_builder.on_successful_encryption(&result.device, used_session, message);
 
                     let user_id = result.device.user_id().to_owned();
                     let device_id = result.device.device_id().to_owned();
-
-                    messages
-                        .entry(user_id.to_owned())
-                        .or_insert_with(BTreeMap::new)
-                        .insert(DeviceIdOrAllDevices::DeviceId(device_id.to_owned()), message);
-
                     share_infos
                         .entry(user_id)
                         .or_insert_with(BTreeMap::new)
                         .insert(device_id, share_info);
                 }
-                MaybeEncryptedRoomKey::Withheld { code } => {
-                    withheld_devices.push((result.device, code));
+                MaybeEncryptedRoomKey::MissingSession => {
+                    result_builder.on_missing_session(result.device);
                 }
             }
         }
 
-        let txn_id = TransactionId::new();
-        let request = ToDeviceRequest {
-            event_type: ToDeviceEventType::RoomEncrypted,
-            txn_id: txn_id.to_owned(),
-            messages,
-        };
-
-        Ok((txn_id, request, share_infos, changed_sessions, withheld_devices))
+        Ok((result_builder.into_result(), share_infos))
     }
 
     /// Given a list of user and an outbound session, return the list of users
@@ -353,21 +335,16 @@ impl GroupSessionManager {
         outbound: OutboundGroupSession,
         sessions: GroupSessionCache,
     ) -> OlmResult<(Vec<Session>, Vec<(DeviceData, WithheldCode)>)> {
-        let (id, request, share_infos, used_sessions, no_olm) =
+        let (result, share_infos) =
             Self::encrypt_session_for(store, outbound.clone(), chunk).await?;
 
-        if !request.messages.is_empty() {
-            trace!(
-                recipient_count = request.message_count(),
-                transaction_id = ?id,
-                "Created a to-device request carrying a room_key"
-            );
-
+        if let Some(request) = result.to_device_request {
+            let id = request.txn_id.clone();
             outbound.add_request(id.clone(), request.into(), share_infos);
             sessions.mark_as_being_shared(id, outbound.clone());
         }
 
-        Ok((used_sessions, no_olm))
+        Ok((result.updated_olm_sessions, result.no_olm_devices))
     }
 
     pub(crate) fn session_cache(&self) -> GroupSessionCache {
@@ -771,6 +748,88 @@ impl GroupSessionManager {
     }
 }
 
+/// Result of [`GroupSessionManager::encrypt_session_for`]
+#[derive(Debug)]
+struct EncryptForDevicesResult {
+    /// The request to send the to-device messages containing the encrypted
+    /// payload, if any devices were found.
+    to_device_request: Option<ToDeviceRequest>,
+
+    /// The devices which lack an Olm session and therefore need a withheld code
+    no_olm_devices: Vec<(DeviceData, WithheldCode)>,
+
+    /// The Olm sessions which were used to encrypt the requests and now need
+    /// persisting to the store.
+    updated_olm_sessions: Vec<Session>,
+}
+
+/// A helper for building [`EncryptForDevicesResult`]
+#[derive(Debug, Default)]
+struct EncryptForDevicesResultBuilder {
+    /// The payloads of the to-device messages
+    messages: BTreeMap<OwnedUserId, BTreeMap<DeviceIdOrAllDevices, Raw<AnyToDeviceEventContent>>>,
+
+    /// The devices which lack an Olm session and therefore need a withheld code
+    no_olm_devices: Vec<(DeviceData, WithheldCode)>,
+
+    /// The Olm sessions which were used to encrypt the requests and now need
+    /// persisting to the store.
+    updated_olm_sessions: Vec<Session>,
+}
+
+impl EncryptForDevicesResultBuilder {
+    /// Record a successful encryption. The encrypted message is added to the
+    /// list to be sent, and the olm session is added to the list of those
+    /// that have been modified.
+    pub fn on_successful_encryption(
+        &mut self,
+        device: &DeviceData,
+        used_session: Session,
+        message: Raw<AnyToDeviceEventContent>,
+    ) {
+        self.updated_olm_sessions.push(used_session);
+
+        self.messages
+            .entry(device.user_id().to_owned())
+            .or_default()
+            .insert(DeviceIdOrAllDevices::DeviceId(device.device_id().to_owned()), message);
+    }
+
+    /// Record a device which didn't have an active Olm session.
+    pub fn on_missing_session(&mut self, device: DeviceData) {
+        self.no_olm_devices.push((device, WithheldCode::NoOlm));
+    }
+
+    /// Transform the accumulated results into an [`EncryptForDevicesResult`],
+    /// wrapping the messages, if any, into a `ToDeviceRequest`.
+    pub fn into_result(self) -> EncryptForDevicesResult {
+        let EncryptForDevicesResultBuilder { updated_olm_sessions, no_olm_devices, messages } =
+            self;
+
+        let mut encrypt_for_devices_result = EncryptForDevicesResult {
+            to_device_request: None,
+            updated_olm_sessions,
+            no_olm_devices,
+        };
+
+        if !messages.is_empty() {
+            let request = ToDeviceRequest {
+                event_type: ToDeviceEventType::RoomEncrypted,
+                txn_id: TransactionId::new(),
+                messages,
+            };
+            trace!(
+                recipient_count = request.message_count(),
+                transaction_id = ?request.txn_id,
+                "Created a to-device request carrying room keys",
+            );
+            encrypt_for_devices_result.to_device_request = Some(request);
+        };
+
+        encrypt_for_devices_result
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::{

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

@@ -156,23 +156,13 @@ pub(crate) async fn collect_session_recipients(
     settings: &EncryptionSettings,
     outbound: &OutboundGroupSession,
 ) -> OlmResult<CollectRecipientsResult> {
-    let users: BTreeSet<&UserId> = users.collect();
-    let mut result = CollectRecipientsResult::default();
-    let mut verified_users_with_new_identities: Vec<OwnedUserId> = Default::default();
-
-    trace!(?users, ?settings, "Calculating group session recipients");
-
-    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();
-
-    // A user left if a user is missing from the set of users that should
-    // get the session but is in the set of users that received the session.
-    let user_left = !users_shared_with.difference(&users).collect::<BTreeSet<_>>().is_empty();
-
-    let visibility_changed = outbound.settings().history_visibility != settings.history_visibility;
-    let algorithm_changed = outbound.settings().algorithm != settings.algorithm;
+    let mut result = collect_recipients_for_share_strategy(
+        store,
+        users,
+        &settings.sharing_strategy,
+        Some(outbound),
+    )
+    .await?;
 
     // To protect the room history we need to rotate the session if either:
     //
@@ -181,13 +171,65 @@ pub(crate) async fn collect_session_recipients(
     // 3. The history visibility changed.
     // 4. The encryption algorithm changed.
     //
-    // This is calculated in the following code and stored in this variable.
-    result.should_rotate = user_left || visibility_changed || algorithm_changed;
+    // `result.should_rotate` is true if the first or second in that list is true;
+    // we now need to check for the other two.
+    let device_removed = result.should_rotate;
+
+    let visibility_changed = outbound.settings().history_visibility != settings.history_visibility;
+    let algorithm_changed = outbound.settings().algorithm != settings.algorithm;
+
+    result.should_rotate = device_removed || visibility_changed || algorithm_changed;
+
+    if result.should_rotate {
+        debug!(
+            device_removed,
+            visibility_changed, algorithm_changed, "Rotating room key to protect room history",
+        );
+    }
+
+    Ok(result)
+}
+
+/// Given a list of users and a [`CollectStrategy`], return the list of devices
+/// that cryptographic keys should be shared with, or that withheld notices
+/// should be sent to.
+///
+/// If an existing [`OutboundGroupSession`] is provided, will also check the
+/// list of devices that the session has been *previously* shared with, and
+/// if that list is too broad, returns a flag indicating that the session should
+/// be rotated (e.g., because a device has been deleted or a user has left the
+/// chat).
+pub(crate) async fn collect_recipients_for_share_strategy(
+    store: &Store,
+    users: impl Iterator<Item = &UserId>,
+    share_strategy: &CollectStrategy,
+    outbound: Option<&OutboundGroupSession>,
+) -> OlmResult<CollectRecipientsResult> {
+    let users: BTreeSet<&UserId> = users.collect();
+    trace!(?users, ?share_strategy, "Calculating group session recipients");
+
+    let mut result = CollectRecipientsResult::default();
+    let mut verified_users_with_new_identities: Vec<OwnedUserId> = Default::default();
+
+    // If we have an outbound session, check if a user is missing from the set of
+    // 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 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");
+            result.should_rotate = true;
+        }
+    }
 
     let own_identity = store.get_user_identity(store.user_id()).await?.and_then(|i| i.into_own());
 
     // Get the recipient and withheld devices, based on the collection strategy.
-    match settings.sharing_strategy {
+    match share_strategy {
         CollectStrategy::AllDevices => {
             for user_id in users {
                 trace!(
@@ -325,15 +367,6 @@ pub(crate) async fn collect_session_recipients(
         ));
     }
 
-    if result.should_rotate {
-        debug!(
-            result.should_rotate,
-            user_left,
-            visibility_changed,
-            algorithm_changed,
-            "Rotating room key to protect room history",
-        );
-    }
     trace!(result.should_rotate, "Done calculating group session recipients");
 
     Ok(result)
@@ -343,17 +376,22 @@ pub(crate) async fn collect_session_recipients(
 /// user.
 fn update_recipients_for_user(
     recipients: &mut CollectRecipientsResult,
-    outbound: &OutboundGroupSession,
+    outbound: Option<&OutboundGroupSession>,
     user_id: &UserId,
     recipient_devices: RecipientDevicesForUser,
 ) {
     // If we haven't already concluded that the session should be
     // rotated for other reasons, we also need to check whether any
     // of the devices in the session got deleted or blacklisted in the
     // meantime. If so, we should also rotate the session.
-    if !recipients.should_rotate {
-        recipients.should_rotate =
-            is_session_overshared_for_user(outbound, user_id, &recipient_devices.allowed_devices)
+    if let Some(outbound) = outbound {
+        if !recipients.should_rotate {
+            recipients.should_rotate = is_session_overshared_for_user(
+                outbound,
+                user_id,
+                &recipient_devices.allowed_devices,
+            )
+        }
     }
 
     recipients