Skip to content

refactor(crypto): Make the room key importing logic more generic #5041

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
May 15, 2025
Merged

refactor(crypto): Make the room key importing logic more generic #5041

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 16 additions & 2 deletions crates/matrix-sdk-crypto/src/olm/group_sessions/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,8 @@ use crate::types::events::forwarded_room_key::ForwardedMegolmV2AesSha2Content;
use crate::types::{
deserialize_curve_key, deserialize_curve_key_vec,
events::forwarded_room_key::{ForwardedMegolmV1AesSha2Content, ForwardedRoomKeyContent},
serialize_curve_key, serialize_curve_key_vec, EventEncryptionAlgorithm, SigningKey,
SigningKeys,
serialize_curve_key, serialize_curve_key_vec, EventEncryptionAlgorithm, RoomKeyExport,
SigningKey, SigningKeys,
};

/// An error type for the creation of group sessions.
Expand Down Expand Up @@ -141,6 +141,20 @@ impl ExportedRoomKey {
}
}

impl RoomKeyExport for &ExportedRoomKey {
fn room_id(&self) -> &ruma::RoomId {
&self.room_id
}

fn session_id(&self) -> &str {
&self.session_id
}

fn sender_key(&self) -> Curve25519PublicKey {
self.sender_key
}
}

/// A backed up version of an [`InboundGroupSession`].
///
/// This can be used to back up the [`InboundGroupSession`] to the server using
Expand Down
97 changes: 56 additions & 41 deletions crates/matrix-sdk-crypto/src/store/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,8 @@ use crate::{
},
types::{
events::room_key_withheld::RoomKeyWithheldEvent, BackupSecrets, CrossSigningSecrets,
EventEncryptionAlgorithm, MegolmBackupV1Curve25519AesSha2Secrets, SecretsBundle,
EventEncryptionAlgorithm, MegolmBackupV1Curve25519AesSha2Secrets, RoomKeyExport,
SecretsBundle,
},
verification::VerificationMachine,
CrossSigningStatus, OwnUserIdentityData, RoomKeyImportResult,
Expand Down Expand Up @@ -1832,6 +1833,54 @@ impl Store {
from_backup_version: Option<&str>,
progress_listener: impl Fn(usize, usize),
) -> Result<RoomKeyImportResult> {
let exported_keys: Vec<&ExportedRoomKey> = exported_keys.iter().collect();
self.import_sessions_impl(exported_keys, from_backup_version, progress_listener).await
}

/// Import the given room keys into our store.
///
/// # Arguments
///
/// * `exported_keys` - A list of previously exported keys that should be
/// imported into our store. If we already have a better version of a key
/// the key will *not* be imported.
///
/// Returns a tuple of numbers that represent the number of sessions that
/// were imported and the total number of sessions that were found in the
/// key export.
///
/// # Examples
///
/// ```no_run
/// # use std::io::Cursor;
/// # use matrix_sdk_crypto::{OlmMachine, decrypt_room_key_export};
/// # use ruma::{device_id, user_id};
/// # let alice = user_id!("@alice:example.org");
/// # async {
/// # let machine = OlmMachine::new(&alice, device_id!("DEVICEID")).await;
/// # let export = Cursor::new("".to_owned());
/// let exported_keys = decrypt_room_key_export(export, "1234").unwrap();
/// machine.store().import_exported_room_keys(exported_keys, |_, _| {}).await.unwrap();
/// # };
/// ```
pub async fn import_exported_room_keys(
&self,
exported_keys: Vec<ExportedRoomKey>,
progress_listener: impl Fn(usize, usize),
) -> Result<RoomKeyImportResult> {
self.import_room_keys(exported_keys, None, progress_listener).await
}

