crypto: Calculate the deleted devices in a key query.

Cargo.toml

@@ -39,6 +39,8 @@ zeroize = { version = "1.1.0", optional = true }
 thiserror = "1.0.13"
 async-trait = { version = "0.1.26", optional = true }
 tracing = "0.1.13"
+atomic = "0.4.5"
+dashmap = "3.9.1"
 
 [dependencies.tracing-futures]
 version = "0.2.3"

src/crypto/device.rs

@@ -13,22 +13,26 @@
 // limitations under the License.
 
 use std::collections::HashMap;
+use std::sync::atomic::AtomicBool;
+use std::sync::Arc;
+
+use atomic::Atomic;
 
 use ruma_client_api::r0::keys::{DeviceKeys, KeyAlgorithm};
 use ruma_events::Algorithm;
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Device {
-    user_id: String,
-    device_id: String,
-    algorithms: Vec<Algorithm>,
-    keys: HashMap<KeyAlgorithm, String>,
-    display_name: Option<String>,
-    deleted: bool,
-    trust_state: TrustState,
+    user_id: Arc<String>,
+    device_id: Arc<String>,
+    algorithms: Arc<Vec<Algorithm>>,
+    keys: Arc<HashMap<KeyAlgorithm, String>>,
+    display_name: Arc<Option<String>>,
+    deleted: Arc<AtomicBool>,
+    trust_state: Arc<Atomic<TrustState>>,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone, Copy)]
 pub enum TrustState {
     Verified,
     BlackListed,
@@ -37,7 +41,7 @@ pub enum TrustState {
 }
 
 impl Device {
-    pub fn id(&self) -> &str {
+    pub fn device_id(&self) -> &str {
         &self.device_id
     }
 
@@ -56,16 +60,18 @@ impl From<&DeviceKeys> for Device {
         }
 
         Device {
-            user_id: device_keys.user_id.to_string(),
-            device_id: device_keys.device_id.clone(),
-            algorithms: device_keys.algorithms.clone(),
-            keys,
-            display_name: device_keys
+            user_id: Arc::new(device_keys.user_id.to_string()),
+            device_id: Arc::new(device_keys.device_id.clone()),
+            algorithms: Arc::new(device_keys.algorithms.clone()),
+            keys: Arc::new(keys),
+            display_name: Arc::new(
+                device_keys
                     .unsigned
                     .as_ref()
                     .map(|d| d.device_display_name.clone()),
-            deleted: false,
-            trust_state: TrustState::Unset,
+            ),
+            deleted: Arc::new(AtomicBool::new(false)),
+            trust_state: Arc::new(Atomic::new(TrustState::Unset)),
         }
     }
 }

src/crypto/machine.rs

@@ -240,17 +240,27 @@ impl OlmMachine {
 
                 let device = self
                     .store
-                    .get_user_device(&user_id_string, device_id)
+                    .get_device(&user_id_string, device_id)
                     .await
                     .expect("Can't load device");
 
                 if let Some(d) = device {
-                    todo!()
+                    // TODO check what and if anything changed for the device.
                 } else {
                     let device = Device::from(device_keys);
                     info!("Found new device {:?}", device);
                 }
             }
+
+            let current_devices: HashSet<&String> = device_map.keys().collect();
+            let stored_devices = self.store.get_user_devices(&user_id_string).await.unwrap();
+            let stored_devices_set: HashSet<&String> = stored_devices.keys().collect();
+
+            let deleted_devices = stored_devices_set.difference(&current_devices);
+
+            for device_id in deleted_devices {
+                // TODO delete devices here.
+            }
         }
         Ok(())
     }

src/crypto/memory_stores.rs

@@ -15,8 +15,10 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
+use dashmap::{DashMap, ReadOnlyView};
 use tokio::sync::Mutex;
 
+use super::device::Device;
 use super::olm::{InboundGroupSession, Session};
 
 #[derive(Debug)]
@@ -96,3 +98,57 @@ impl GroupSessionStore {
             .and_then(|m| m.get(sender_key).and_then(|m| m.get(session_id).cloned()))
     }
 }
