Generate one-time keys when we receive new one-time key counts (#2249)

Previously we would generate one-time keys, if needed, whenever we tried
to upload them. This code-path is critically missing a `save_account()`
call and we would only persist the account once we uploaded the one-time
keys.

This patch changes things up to generate one-time keys whenever we
receive new one-time key counts from the sync response. This aligns
with the way we generate fallback keys and removes the need to introduce
a new place where we persist the `Account`.

It's still possible to re-upload the same one-time keys, in the case
where the upload process succeeds on the server side but we fail to
receive the response.

Co-authored-by: Denis Kasak <dkasak@termina.org.uk>

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

@@ -1987,6 +1987,7 @@ pub(crate) mod tests {
         let machine = OlmMachine::new(user_id(), bob_device_id()).await;
         machine.account().generate_fallback_key_helper().await;
         machine.account().update_uploaded_key_count(0);
+        machine.account().generate_one_time_keys().await;
         let request = machine.keys_for_upload().await.expect("Can't prepare initial key upload");
         let response = keys_upload_response();
         machine.receive_keys_upload_response(&response).await.unwrap();
@@ -2146,6 +2147,7 @@ pub(crate) mod tests {
     async fn one_time_key_signing() {
         let machine = OlmMachine::new(user_id(), alice_device_id()).await;
         machine.account().update_uploaded_key_count(49);
+        machine.account().generate_one_time_keys().await;
 
         let mut one_time_keys = machine.account().signed_one_time_keys().await;
         let ed25519_key = machine.account().identity_keys().ed25519;
@@ -2169,7 +2171,11 @@ pub(crate) mod tests {
     #[async_test]
     async fn test_keys_for_upload() {
         let machine = OlmMachine::new(user_id(), alice_device_id()).await;
-        machine.account().update_uploaded_key_count(0);
+        let key_counts = BTreeMap::from([(DeviceKeyAlgorithm::SignedCurve25519, 49u8.into())]);
+        machine
+            .receive_sync_changes(Vec::new(), &Default::default(), &key_counts, None)
+            .await
+            .expect("We should be able to update our one-time key counts");
 
         let ed25519_key = machine.account().identity_keys().ed25519;
 
@@ -2203,7 +2209,7 @@ pub(crate) mod tests {
         let mut response = keys_upload_response();
         response.one_time_key_counts.insert(
             DeviceKeyAlgorithm::SignedCurve25519,
-            (request.one_time_keys.len() as u64).try_into().unwrap(),
+            (machine.account().max_one_time_keys().await).try_into().unwrap(),
         );
 
         machine.receive_keys_upload_response(&response).await.unwrap();

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

@@ -242,8 +242,10 @@ impl Account {
         self.inner.mark_as_shared();
 
         debug!("Marking one-time keys as published");
-        self.update_key_counts(&response.one_time_key_counts, None).await;
+        // First mark the current keys as published, as updating the key counts might
+        // generate some new keys if we're still below the limit.
         self.inner.mark_keys_as_published().await;
+        self.update_key_counts(&response.one_time_key_counts, None).await;
         self.store.save_account(self.inner.clone()).await?;
 
         Ok(())
@@ -558,9 +560,22 @@ impl ReadOnlyAccount {
     /// Create a fresh new account, this will generate the identity key-pair.
     #[allow(clippy::ptr_arg)]
     pub fn new(user_id: &UserId, device_id: &DeviceId) -> Self {
-        let account = InnerAccount::new();
+        let mut account = InnerAccount::new();
         let identity_keys = account.identity_keys();
 
+        // Let's generate some initial one-time keys while we're here. Since we know
+        // that this is a completely new [`Account`] we're certain that the
+        // server does not yet have any one-time keys of ours.
+        //
+        // This ensures we upload one-time keys along with our device keys right
+        // away, rather than waiting for the key counts to be echoed back to us
+        // from the server.
+        //
+        // It would be nice to do this for the fallback key as well but we can't assume
+        // that the server supports fallback keys. Maybe one of these days we
+        // will be able to do so.
+        account.generate_one_time_keys(account.max_number_of_one_time_keys());
+
         Self {
             user_id: user_id.into(),
             device_id: device_id.into(),
@@ -661,6 +676,7 @@ impl ReadOnlyAccount {
             }
 
             self.update_uploaded_key_count(count);
+            self.generate_one_time_keys().await;
         }
 
         if let Some(unused) = unused_fallback_keys {
@@ -929,8 +945,6 @@ impl ReadOnlyAccount {
     pub async fn signed_one_time_keys(
         &self,
     ) -> BTreeMap<OwnedDeviceKeyId, Raw<ruma::encryption::OneTimeKey>> {
-        let _ = self.generate_one_time_keys().await;
-
         let one_time_keys = self.one_time_keys().await;
 
         if one_time_keys.is_empty() {
@@ -1332,11 +1346,13 @@ mod tests {
 
         account.mark_keys_as_published().await;
         account.update_uploaded_key_count(50);
+        account.generate_one_time_keys().await;
 
         let (_, third_one_time_keys, _) = account.keys_for_upload().await;
         assert!(third_one_time_keys.is_empty());
 
         account.update_uploaded_key_count(0);
+        account.generate_one_time_keys().await;
 
         let (_, fourth_one_time_keys, _) = account.keys_for_upload().await;
         assert!(!fourth_one_time_keys.is_empty());

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

@@ -111,9 +111,6 @@ pub(crate) mod tests {
     #[async_test]
     async fn one_time_keys_creation() {
         let account = ReadOnlyAccount::new(alice_id(), alice_device_id());
-        let one_time_keys = account.one_time_keys().await;
-
-        assert!(one_time_keys.is_empty());
         assert_ne!(account.max_one_time_keys().await, 0);
 
         account.generate_one_time_keys_helper(10).await;