feat(sdk): Poll encryption state on sync event for up to 3 seconds.

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

@@ -1908,29 +1908,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(())