Merge pull request #5559 from matrix-org/kaylendog/encrypt-poll-sync

This fixes a small issue encountered while developing the integration test for encrypted state events. If the first sync response received from the server after enabling encryption does not contain the m.room.encryption event, the client will report the encryption state as unknown, even if the next sync response does contain the event.

 - [x] Modify Room::enable_encryption_inner to spin on the room encryption state, rather than bailing after the first sync response is received.

Signed-off-by: Skye Elliot

Author: richvdh

crates/matrix-sdk/CHANGELOG.md

@@ -8,6 +8,10 @@ All notable changes to this project will be documented in this file.
 
 ### Features
 
+- `Room::enable_encryption` and `Room::enable_encryption_with_state_event_encryption` will poll
+  the encryption state for up to 3 seconds, rather than checking once after a single sync has
+  completed.
+  ([#5559](https://github.com/matrix-org/matrix-rust-sdk/pull/5559))
 - Add `Room::enable_encryption_with_state` to enable E2E encryption with encrypted state event
   support, gated behind the `experimental-encrypted-state-events` feature.
   ([#5557](https://github.com/matrix-org/matrix-rust-sdk/pull/5557))

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

@@ -1906,29 +1906,62 @@ impl Room {
             }
             self.send_state_event(content).await?;
 
+            // Spin on the sync beat event, since the first sync we receive might not
+            // include the encryption event.
+            //
             // TODO do we want to return an error here if we time out? This
             // could be quite useful if someone wants to enable encryption and
             // send a message right after it's enabled.
-            _ = timeout(self.client.inner.sync_beat.listen(), SYNC_WAIT_TIME).await;
+            let res = timeout(
+                async {
+                    loop {
+                        // Listen for sync events, then check if the encryption state is known.
+                        self.client.inner.sync_beat.listen().await;
+                        let _sync_lock = self.client.base_client().sync_lock().lock().await;
+                        if !self.inner.encryption_state().is_unknown() {
+                            break;
+                        }
+                    }
+                },
+                SYNC_WAIT_TIME,
+            )
+            .await;
 
-            // If after waiting for a sync, we don't have the encryption state we expect,
-            // assume the local encryption state is incorrect; this will cause
-            // the SDK to re-request it later for confirmation, instead of
-            // assuming it's sync'd and correct (and not encrypted).
             let _sync_lock = self.client.base_client().sync_lock().lock().await;
-            if !self.inner.encryption_state().is_encrypted() {
-                debug!("still not marked as encrypted, marking encryption state as missing");
 
-                let mut room_info = self.clone_info();
-                room_info.mark_encryption_state_missing();
-                let mut changes = StateChanges::default();
-                changes.add_room(room_info.clone());
+            // If encryption was enabled, return.
+            #[cfg(not(feature = "experimental-encrypted-state-events"))]
+            if res.is_ok() && self.inner.encryption_state().is_encrypted() {
+                debug!("room successfully marked as encrypted");
+                return Ok(());
+            }
 
-                self.client.state_store().save_changes(&changes).await?;
-                self.set_room_info(room_info, RoomInfoNotableUpdateReasons::empty());
-            } else {
+            // If encryption with state event encryption was enabled, return.
+            #[cfg(feature = "experimental-encrypted-state-events")]
+            if res.is_ok() && {
+                if encrypted_state_events {
+                    self.inner.encryption_state().is_state_encrypted()
+                } else {
+                    self.inner.encryption_state().is_encrypted()
+                }
+            } {
                 debug!("room successfully marked as encrypted");
+                return Ok(());
             }
+
+            // If after waiting for multiple syncs, we don't have the encryption state we
+            // expect, assume the local encryption state is incorrect; this will
+            // cause the SDK to re-request it later for confirmation, instead of
+            // assuming it's sync'd and correct (and not encrypted).
+            debug!("still not marked as encrypted, marking encryption state as missing");
+
+            let mut room_info = self.clone_info();
+            room_info.mark_encryption_state_missing();
+            let mut changes = StateChanges::default();
+            changes.add_room(room_info.clone());
+
+            self.client.state_store().save_changes(&changes).await?;
+            self.set_room_info(room_info, RoomInfoNotableUpdateReasons::empty());
         }
 
         Ok(())