feat(widgetDriver): add to-device encryption support

Author: toger5

crates/matrix-sdk/src/widget/machine/driver_req.rs

@@ -273,6 +273,7 @@ impl FromMatrixDriverResponse for SendEventResponse {
 
 /// Ask the client to send a to-device message that corresponds to the given
 /// description.
+/// see [MSC4140](https://github.com/matrix-org/matrix-spec-proposals/pull/4140)) as a response.
 #[derive(Clone, Debug, Deserialize)]
 pub(crate) struct SendToDeviceRequest {
     /// The type of the to-device message.

crates/matrix-sdk/src/widget/matrix.rs

@@ -17,6 +17,7 @@
 
 use std::collections::BTreeMap;
 
+use futures_util::future::join_all;
 use matrix_sdk_base::deserialized_responses::{EncryptionInfo, RawAnySyncOrStrippedState};
 use ruma::{
     api::client::{
@@ -197,10 +198,11 @@ impl MatrixDriver {
                 async {}
             });
         let drop_guard_msg_like = self.room.client().event_handler_drop_guard(handle_msg_like);
-
+        let _room_id = room_id;
+        let _tx = tx;
         // Get only all state events from the state section of the sync.
         let handle_state = self.room.add_event_handler(move |raw: Raw<AnySyncStateEvent>| {
-            let _ = tx.send(attach_room_id(raw.cast_ref(), &room_id));
+            let _ = _tx.send(attach_room_id(raw.cast_ref(), &_room_id));
             async {}
         });
         let drop_guard_state = self.room.client().event_handler_drop_guard(handle_state);
@@ -224,8 +226,11 @@ impl MatrixDriver {
             // EncryptionInfo. The widgetAPI expects a boolean `encrypted` to be added
             // (!) to the raw content to know if the to-device message was encrypted or
             // not (as per MSC3819).
-            move |raw: Raw<AnyToDeviceEvent>, _: Option<EncryptionInfo>| {
-                let _ = tx.send(raw);
+            move |raw: Raw<AnyToDeviceEvent>, encryption_info: Option<EncryptionInfo>| {
+                // Only sent encrypted to-device events
+                if encryption_info.is_some() {
+                    let _ = tx.send(raw);
+                }
                 async {}
             },
         );
@@ -248,9 +253,36 @@ impl MatrixDriver {
         let client = self.room.client();
 
         let request = if encrypted {
-            return Err(Error::UnknownError(
-                "Sending encrypted to-device events is not supported by the widget driver.".into(),
-            ));
+            // We first want to get all missing session before we start any to device
+            // sending!
+            client.claim_one_time_keys(messages.keys().map(|u| u.as_ref())).await?;
+            let encrypted_content: BTreeMap<
+                OwnedUserId,
+                BTreeMap<DeviceIdOrAllDevices, Raw<AnyToDeviceEventContent>>,
+            > = join_all(messages.into_iter().map(|(user_id, device_content_map)| {
+                let event_type = event_type.clone();
+                async move {
+                    (
+                        user_id.clone(),
+                        to_device_crypto::encrypted_device_content_map(
+                            &self.room.client(),
+                            &user_id,
+                            &event_type,
+                            device_content_map,
+                        )
+                        .await,
+                    )
+                }
+            }))
+            .await
+            .into_iter()
+            .collect();
+
+            RumaToDeviceRequest::new_raw(
+                ToDeviceEventType::RoomEncrypted,
+                TransactionId::new(),
+                encrypted_content,
+            )
         } else {
             RumaToDeviceRequest::new_raw(event_type, TransactionId::new(), messages)
         };
@@ -280,13 +312,153 @@ fn attach_room_id(raw_ev: &Raw<AnySyncTimelineEvent>, room_id: &RoomId) -> Raw<A
     Raw::new(&ev_obj).unwrap().cast()
 }
 
+/// Move this into the `matrix_crypto` crate!
+/// This module contains helper functions to encrypt to device events.
+mod to_device_crypto {
+    use std::collections::BTreeMap;
+
+    use futures_util::future::join_all;
+    use ruma::{
+        events::{AnyToDeviceEventContent, ToDeviceEventType},
+        serde::Raw,
+        to_device::DeviceIdOrAllDevices,
+        UserId,
+    };
+    use serde_json::Value;
+    use tracing::{info, warn};
+
+    use crate::{encryption::identities::Device, executor::spawn, Client, Error, Result};
+
+    /// This encrypts to device content for a collection of devices.
+    /// It will ignore all devices where errors occurred or where the device
+    /// is not verified or where th user has a has_verification_violation.
+    async fn encrypted_content_for_devices(
+        unencrypted_content: &Raw<AnyToDeviceEventContent>,
+        devices: Vec<Device>,
+        event_type: &ToDeviceEventType,
+    ) -> Result<impl Iterator<Item = (DeviceIdOrAllDevices, Raw<AnyToDeviceEventContent>)>> {
+        let content: Value = unencrypted_content.deserialize_as().map_err(Into::<Error>::into)?;
+        let event_type = event_type.clone();
+        let device_content_tasks = devices.into_iter().map(|device| spawn({
+                let event_type = event_type.clone();
+                let content = content.clone();
+
+                async move {
+                    // This is not yet used. It is incompatible with the spa guest mode (the spa will not verify its crypto identity)
+                    // if !device.is_cross_signed_by_owner() {
+                    //     info!("Device {} is not verified, skipping encryption", device.device_id());
+                    //     return None;
+                    // }
+                    match device
+                        .inner
+                        .encrypt_event_raw(&event_type.to_string(), &content)
+                        .await {
+                            Ok(encrypted) => Some((device.device_id().to_owned().into(), encrypted.cast())),
+                            Err(e) =>{ info!("Failed to encrypt to_device event from widget for device: {} because, {}", device.device_id(), e); None},
+                        }
+                }
+            }));
+        let device_encrypted_content_map =
+            join_all(device_content_tasks).await.into_iter().flatten().flatten();
+        Ok(device_encrypted_content_map)
+    }
+
+    /// Convert the device content map for one user into the same content
+    /// map with encrypted content This needs to flatten the vectors
+    /// we get from `encrypted_content_for_devices`
+    /// since one `DeviceIdOrAllDevices` id can be multiple devices.
+    pub(super) async fn encrypted_device_content_map(
+        client: &Client,
+        user_id: &UserId,
+        event_type: &ToDeviceEventType,
+        device_content_map: BTreeMap<DeviceIdOrAllDevices, Raw<AnyToDeviceEventContent>>,
+    ) -> BTreeMap<DeviceIdOrAllDevices, Raw<AnyToDeviceEventContent>> {
+        let device_map_futures =
+                device_content_map.into_iter().map(|(device_or_all_id, content)| spawn({
+                    let client = client.clone();
+                    let user_id = user_id.to_owned();
+                    let event_type = event_type.clone();
+                    async move {
+                        let Ok(user_devices) = client.encryption().get_user_devices(&user_id).await else {
+                            warn!("Failed to get user devices for user: {}", user_id);
+                            return None;
+                        };
+                        // This is not yet used. It is incompatible with the spa guest mode (the spa will not verify its crypto identity)
+                        // let Ok(user_identity) = client.encryption().get_user_identity(&user_id).await else{
+                        //     warn!("Failed to get user identity for user: {}", user_id);
+                        //     return None;
+                        // };
+                        // if user_identity.map(|i|i.has_verification_violation()).unwrap_or(false) {
+                        //     info!("User {} has a verification violation, skipping encryption", user_id);
+                        //     return None;
+                        // }
+                        let devices: Vec<Device> = match device_or_all_id {
+                            DeviceIdOrAllDevices::DeviceId(device_id) => {
+                                vec![user_devices.get(&device_id)].into_iter().flatten().collect()
+                            }
+                            DeviceIdOrAllDevices::AllDevices => user_devices.devices().collect(),
+                        };
+                        encrypted_content_for_devices(
+                            &content,
+                            devices,
+                            &event_type,
+                        )
+                        .await
+                        .map_err(|e| info!("WidgetDriver: could not encrypt content for to device widget event content: {}. because, {}", content.json(), e))
+                        .ok()
+                }}));
+        let content_map_iterator = join_all(device_map_futures).await.into_iter();
+
+        // The first flatten takes the iterator over Result<Option<impl Iterator<Item =
+        // (DeviceIdOrAllDevices, Raw<AnyToDeviceEventContent>)>>, JoinError>>
+        // and flattens the Result (drops Err() items)
+        // The second takes the iterator over: Option<impl Iterator<Item =
+        // (DeviceIdOrAllDevices, Raw<AnyToDeviceEventContent>)>>
+        // and flattens the Option (drops None items)
+        // The third takes the iterator over iterators: impl Iterator<Item =
+        // (DeviceIdOrAllDevices, Raw<AnyToDeviceEventContent>)>
+        // and flattens it to just an iterator over (DeviceIdOrAllDevices,
+        // Raw<AnyToDeviceEventContent>)
+        content_map_iterator.flatten().flatten().flatten().collect()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use insta;
-    use ruma::{events::AnyTimelineEvent, room_id, serde::Raw};
+    use ruma::{
+        events::{AnySyncTimelineEvent, AnyTimelineEvent},
+        room_id,
+        serde::Raw,
+    };
     use serde_json::{json, Value};
 
     use super::attach_room_id;
+    #[test]
+    fn test_app_props_to_raw() {
+        let raw = Raw::new(&json!({
+            "encrypted": true,
+            "type": "m.room.message",
+            "content": {
+                "body": "Hello world"
+            }
+        }))
+        .unwrap()
+        .cast::<AnySyncTimelineEvent>();
+        let room_id = room_id!("!my_id:example.org");
+        let new = attach_room_id(&raw, room_id);
+        assert_eq!(
+            serde_json::to_value(new).unwrap(),
+            json!({
+                "encrypted": true,
+                "room_id": "!my_id:example.org",
+                "type": "m.room.message",
+                "content": {
+                    "body": "Hello world"
+                }
+            })
+        );
+    }
 
     #[test]
     fn test_add_room_id_to_raw() {