From 758171931d3c375e03ee8bcd77f78712f67514f3 Mon Sep 17 00:00:00 2001 From: Hubert Chathi Date: Tue, 27 Aug 2024 06:59:38 -0400 Subject: [PATCH] refactor collect_session_recipients in advance of some behavioural changes (#3884) This is part of https://github.com/matrix-org/matrix-rust-sdk/pull/3662, pulled out to into a separate PR. Recent changes in `main` made it pretty much impossible to merge this section of code from `main` into that PR, and Rich wanted to see the refactoring bits separate from the behavioural changes. So I've re-written the refactoring. Pulls the `match` on `sharing_strategy` outside of the `for` loop, and moves any code that is specific to one strategy into the appropriate branch. --- .../group_sessions/share_strategy.rs | 189 +++++++++++------- 1 file changed, 112 insertions(+), 77 deletions(-) diff --git a/crates/matrix-sdk-crypto/src/session_manager/group_sessions/share_strategy.rs b/crates/matrix-sdk-crypto/src/session_manager/group_sessions/share_strategy.rs index 99fcf534c68..5c426d0d425 100644 --- a/crates/matrix-sdk-crypto/src/session_manager/group_sessions/share_strategy.rs +++ b/crates/matrix-sdk-crypto/src/session_manager/group_sessions/share_strategy.rs @@ -120,9 +120,6 @@ pub(crate) async fn collect_session_recipients( let users: BTreeSet<&UserId> = users.collect(); let mut devices: BTreeMap> = Default::default(); let mut withheld_devices: Vec<(DeviceData, WithheldCode)> = Default::default(); - let mut unsigned_devices_of_verified_users: BTreeMap> = - Default::default(); - let mut verified_users_with_new_identities: Vec = Default::default(); trace!(?users, ?settings, "Calculating group session recipients"); @@ -148,17 +145,23 @@ pub(crate) async fn collect_session_recipients( // This is calculated in the following code and stored in this variable. let mut should_rotate = user_left || visibility_changed || algorithm_changed; - 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 { + CollectStrategy::DeviceBasedStrategy { + only_allow_trusted_devices, + error_on_verified_user_problem, + } => { + let mut unsigned_devices_of_verified_users: BTreeMap> = + Default::default(); + let mut verified_users_with_new_identities: Vec = Default::default(); - for user_id in users { - trace!("Considering recipient devices for user {}", user_id); - let user_devices = store.get_device_data_for_user_filtered(user_id).await?; + let own_identity = + store.get_user_identity(store.user_id()).await?.and_then(|i| i.into_own()); + + for user_id in users { + trace!("Considering recipient devices for user {}", user_id); + let user_devices = store.get_device_data_for_user_filtered(user_id).await?; - let recipient_devices = match settings.sharing_strategy { - CollectStrategy::DeviceBasedStrategy { - only_allow_trusted_devices, - error_on_verified_user_problem, - } => { // We only need the user identity if `only_allow_trusted_devices` or // `error_on_verified_user_problem` is set. let device_owner_identity = @@ -179,73 +182,99 @@ pub(crate) async fn collect_session_recipients( continue; } - split_devices_for_user( + let recipient_devices = split_devices_for_user( user_devices, &own_identity, &device_owner_identity, only_allow_trusted_devices, error_on_verified_user_problem, - ) - } - CollectStrategy::IdentityBasedStrategy => { - let device_owner_identity = store.get_user_identity(user_id).await?; - split_recipients_withhelds_for_user_based_on_identity( - user_devices, - &device_owner_identity, - ) - } - }; + ); + + // If `error_on_verified_user_problem` is set, then + // `unsigned_of_verified_user` may be populated. If so, add an entry to the + // list of users with unsigned devices. + if !recipient_devices.unsigned_of_verified_user.is_empty() { + unsigned_devices_of_verified_users.insert( + user_id.to_owned(), + recipient_devices + .unsigned_of_verified_user + .into_iter() + .map(|d| d.device_id().to_owned()) + .collect(), + ); + } - // If we're using a `DeviceBasedStrategy` with - // `error_on_verified_user_problem` set, then - // `unsigned_of_verified_user` may be populated. If so, add an entry to the - // list of users with unsigned devices. - if !recipient_devices.unsigned_of_verified_user.is_empty() { - unsigned_devices_of_verified_users.insert( - user_id.to_owned(), - recipient_devices - .unsigned_of_verified_user - .into_iter() - .map(|d| d.device_id().to_owned()) - .collect(), - ); - } + // 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 !should_rotate { + should_rotate = is_session_overshared_for_user( + outbound, + user_id, + &recipient_devices.allowed_devices, + ) + } + + devices + .entry(user_id.to_owned()) + .or_default() + .extend(recipient_devices.allowed_devices); + withheld_devices.extend(recipient_devices.denied_devices_with_code); + } - let recipients = recipient_devices.allowed_devices; - let withheld_recipients = recipient_devices.denied_devices_with_code; + // If `error_on_verified_user_problem` is set, then + // `unsigned_devices_of_verified_users` may be populated. If so, we need to bail + // out with an error. + if !unsigned_devices_of_verified_users.is_empty() { + return Err(OlmError::SessionRecipientCollectionError( + SessionRecipientCollectionError::VerifiedUserHasUnsignedDevice( + unsigned_devices_of_verified_users, + ), + )); + } - // 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 !should_rotate { - should_rotate = is_session_overshared_for_user(outbound, user_id, &recipients) + // Alternatively, we may have encountered previously-verified users who have + // changed their identities. We bail out for that, too. + if !verified_users_with_new_identities.is_empty() { + return Err(OlmError::SessionRecipientCollectionError( + SessionRecipientCollectionError::VerifiedUserChangedIdentity( + verified_users_with_new_identities, + ), + )); + } } + CollectStrategy::IdentityBasedStrategy => { + for user_id in users { + trace!("Considering recipient devices for user {}", user_id); + let user_devices = store.get_device_data_for_user_filtered(user_id).await?; - devices.entry(user_id.to_owned()).or_default().extend(recipients); - withheld_devices.extend(withheld_recipients); - } + let device_owner_identity = store.get_user_identity(user_id).await?; - // If we're using a `DeviceBasedStrategy` with - // `error_on_verified_user_problem` set, then - // `unsigned_devices_of_verified_users` may be populated. If so, we need to bail - // out with an error. - if !unsigned_devices_of_verified_users.is_empty() { - return Err(OlmError::SessionRecipientCollectionError( - SessionRecipientCollectionError::VerifiedUserHasUnsignedDevice( - unsigned_devices_of_verified_users, - ), - )); - } + let recipient_devices = split_recipients_withhelds_for_user_based_on_identity( + user_devices, + &device_owner_identity, + ); + + // 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 !should_rotate { + should_rotate = is_session_overshared_for_user( + outbound, + user_id, + &recipient_devices.allowed_devices, + ) + } - // Alternatively, we may have encountered previously-verified users who have - // changed their identities. We bail out for that, too. - if !verified_users_with_new_identities.is_empty() { - return Err(OlmError::SessionRecipientCollectionError( - SessionRecipientCollectionError::VerifiedUserChangedIdentity( - verified_users_with_new_identities, - ), - )); + devices + .entry(user_id.to_owned()) + .or_default() + .extend(recipient_devices.allowed_devices); + withheld_devices.extend(recipient_devices.denied_devices_with_code); + } + } } if should_rotate { @@ -315,10 +344,9 @@ fn is_session_overshared_for_user( should_rotate } -/// Result type for [`split_devices_for_user`] and -/// [`split_recipients_withhelds_for_user_based_on_identity`]. +/// Result type for [`split_devices_for_user`]. #[derive(Default)] -struct RecipientDevices { +struct DeviceBasedRecipientDevices { /// Devices that should receive the room key. allowed_devices: Vec, /// Devices that should receive a withheld code. @@ -349,8 +377,8 @@ fn split_devices_for_user( device_owner_identity: &Option, only_allow_trusted_devices: bool, error_on_verified_user_problem: bool, -) -> RecipientDevices { - let mut recipient_devices: RecipientDevices = Default::default(); +) -> DeviceBasedRecipientDevices { + let mut recipient_devices: DeviceBasedRecipientDevices = Default::default(); for d in user_devices.into_values() { if d.is_blacklisted() { recipient_devices.denied_devices_with_code.push((d, WithheldCode::Blacklisted)); @@ -375,21 +403,29 @@ fn split_devices_for_user( recipient_devices } +/// Result type for [`split_recipients_withhelds_for_user_based_on_identity`]. +#[derive(Default)] +struct IdentityBasedRecipientDevices { + /// Devices that should receive the room key. + allowed_devices: Vec, + /// Devices that should receive a withheld code. + denied_devices_with_code: Vec<(DeviceData, WithheldCode)>, +} + fn split_recipients_withhelds_for_user_based_on_identity( user_devices: HashMap, device_owner_identity: &Option, -) -> RecipientDevices { +) -> IdentityBasedRecipientDevices { match device_owner_identity { None => { // withheld all the users devices, we need to have an identity for this // distribution mode - RecipientDevices { + IdentityBasedRecipientDevices { allowed_devices: Vec::default(), denied_devices_with_code: user_devices .into_values() .map(|d| (d, WithheldCode::Unauthorised)) .collect(), - unsigned_of_verified_user: Vec::default(), } } Some(device_owner_identity) => { @@ -404,10 +440,9 @@ fn split_recipients_withhelds_for_user_based_on_identity( Either::Right((d, WithheldCode::Unauthorised)) } }); - RecipientDevices { + IdentityBasedRecipientDevices { allowed_devices: recipients, denied_devices_with_code: withheld_recipients, - unsigned_of_verified_user: Vec::default(), } } }