bitwarden_state/
registry.rs

1use std::sync::Arc;
2
3use bitwarden_error::bitwarden_error;
4use thiserror::Error;
5
6use crate::{
7    any_map::AnyMap,
8    persisted_value::{PersistedValue, PersistedValueKey},
9    repository::{Repository, RepositoryItem, RepositoryMigrations},
10    sdk_managed::{Database, DatabaseConfiguration, DatabaseError, MemoryDatabase, SystemDatabase},
11};
12
13/// A registry that contains repositories for different types of items.
14/// These repositories can be either managed by the client or by the SDK itself.
15pub struct StateRegistry {
16    database: SystemDatabase,
17    client_managed: AnyMap,
18    /// Dev-only index of per-type get/set/list shims backing the generic debug
19    /// browse. Populated at registration (both client- and SDK-managed), where
20    /// the concrete type is known. See [`crate::debug`].
21    #[cfg(feature = "debug-capabilities")]
22    pub(crate) debug: crate::debug::DebugRegistry,
23}
24
25impl std::fmt::Debug for StateRegistry {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        f.debug_struct("StateRegistry").finish()
28    }
29}
30
31#[allow(missing_docs)]
32#[bitwarden_error(flat)]
33#[derive(Debug, Error)]
34pub enum StateRegistryError {
35    #[error("Database is not initialized")]
36    DatabaseNotInitialized,
37
38    #[error(transparent)]
39    Database(#[from] DatabaseError),
40}
41
42impl StateRegistry {
43    /// Creates a new `StateRegistry` backed by an in-memory database.
44    pub fn new_with_memory_db() -> Self {
45        StateRegistry {
46            database: SystemDatabase::Memory(MemoryDatabase::new()),
47            client_managed: AnyMap::new(),
48            #[cfg(feature = "debug-capabilities")]
49            debug: crate::debug::DebugRegistry::new(),
50        }
51    }
52
53    /// Creates a new `StateRegistry` backed by a database.
54    pub async fn new_with_db(
55        configuration: DatabaseConfiguration,
56        migrations: RepositoryMigrations,
57    ) -> Result<Self, DatabaseError> {
58        let database = SystemDatabase::initialize(configuration, migrations.clone()).await?;
59        let registry = StateRegistry {
60            database,
61            client_managed: AnyMap::new(),
62            #[cfg(feature = "debug-capabilities")]
63            debug: crate::debug::DebugRegistry::new(),
64        };
65        // Capture debug shims for every SDK-managed type declared in the migrations.
66        #[cfg(feature = "debug-capabilities")]
67        registry.debug.register_migrations(&migrations);
68        Ok(registry)
69    }
70
71    /// Get a handle to a persisted value by its type-safe key.
72    pub fn value<K: PersistedValueKey>(&self) -> PersistedValue<K> {
73        PersistedValue::new(self.database.get_persisted_value::<K::Value>(K::NAME))
74    }
75
76    /// Registers a client-managed repository into the map, associating it with its type.
77    pub fn register_client_managed<T: RepositoryItem>(&self, value: Arc<dyn Repository<T>>) {
78        self.client_managed.insert(value);
79        #[cfg(feature = "debug-capabilities")]
80        self.debug.register::<T>();
81    }
82
83    /// Retrieves a client-managed repository from the map given its type.
84    fn get_client_managed<T: RepositoryItem>(&self) -> Option<Arc<dyn Repository<T>>> {
85        self.client_managed.get()
86    }
87
88    /// Retrieves a SDK-managed repository from the database.
89    fn get_sdk_managed<T: RepositoryItem>(
90        &self,
91    ) -> Result<Arc<dyn Repository<T>>, StateRegistryError> {
92        Ok(self.database.get_repository::<T>())
93    }
94
95    /// Get a repository with fallback: prefer client-managed, fall back to SDK-managed.
96    ///
97    /// This method first attempts to retrieve a client-managed repository. If not found,
98    /// it falls back to an SDK-managed repository. Both are returned as `Arc<dyn Repository<T>>`.
99    ///
100    /// # Errors
101    /// This method never fails, but returns a Result for backwards compatibility.
102    pub fn get<T>(&self) -> Result<Arc<dyn Repository<T>>, StateRegistryError>
103    where
104        T: RepositoryItem,
105    {
106        if let Some(repo) = self.get_client_managed::<T>() {
107            return Ok(repo);
108        }
109
110        self.get_sdk_managed::<T>()
111    }
112
113    /// Wipes all state from this registry, and deletes any files or databases associated with it.
114    /// Intended to be used during logout, where the Client will be dropped right after.
115    ///
116    /// # Warning
117    ///
118    /// This closes the SDK-managed database and deletes persistent storage (SQLite file + WAL/SHM,
119    /// IndexedDB database). Outstanding [`Repository`] handles will return
120    /// [`DatabaseError::Closed`] on subsequent operations. Client-managed repositories are also
121    /// cleared.
122    pub async fn wipe(&self) -> Result<(), DatabaseError> {
123        // Clear client-managed first so a failure in the persistent-store wipe
124        // still releases the in-memory Arc references.
125        self.client_managed.clear();
126        self.database.wipe().await
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133    use crate::{
134        register_repository_item,
135        repository::{RepositoryError, RepositoryItem},
136        sdk_managed::DatabaseError,
137    };
138
139    macro_rules! impl_repository {
140        ($name:ident, $ty:ty) => {
141            #[async_trait::async_trait]
142            impl Repository<$ty> for $name {
143                async fn get(&self, _key: String) -> Result<Option<$ty>, RepositoryError> {
144                    Ok(Some(TestItem(self.0.clone())))
145                }
146                async fn list(&self) -> Result<Vec<$ty>, RepositoryError> {
147                    unimplemented!()
148                }
149                async fn set(&self, _key: String, _value: $ty) -> Result<(), RepositoryError> {
150                    unimplemented!()
151                }
152                async fn set_bulk(
153                    &self,
154                    _values: Vec<(String, $ty)>,
155                ) -> Result<(), RepositoryError> {
156                    unimplemented!()
157                }
158                async fn remove(&self, _key: String) -> Result<(), RepositoryError> {
159                    unimplemented!()
160                }
161                async fn remove_bulk(&self, _keys: Vec<String>) -> Result<(), RepositoryError> {
162                    unimplemented!()
163                }
164                async fn remove_all(&self) -> Result<(), RepositoryError> {
165                    unimplemented!()
166                }
167            }
168        };
169    }
170
171    use serde::{Deserialize, Serialize};
172
173    #[derive(PartialEq, Eq, Debug)]
174    struct TestA(usize);
175    #[derive(PartialEq, Eq, Debug)]
176    struct TestB(String);
177    #[derive(PartialEq, Eq, Debug)]
178    struct TestC(Vec<u8>);
179    /// A second implementation for the same item type as [`TestA`].
180    #[derive(PartialEq, Eq, Debug)]
181    struct TestD(usize);
182    #[derive(PartialEq, Eq, Debug, Serialize, Deserialize)]
183    struct TestItem<T>(T);
184
185    register_repository_item!(String => TestItem<usize>, "TestItem_usize");
186    register_repository_item!(String => TestItem<String>, "TestItem_String");
187    register_repository_item!(String => TestItem<Vec<u8>>, "TestItem_Vec");
188
189    impl_repository!(TestA, TestItem<usize>);
190    impl_repository!(TestB, TestItem<String>);
191    impl_repository!(TestC, TestItem<Vec<u8>>);
192    impl_repository!(TestD, TestItem<usize>);
193
194    #[tokio::test]
195    async fn test_state_registry() {
196        let a = Arc::new(TestA(145832));
197        let b = Arc::new(TestB("test".to_string()));
198        let c = Arc::new(TestC(vec![1, 2, 3, 4, 5, 6, 7, 8, 9]));
199
200        let map = StateRegistry::new_with_memory_db();
201
202        async fn get<T: RepositoryItem>(map: &StateRegistry) -> Option<T>
203        where
204            T::Key: Default,
205        {
206            map.get_client_managed::<T>()
207                .unwrap()
208                .get(Default::default())
209                .await
210                .unwrap()
211        }
212
213        assert!(map.get_client_managed::<TestItem<usize>>().is_none());
214        assert!(map.get_client_managed::<TestItem<String>>().is_none());
215        assert!(map.get_client_managed::<TestItem<Vec<u8>>>().is_none());
216
217        map.register_client_managed(a.clone());
218        assert_eq!(get(&map).await, Some(TestItem(a.0)));
219        assert!(map.get_client_managed::<TestItem<String>>().is_none());
220        assert!(map.get_client_managed::<TestItem<Vec<u8>>>().is_none());
221
222        map.register_client_managed(b.clone());
223        assert_eq!(get(&map).await, Some(TestItem(a.0)));
224        assert_eq!(get(&map).await, Some(TestItem(b.0.clone())));
225        assert!(map.get_client_managed::<TestItem<Vec<u8>>>().is_none());
226
227        map.register_client_managed(c.clone());
228        assert_eq!(get(&map).await, Some(TestItem(a.0)));
229        assert_eq!(get(&map).await, Some(TestItem(b.0.clone())));
230        assert_eq!(get(&map).await, Some(TestItem(c.0.clone())));
231    }
232
233    #[tokio::test]
234    async fn test_fallback_client_managed_found() {
235        let registry = StateRegistry::new_with_memory_db();
236        let test_repo = Arc::new(TestA(12345));
237
238        registry.register_client_managed(test_repo.clone());
239
240        let repo = registry.get::<TestItem<usize>>().unwrap();
241        let result = repo.get(String::new()).await.unwrap();
242
243        assert_eq!(result, Some(TestItem(12345)));
244    }
245
246    #[tokio::test]
247    async fn test_new_with_memory_db_sync() {
248        // Construct in sync context (no .await on the constructor itself)
249        let registry = StateRegistry::new_with_memory_db();
250        // Database must be accessible via async get after sync construction
251        let repo = registry.get::<TestItem<usize>>().unwrap();
252        let result = repo.get(String::new()).await;
253        // Should return Ok(None) — key not found, not an error
254        // (Note: TestItem<usize> is registered in this test module already)
255        assert!(result.is_ok());
256    }
257
258    #[tokio::test]
259    async fn test_wipe_disconnects_outstanding_repository_handles() {
260        let registry = StateRegistry::new_with_memory_db();
261        let repo = registry.get::<TestItem<usize>>().unwrap();
262        repo.set(String::new(), TestItem(42usize)).await.unwrap();
263
264        registry.wipe().await.unwrap();
265
266        assert!(matches!(
267            repo.get(String::new()).await,
268            Err(RepositoryError::Database(DatabaseError::Closed))
269        ));
270        assert!(matches!(
271            repo.list().await,
272            Err(RepositoryError::Database(DatabaseError::Closed))
273        ));
274    }
275
276    #[tokio::test]
277    async fn test_wipe_clears_client_managed() {
278        let registry = StateRegistry::new_with_memory_db();
279        registry.register_client_managed(Arc::new(TestA(99)));
280
281        registry.wipe().await.unwrap();
282
283        // Client-managed is gone; falls through to SDK-managed (now closed).
284        let repo = registry.get::<TestItem<usize>>().unwrap();
285        assert!(matches!(
286            repo.get(String::new()).await,
287            Err(RepositoryError::Database(DatabaseError::Closed))
288        ));
289    }
290
291    #[tokio::test]
292    async fn test_wipe_is_idempotent() {
293        let registry = StateRegistry::new_with_memory_db();
294        registry.wipe().await.unwrap();
295        registry.wipe().await.unwrap();
296    }
297
298    #[tokio::test]
299    async fn test_persisted_value_on_memory_db() {
300        use crate::register_persisted_value;
301        register_persisted_value!(TEST_VALUE: String = "test_registry_persisted_value_key");
302
303        let registry = StateRegistry::new_with_memory_db();
304        let value = registry.value::<TEST_VALUE>();
305
306        // Value must not exist initially
307        assert_eq!(value.try_get().await.unwrap(), None::<String>);
308
309        // Update and read back
310        value.set("hello".to_string()).await.unwrap();
311        assert_eq!(value.try_get().await.unwrap(), Some("hello".to_string()));
312
313        // Delete and confirm gone
314        value.remove().await.unwrap();
315        assert_eq!(value.try_get().await.unwrap(), None::<String>);
316    }
317
318    #[tokio::test]
319    async fn test_persisted_values_are_isolated_by_key() {
320        use crate::register_persisted_value;
321        register_persisted_value!(THEME: String = "test_theme");
322        register_persisted_value!(LOCALE: String = "test_locale");
323
324        let registry = StateRegistry::new_with_memory_db();
325        let theme = registry.value::<THEME>();
326        let locale = registry.value::<LOCALE>();
327
328        theme.set("dark".to_string()).await.unwrap();
329        locale.set("en-US".to_string()).await.unwrap();
330
331        theme.remove().await.unwrap();
332        assert_eq!(theme.try_get().await.unwrap(), None::<String>);
333        assert_eq!(locale.try_get().await.unwrap(), Some("en-US".to_string()));
334    }
335
336    #[tokio::test]
337    async fn test_persisted_value_is_stored_in_the_persisted_value_table_as_bare_json() {
338        use crate::{persisted_value::PersistedValueItem, register_persisted_value};
339
340        #[derive(Debug, PartialEq, serde::Serialize, serde::Deserialize)]
341        struct Config {
342            theme: String,
343        }
344        register_persisted_value!(CONFIG: Config = "test_config");
345
346        let registry = StateRegistry::new_with_memory_db();
347        let value = Config {
348            theme: "dark".to_string(),
349        };
350        let expected = serde_json::to_value(&value).unwrap();
351        registry.value::<CONFIG>().set(value).await.unwrap();
352
353        // Storage contract for existing databases: persisted values live in the `PersistedValue`
354        // table, addressed by key name, holding the bare serialized value.
355        assert_eq!(PersistedValueItem::NAME, "PersistedValue");
356        let raw: PersistedValueItem = registry
357            .database
358            .get::<PersistedValueItem>("test_config")
359            .await
360            .unwrap()
361            .expect("persisted value is present");
362        assert_eq!(raw.0, expected);
363    }
364
365    #[tokio::test]
366    async fn test_persisted_value_reports_closed_after_wipe() {
367        use crate::{persisted_value::PersistedValueError, register_persisted_value};
368        register_persisted_value!(TEST_VALUE: String = "test_wiped_persisted_value");
369
370        let registry = StateRegistry::new_with_memory_db();
371        let value = registry.value::<TEST_VALUE>();
372        value.set("hello".to_string()).await.unwrap();
373
374        registry.wipe().await.unwrap();
375
376        assert!(matches!(
377            value.try_get().await,
378            Err(PersistedValueError::Closed)
379        ));
380        assert!(matches!(
381            value.set("bye".to_string()).await,
382            Err(PersistedValueError::Closed)
383        ));
384    }
385
386    /// The concrete implementation is erased by the coercion to `Arc<dyn Repository<T>>`, so two
387    /// implementations of the same item type share a slot and the later registration wins.
388    #[tokio::test]
389    async fn test_register_client_managed_replaces_previous_implementation() {
390        let registry = StateRegistry::new_with_memory_db();
391
392        registry.register_client_managed(Arc::new(TestA(1)));
393        registry.register_client_managed(Arc::new(TestD(2)));
394
395        let repo = registry.get_client_managed::<TestItem<usize>>().unwrap();
396        assert_eq!(repo.get(String::new()).await.unwrap(), Some(TestItem(2)));
397    }
398}
bitwarden_state/registry.rs

bitwarden_state/
registry.rs
