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

…ative transaction

We're using application-level 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 application-level 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.

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 {