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

View File

@ -14,11 +14,11 @@
use std::{ use std::{
collections::BTreeMap, collections::BTreeMap,
convert::{TryFrom, TryInto}, convert::TryInto,
fmt, fmt,
ops::Deref, ops::Deref,
sync::{ sync::{
atomic::{AtomicBool, AtomicI64, Ordering}, atomic::{AtomicBool, AtomicU64, Ordering},
Arc, 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) { if let Some(count) = key_count.get(&DeviceKeyAlgorithm::SignedCurve25519) {
let count: u64 = (*count).into(); 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); self.inner.update_uploaded_key_count(count);
} }
} }
@ -197,16 +210,8 @@ impl Account {
} }
self.inner.mark_as_shared(); self.inner.mark_as_shared();
let one_time_key_count = debug!("Marking one-time keys as published");
response.one_time_key_counts.get(&DeviceKeyAlgorithm::SignedCurve25519); self.update_uploaded_key_count(&response.one_time_key_counts);
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);
self.inner.mark_keys_as_published().await; self.inner.mark_keys_as_published().await;
self.store.save_account(self.inner.clone()).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 /// 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 /// needs to set this for us, depending on the count we will suggest the
/// client to upload new keys. /// 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 /// A typed representation of a base64 encoded string containing the account
@ -468,7 +473,7 @@ pub struct PickledAccount {
/// Was the account shared. /// Was the account shared.
pub shared: bool, pub shared: bool,
/// The number of uploaded one-time keys we have on the server. /// 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))] #[cfg(not(tarpaulin_include))]
@ -499,7 +504,7 @@ impl ReadOnlyAccount {
inner: Arc::new(Mutex::new(account)), inner: Arc::new(Mutex::new(account)),
identity_keys: Arc::new(identity_keys), identity_keys: Arc::new(identity_keys),
shared: Arc::new(AtomicBool::new(false)), 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. /// * `new_count` - The new count that was reported by the server.
pub(crate) fn update_uploaded_key_count(&self, new_count: u64) { 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(new_count, Ordering::SeqCst);
self.uploaded_signed_key_count.store(key_count, Ordering::Relaxed);
} }
/// Get the currently known uploaded key count. /// Get the currently known uploaded key count.
pub fn uploaded_key_count(&self) -> i64 { pub fn uploaded_key_count(&self) -> u64 {
self.uploaded_signed_key_count.load(Ordering::Relaxed) self.uploaded_signed_key_count.load(Ordering::SeqCst)
} }
/// Has the account been shared with the server. /// Has the account been shared with the server.
pub fn shared(&self) -> bool { pub fn shared(&self) -> bool {
self.shared.load(Ordering::Relaxed) self.shared.load(Ordering::SeqCst)
} }
/// Mark the account as shared. /// Mark the account as shared.
@ -543,7 +547,7 @@ impl ReadOnlyAccount {
/// Messages shouldn't be encrypted with the session before it has been /// Messages shouldn't be encrypted with the session before it has been
/// shared. /// shared.
pub(crate) fn mark_as_shared(&self) { 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. /// 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. /// Returns an empty error if no keys need to be uploaded.
pub(crate) async fn generate_one_time_keys(&self) -> Result<u64, ()> { 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_keys = self.max_one_time_keys().await;
let max_on_server = (max_keys as u64) / 2; let max_on_server = (max_keys as u64) / 2;
@ -588,7 +592,7 @@ impl ReadOnlyAccount {
return true; 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 // 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. // max_one_time_Keys() / 2, otherwise tell the client to upload more.
@ -673,7 +677,7 @@ impl ReadOnlyAccount {
inner: Arc::new(Mutex::new(account)), inner: Arc::new(Mutex::new(account)),
identity_keys: Arc::new(identity_keys), identity_keys: Arc::new(identity_keys),
shared: Arc::new(AtomicBool::from(pickle.shared)), 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)),
}) })
} }