crypto: Streamline how we update one-time key counts

There is no reason to use different logic when updating one-time key
counts coming from an upload response vs when coming from a sync
response.

Switch to an AtomicU64 for the count as well, the i64 is probably a
remnant of the SQLite based crypto store.
This commit is contained in:
Damir Jelić 2021-07-22 16:02:32 +02:00
parent e919a82b2c
commit 46f9c292ab
2 changed files with 36 additions and 28 deletions

View file

@ -736,8 +736,8 @@ impl OlmMachine {
self.verification_machine.get_requests(user_id)
}
async fn update_one_time_key_count(&self, key_count: &BTreeMap<DeviceKeyAlgorithm, UInt>) {
self.account.update_uploaded_key_count(key_count).await;
fn update_one_time_key_count(&self, key_count: &BTreeMap<DeviceKeyAlgorithm, UInt>) {
self.account.update_uploaded_key_count(key_count);
}
async fn handle_to_device_event(&self, event: &AnyToDeviceEvent) {
@ -796,7 +796,7 @@ impl OlmMachine {
let mut changes =
Changes { account: Some(self.account.inner.clone()), ..Default::default() };
self.update_one_time_key_count(one_time_keys_counts).await;
self.update_one_time_key_count(one_time_keys_counts);
for user_id in &changed_devices.changed {
if let Err(e) = self.identity_manager.mark_user_as_changed(user_id).await {
@ -1410,6 +1410,10 @@ pub(crate) mod test {
response.one_time_key_counts.insert(DeviceKeyAlgorithm::SignedCurve25519, uint!(50));
machine.receive_keys_upload_response(&response).await.unwrap();
assert!(!machine.should_upload_keys().await);
response.one_time_key_counts.remove(&DeviceKeyAlgorithm::SignedCurve25519);
machine.receive_keys_upload_response(&response).await.unwrap();
assert!(!machine.should_upload_keys().await);
}
#[tokio::test]

View file

@ -14,11 +14,11 @@
use std::{
collections::BTreeMap,
convert::{TryFrom, TryInto},
convert::TryInto,
fmt,
ops::Deref,
sync::{
atomic::{AtomicBool, AtomicI64, Ordering},
atomic::{AtomicBool, AtomicU64, Ordering},
Arc,
},
};
@ -181,9 +181,22 @@ impl Account {
}
}
pub async fn update_uploaded_key_count(&self, key_count: &BTreeMap<DeviceKeyAlgorithm, UInt>) {
pub fn update_uploaded_key_count(&self, key_count: &BTreeMap<DeviceKeyAlgorithm, UInt>) {
if let Some(count) = key_count.get(&DeviceKeyAlgorithm::SignedCurve25519) {
let count: u64 = (*count).into();
let old_count = self.inner.uploaded_key_count();
// Some servers might always return the key counts in the sync
// response, we don't want to the logs with noop changes if they do
// so.
if count != old_count {
debug!(
"Updated uploaded one-time key count {} -> {}.",
self.inner.uploaded_key_count(),
count
);
}
self.inner.update_uploaded_key_count(count);
}
}
@ -197,16 +210,8 @@ impl Account {
}
self.inner.mark_as_shared();
let one_time_key_count =
response.one_time_key_counts.get(&DeviceKeyAlgorithm::SignedCurve25519);
let count: u64 = one_time_key_count.map_or(0, |c| (*c).into());
debug!(
"Updated uploaded one-time key count {} -> {}, marking keys as published",
self.inner.uploaded_key_count(),
count
);
self.inner.update_uploaded_key_count(count);
debug!("Marking one-time keys as published");
self.update_uploaded_key_count(&response.one_time_key_counts);
self.inner.mark_keys_as_published().await;
self.store.save_account(self.inner.clone()).await?;
@ -432,7 +437,7 @@ pub struct ReadOnlyAccount {
/// this is None, no action will be taken. After a sync request the client
/// needs to set this for us, depending on the count we will suggest the
/// client to upload new keys.
uploaded_signed_key_count: Arc<AtomicI64>,
uploaded_signed_key_count: Arc<AtomicU64>,
}
/// A typed representation of a base64 encoded string containing the account
@ -468,7 +473,7 @@ pub struct PickledAccount {
/// Was the account shared.
pub shared: bool,
/// The number of uploaded one-time keys we have on the server.
pub uploaded_signed_key_count: i64,
pub uploaded_signed_key_count: u64,
}
#[cfg(not(tarpaulin_include))]
@ -499,7 +504,7 @@ impl ReadOnlyAccount {
inner: Arc::new(Mutex::new(account)),
identity_keys: Arc::new(identity_keys),
shared: Arc::new(AtomicBool::new(false)),
uploaded_signed_key_count: Arc::new(AtomicI64::new(0)),
uploaded_signed_key_count: Arc::new(AtomicU64::new(0)),
}
}
@ -524,18 +529,17 @@ impl ReadOnlyAccount {
///
/// * `new_count` - The new count that was reported by the server.
pub(crate) fn update_uploaded_key_count(&self, new_count: u64) {
let key_count = i64::try_from(new_count).unwrap_or(i64::MAX);
self.uploaded_signed_key_count.store(key_count, Ordering::Relaxed);
self.uploaded_signed_key_count.store(new_count, Ordering::SeqCst);
}
/// Get the currently known uploaded key count.
pub fn uploaded_key_count(&self) -> i64 {
self.uploaded_signed_key_count.load(Ordering::Relaxed)
pub fn uploaded_key_count(&self) -> u64 {
self.uploaded_signed_key_count.load(Ordering::SeqCst)
}
/// Has the account been shared with the server.
pub fn shared(&self) -> bool {
self.shared.load(Ordering::Relaxed)
self.shared.load(Ordering::SeqCst)
}
/// Mark the account as shared.
@ -543,7 +547,7 @@ impl ReadOnlyAccount {
/// Messages shouldn't be encrypted with the session before it has been
/// shared.
pub(crate) fn mark_as_shared(&self) {
self.shared.store(true, Ordering::Relaxed);
self.shared.store(true, Ordering::SeqCst);
}
/// Get the one-time keys of the account.
@ -567,7 +571,7 @@ impl ReadOnlyAccount {
///
/// Returns an empty error if no keys need to be uploaded.
pub(crate) async fn generate_one_time_keys(&self) -> Result<u64, ()> {
let count = self.uploaded_key_count() as u64;
let count = self.uploaded_key_count();
let max_keys = self.max_one_time_keys().await;
let max_on_server = (max_keys as u64) / 2;
@ -588,7 +592,7 @@ impl ReadOnlyAccount {
return true;
}
let count = self.uploaded_key_count() as u64;
let count = self.uploaded_key_count();
// If we have a known key count, check that we have more than
// max_one_time_Keys() / 2, otherwise tell the client to upload more.
@ -673,7 +677,7 @@ impl ReadOnlyAccount {
inner: Arc::new(Mutex::new(account)),
identity_keys: Arc::new(identity_keys),
shared: Arc::new(AtomicBool::from(pickle.shared)),
uploaded_signed_key_count: Arc::new(AtomicI64::new(pickle.uploaded_signed_key_count)),
uploaded_signed_key_count: Arc::new(AtomicU64::new(pickle.uploaded_signed_key_count)),
})
}