Use generic data store trait and create a macro to do the messaging

Author: dani-garcia

crates/bitwarden-core/src/client/client.rs

@@ -8,6 +8,7 @@ use super::internal::InternalClient;
 use crate::client::flags::Flags;
 use crate::client::{
     client_settings::ClientSettings,
+    data_store::DataStoreMap,
     internal::{ApiConfigurations, Tokens},
 };
 
@@ -82,8 +83,8 @@ impl Client {
                 })),
                 external_client,
                 key_store: KeyStore::default(),
-
-                cipher_store: RwLock::new(None),
+                #[cfg(feature = "internal")]
+                data_stores: RwLock::new(DataStoreMap::default()),
             },
         }
     }

crates/bitwarden-core/src/client/data_store.rs

@@ -0,0 +1,113 @@
+use std::{
+    any::{Any, TypeId},
+    collections::HashMap,
+    sync::Arc,
+};
+
+#[async_trait::async_trait]
+pub trait DataStore<T>: Send + Sync {
+    async fn get(&self, key: String) -> Option<T>;
+    async fn list(&self) -> Vec<T>;
+    async fn set(&self, key: String, value: T);
+    async fn remove(&self, key: String);
+}
+
+#[derive(Default)]
+pub struct DataStoreMap {
+    stores: HashMap<TypeId, Box<dyn Any + Send + Sync>>,
+}
+
+impl std::fmt::Debug for DataStoreMap {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("DataStoreMap")
+            .field("stores", &self.stores.keys())
+            .finish()
+    }
+}
+
+impl DataStoreMap {
+    pub fn new() -> Self {
+        DataStoreMap {
+            stores: HashMap::new(),
+        }
+    }
+
+    pub fn insert<T: 'static>(&mut self, value: Arc<dyn DataStore<T>>) {
+        self.stores.insert(TypeId::of::<T>(), Box::new(value));
+    }
+
+    pub fn get<T: 'static>(&self) -> Option<Arc<dyn DataStore<T>>> {
+        self.stores
+            .get(&TypeId::of::<T>())
+            .and_then(|boxed| boxed.downcast_ref::<Arc<dyn DataStore<T>>>())
+            .map(Arc::clone)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    macro_rules! impl_data_store {
+        ($name:ident, $ty:ty) => {
+            #[async_trait::async_trait]
+            impl DataStore<$ty> for $name {
+                async fn get(&self, _key: String) -> Option<$ty> {
+                    Some(self.0.clone())
+                }
+                async fn list(&self) -> Vec<$ty> {
+                    unimplemented!()
+                }
+                async fn set(&self, _key: String, _value: $ty) {
+                    unimplemented!()
+                }
+                async fn remove(&self, _key: String) {
+                    unimplemented!()
+                }
+            }
+        };
+    }
+
+    #[derive(PartialEq, Eq, Debug)]
+    struct TestA(usize);
+    #[derive(PartialEq, Eq, Debug)]
+    struct TestB(String);
+    #[derive(PartialEq, Eq, Debug)]
+    struct TestC(Vec<u8>);
+
+    impl_data_store!(TestA, usize);
+    impl_data_store!(TestB, String);
+    impl_data_store!(TestC, Vec<u8>);
+
+    #[tokio::test]
+    async fn test_data_stores_map() {
+        let a = Arc::new(TestA(145832));
+        let b = Arc::new(TestB("test".to_string()));
+        let c = Arc::new(TestC(vec![1, 2, 3, 4, 5, 6, 7, 8, 9]));
+
+        let mut map = DataStoreMap::new();
+
+        async fn get<T: 'static>(map: &DataStoreMap) -> Option<T> {
+            map.get::<T>().unwrap().get(String::new()).await
+        }
+
+        assert!(map.get::<usize>().is_none());
+        assert!(map.get::<String>().is_none());
+        assert!(map.get::<Vec<u8>>().is_none());
+
+        map.insert(a.clone());
+        assert_eq!(get(&map).await, Some(a.0));
+        assert!(map.get::<String>().is_none());
+        assert!(map.get::<Vec<u8>>().is_none());
+
+        map.insert(b.clone());
+        assert_eq!(get(&map).await, Some(a.0));
+        assert_eq!(get(&map).await, Some(b.0.clone()));
+        assert!(map.get::<Vec<u8>>().is_none());
+
+        map.insert(c.clone());
+        assert_eq!(get(&map).await, Some(a.0));
+        assert_eq!(get(&map).await, Some(b.0.clone()));
+        assert_eq!(get(&map).await, Some(c.0.clone()));
+    }
+}

