Ignore to-device messages from dehydrated devices

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

@@ -852,9 +852,14 @@ impl OlmMachine {
         let mut decrypted =
             transaction.account().await?.decrypt_to_device_event(&self.inner.store, event).await?;
 
-        // Handle the decrypted event, e.g. fetch out Megolm sessions out of
-        // the event.
-        self.handle_decrypted_to_device_event(transaction.cache(), &mut decrypted, changes).await?;
+        // We ignore all to-device events from dehydrated devices - we should not
+        // receive any
+        if !self.event_is_from_dehydrated_device(&decrypted, &event.sender).await {
+            // Handle the decrypted event, e.g. fetch out Megolm sessions out of
+            // the event.
+            self.handle_decrypted_to_device_event(transaction.cache(), &mut decrypted, changes)
+                .await?;
+        }
 
         Ok(decrypted)
     }
@@ -1259,13 +1264,20 @@ impl OlmMachine {
         }
     }
 
+    /// Decrypt the supplied to-device event (if needed, and if we can) and
+    /// handle it.
+    ///
+    /// Return the same event, decrypted if possible and needed.
+    ///
+    /// If we can identify that this to-device event came from a dehydrated
+    /// device, this method does not process it, and returns `None`.
     #[instrument(skip_all, fields(sender, event_type, message_id))]
     async fn receive_to_device_event(
         &self,
         transaction: &mut StoreTransaction,
         changes: &mut Changes,
         mut raw_event: Raw<AnyToDeviceEvent>,
-    ) -> Raw<AnyToDeviceEvent> {
+    ) -> Option<Raw<AnyToDeviceEvent>> {
         Self::record_message_id(&raw_event);
 
         let event: ToDeviceEvents = match raw_event.deserialize_as() {
@@ -1274,7 +1286,7 @@ impl OlmMachine {
                 // Skip invalid events.
                 warn!("Received an invalid to-device event: {e}");
 
-                return raw_event;
+                return Some(raw_event);
             }
         };
 
@@ -1299,10 +1311,21 @@ impl OlmMachine {
                             }
                         }
 
-                        return raw_event;
+                        return Some(raw_event);
                     }
                 };
 
+                // We ignore all to-device events from dehydrated devices - we should not
+                // receive any
+                if self.event_is_from_dehydrated_device(&decrypted, &e.sender).await {
+                    warn!(
+                        sender = ?e.sender,
+                        session = ?decrypted.session,
+                        "Received a to-device event from a dehydrated device. This is unexpected"
+                    );
+                    return None;
+                }
+
                 // New sessions modify the account so we need to save that
                 // one as well.
                 match decrypted.session {
@@ -1336,7 +1359,54 @@ impl OlmMachine {
             e => self.handle_to_device_event(changes, &e).await,
         }
 
-        raw_event
+        Some(raw_event)
+    }
+
+    async fn event_is_from_dehydrated_device(
+        &self,
+        decrypted: &OlmDecryptionInfo,
+        sender_user_id: &UserId,
+    ) -> bool {
+        // Does the to-device message include device info?
+        if let Some(device_keys) = decrypted.result.event.sender_device_keys() {
+            // There is no need to check whether the device keys are signed correctly - any
+            // to-device message that claims to be from a dehydrated device is weird, so we
+            // will drop it.
+
+            // Does the included device info say the device is dehydrated?
+            if device_keys.dehydrated.unwrap_or(false) {
+                return true;
+            }
+            // If not, fall through and check our existing list of devices
+            // below, just in case the sender is sending us
+            // incorrect information embedded in the to-device message, but we
+            // know better.
+        }
+
+        // Do we already know about this device?
+        match self
+            .store()
+            .get_device_from_curve_key(sender_user_id, decrypted.result.sender_key)
+            .await
+        {
+            Ok(device) => {
+                // Yes - is it dehydrated?
+                if let Some(device) = device {
+                    device.is_dehydrated()
+                } else {
+                    false
+                }
+            }
+            Err(e) => {
+                // No - we received a to-device message from a device we don't know about. We
+                // continue to process it, but this is unexpected so we log an error.
+                warn!(
+                    error = ?e,
+                    "Received a to-device message from a device we don't recognise",
+                );
+                false
+            }
+        }
     }
 
     /// Handle a to-device and one-time key counts from a sync response.
@@ -1377,6 +1447,14 @@ impl OlmMachine {
         Ok((events, room_key_updates))
     }
 
