e2ee: save the account after generating new keys in a separate transaction

We're using applicative transactions to make sure that the account is properly synchronized in the cache vs in the database.

Before this commit, the transaction would be committed only when *all* the operations in it succeeded. This was based on the
assumption that most encryption requests could be replayed, by re-sending them to the server. Unfortunately, this assumption
doesn't hold for when generating one-time keys: it could be that one time-keys would be generated by the client, then
the applicative transaction would fail, resulting in the client "forgetting" about the one time keys it uploaded. The server
rejects reuploads of existing one-time keys, so that would end up wedging a device, causing unable-to-decrypt events, without
a proper way out.

Here, we propose to save the account just after one-time keys have been generated.

Author: bnjbvr

crates/matrix-sdk-crypto/src/dehydrated_devices.rs

@@ -240,12 +240,8 @@ impl RehydratedDevice {
 
         // Let us first give the events to the rehydrated device, this will decrypt any
         // encrypted to-device events and fetch out the room keys.
-        let mut rehydrated_transaction = self.rehydrated.store().transaction().await;
-
-        let (_, changes) = self
-            .rehydrated
-            .preprocess_sync_changes(&mut rehydrated_transaction, sync_changes)
-            .await?;
+        let (rehydrated_transaction, _, changes) =
+            self.rehydrated.preprocess_sync_changes(sync_changes).await?;
 
         // Now take the room keys and persist them in our original `OlmMachine`.
         let room_keys = &changes.inbound_group_sessions;

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

@@ -1197,10 +1197,8 @@ impl OlmMachine {
         &self,
         sync_changes: EncryptionSyncChanges<'_>,
     ) -> OlmResult<(Vec<Raw<AnyToDeviceEvent>>, Vec<RoomKeyInfo>)> {
-        let mut store_transaction = self.inner.store.transaction().await;
-
-        let (events, changes) =
-            self.preprocess_sync_changes(&mut store_transaction, sync_changes).await?;
+        let (store_transaction, events, changes) =
+            self.preprocess_sync_changes(sync_changes).await?;
 
         // Technically save_changes also does the same work, so if it's slow we could
         // refactor this to do it only once.
@@ -1215,23 +1213,27 @@ impl OlmMachine {
 
     pub(crate) async fn preprocess_sync_changes(
         &self,
-        transaction: &mut StoreTransaction,
         sync_changes: EncryptionSyncChanges<'_>,
-    ) -> OlmResult<(Vec<Raw<AnyToDeviceEvent>>, Changes)> {
+    ) -> OlmResult<(StoreTransaction, Vec<Raw<AnyToDeviceEvent>>, Changes)> {
         // Remove verification objects that have expired or are done.
         let mut events = self.inner.verification_machine.garbage_collect();
 
-        // The account is automatically saved by the store transaction created by the
-        // caller.
-        let mut changes = Default::default();
+        // Updating key counts may generate new one-time keys, and we should *never*
+        // forget about them once generated (otherwise this leads to issues like
+        // #1415), so save them in a separate transaction here.
+        self.inner
+            .store
+            .with_transaction(|mut transaction| async {
+                let account = transaction.account().await.map_err(OlmError::Store)?;
+                account.update_key_counts(
+                    sync_changes.one_time_keys_counts,
+                    sync_changes.unused_fallback_keys,
+                );
+                Ok((transaction, ()))
+            })
+            .await?;
 
-        {
-            let account = transaction.account().await?;
-            account.update_key_counts(
-                sync_changes.one_time_keys_counts,
-                sync_changes.unused_fallback_keys,
-            )
-        }
+        let mut transaction = self.inner.store.transaction().await;
 
         if let Err(e) = self
             .inner
@@ -1245,9 +1247,10 @@ impl OlmMachine {
             error!(error = ?e, "Error marking a tracked user as changed");
         }
 
+        let mut changes = Default::default();
         for raw_event in sync_changes.to_device_events {
             let raw_event =
-                Box::pin(self.receive_to_device_event(transaction, &mut changes, raw_event))
+                Box::pin(self.receive_to_device_event(&mut transaction, &mut changes, raw_event))
                     .await?;
             events.push(raw_event);
         }
@@ -1261,7 +1264,7 @@ impl OlmMachine {
         changes.sessions.extend(changed_sessions);
         changes.next_batch_token = sync_changes.next_batch_token;
 
-        Ok((events, changes))
+        Ok((transaction, events, changes))
     }
 
     /// Request a room key from our devices.

crates/matrix-sdk-crypto/src/olm/account.rs

@@ -473,6 +473,12 @@ impl Account {
         self.inner.max_number_of_one_time_keys()
     }
 
+    /// Generate new one-time keys if needed, and update the key counts
+    /// accordingly.
+    ///
+    /// Note: the `Account` *must* be persisted to disk after this has been
+    /// called, otherwise it's possible that one-time keys get uploaded and
+    /// the client forgets about those later.
     pub(crate) fn update_key_counts(
         &mut self,
         one_time_key_counts: &BTreeMap<DeviceKeyAlgorithm, UInt>,

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

@@ -449,15 +449,17 @@ impl StoreTransaction {
 
     /// Commits all dirty fields to the store, and maintains the cache so it
     /// reflects the current state of the database.
-    pub async fn commit(self) -> Result<()> {
+    pub async fn commit(mut self) -> Result<()> {
         if self.changes.is_empty() {
             return Ok(());
         }
 
         // Save changes in the database.
         let account = self.changes.account.as_ref().map(|acc| acc.deep_clone());
 
-        self.store.save_pending_changes(self.changes).await?;
+        // Note: after this `self.changes` is empty, and will be refilled as usual on
+        // the next attempt to read data from it.
+        self.store.save_pending_changes(std::mem::take(&mut self.changes)).await?;
 
         // Make the cache coherent with the database.
         if let Some(account) = account {