Skip to content

Commit

Permalink
refactor collect_session_recipients in advance of some behavioural ch…
Browse files Browse the repository at this point in the history
…anges (#3884)

This is part of #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.
  • Loading branch information
uhoreg authored Aug 27, 2024
1 parent 47671a1 commit 7581719
Showing 1 changed file with 112 additions and 77 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -120,9 +120,6 @@ pub(crate) async fn collect_session_recipients(
let users: BTreeSet<&UserId> = users.collect();
let mut devices: BTreeMap<OwnedUserId, Vec<DeviceData>> = Default::default();
let mut withheld_devices: Vec<(DeviceData, WithheldCode)> = Default::default();
let mut unsigned_devices_of_verified_users: BTreeMap<OwnedUserId, Vec<OwnedDeviceId>> =
Default::default();
let mut verified_users_with_new_identities: Vec<OwnedUserId> = Default::default();

trace!(?users, ?settings, "Calculating group session recipients");

Expand All @@ -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<OwnedUserId, Vec<OwnedDeviceId>> =
Default::default();
let mut verified_users_with_new_identities: Vec<OwnedUserId> = 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 =
Expand All @@ -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 {
Expand Down Expand Up @@ -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<DeviceData>,
/// Devices that should receive a withheld code.
Expand Down Expand Up @@ -349,8 +377,8 @@ fn split_devices_for_user(
device_owner_identity: &Option<UserIdentityData>,
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));
Expand All @@ -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<DeviceData>,
/// 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<OwnedDeviceId, DeviceData>,
device_owner_identity: &Option<UserIdentityData>,
) -> 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) => {
Expand All @@ -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(),
}
}
}
Expand Down

0 comments on commit 7581719

Please sign in to comment.