async fn import_sessions_impl<T>(
&self,
room_keys: Vec<T>,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can this be IntoIter<T> or something, to avoid the collect() above when we call it?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Kinda, though it can't be just that, because we have a room_keys.len() call.

Which then makes all of the trait bounds even more complex.

The other option is to remove the len() call, but then our progress listener doesn't work anymore.

Don't think it's worth it to avoid this collect() call, after all, allocating a vector of references isn't that big of a deal.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep, agree it's fine as-is.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(Or pass in the length?)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(Or pass in the length?)

That sounds like it would be a bit of a clunky API, not terrible but clunky enough for me to rather take the allocation.

The only reason why we take a Vec<&Foo> vs a Vec<Foo> is so we can gather info about the specific key in the error branch if creating a InboundGroupSession fails.

Perhaps gathering the info before the try_into() call consumes the key would be the wiser choice if we really want to avoid the allocation.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't mind either way.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Alright, let's keep the allocation, if we change our mind I think gathering the info before the consumption of the key would be the way to go.

from_backup_version: Option<&str>,
progress_listener: impl Fn(usize, usize),
) -> Result<RoomKeyImportResult>
where
T: TryInto<InboundGroupSession> + RoomKeyExport + Copy,
T::Error: Debug,
{
let mut sessions = Vec::new();

async fn new_session_better(
Expand All @@ -1845,11 +1894,11 @@ impl Store {
}
}

let total_count = exported_keys.len();
let total_count = room_keys.len();
let mut keys = BTreeMap::new();

for (i, key) in exported_keys.into_iter().enumerate() {
match InboundGroupSession::from_export(&key) {
for (i, key) in room_keys.into_iter().enumerate() {
match key.try_into() {
Ok(session) => {
let old_session = self
.inner
Expand All @@ -1875,9 +1924,9 @@ impl Store {
}
Err(e) => {
warn!(
sender_key= key.sender_key.to_base64(),
room_id = ?key.room_id,
session_id = key.session_id,
sender_key = key.sender_key().to_base64(),
room_id = ?key.room_id(),
session_id = key.session_id(),
error = ?e,
"Couldn't import a room key from a file export."
);
Expand All @@ -1896,40 +1945,6 @@ impl Store {
Ok(RoomKeyImportResult::new(imported_count, total_count, keys))
}

/// Import the given room keys into our store.
///
/// # Arguments
///
/// * `exported_keys` - A list of previously exported keys that should be
/// imported into our store. If we already have a better version of a key
/// the key will *not* be imported.
///
/// Returns a tuple of numbers that represent the number of sessions that
/// were imported and the total number of sessions that were found in the
/// key export.
///
/// # Examples
///
/// ```no_run
/// # use std::io::Cursor;
/// # use matrix_sdk_crypto::{OlmMachine, decrypt_room_key_export};
/// # use ruma::{device_id, user_id};
/// # let alice = user_id!("@alice:example.org");
/// # async {
/// # let machine = OlmMachine::new(&alice, device_id!("DEVICEID")).await;
/// # let export = Cursor::new("".to_owned());
/// let exported_keys = decrypt_room_key_export(export, "1234").unwrap();
/// machine.store().import_exported_room_keys(exported_keys, |_, _| {}).await.unwrap();
/// # };
/// ```
pub async fn import_exported_room_keys(
&self,
exported_keys: Vec<ExportedRoomKey>,
progress_listener: impl Fn(usize, usize),
) -> Result<RoomKeyImportResult> {
self.import_room_keys(exported_keys, None, progress_listener).await
}

pub(crate) fn crypto_store(&self) -> Arc<CryptoStoreWrapper> {
self.inner.store.clone()
}
Expand Down
14 changes: 13 additions & 1 deletion crates/matrix-sdk-crypto/src/types/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ use matrix_sdk_common::deserialized_responses::PrivOwnedStr;
use ruma::{
events::AnyToDeviceEvent,
serde::{Raw, StringEnum},
DeviceKeyAlgorithm, DeviceKeyId, OwnedDeviceKeyId, OwnedUserId, UserId,
DeviceKeyAlgorithm, DeviceKeyId, OwnedDeviceKeyId, OwnedUserId, RoomId, UserId,
};
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use vodozemac::{Curve25519PublicKey, Ed25519PublicKey, Ed25519Signature, KeyError};
Expand Down Expand Up @@ -660,3 +660,15 @@ impl ProcessedToDeviceEvent {
}
}
}

/// Trait to express the various room key export formats we have in a unified
/// manner.
pub trait RoomKeyExport {
/// The ID of the room where the exported room key was used.
fn room_id(&self) -> &RoomId;
/// The unique ID of the exported room key.
fn session_id(&self) -> &str;
/// The [Curve25519PublicKey] long-term identity key of the sender of this
/// room key.
fn sender_key(&self) -> Curve25519PublicKey;
}
Loading