feat(crypto): Implement OldMachine::mark_all_tracked_users_as_dirty.

This patch adds the `OldMachine::mark_all_tracked_users_as_dirty`.

This patch rewrites a bit `OlmMachine::new_helper` by extracting some
piece of it inside `OlmMachine::new_helper_prelude`. With that, we
can rewrite `OlmMachine::migration_post_verified_latch_support` to use
`IdentityManager::mark_all_tracked_users_as_dirty`.
This latter is the shared implementation with
`OlmMachine::mark_all_tracked_users_as_dirty`.

This patch adds a test for `OlmMachine:mark_all_tracked_users_as_dirty`.

crates/matrix-sdk-crypto/src/identities/manager.rs

@@ -36,7 +36,7 @@ use crate::{
     requests::KeysQueryRequest,
     store::{
         caches::SequenceNumber, Changes, DeviceChanges, IdentityChanges, KeyQueryManager,
-        Result as StoreResult, Store, StoreCache, UserKeyQueryResult,
+        Result as StoreResult, Store, StoreCache, StoreCacheGuard, UserKeyQueryResult,
     },
     types::{CrossSigningKey, DeviceKeys, MasterPubkey, SelfSigningPubkey, UserSigningPubkey},
     CryptoStoreError, LocalTrust, OwnUserIdentity, SignatureError, UserIdentities,
@@ -1147,6 +1147,28 @@ impl IdentityManager {
 
         Ok(())
     }
+
+    /// Mark all tracked users as dirty.
+    ///
+    /// See `SyncedKeyQueryManager::mark_tracked_users_as_changed()` to learn
+    /// more.
+    pub(crate) async fn mark_all_tracked_users_as_dirty(
+        &self,
+        store_cache: StoreCacheGuard,
+    ) -> StoreResult<()> {
+        let store_wrapper = store_cache.store_wrapper();
+        let tracked_users = store_wrapper.load_tracked_users().await?;
+
+        self.key_query_manager
+            .synced(&store_cache)
+            .await?
+            .mark_tracked_users_as_changed(
+                tracked_users.iter().map(|tracked_user| tracked_user.user_id.as_ref()),
+            )
+            .await?;
+
+        Ok(())
+    }
 }
 
 /// Log information about what changed after processing a /keys/query response.

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

@@ -185,30 +185,43 @@ impl OlmMachine {
             })
             .await?;
 
+        let (verification_machine, store, identity_manager) =
+            Self::new_helper_prelude(store, static_account, self.store().private_identity());
+
         Ok(Self::new_helper(
             device_id,
             store,
-            static_account,
+            verification_machine,
+            identity_manager,
             self.store().private_identity(),
             None,
         ))
     }
 