crates/bitwarden-core/src/client/internal.rs

@@ -18,6 +18,7 @@ use crate::{
 };
 #[cfg(feature = "internal")]
 use crate::{
+    client::data_store::{DataStore, DataStoreMap},
     client::encryption_settings::EncryptionSettingsError,
     client::{flags::Flags, login_method::UserLoginMethod},
     error::NotAuthenticatedError,
@@ -61,7 +62,8 @@ pub struct InternalClient {
 
     pub(super) key_store: KeyStore<KeyIds>,
 
-    pub(super) cipher_store: RwLock<Option<Arc<dyn CipherStore>>>,
+    #[cfg(feature = "internal")]
+    pub(super) data_stores: RwLock<DataStoreMap>,
 }
 
 impl InternalClient {
@@ -222,26 +224,19 @@ impl InternalClient {
         EncryptionSettings::set_org_keys(org_keys, &self.key_store)
     }
 
-    pub fn register_cipher_store(&self, store: Arc<dyn CipherStore>) {
-        self.cipher_store
+    #[cfg(feature = "internal")]
+    pub fn register_data_store<T: 'static + DataStore<V>, V: 'static>(&self, store: Arc<T>) {
+        self.data_stores
             .write()
             .expect("RwLock is not poisoned")
-            .replace(store);
+            .insert(store);
     }
 
-    pub fn get_cipher_store(&self) -> Option<Arc<dyn CipherStore>> {
-        self.cipher_store
+    #[cfg(feature = "internal")]
+    pub fn get_data_store<T: 'static>(&self) -> Option<Arc<dyn DataStore<T>>> {
+        self.data_stores
             .read()
             .expect("RwLock is not poisoned")
-            .clone()
+            .get()
     }
 }
-
-// TODO: We can't expose the store as returning a bitwarden_vault::Cipher to avoid a circular dependency, we'll fix that at some point somehow
-#[async_trait::async_trait]
-pub trait CipherStore: std::fmt::Debug + Send + Sync {
-    async fn get(&self, key: &str) -> Option<String>;
-    async fn list(&self) -> Vec<String>;
-    async fn set(&self, key: &str, value: String);
-    async fn remove(&self, key: &str);
-}

crates/bitwarden-core/src/client/mod.rs

@@ -18,3 +18,6 @@ pub use client_settings::{ClientSettings, DeviceType};
 
 #[cfg(feature = "internal")]
 pub mod test_accounts;
+
+#[cfg(feature = "internal")]
+pub mod data_store;

crates/bitwarden-error-macro/src/lib.rs

@@ -5,6 +5,7 @@ mod attribute;
 mod basic;
 mod flat;
 mod full;
+mod wasm_ipc;
 
 /// A procedural macro for generating error types with customizable serialization behavior.
 ///
@@ -120,3 +121,14 @@ pub fn bitwarden_error(
 ) -> proc_macro::TokenStream {
     attribute::bitwarden_error(args, item)
 }
+
+#[proc_macro_attribute]
+pub fn bitwarden_wasm_ipc_channel(
+    args: proc_macro::TokenStream,
+    item: proc_macro::TokenStream,
+) -> proc_macro::TokenStream {
+    match wasm_ipc::bitwarden_wasm_ipc_channel_internal(args, item) {
+        Ok(v) => v,
+        Err(e) => proc_macro::TokenStream::from(e.to_compile_error()),
+    }
+}