crypto: withhold outgoing messages to unsigned dehydrated devices

Per #4313, we should not
send outgoing messages to dehydrated devices that are not signed by the current
pinned/verified identity.

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

@@ -160,15 +160,7 @@ pub(crate) async fn collect_session_recipients(
             for user_id in users {
                 trace!("Considering recipient devices for user {}", user_id);
                 let user_devices = store.get_device_data_for_user_filtered(user_id).await?;
-
-                // We only need the user identity if `only_allow_trusted_devices` or
-                // `error_on_verified_user_problem` is set.
-                let device_owner_identity =
-                    if only_allow_trusted_devices || error_on_verified_user_problem {
-                        store.get_user_identity(user_id).await?
-                    } else {
-                        None
-                    };
+                let device_owner_identity = store.get_user_identity(user_id).await?;
 
                 if error_on_verified_user_problem
                     && has_identity_verification_violation(
@@ -406,27 +398,88 @@ fn split_devices_for_user(
     for d in user_devices.into_values() {
         if d.is_blacklisted() {
             recipient_devices.denied_devices_with_code.push((d, WithheldCode::Blacklisted));
-        } else if d.local_trust_state() == LocalTrust::Ignored {
+            continue;
+        }
+
+        if d.local_trust_state() == LocalTrust::Ignored {
             // Ignore the trust state of that device and share
             recipient_devices.allowed_devices.push(d);
-        } else if only_allow_trusted_devices && !d.is_verified(own_identity, device_owner_identity)
-        {
+            continue;
+        }
+
+        if only_allow_trusted_devices && !d.is_verified(own_identity, device_owner_identity) {
             recipient_devices.denied_devices_with_code.push((d, WithheldCode::Unverified));
-        } else if error_on_verified_user_problem
+            continue;
+        }
+
+        // If `error_on_verified_user_problem` is set, and the user is verified, but the
+        // device is not signed, we fail the request by adding an entry to
+        // `unsigned_of_verified_user`.
+        if error_on_verified_user_problem
             && is_unsigned_device_of_verified_user(
                 own_identity.as_ref(),
                 device_owner_identity.as_ref(),
                 &d,
             )
         {
-            recipient_devices.unsigned_of_verified_user.push(d)
-        } else {
-            recipient_devices.allowed_devices.push(d);
+            recipient_devices.unsigned_of_verified_user.push(d);
+            continue;
         }
+
+        // If the device is dehydrated, it must be signed by the user, even if that user
+        // is not verified; further, that user must not have an identity violation.
+        let dehydrated = d.device_keys.dehydrated.unwrap_or(false);
+        if dehydrated {
+            if should_withhold_to_verified_device(
+                own_identity.as_ref(),
+                device_owner_identity.as_ref(),
+                &d,
+            ) {
+                recipient_devices.denied_devices_with_code.push((d, WithheldCode::Unverified));
+                continue;
+            }
+        }
+
+        recipient_devices.allowed_devices.push(d);
     }
     recipient_devices
 }
 
+fn should_withhold_to_verified_device(
+    own_identity: Option<&OwnUserIdentityData>,
+    device_owner_identity: Option<&UserIdentityData>,
+    d: &DeviceData,
+) -> bool {
+    let Some(owner_id) = device_owner_identity else {
+        // Dehydrated devices must be signed by their owners, and if we don't have the
+        // owner's identity, that can't be the case.
+        return true;
+    };
+
+    // If the user has changed identity since we first saw them, withhold.
+    //
+    // A pin violation is not well-defined for our own user, and in any case, there
+    // is no need to check it: such a failure will result in a verification
+    // violation.
+    if let UserIdentityData::Other(other) = owner_id {
+        if other.has_pin_violation() {
+            return true;
+        }
+    }
+
+    // If the user has changed identity since we verified them, withhold the message
+    if owner_id.was_previously_verified() && !is_user_verified(own_identity, owner_id) {
+        return true;
+    }
+
+    // Dehydrated devices must be signed by their owners
+    if !d.is_cross_signed_by_owner(owner_id) {
+        return true;
+    }
+
+    false
+}
+
 /// Result type for [`split_recipients_withhelds_for_user_based_on_identity`].
 #[derive(Default)]
 struct IdentityBasedRecipientDevices {
@@ -453,6 +506,9 @@ fn split_recipients_withhelds_for_user_based_on_identity(
             }
         }
         Some(device_owner_identity) => {
+            // TODO: if it is dehydrated, check that the current identity doesn't have a pin
+            //   violation?
+
             // Only accept devices signed by the current identity
             let (recipients, withheld_recipients): (
                 Vec<DeviceData>,
@@ -540,30 +596,33 @@ mod tests {
         CrossSigningKeyExport, EncryptionSettings, LocalTrust, OlmError, OlmMachine,
     };
 
-    async fn set_up_test_machine() -> OlmMachine {
-        let machine = OlmMachine::new(
-            KeyDistributionTestData::me_id(),
-            KeyDistributionTestData::me_device_id(),
-        )
-        .await;
+    /// Returns an `OlmMachine` set up for the test user in
+    /// [`KeyDistributionTestData`], with cross-signing set up and the
+    /// private cross-signing keys imported.
+    async fn test_machine() -> OlmMachine {
+        use KeyDistributionTestData as DataSet;
 
-        let keys_query = KeyDistributionTestData::me_keys_query_response();
-        let txn_id = TransactionId::new();
-        machine.mark_request_as_sent(&txn_id, &keys_query).await.unwrap();
+        // Create the local user (`@me`), and import the public identity keys
+        let machine = OlmMachine::new(DataSet::me_id(), DataSet::me_device_id()).await;
+        let keys_query = DataSet::me_keys_query_response();
+        machine.mark_request_as_sent(&TransactionId::new(), &keys_query).await.unwrap();
 
+        // Also import the private cross signing keys
         machine
             .import_cross_signing_keys(CrossSigningKeyExport {
-                master_key: KeyDistributionTestData::MASTER_KEY_PRIVATE_EXPORT.to_owned().into(),
-                self_signing_key: KeyDistributionTestData::SELF_SIGNING_KEY_PRIVATE_EXPORT
-                    .to_owned()
-                    .into(),
-                user_signing_key: KeyDistributionTestData::USER_SIGNING_KEY_PRIVATE_EXPORT
-                    .to_owned()
-                    .into(),
+                master_key: DataSet::MASTER_KEY_PRIVATE_EXPORT.to_owned().into(),
+                self_signing_key: DataSet::SELF_SIGNING_KEY_PRIVATE_EXPORT.to_owned().into(),
+                user_signing_key: DataSet::USER_SIGNING_KEY_PRIVATE_EXPORT.to_owned().into(),
             })
             .await
             .unwrap();
 
+        machine
+    }
+
+    /// Import device data for `@dan`, `@dave`, and `@good`, as referenced in
+    /// [`KeyDistributionTestData`], into the given OlmMachine
+    async fn import_known_users_to_test_machine(machine: &OlmMachine) {
         let keys_query = KeyDistributionTestData::dan_keys_query_response();
         let txn_id = TransactionId::new();
         machine.mark_request_as_sent(&txn_id, &keys_query).await.unwrap();
@@ -575,13 +634,12 @@ mod tests {
         let txn_id_good = TransactionId::new();
         let keys_query_good = KeyDistributionTestData::good_keys_query_response();
         machine.mark_request_as_sent(&txn_id_good, &keys_query_good).await.unwrap();
-
-        machine
     }
 
     #[async_test]
     async fn test_share_with_per_device_strategy_to_all() {
-        let machine = set_up_test_machine().await;
+        let machine = test_machine().await;
+        import_known_users_to_test_machine(&machine).await;
 
         let encryption_settings = EncryptionSettings {
             sharing_strategy: CollectStrategy::DeviceBasedStrategy {
@@ -643,7 +701,8 @@ mod tests {
     /// Common helper for [`test_share_with_per_device_strategy_only_trusted`]
     /// and [`test_share_with_per_device_strategy_only_trusted_error_on_unsigned_of_verified`].
     async fn test_share_only_trusted_helper(error_on_verified_user_problem: bool) {
-        let machine = set_up_test_machine().await;
+        let machine = test_machine().await;
+        import_known_users_to_test_machine(&machine).await;
 
         let encryption_settings = EncryptionSettings {
             sharing_strategy: CollectStrategy::DeviceBasedStrategy {
@@ -1086,9 +1145,177 @@ mod tests {
         .unwrap();
     }
 
+    /// A set of tests for the behaviour of
+    /// [`CollectStrategy::DeviceBasedStrategy`] with a dehydrated device
+    mod dehydrated_device {
+        use std::iter;
+
+        use insta::{assert_json_snapshot, with_settings};
+        use matrix_sdk_common::deserialized_responses::WithheldCode;
+        use matrix_sdk_test::{
+            async_test, ruma_response_to_json,
+            test_json::keys_query_sets::{
+                KeyQueryResponseTemplate, KeyQueryResponseTemplateDeviceOptions,
+            },
+        };
+        use ruma::{device_id, user_id, DeviceId, TransactionId, UserId};
+        use vodozemac::{Curve25519PublicKey, Ed25519SecretKey};
+
+        use super::{create_test_outbound_group_session, test_machine};
+        use crate::{
+            session_manager::group_sessions::{
+                share_strategy::collect_session_recipients, CollectRecipientsResult,
+            },
+            EncryptionSettings, OlmMachine,
+        };
+
+        /// Common helper for these tests: start a [`KeysQueryResponseTemplate`]
+        /// for the given user, with cross-signing keys.
+        fn cross_signed_key_query_response_template(user_id: &UserId) -> KeyQueryResponseTemplate {
+            KeyQueryResponseTemplate::new(user_id.to_owned()).with_cross_signing_keys(
+                Ed25519SecretKey::from_slice(b"master12master12master12master12"),
+                Ed25519SecretKey::from_slice(b"self1234self1234self1234self1234"),
+                Ed25519SecretKey::from_slice(b"user1234user1234user1234user1234"),
+            )
+        }
+
+        #[async_test]
+        async fn test_should_share_with_verified_dehydrated_device() {
+            let machine = test_machine().await;
+
+            // Bob is a user with cross-signing, who has a single dehydrated device.
+            let bob_user_id = user_id!("@bob:localhost");
+            let bob_dehydrated_device_id = device_id!("DEHYDRATED_DEVICE");
+            let keys_query = cross_signed_key_query_response_template(bob_user_id)
+                .with_device(
+                    bob_dehydrated_device_id,
+                    &Curve25519PublicKey::from(b"curvepubcurvepubcurvepubcurvepub".clone()),
+                    &Ed25519SecretKey::from_slice(b"device12device12device12device12"),
+                    KeyQueryResponseTemplateDeviceOptions::new().dehydrated(true).verified(true),
+                )
+                .build_response();
+            with_settings!({sort_maps => true}, { assert_json_snapshot!(ruma_response_to_json(keys_query.clone())) });
+            machine.mark_request_as_sent(&TransactionId::new(), &keys_query).await.unwrap();
+
+            // When we collect the recipients ...
+            let recips = share_test_session_and_collect_recipients(&machine, bob_user_id).await;
+
+            // ... then the dehydrated device should be included
+            let bob_devices_shared: Vec<_> = recips
+                .devices
+                .get(bob_user_id)
+                .expect(&format!("session not shared with {bob_user_id}"))
+                .iter()
+                .map(|d| d.device_id().to_owned())
+                .collect();
+            assert_eq!(bob_devices_shared, vec![bob_dehydrated_device_id]);
+        }
+
+        #[async_test]
+        async fn test_should_not_share_with_unverified_dehydrated_device() {
+            let machine = test_machine().await;
+
+            // Bob is a user with cross-signing, who has a single dehydrated device.
+            let bob_user_id = user_id!("@bob:localhost");
+            let bob_dehydrated_device_id = device_id!("DEHYDRATED_DEVICE");
+            let keys_query = cross_signed_key_query_response_template(bob_user_id)
+                .with_device(
+                    bob_dehydrated_device_id,
+                    &Curve25519PublicKey::from(b"curvepubcurvepubcurvepubcurvepub".clone()),
+                    &Ed25519SecretKey::from_slice(b"device12device12device12device12"),
+                    KeyQueryResponseTemplateDeviceOptions::new().dehydrated(true).verified(false),
+                )
+                .build_response();
+            with_settings!({sort_maps => true}, { assert_json_snapshot!(ruma_response_to_json(keys_query.clone())) });
+            machine.mark_request_as_sent(&TransactionId::new(), &keys_query).await.unwrap();
+
+            // When we collect the recipients ...
+            let recips = share_test_session_and_collect_recipients(&machine, bob_user_id).await;
+
+            // ... it shouldn't be shared with anyone, and there should be a withheld
+            // message for the dehydrated device.
+            assert_withheld_to(recips, &bob_user_id, bob_dehydrated_device_id);
+        }
+
+        #[async_test]
+        async fn test_should_not_share_with_verified_device_of_pin_violation_user() {
+            let machine = test_machine().await;
+
+            // Bob starts out with one identity
+            let bob_user_id = user_id!("@bob:localhost");
+            let keys_query = cross_signed_key_query_response_template(bob_user_id).build_response();
+            machine.mark_request_as_sent(&TransactionId::new(), &keys_query).await.unwrap();
+
+            // He then changes identity, and adds a dehydrated device (signed with his new
+            // identity)
+            let bob_dehydrated_device_id = device_id!("DEHYDRATED_DEVICE");
+            let keys_query = KeyQueryResponseTemplate::new(bob_user_id.to_owned())
+                .with_cross_signing_keys(
+                    Ed25519SecretKey::from_slice(b"newmaster__newmaster__newmaster_"),
+                    Ed25519SecretKey::from_slice(b"self1234self1234self1234self1234"),
+                    Ed25519SecretKey::from_slice(b"user1234user1234user1234user1234"),
+                )
+                .with_device(
+                    bob_dehydrated_device_id,
+                    &Curve25519PublicKey::from(b"curvepubcurvepubcurvepubcurvepub".clone()),
+                    &Ed25519SecretKey::from_slice(b"device12device12device12device12"),
+                    KeyQueryResponseTemplateDeviceOptions::new().dehydrated(true).verified(true),
+                )
+                .build_response();
+            with_settings!({sort_maps => true}, { assert_json_snapshot!(ruma_response_to_json(keys_query.clone())) });
+            machine.mark_request_as_sent(&TransactionId::new(), &keys_query).await.unwrap();
+
+            // When we collect the recipients ...
+            let recips = share_test_session_and_collect_recipients(&machine, bob_user_id).await;
+
+            // ... it shouldn't be shared with anyone, and there should be a withheld
+            // message for the dehydrated device.
+            assert_withheld_to(recips, &bob_user_id, bob_dehydrated_device_id);
+        }
+
+        /// Create a test megolm session and prepare to share it with the given
+        /// users. Returns a list of the devices belonging to that user
+        /// that received the session.
+        async fn share_test_session_and_collect_recipients(
+            machine: &OlmMachine,
+            target_user_id: &UserId,
+        ) -> CollectRecipientsResult {
+            let group_session =
+                create_test_outbound_group_session(&machine, &EncryptionSettings::default());
+            collect_session_recipients(
+                machine.store(),
+                iter::once(target_user_id),
+                &EncryptionSettings::default(),
+                &group_session,
+            )
+            .await
+            .unwrap()
+        }
+
+        /// Assert that the session is not shared with any devices, and that
+        /// there is a withheld code for the given device.
+        fn assert_withheld_to(
+            recips: CollectRecipientsResult,
+            bob_user_id: &UserId,
+            bob_dehydrated_device_id: &DeviceId,
+        ) {
+            // The share list should be empty
+            for (user, device_list) in recips.devices {
+                assert_eq!(device_list.len(), 0, "session shared with {}", user);
+            }
+
+            // ... and there should be one withheld message
+            assert_eq!(recips.withheld_devices.len(), 1);
+            assert_eq!(recips.withheld_devices[0].0.user_id(), bob_user_id);
+            assert_eq!(recips.withheld_devices[0].0.device_id(), bob_dehydrated_device_id);
+            assert_eq!(recips.withheld_devices[0].1, WithheldCode::Unverified);
+        }
+    }
+
     #[async_test]
     async fn test_share_with_identity_strategy() {
-        let machine = set_up_test_machine().await;
+        let machine = test_machine().await;
+        import_known_users_to_test_machine(&machine).await;
 
         let strategy = CollectStrategy::new_identity_based();
 
@@ -1388,7 +1615,8 @@ mod tests {
 
     #[async_test]
     async fn test_should_rotate_based_on_visibility() {
-        let machine = set_up_test_machine().await;
+        let machine = test_machine().await;
+        import_known_users_to_test_machine(&machine).await;
 
         let strategy = CollectStrategy::DeviceBasedStrategy {
             only_allow_trusted_devices: false,
@@ -1436,7 +1664,8 @@ mod tests {
     /// his devices.
     #[async_test]
     async fn test_should_rotate_based_on_device_excluded() {
-        let machine = set_up_test_machine().await;
+        let machine = test_machine().await;
+        import_known_users_to_test_machine(&machine).await;
 
         let fake_room_id = room_id!("!roomid:localhost");