Skip to content

Commit

Permalink
refactor(sdk-crypto): Room key sharing, introduce extensible strategy
Browse files Browse the repository at this point in the history 
This patch set does two things:

1. Extracted the logic to collect the devices that should receive a room key.
2. Introduce a new CollectStrategy enum which defines which rules are
   used to collect recipient devices for a room key. Currently only the
   existing rules have beenmoved under this enum.
  • Loading branch information
@poljar
poljar authored Jun 25, 2024
2 parents 1221d15 + ca6537b commit 3be84a5
Show file tree
Hide file tree
Showing 9 changed files with 1,001 additions and 150 deletions.
4 changes: 2 additions & 2 deletions bindings/matrix-sdk-crypto-ffi/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ use matrix_sdk_crypto::{
olm::{IdentityKeys, InboundGroupSession, Session},
store::{Changes, CryptoStore, PendingChanges, RoomSettings as RustRoomSettings},
types::{EventEncryptionAlgorithm as RustEventEncryptionAlgorithm, SigningKey},
EncryptionSettings as RustEncryptionSettings,
CollectStrategy, EncryptionSettings as RustEncryptionSettings,
};
use matrix_sdk_sqlite::SqliteCryptoStore;
pub use responses::{
Expand Down Expand Up @@ -650,7 +650,7 @@ impl From<EncryptionSettings> for RustEncryptionSettings {
rotation_period: Duration::from_secs(v.rotation_period),
rotation_period_msgs: v.rotation_period_msgs,
history_visibility: v.history_visibility.into(),
only_allow_trusted_devices: v.only_allow_trusted_devices,
sharing_strategy: CollectStrategy::new_device_based(v.only_allow_trusted_devices),
}
}
}
Expand Down
1 change: 1 addition & 0 deletions crates/matrix-sdk-crypto/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,7 @@ pub use requests::{
IncomingResponse, KeysBackupRequest, KeysQueryRequest, OutgoingRequest, OutgoingRequests,
OutgoingVerificationRequest, RoomMessageRequest, ToDeviceRequest, UploadSigningKeysRequest,
};
pub use session_manager::CollectStrategy;
pub use store::{
CrossSigningKeyExport, CryptoStoreError, SecretImportError, SecretInfo, TrackedUser,
};
Expand Down
7 changes: 5 additions & 2 deletions crates/matrix-sdk-crypto/src/machine.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2412,6 +2412,7 @@ pub(crate) mod tests {
olm::{
BackedUpRoomKey, ExportedRoomKey, InboundGroupSession, OutboundGroupSession, VerifyJson,
},
session_manager::CollectStrategy,
store::{BackupDecryptionKey, Changes, CryptoStore, MemoryStore, RoomSettings},
types::{
events::{
Expand Down Expand Up @@ -3230,8 +3231,10 @@ pub(crate) mod tests {
let room_id = room_id!("!test:example.org");

let encryption_settings = EncryptionSettings::default();
let encryption_settings =
EncryptionSettings { only_allow_trusted_devices: true, ..encryption_settings };
let encryption_settings = EncryptionSettings {
sharing_strategy: CollectStrategy::new_device_based(true),
..encryption_settings
};

let to_device_requests = alice
.share_room_key(room_id, iter::once(bob.user_id()), encryption_settings)
Expand Down
11 changes: 6 additions & 5 deletions crates/matrix-sdk-crypto/src/olm/group_sessions/outbound.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ use super::SessionCreationError;
#[cfg(feature = "experimental-algorithms")]
use crate::types::events::room::encrypted::MegolmV2AesSha2Content;
use crate::{
session_manager::CollectStrategy,
types::{
events::{
room::encrypted::{
Expand Down Expand Up @@ -85,10 +86,10 @@ pub struct EncryptionSettings {
pub rotation_period_msgs: u64,
/// The history visibility of the room when the session was created.
pub history_visibility: HistoryVisibility,
/// Should untrusted devices receive the room key, or should they be
/// excluded from the conversation.
/// The strategy used to distribute the room keys to participant.
/// Default will send to all devices.
#[serde(default)]
pub only_allow_trusted_devices: bool,
pub sharing_strategy: CollectStrategy,
}

impl Default for EncryptionSettings {
Expand All @@ -98,7 +99,7 @@ impl Default for EncryptionSettings {
rotation_period: ROTATION_PERIOD,
rotation_period_msgs: ROTATION_MESSAGES,
history_visibility: HistoryVisibility::Shared,
only_allow_trusted_devices: false,
sharing_strategy: CollectStrategy::default(),
}
}
}
Expand All @@ -122,7 +123,7 @@ impl EncryptionSettings {
rotation_period,
rotation_period_msgs,
history_visibility,
only_allow_trusted_devices,
sharing_strategy: CollectStrategy::new_device_based(only_allow_trusted_devices),
}
}
}
Expand Down
161 changes: 20 additions & 141 deletions ...pto/src/session_manager/group_sessions.rs → ...src/session_manager/group_sessions/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,23 +12,25 @@
// See the License for the specific language governing permissions and
// limitations under the License.

mod share_strategy;

use std::{
collections::{BTreeMap, BTreeSet},
fmt::Debug,
ops::Deref,
sync::{Arc, RwLock as StdRwLock},
};

use futures_util::future::join_all;
use itertools::{Either, Itertools};
use itertools::Itertools;
use matrix_sdk_common::executor::spawn;
use ruma::{
events::{AnyMessageLikeEventContent, ToDeviceEventType},
serde::Raw,
to_device::DeviceIdOrAllDevices,
DeviceId, OwnedDeviceId, OwnedRoomId, OwnedTransactionId, OwnedUserId, RoomId, TransactionId,
UserId,
OwnedDeviceId, OwnedRoomId, OwnedTransactionId, OwnedUserId, RoomId, TransactionId, UserId,
};
pub(crate) use share_strategy::CollectRecipientsResult;
pub use share_strategy::CollectStrategy;
use tracing::{debug, error, info, instrument, trace};

use crate::{
Expand Down Expand Up @@ -116,23 +118,6 @@ impl GroupSessionCache {
}
}

/// Returned by `collect_session_recipients`.
///
/// Information indicating whether the session needs to be rotated
/// (`should_rotate`) and the list of users/devices that should receive
/// (`devices`) or not the session, including withheld reason
/// `withheld_devices`.
#[derive(Debug)]
pub(crate) struct CollectRecipientsResult {
/// If true the outbound group session should be rotated
pub should_rotate: bool,
/// The map of user|device that should receive the session
pub devices: BTreeMap<OwnedUserId, Vec<ReadOnlyDevice>>,
/// The map of user|device that won't receive the key with the withheld
/// code.
pub withheld_devices: Vec<(ReadOnlyDevice, WithheldCode)>,
}

#[derive(Debug, Clone)]
pub(crate) struct GroupSessionManager {
/// Store for the encryption keys.
Expand Down Expand Up @@ -350,118 +335,7 @@ impl GroupSessionManager {
settings: &EncryptionSettings,
outbound: &OutboundGroupSession,
) -> OlmResult<CollectRecipientsResult> {
let users: BTreeSet<&UserId> = users.collect();
let mut devices: BTreeMap<OwnedUserId, Vec<ReadOnlyDevice>> = Default::default();
let mut withheld_devices: Vec<(ReadOnlyDevice, WithheldCode)> = Default::default();

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

let users_shared_with: BTreeSet<OwnedUserId> =
outbound.shared_with_set.read().unwrap().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;

// To protect the room history we need to rotate the session if either:
//
// 1. Any user left the room.
// 2. Any of the users' devices got deleted or blacklisted.
// 3. The history visibility changed.
// 4. The encryption algorithm changed.
//
// 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 =
self.store.get_user_identity(self.store.user_id()).await?.and_then(|i| i.into_own());

for user_id in users {
let user_devices = self.store.get_readonly_devices_filtered(user_id).await?;

// We only need the user identity if settings.only_allow_trusted_devices is set.
let device_owner_identity = if settings.only_allow_trusted_devices {
self.store.get_user_identity(user_id).await?
} else {
None
};

// From all the devices a user has, we're splitting them into two
// buckets, a bucket of devices that should receive the
// room key and a bucket of devices that should receive
// a withheld code.
let (recipients, withheld_recipients): (
Vec<ReadOnlyDevice>,
Vec<(ReadOnlyDevice, WithheldCode)>,
) = user_devices.into_values().partition_map(|d| {
if d.is_blacklisted() {
Either::Right((d, WithheldCode::Blacklisted))
} else if settings.only_allow_trusted_devices
&& !d.is_verified(&own_identity, &device_owner_identity)
{
Either::Right((d, WithheldCode::Unverified))
} else {
Either::Left(d)
}
});

// 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 {
// Device IDs that should receive this session
let recipient_device_ids: BTreeSet<&DeviceId> =
recipients.iter().map(|d| d.device_id()).collect();

if let Some(shared) = outbound.shared_with_set.read().unwrap().get(user_id) {
// Devices that received this session
let shared: BTreeSet<OwnedDeviceId> = shared.keys().cloned().collect();
let shared: BTreeSet<&DeviceId> = shared.iter().map(|d| d.as_ref()).collect();

// The set difference between
//
// 1. Devices that had previously received 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<_>>();

should_rotate = !newly_deleted_or_blacklisted.is_empty();
if should_rotate {
debug!(
"Rotating a room key due to these devices being deleted/blacklisted {:?}",
newly_deleted_or_blacklisted,
);
}
};
}

devices.entry(user_id.to_owned()).or_default().extend(recipients);
withheld_devices.extend(withheld_recipients);
}

if should_rotate {
debug!(
should_rotate,
user_left,
visibility_changed,
algorithm_changed,
"Rotating room key to protect room history",
);
}
trace!(should_rotate, "Done calculating group session recipients");

Ok(CollectRecipientsResult { should_rotate, devices, withheld_devices })
share_strategy::collect_session_recipients(&self.store, users, settings, outbound).await
}

async fn encrypt_request(
Expand Down Expand Up @@ -871,10 +745,11 @@ mod tests {
use serde_json::{json, Value};

use crate::{
session_manager::group_sessions::CollectRecipientsResult,
session_manager::{group_sessions::CollectRecipientsResult, CollectStrategy},
types::{
events::room_key_withheld::{
RoomKeyWithheldContent, RoomKeyWithheldContent::MegolmV1AesSha2, WithheldCode,
RoomKeyWithheldContent::{self, MegolmV1AesSha2},
WithheldCode,
},
EventEncryptionAlgorithm,
},
Expand All @@ -891,7 +766,7 @@ mod tests {

/// Returns a /keys/query response for user "@example:localhost"
fn keys_query_response() -> get_keys::v3::Response {
let data = include_bytes!("../../../../benchmarks/benches/crypto_bench/keys_query.json");
let data = include_bytes!("../../../../../benchmarks/benches/crypto_bench/keys_query.json");
let data: Value = serde_json::from_slice(data).unwrap();
let data = response_from_file(&data);
get_keys::v3::Response::try_from_http_response(data)
Expand Down Expand Up @@ -958,7 +833,7 @@ mod tests {
/// Returns a key claim response for device `NMMBNBUSNR` of user
/// `@example2:localhost`
fn keys_claim_response() -> claim_keys::v3::Response {
let data = include_bytes!("../../../../benchmarks/benches/crypto_bench/keys_claim.json");
let data = include_bytes!("../../../../../benchmarks/benches/crypto_bench/keys_claim.json");
let data: Value = serde_json::from_slice(data).unwrap();
let data = response_from_file(&data);
claim_keys::v3::Response::try_from_http_response(data)
Expand Down Expand Up @@ -1237,8 +1112,10 @@ mod tests {
.iter()
.any(|d| d.user_id() == user_id && d.device_id() == device_id));

let settings =
EncryptionSettings { only_allow_trusted_devices: true, ..Default::default() };
let settings = EncryptionSettings {
sharing_strategy: CollectStrategy::new_device_based(true),
..Default::default()
};
let users = [user_id].into_iter();

let CollectRecipientsResult { devices: recipients, .. } = machine
Expand Down Expand Up @@ -1295,8 +1172,10 @@ mod tests {
let keys_claim = keys_claim_response();

let users = keys_claim.one_time_keys.keys().map(Deref::deref);
let settings =
EncryptionSettings { only_allow_trusted_devices: true, ..Default::default() };
let settings = EncryptionSettings {
sharing_strategy: CollectStrategy::new_device_based(true),
..Default::default()
};

// Trust only one
let user_id = user_id!("@example:localhost");
Expand Down
Loading

0 comments on commit 3be84a5

Please sign in to comment.