+
+#[derive(Debug)]
+pub struct DeviceStore {
+    entries: DashMap<String, DashMap<String, Device>>,
+}
+
+pub struct UserDevices {
+    entries: ReadOnlyView<String, Device>,
+}
+
+impl UserDevices {
+    pub fn get(&self, device_id: &str) -> Option<Device> {
+        self.entries.get(device_id).cloned()
+    }
+
+    pub fn keys(&self) -> impl Iterator<Item = &String> {
+        self.entries.keys()
+    }
+}
+
+impl DeviceStore {
+    pub fn new() -> Self {
+        DeviceStore {
+            entries: DashMap::new(),
+        }
+    }
+
+    pub fn add(&self, device: Device) -> bool {
+        if !self.entries.contains_key(device.user_id()) {
+            self.entries
+                .insert(device.user_id().to_owned(), DashMap::new());
+        }
+        let mut device_map = self.entries.get_mut(device.user_id()).unwrap();
+
+        device_map
+            .insert(device.device_id().to_owned(), device)
+            .is_some()
+    }
+
+    pub fn get(&self, user_id: &str, device_id: &str) -> Option<Device> {
+        self.entries
+            .get(user_id)
+            .and_then(|m| m.get(device_id).map(|d| d.value().clone()))
+    }
+
+    pub fn user_devices(&self, user_id: &str) -> UserDevices {
+        if !self.entries.contains_key(user_id) {
+            self.entries.insert(user_id.to_owned(), DashMap::new());
+        }
+        UserDevices {
+            entries: self.entries.get(user_id).unwrap().clone().into_read_only(),
+        }
+    }
+}

src/crypto/store/memorystore.rs

@@ -20,13 +20,14 @@ use tokio::sync::Mutex;
 
 use super::{Account, CryptoStore, CryptoStoreError, InboundGroupSession, Result, Session};
 use crate::crypto::device::Device;
-use crate::crypto::memory_stores::{GroupSessionStore, SessionStore};
+use crate::crypto::memory_stores::{DeviceStore, GroupSessionStore, SessionStore, UserDevices};
 
 #[derive(Debug)]
 pub struct MemoryStore {
     sessions: SessionStore,
     inbound_group_sessions: GroupSessionStore,
     tracked_users: HashSet<String>,
+    devices: DeviceStore,
 }
 
 impl MemoryStore {
@@ -35,6 +36,7 @@ impl MemoryStore {
             sessions: SessionStore::new(),
             inbound_group_sessions: GroupSessionStore::new(),
             tracked_users: HashSet::new(),
+            devices: DeviceStore::new(),
         }
     }
 }
@@ -88,7 +90,11 @@ impl CryptoStore for MemoryStore {
         Ok(self.tracked_users.insert(user.to_string()))
     }
 
-    async fn get_user_device(&self, user_id: &str, device_id: &str) -> Result<Option<Device>> {
-        Ok(None)
+    async fn get_device(&self, user_id: &str, device_id: &str) -> Result<Option<Device>> {
+        Ok(self.devices.get(user_id, device_id))
+    }
+
+    async fn get_user_devices(&self, user_id: &str) -> Result<UserDevices> {
+        Ok(self.devices.user_devices(user_id))
     }
 }

src/crypto/store/mod.rs

@@ -25,6 +25,7 @@ use thiserror::Error;
 use tokio::sync::Mutex;
 
 use super::device::Device;
+use super::memory_stores::UserDevices;
 use super::olm::{Account, InboundGroupSession, Session};
 use olm_rs::errors::{OlmAccountError, OlmGroupSessionError, OlmSessionError};
 use olm_rs::PicklingMode;
@@ -82,5 +83,6 @@ pub trait CryptoStore: Debug + Send + Sync {
     ) -> Result<Option<Arc<Mutex<InboundGroupSession>>>>;
     fn tracked_users(&self) -> &HashSet<String>;
     async fn add_user_for_tracking(&mut self, user: &str) -> Result<bool>;
-    async fn get_user_device(&self, user_id: &str, device_id: &str) -> Result<Option<Device>>;
+    async fn get_device(&self, user_id: &str, device_id: &str) -> Result<Option<Device>>;
+    async fn get_user_devices(&self, user_id: &str) -> Result<UserDevices>;
 }

src/crypto/store/sqlite.rs

@@ -28,7 +28,7 @@ use zeroize::Zeroizing;
 
 use super::{Account, CryptoStore, CryptoStoreError, InboundGroupSession, Result, Session};
 use crate::crypto::device::Device;
-use crate::crypto::memory_stores::{GroupSessionStore, SessionStore};
+use crate::crypto::memory_stores::{GroupSessionStore, SessionStore, UserDevices};
 
 pub struct SqliteStore {
     user_id: Arc<String>,
@@ -410,7 +410,11 @@ impl CryptoStore for SqliteStore {
         Ok(self.tracked_users.insert(user.to_string()))
     }
 
-    async fn get_user_device(&self, user_id: &str, device_id: &str) -> Result<Option<Device>> {
+    async fn get_device(&self, user_id: &str, device_id: &str) -> Result<Option<Device>> {
+        todo!()
+    }
+
+    async fn get_user_devices(&self, user_id: &str) -> Result<UserDevices> {
         todo!()
     }
 }