feat(crypto): Add new encrypt_and_send_custom_to_device to the client

Author: BillCarsonFr

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

@@ -23,7 +23,7 @@ use matrix_sdk_common::{
     deserialized_responses::{
         AlgorithmInfo, DecryptedRoomEvent, DeviceLinkProblem, EncryptionInfo, UnableToDecryptInfo,
         UnableToDecryptReason, UnsignedDecryptionResult, UnsignedEventLocation, VerificationLevel,
-        VerificationState,
+        VerificationState, WithheldCode,
     },
     locks::RwLock as StdRwLock,
     BoxFuture,
@@ -53,7 +53,7 @@ use tokio::sync::Mutex;
 use tracing::{
     debug, error,
     field::{debug, display},
-    info, instrument, warn, Span,
+    info, instrument, trace, warn, Span,
 };
 use vodozemac::{
     megolm::{DecryptionError, SessionOrdering},
@@ -1111,6 +1111,49 @@ impl OlmMachine {
         self.inner.group_session_manager.share_room_key(room_id, users, encryption_settings).await
     }
 
+    /// Encrypts the given content using Olm for each of the given devices.
+    ///
+    /// The 1-to-1 session must be established prior to this
+    /// call by using the [`OlmMachine::get_missing_sessions`] method or the
+    /// encryption will fail.
+    ///
+    /// The caller is responsible for sending the encrypted
+    /// event to the target device, and should do it ASAP to avoid out-of-order
+    /// messages.
+    ///
+    /// # Returns
+    /// A list of `ToDeviceRequest` to send out the event, and the list of
+    /// devices where encryption did not succeed (device excluded or no olm)
+    pub async fn encrypt_content_for_devices(
+        &self,
+        devices: Vec<DeviceData>,
+        event_type: &str,
+        content: &Value,
+    ) -> OlmResult<(Vec<ToDeviceRequest>, Vec<(DeviceData, WithheldCode)>)> {
+        // TODO: Use a `CollectStrategy` arguments to filter our devices depending on
+        // safety settings (like not sending to insecure devices).
+        let mut changes = Changes::default();
+
+        let result = self
+            .inner
+            .group_session_manager
+            .encrypt_content_for_devices(devices, event_type, content.clone(), &mut changes)
+            .await;
+
+        // Persist any changes we might have collected.
+        if !changes.is_empty() {
+            let session_count = changes.sessions.len();
+
+            self.inner.store.save_changes(changes).await?;
+
+            trace!(
+                session_count = session_count,
+                "Stored the changed sessions after encrypting a custom to device event"
+            );
+        }
+
+        result
+    }
     /// Collect the devices belonging to the given user, and send the details of
     /// a room key bundle to those devices.
     ///

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

@@ -827,7 +827,7 @@ impl GroupSessionManager {
     /// Returns a tuple containing (1) the list of to-device requests, and (2)
     /// the list of devices that we could not find an olm session for (so
     /// need a withheld message).
-    async fn encrypt_content_for_devices(
+    pub(crate) async fn encrypt_content_for_devices(
         &self,
         recipient_devices: Vec<DeviceData>,
         event_type: &str,

crates/matrix-sdk/src/client/mod.rs

@@ -14,6 +14,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#[cfg(feature = "e2e-encryption")]
+use std::ops::Deref;
 use std::{
     collections::{btree_map, BTreeMap},
     fmt::{self, Debug},
@@ -37,8 +39,6 @@ use matrix_sdk_base::{
     StateStoreDataKey, StateStoreDataValue, SyncOutsideWasm,
 };
 use matrix_sdk_common::ttl_cache::TtlCache;
-#[cfg(feature = "e2e-encryption")]
-use ruma::events::{room::encryption::RoomEncryptionEventContent, InitialStateEvent};
 use ruma::{
     api::{
         client::{
@@ -69,6 +69,15 @@ use ruma::{
     DeviceId, OwnedDeviceId, OwnedEventId, OwnedRoomId, OwnedRoomOrAliasId, OwnedServerName,
     RoomAliasId, RoomId, RoomOrAliasId, ServerName, UInt, UserId,
 };
+#[cfg(feature = "e2e-encryption")]
+use ruma::{
+    events::{
+        room::encryption::RoomEncryptionEventContent, AnyToDeviceEventContent, InitialStateEvent,
+    },
+    serde::Raw,
+    to_device::DeviceIdOrAllDevices,
+    OwnedUserId,
+};
 use serde::de::DeserializeOwned;
 use tokio::sync::{broadcast, Mutex, OnceCell, RwLock, RwLockReadGuard};
 use tracing::{debug, error, instrument, trace, warn, Instrument, Span};
@@ -99,7 +108,9 @@ use crate::{
 };
 #[cfg(feature = "e2e-encryption")]
 use crate::{
-    encryption::{Encryption, EncryptionData, EncryptionSettings, VerificationState},
+    encryption::{
+        identities::Device, Encryption, EncryptionData, EncryptionSettings, VerificationState,
+    },
     store_locks::CrossProcessStoreLock,
 };
 
@@ -2513,6 +2524,74 @@ impl Client {
         let base_room = self.inner.base_client.room_knocked(&response.room_id).await?;
         Ok(Room::new(self.clone(), base_room))
     }
+
+    /// Encrypts then send the given content via the `sendToDevice` end-point
+    /// using olm encryption.
+    ///
+    /// If there are a lot of targets this will be break down by chunks.
+    ///
+    /// # Returns
+    /// A list of `ToDeviceRequest` to send out the event, and the list of
+    /// devices where encryption did not succeed (device excluded or no olm)
+    #[cfg(feature = "e2e-encryption")]
+    pub async fn encrypt_and_send_custom_to_device(
+        &self,
+        targets: Vec<&Device>,
+        event_type: &str,
+        content: Raw<AnyToDeviceEventContent>,
+    ) -> Result<Vec<(OwnedUserId, OwnedDeviceId)>> {
+        let users = targets.iter().map(|device| device.user_id());
+
+        // Will claim one-time-key for users that needs it
+        // TODO: For later optimisation: This will establish missing olm sessions with
+        // all this users devices, but we just want for some devices.
+        self.claim_one_time_keys(users).await?;
+
+        let olm = self.olm_machine().await;
+        let olm = olm.as_ref().expect("Olm machine wasn't started");
+
+        let (requests, withhelds) = olm
+            .encrypt_content_for_devices(
+                targets.into_iter().map(|d| d.deref().clone()).collect(),
+                event_type,
+                &content
+                    .deserialize_as::<serde_json::Value>()
+                    .expect("Deserialize as Value will always work"),
+            )
+            .await?;
+
+        let mut failures: Vec<(OwnedUserId, OwnedDeviceId)> = Default::default();
+
+        // Push the withhelds in the failures
+        withhelds.iter().for_each(|(d, _)| {
+            failures.push((d.user_id().to_owned(), d.device_id().to_owned()));
+        });
+
+        // TODO: parallelize that? it's already grouping 250 devices per chunk.
+        for request in requests {
+            let send_result =
+                self.send_to_device_with_config(&request, RequestConfig::short_retry()).await;
+
+            // If the sending failed we need to collect the failures to report them
+            if send_result.is_err() {
+                // Mark the sending as failed
+                for (user_id, device_map) in request.messages {
+                    for device_id in device_map.keys() {
+                        match device_id {
+                            DeviceIdOrAllDevices::DeviceId(device_id) => {
+                                failures.push((user_id.clone(), device_id.to_owned()));
+                            }
+                            DeviceIdOrAllDevices::AllDevices => {
+                                // Cannot happen in this case
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        Ok(failures)
+    }
 }
 
 /// A weak reference to the inner client, useful when trying to get a handle

crates/matrix-sdk/src/encryption/mod.rs

@@ -99,6 +99,7 @@ pub use matrix_sdk_base::crypto::{
     SessionCreationError, SignatureError, VERSION,
 };
 
+use crate::config::RequestConfig;
 pub use crate::error::RoomKeyImportError;
 
 /// All the data related to the encryption state.
@@ -572,6 +573,20 @@ impl Client {
         self.send(request).await
     }
 
+    pub(crate) async fn send_to_device_with_config(
+        &self,
+        request: &ToDeviceRequest,
+        config: RequestConfig,
+    ) -> HttpResult<ToDeviceResponse> {
+        let request = RumaToDeviceRequest::new_raw(
+            request.event_type.clone(),
+            request.txn_id.clone(),
+            request.messages.clone(),
+        );
+
+        self.send_inner(request, Some(config), Default::default()).await
+    }
+
     pub(crate) async fn send_verification_request(
         &self,
         request: OutgoingVerificationRequest,

crates/matrix-sdk/tests/integration/encryption.rs

@@ -2,6 +2,7 @@ mod backups;
 mod cross_signing;
 mod recovery;
 mod secret_storage;
+mod to_device;
 mod verification;
 
 /// The backup key, which is also returned (encrypted) as part of the secret