diff --git a/crates/matrix-sdk-crypto/src/machine.rs b/crates/matrix-sdk-crypto/src/machine.rs index e0cc704f5..3fd26f7cc 100644 --- a/crates/matrix-sdk-crypto/src/machine.rs +++ b/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(); diff --git a/crates/matrix-sdk-crypto/src/olm/account.rs b/crates/matrix-sdk-crypto/src/olm/account.rs index 5ff298ed0..9212185b4 100644 --- a/crates/matrix-sdk-crypto/src/olm/account.rs +++ b/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> { - 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()); diff --git a/crates/matrix-sdk-crypto/src/olm/mod.rs b/crates/matrix-sdk-crypto/src/olm/mod.rs index b144101f5..16a3878ba 100644 --- a/crates/matrix-sdk-crypto/src/olm/mod.rs +++ b/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;