+    fn new_helper_prelude(
+        store_wrapper: Arc<CryptoStoreWrapper>,
+        account: StaticAccountData,
+        user_identity: Arc<Mutex<PrivateCrossSigningIdentity>>,
+    ) -> (VerificationMachine, Store, IdentityManager) {
+        let verification_machine =
+            VerificationMachine::new(account.clone(), user_identity.clone(), store_wrapper.clone());
+        let store = Store::new(account, user_identity, store_wrapper, verification_machine.clone());
+
+        let identity_manager = IdentityManager::new(store.clone());
+
+        (verification_machine, store, identity_manager)
+    }
+
     fn new_helper(
         device_id: &DeviceId,
-        store: Arc<CryptoStoreWrapper>,
-        account: StaticAccountData,
+        store: Store,
+        verification_machine: VerificationMachine,
+        identity_manager: IdentityManager,
         user_identity: Arc<Mutex<PrivateCrossSigningIdentity>>,
         maybe_backup_key: Option<MegolmV1BackupKey>,
     ) -> Self {
-        let verification_machine =
-            VerificationMachine::new(account.clone(), user_identity.clone(), store.clone());
-        let store = Store::new(account, user_identity.clone(), store, verification_machine.clone());
-
         let group_session_manager = GroupSessionManager::new(store.clone());
 
-        let identity_manager = IdentityManager::new(store.clone());
-
         let users_for_key_claim = Arc::new(StdRwLock::new(BTreeMap::new()));
         let key_request_machine = GossipMachine::new(
             store.clone(),
@@ -360,11 +373,21 @@ impl OlmMachine {
         let identity = Arc::new(Mutex::new(identity));
         let store = Arc::new(CryptoStoreWrapper::new(user_id, device_id, store));
 
+        let (verification_machine, store, identity_manager) =
+            Self::new_helper_prelude(store, static_account, identity.clone());
+
         // FIXME: We might want in the future a more generic high-level data migration
         // mechanism (at the store wrapper layer).
-        Self::migration_post_verified_latch_support(&store).await?;
+        Self::migration_post_verified_latch_support(&store, &identity_manager).await?;
 
-        Ok(OlmMachine::new_helper(device_id, store, static_account, identity, maybe_backup_key))
+        Ok(Self::new_helper(
+            device_id,
+            store,
+            verification_machine,
+            identity_manager,
+            identity,
+            maybe_backup_key,
+        ))
     }
 
     // The sdk now support verified identity change detection.
@@ -375,19 +398,15 @@ impl OlmMachine {
     //
     // pub(crate) visibility for testing.
     pub(crate) async fn migration_post_verified_latch_support(
-        store: &CryptoStoreWrapper,
+        store: &Store,
+        identity_manager: &IdentityManager,
     ) -> Result<(), CryptoStoreError> {
         let maybe_migrate_for_identity_verified_latch =
             store.get_custom_value(Self::HAS_MIGRATED_VERIFICATION_LATCH).await?.is_none();
+
         if maybe_migrate_for_identity_verified_latch {
-            // We want to mark all tracked users as dirty to ensure the verified latch is
-            // set up correctly.
-            let tracked_user = store.load_tracked_users().await?;
-            let mut store_updates = Vec::with_capacity(tracked_user.len());
-            tracked_user.iter().for_each(|tu| {
-                store_updates.push((tu.user_id.as_ref(), true));
-            });
-            store.save_tracked_users(&store_updates).await?;
+            identity_manager.mark_all_tracked_users_as_dirty(store.cache().await?).await?;
+
             store.set_custom_value(Self::HAS_MIGRATED_VERIFICATION_LATCH, vec![0]).await?
         }
         Ok(())
@@ -1992,6 +2011,17 @@ impl OlmMachine {
         self.inner.identity_manager.update_tracked_users(users).await
     }
 
+    /// Mark all tracked users as dirty.
+    ///
+    /// See `IdentityManager::mark_all_tracked_users_as_dirty()` to learn
+    /// more.
+    pub async fn mark_all_tracked_users_as_dirty(&self) -> StoreResult<()> {
+        self.inner
+            .identity_manager
+            .mark_all_tracked_users_as_dirty(self.inner.store.cache().await?)
+            .await
+    }
+
     async fn wait_if_user_pending(
         &self,
         user_id: &UserId,
@@ -2404,10 +2434,10 @@ impl OlmMachine {
         Ok(())
     }
 
-    #[cfg(any(feature = "testing", test))]
     /// Returns whether this `OlmMachine` is the same another one.
     ///
     /// Useful for testing purposes only.
+    #[cfg(any(feature = "testing", test))]
     pub fn same_as(&self, other: &OlmMachine) -> bool {
         Arc::ptr_eq(&self.inner, &other.inner)
     }
@@ -2419,6 +2449,18 @@ impl OlmMachine {
         let account = cache.account().await?;
         Ok(account.uploaded_key_count())
     }
+
+    /// Returns the identity manager.
+    #[cfg(test)]
+    pub(crate) fn identity_manager(&self) -> &IdentityManager {
+        &self.inner.identity_manager
+    }
+
+    /// Returns a store key, only useful for testing purposes.
+    #[cfg(test)]
+    pub(crate) fn key_for_has_migrated_verification_latch() -> &'static str {
+        Self::HAS_MIGRATED_VERIFICATION_LATCH
+    }
 }
 
 fn sender_data_to_verification_state(

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

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::{collections::BTreeMap, iter, sync::Arc, time::Duration};
+use std::{collections::BTreeMap, iter, ops::Not, sync::Arc, time::Duration};
 
 use assert_matches2::assert_matches;
 use futures_util::{pin_mut, FutureExt, StreamExt};
@@ -1528,6 +1528,43 @@ async fn test_unsigned_decryption() {
     assert_matches!(thread_encryption_result, UnsignedDecryptionResult::Decrypted(_));
 }
 
+#[async_test]
+async fn test_mark_all_tracked_users_as_dirty() {
+    let store = MemoryStore::new();
+    let account = vodozemac::olm::Account::new();
+
+    // Put some tracked users
+    let damir = user_id!("@damir:localhost");
+    let ben = user_id!("@ben:localhost");
+    let ivan = user_id!("@ivan:localhost");
+
+    // Mark them as not dirty.
+    store.save_tracked_users(&[(damir, false), (ben, false), (ivan, false)]).await.unwrap();
+
+    // Let's imagine the migration has been done: this is useful so that tracked
+    // users are not marked as dirty when creating the `OlmMachine`.
+    store
+        .set_custom_value(OlmMachine::key_for_has_migrated_verification_latch(), vec![0])
+        .await
+        .unwrap();
+
+    let alice =
+        OlmMachine::with_store(user_id(), alice_device_id(), store, Some(account)).await.unwrap();
+
+    // All users are marked as not dirty.
+    alice.store().load_tracked_users().await.unwrap().iter().for_each(|tracked_user| {
+        assert!(tracked_user.dirty.not());
+    });
+
+    // Now, mark all tracked users as dirty.
+    alice.mark_all_tracked_users_as_dirty().await.unwrap();
+
+    // All users are now marked as dirty.
+    alice.store().load_tracked_users().await.unwrap().iter().for_each(|tracked_user| {
+        assert!(tracked_user.dirty);
+    });
+}
+
 #[async_test]
 async fn test_verified_latch_migration() {
     let store = MemoryStore::new();
@@ -1544,20 +1581,22 @@ async fn test_verified_latch_migration() {
     let alice =
         OlmMachine::with_store(user_id(), alice_device_id(), store, Some(account)).await.unwrap();
 
+    let alice_store = alice.store();
+
     // A migration should have occurred and all users should be marked as dirty
-    alice.store().load_tracked_users().await.unwrap().iter().for_each(|tu| {
+    alice_store.load_tracked_users().await.unwrap().iter().for_each(|tu| {
         assert!(tu.dirty);
     });
 
     // Ensure it does so only once
-    alice.store().save_tracked_users(&to_track_not_dirty).await.unwrap();
+    alice_store.save_tracked_users(&to_track_not_dirty).await.unwrap();
 
-    OlmMachine::migration_post_verified_latch_support(alice.store().crypto_store().as_ref())
+    OlmMachine::migration_post_verified_latch_support(alice_store, alice.identity_manager())
         .await
         .unwrap();
 
     // Migration already done, so user should not be marked as dirty
-    alice.store().load_tracked_users().await.unwrap().iter().for_each(|tu| {
+    alice_store.load_tracked_users().await.unwrap().iter().for_each(|tu| {
         assert!(!tu.dirty);
     });
 }

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

@@ -351,6 +351,10 @@ pub(crate) struct StoreCache {
 }
 
 impl StoreCache {
+    pub(crate) fn store_wrapper(&self) -> &CryptoStoreWrapper {
+        self.store.as_ref()
+    }
+
     /// Returns a reference to the `Account`.
     ///
     /// Either load the account from the cache, or the store if missing from