Client-managed validation

Author: dani-garcia

Cargo.lock

@@ -2,7 +2,7 @@ use std::sync::{Arc, OnceLock, RwLock};
 
 use bitwarden_crypto::KeyStore;
 #[cfg(feature = "internal")]
-use bitwarden_state::repository::RepositoryMap;
+use bitwarden_state::registry::StateRegistry;
 use reqwest::header::{self, HeaderValue};
 
 use super::internal::InternalClient;
@@ -91,7 +91,7 @@ impl Client {
                 external_client,
                 key_store: KeyStore::default(),
                 #[cfg(feature = "internal")]
-                repository_map: RwLock::new(RepositoryMap::default()),
+                repository_map: RwLock::new(StateRegistry::new()),
             }),
         }
     }

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

@@ -6,7 +6,7 @@ use bitwarden_crypto::SymmetricCryptoKey;
 #[cfg(feature = "internal")]
 use bitwarden_crypto::{EncString, Kdf, MasterKey, PinKey, UnsignedSharedKey};
 #[cfg(feature = "internal")]
-use bitwarden_state::repository::RepositoryMap;
+use bitwarden_state::registry::StateRegistry;
 use chrono::Utc;
 use uuid::Uuid;
 
@@ -67,7 +67,7 @@ pub struct InternalClient {
     pub(super) key_store: KeyStore<KeyIds>,
 
     #[cfg(feature = "internal")]
-    pub(crate) repository_map: RwLock<RepositoryMap>,
+    pub(crate) repository_map: RwLock<StateRegistry>,
 }
 
 impl InternalClient {

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

@@ -1,6 +1,9 @@
 use std::sync::Arc;
 
-use bitwarden_state::repository::Repository;
+use bitwarden_state::{
+    registry::StateRegistryError,
+    repository::{Repository, RepositoryItem},
+};
 
 use crate::Client;
 
@@ -11,22 +14,25 @@ pub struct StateClient {
 
 impl StateClient {
     /// Register a client managed state repository for a specific type.
-    pub fn register_repository<T: 'static + Repository<V>, V: 'static>(&self, store: Arc<T>) {
+    pub fn register_client_managed<T: 'static + Repository<V>, V: RepositoryItem>(
+        &self,
+        store: Arc<T>,
+    ) -> Result<(), StateRegistryError> {
         self.client
             .internal
             .repository_map
             .write()
             .expect("RwLock is not poisoned")
-            .insert(store);
+            .register_client_managed(store)
     }
 
     /// Get a client managed state repository for a specific type, if it exists.
-    pub fn get_repository<T: 'static>(&self) -> Option<Arc<dyn Repository<T>>> {
+    pub fn get_client_managed<T: RepositoryItem>(&self) -> Option<Arc<dyn Repository<T>>> {
         self.client
             .internal
             .repository_map
             .read()
             .expect("RwLock is not poisoned")
-            .get()
+            .get_client_managed()
     }
 }

crates/bitwarden-core/src/platform/state_client.rs

@@ -16,6 +16,7 @@ wasm = []
 [dependencies]
 async-trait = { workspace = true }
 bitwarden-error = { workspace = true }
+inventory = "0.3.20"
 log = { workspace = true }
 serde = { workspace = true }
 thiserror = { workspace = true }

crates/bitwarden-state/Cargo.toml

@@ -2,3 +2,6 @@
 
 /// This module provides a generic repository interface for storing and retrieving items.
 pub mod repository;
+
+/// This module provides a registry for managing repositories of different types.
+pub mod registry;

crates/bitwarden-state/src/lib.rs

@@ -0,0 +1,183 @@
+use std::{
+    any::{Any, TypeId},
+    collections::HashMap,
+    sync::Arc,
+};
+
+use bitwarden_error::bitwarden_error;
+use thiserror::Error;
+
+use super::repository::RepositoryItemRegistration;
+use crate::repository::{Repository, RepositoryItem};
+
+/// A registry that contains repositories for different types of items.
+/// These repositories can be either managed by the client or by the SDK itself.
+pub struct StateRegistry {
+    client_managed: HashMap<TypeId, Box<dyn Any + Send + Sync>>,
+}
+
+impl std::fmt::Debug for StateRegistry {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("StateRegistry")
+            .field("client_managed", &self.client_managed.keys())
+            .finish()
+    }
+}
+
+#[allow(missing_docs)]
+#[bitwarden_error(flat)]
+#[derive(Debug, Error)]
+pub enum StateRegistryError {
+    #[error("Repository for type {0} is not registered as client-managed")]
+    RepositoryNotClientManaged(&'static str),
+
+    #[error("Not all client-managed repositories are registered")]
+    NotAllRepositoriesRegistered,
+}
+
+impl StateRegistry {
+    /// Creates a new empty `StateRegistry`.
+    #[allow(clippy::new_without_default)]
+    pub fn new() -> Self {
+        StateRegistry {
+            client_managed: HashMap::new(),
+        }
+    }
+
+    /// Registers a client-managed repository into the map, associating it with its type.
+    pub fn register_client_managed<T: RepositoryItem>(
+        &mut self,
+        value: Arc<dyn Repository<T>>,
+    ) -> Result<(), StateRegistryError> {
+        let mut possible_registrations = RepositoryItemRegistration::iter();
+        match possible_registrations.find(|reg| reg.type_id() == TypeId::of::<T>()) {
+            Some(reg) => {
+                if !reg.rtype.is_client_managed() {
+                    return Err(StateRegistryError::RepositoryNotClientManaged(reg.name));
+                }
+            }
+            // This should never happen, as we have tests to ensure all repositories are registered.
+            _ => {
+                return Err(StateRegistryError::NotAllRepositoriesRegistered);
+            }
+        }
+
+        self.client_managed
+            .insert(TypeId::of::<T>(), Box::new(value));
+
+        Ok(())
+    }
+
+    /// Retrieves a client-managed repository from the map given its type.
+    pub fn get_client_managed<T: RepositoryItem>(&self) -> Option<Arc<dyn Repository<T>>> {
+        self.client_managed
+            .get(&TypeId::of::<T>())
+            .and_then(|boxed| boxed.downcast_ref::<Arc<dyn Repository<T>>>())
+            .map(Arc::clone)
+    }
+
+    /// Validates that all repositories registered in the client-managed state registry.
+    /// This should only be called after all the repositories have been registered by the clients.
+    pub fn validate_repositories(&self) -> Result<(), StateRegistryError> {
+        let possible_registrations = RepositoryItemRegistration::iter();
+        let mut missing_repository = false;
+
+        for reg in possible_registrations {
+            if reg.rtype.is_client_managed() && !self.client_managed.contains_key(&reg.type_id()) {
+                log::error!(
+                    "Repository for type {} is not registered in the client-managed state registry",
+                    reg.name
+                );
+                missing_repository = true;
+            }
+        }
+
+        if missing_repository {
+            return Err(StateRegistryError::NotAllRepositoriesRegistered);
+        }
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{
+        register_repository_item,
+        repository::{RepositoryError, RepositoryItem},
+    };
+
+    macro_rules! impl_repository {
+        ($name:ident, $ty:ty) => {
+            #[async_trait::async_trait]
+            impl Repository<$ty> for $name {
+                async fn get(&self, _key: String) -> Result<Option<$ty>, RepositoryError> {
+                    Ok(Some(TestItem(self.0.clone())))
+                }
+                async fn list(&self) -> Result<Vec<$ty>, RepositoryError> {
+                    unimplemented!()
+                }
+                async fn set(&self, _key: String, _value: $ty) -> Result<(), RepositoryError> {
+                    unimplemented!()
+                }
+                async fn remove(&self, _key: String) -> Result<(), RepositoryError> {
+                    unimplemented!()
+                }
+            }
+        };
+    }
+
+    #[derive(PartialEq, Eq, Debug)]
+    struct TestA(usize);
+    #[derive(PartialEq, Eq, Debug)]
+    struct TestB(String);
+    #[derive(PartialEq, Eq, Debug)]
+    struct TestC(Vec<u8>);
+    #[derive(PartialEq, Eq, Debug)]
+    struct TestItem<T>(T);
+
+    register_repository_item!(TestItem<usize>, "TestItem<usize>", ClientManaged);
+    register_repository_item!(TestItem<String>, "TestItem<String>", ClientManaged);
+    register_repository_item!(TestItem<Vec<u8>>, "TestItem<Vec<u8>>", ClientManaged);
+
+    impl_repository!(TestA, TestItem<usize>);
+    impl_repository!(TestB, TestItem<String>);
+    impl_repository!(TestC, TestItem<Vec<u8>>);
+
+    #[tokio::test]
+    async fn test_repository_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 = StateRegistry::new();
+
+        async fn get<T: RepositoryItem>(map: &StateRegistry) -> Option<T> {
+            map.get_client_managed::<T>()
+                .unwrap()
+                .get(String::new())
+                .await
+                .unwrap()
+        }
+
+        assert!(map.get_client_managed::<TestItem<usize>>().is_none());
+        assert!(map.get_client_managed::<TestItem<String>>().is_none());
+        assert!(map.get_client_managed::<TestItem<Vec<u8>>>().is_none());
+
+        map.register_client_managed(a.clone()).unwrap();
+        assert_eq!(get(&map).await, Some(TestItem(a.0)));
+        assert!(map.get_client_managed::<TestItem<String>>().is_none());
+        assert!(map.get_client_managed::<TestItem<Vec<u8>>>().is_none());
+
+        map.register_client_managed(b.clone()).unwrap();
+        assert_eq!(get(&map).await, Some(TestItem(a.0)));
+        assert_eq!(get(&map).await, Some(TestItem(b.0.clone())));
+        assert!(map.get_client_managed::<TestItem<Vec<u8>>>().is_none());
+
+        map.register_client_managed(c.clone()).unwrap();
+        assert_eq!(get(&map).await, Some(TestItem(a.0)));
+        assert_eq!(get(&map).await, Some(TestItem(b.0.clone())));
+        assert_eq!(get(&map).await, Some(TestItem(c.0.clone())));
+    }
+}