+    /// Initial processing of the changes specified within a sync response.
+    ///
+    /// Returns the to-device events (decrypted where needed and where possible)
+    /// and the processed set of changes.
+    ///
+    /// If any of the to-device events in the supplied changes were sent from
+    /// dehydrated devices, these are not processed, and are omitted from
+    /// the returned list, as per MSC3814.
     pub(crate) async fn preprocess_sync_changes(
         &self,
         transaction: &mut StoreTransaction,
@@ -1412,7 +1490,10 @@ impl OlmMachine {
         for raw_event in sync_changes.to_device_events {
             let raw_event =
                 Box::pin(self.receive_to_device_event(transaction, &mut changes, raw_event)).await;
-            events.push(raw_event);
+
+            if let Some(raw_event) = raw_event {
+                events.push(raw_event);
+            }
         }
 
         let changed_sessions = self

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

@@ -34,7 +34,7 @@ use ruma::{
 use serde_json::json;
 
 use crate::{
-    store::Changes,
+    store::{Changes, MemoryStore},
     types::{events::ToDeviceEvent, requests::AnyOutgoingRequest},
     CrossSigningBootstrapRequests, DeviceData, OlmMachine,
 };
@@ -102,6 +102,23 @@ pub async fn get_machine_after_query_test_helper() -> (OlmMachine, OneTimeKeys)
     (machine, otk)
 }
 
+pub async fn get_machine_pair_using_store(
+    alice: &UserId,
+    bob: &UserId,
+    use_fallback_key: bool,
+    alice_store: MemoryStore,
+    alice_device_id: &DeviceId,
+) -> (OlmMachine, OlmMachine, OneTimeKeys) {
+    let (bob, otk) = get_prepared_machine_test_helper(bob, use_fallback_key).await;
+
+    let alice = OlmMachine::with_store(alice, alice_device_id, alice_store, None)
+        .await
+        .expect("Failed to create OlmMachine from supplied store");
+
+    store_each_others_device_data(&alice, &bob).await;
+    (alice, bob, otk)
+}
+
 pub async fn get_machine_pair(
     alice: &UserId,
     bob: &UserId,
@@ -112,12 +129,34 @@ pub async fn get_machine_pair(
     let alice_device = alice_device_id();
     let alice = OlmMachine::new(alice, alice_device).await;
 
-    let alice_device = DeviceData::from_machine_test_helper(&alice).await.unwrap();
-    let bob_device = DeviceData::from_machine_test_helper(&bob).await.unwrap();
+    store_each_others_device_data(&alice, &bob).await;
+    (alice, bob, otk)
+}
+
+/// Store alice's device data in bob's store and vice versa
+async fn store_each_others_device_data(alice: &OlmMachine, bob: &OlmMachine) {
+    let alice_device = DeviceData::from_machine_test_helper(alice).await.unwrap();
+    let bob_device = DeviceData::from_machine_test_helper(bob).await.unwrap();
     alice.store().save_device_data(&[bob_device]).await.unwrap();
     bob.store().save_device_data(&[alice_device]).await.unwrap();
+}
 
-    (alice, bob, otk)
+/// Return a pair of [`OlmMachine`]s, with an olm session created on Alice's
+/// side, but with no message yet sent.
+///
+/// Create Alice's `OlmMachine` using the [`MemoryStore`] provided
+pub async fn get_machine_pair_with_session_using_store(
+    alice: &UserId,
+    bob: &UserId,
+    use_fallback_key: bool,
+    alice_store: MemoryStore,
+    alice_device_id: &DeviceId,
+) -> (OlmMachine, OlmMachine) {
+    let (alice, bob, one_time_keys) =
+        get_machine_pair_using_store(alice, bob, use_fallback_key, alice_store, alice_device_id)
+            .await;
+
+    build_session_for_pair(alice, bob, one_time_keys).await
 }
 
 /// Return a pair of [`OlmMachine`]s, with an olm session created on Alice's
@@ -127,8 +166,20 @@ pub async fn get_machine_pair_with_session(
     bob: &UserId,
     use_fallback_key: bool,
 ) -> (OlmMachine, OlmMachine) {
-    let (alice, bob, mut one_time_keys) = get_machine_pair(alice, bob, use_fallback_key).await;
+    let (alice, bob, one_time_keys) = get_machine_pair(alice, bob, use_fallback_key).await;
 
+    build_session_for_pair(alice, bob, one_time_keys).await
+}
+
+/// Create a session for the two supplied Olm machines to communicate.
+async fn build_session_for_pair(
+    alice: OlmMachine,
+    bob: OlmMachine,
+    mut one_time_keys: BTreeMap<
+        ruma::OwnedKeyId<ruma::OneTimeKeyAlgorithm, ruma::OneTimeKeyName>,
+        Raw<OneTimeKey>,
+    >,
+) -> (OlmMachine, OlmMachine) {
     let (device_key_id, one_time_key) = one_time_keys.pop_first().unwrap();
 
     let one_time_keys = BTreeMap::from([(