Improve performance of share_room_key

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

@@ -431,28 +431,6 @@ impl Device {
         self.inner.encrypt(self.verification_machine.store.inner(), event_type, content).await
     }
 
-    pub(crate) async fn maybe_encrypt_room_key(
-        &self,
-        session: OutboundGroupSession,
-    ) -> OlmResult<MaybeEncryptedRoomKey> {
-        let content = session.as_content().await;
-        let message_index = session.message_index().await;
-        let event_type = content.event_type();
-
-        match self.encrypt(event_type, content).await {
-            Ok((session, encrypted)) => Ok(MaybeEncryptedRoomKey::Encrypted {
-                share_info: ShareInfo::new_shared(session.sender_key().to_owned(), message_index),
-                used_session: session,
-                message: encrypted.cast(),
-            }),
-
-            Err(OlmError::MissingSession | OlmError::EventError(EventError::MissingSenderKey)) => {
-                Ok(MaybeEncryptedRoomKey::Withheld { code: WithheldCode::NoOlm })
-            }
-            Err(e) => Err(e),
-        }
-    }
-
     /// Encrypt the given inbound group session as a forwarded room key for this
     /// device.
     pub async fn encrypt_room_key_for_forwarding(
@@ -788,6 +766,29 @@ impl ReadOnlyDevice {
         }
     }
 
+    pub(crate) async fn maybe_encrypt_room_key(
+        &self,
+        store: &CryptoStoreWrapper,
+        session: OutboundGroupSession,
+    ) -> OlmResult<MaybeEncryptedRoomKey> {
+        let content = session.as_content().await;
+        let message_index = session.message_index().await;
+        let event_type = content.event_type();
+
+        match self.encrypt(store, event_type, content).await {
+            Ok((session, encrypted)) => Ok(MaybeEncryptedRoomKey::Encrypted {
+                share_info: ShareInfo::new_shared(session.sender_key().to_owned(), message_index),
+                used_session: session,
+                message: encrypted.cast(),
+            }),
+
+            Err(OlmError::MissingSession | OlmError::EventError(EventError::MissingSenderKey)) => {
+                Ok(MaybeEncryptedRoomKey::Withheld { code: WithheldCode::NoOlm })
+            }
+            Err(e) => Err(e),
+        }
+    }
+
     /// Update a device with a new device keys struct.
     pub(crate) fn update_device(&mut self, device_keys: &DeviceKeys) -> Result<(), SignatureError> {
         self.verify_device_keys(device_keys)?;

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

@@ -35,7 +35,7 @@ use crate::{
     error::{EventError, MegolmResult, OlmResult},
     identities::device::MaybeEncryptedRoomKey,
     olm::{InboundGroupSession, OutboundGroupSession, Session, ShareInfo, ShareState},
-    store::{Changes, Result as StoreResult, Store},
+    store::{Changes, CryptoStoreWrapper, Result as StoreResult, Store},
     types::events::{room::encrypted::RoomEncryptedEventContent, room_key_withheld::WithheldCode},
     Device, EncryptionSettings, OlmError, ToDeviceRequest,
 };
@@ -263,6 +263,7 @@ impl GroupSessionManager {
     /// Encrypt the given content for the given devices and create a to-device
     /// requests that sends the encrypted content to them.
     async fn encrypt_session_for(
+        store: Arc<CryptoStoreWrapper>,
         group_session: OutboundGroupSession,
         devices: Vec<Device>,
     ) -> OlmResult<(
@@ -283,14 +284,21 @@ impl GroupSessionManager {
         let mut share_infos = BTreeMap::new();
         let mut withheld_devices = Vec::new();
 
-        let encrypt = |device: Device, session: OutboundGroupSession| async move {
-            let encryption_result = device.maybe_encrypt_room_key(session).await?;
+        // XXX is there a way to do this that doesn't involve cloning the
+        // `Arc<CryptoStoreWrapper>` fore each device?
+        let encrypt = |store: Arc<CryptoStoreWrapper>,
+                       device: Device,
+                       session: OutboundGroupSession| async move {
+            let encryption_result =
+                device.inner.maybe_encrypt_room_key(store.as_ref(), session).await?;
 
             Ok::<_, OlmError>(DeviceResult { device, maybe_encrypted_room_key: encryption_result })
         };
 
-        let tasks: Vec<_> =
-            devices.iter().map(|d| spawn(encrypt(d.clone(), group_session.clone()))).collect();
+        let tasks: Vec<_> = devices
+            .iter()
+            .map(|d| spawn(encrypt(store.clone(), d.clone(), group_session.clone())))
+            .collect();
 
         let results = join_all(tasks).await;
 
@@ -428,12 +436,13 @@ impl GroupSessionManager {
     }
 
     async fn encrypt_request(
+        store: Arc<CryptoStoreWrapper>,
         chunk: Vec<Device>,
         outbound: OutboundGroupSession,
         sessions: GroupSessionCache,
     ) -> OlmResult<(Vec<Session>, Vec<(Device, WithheldCode)>)> {
         let (id, request, share_infos, used_sessions, no_olm) =
-            Self::encrypt_session_for(outbound.clone(), chunk).await?;
+            Self::encrypt_session_for(store, outbound.clone(), chunk).await?;
 
         if !request.messages.is_empty() {
             trace!(
@@ -520,6 +529,7 @@ impl GroupSessionManager {
             .chunks(Self::MAX_TO_DEVICE_MESSAGES)
             .map(|chunk| {
                 spawn(Self::encrypt_request(
+                    self.store.crypto_store(),
                     chunk.to_vec(),
                     group_session.clone(),
                     self.sessions.clone(),

crates/matrix-sdk-crypto/src/store/mod.rs

@@ -1591,6 +1591,10 @@ impl Store {
     ) -> Result<RoomKeyImportResult> {
         self.import_room_keys(exported_keys, false, progress_listener).await
     }
+
+    pub(crate) fn crypto_store(&self) -> Arc<CryptoStoreWrapper> {
+        self.inner.store.clone()
+    }
 }
 
 impl Deref for Store {