[PM-37190] Use state bridge instead of platform state

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

@@ -11,7 +11,7 @@ use bitwarden_test::MemoryRepository;
 use crate::{
     Client, ClientSettings, UserId,
     key_management::{
-        LocalUserDataKeyState, MasterPasswordUnlockData, UserKeyState,
+        LocalUserDataKeyState, MasterPasswordUnlockData,
         account_cryptographic_state::WrappedAccountCryptographicState,
         crypto::{InitOrgCryptoRequest, InitUserCryptoMethod, InitUserCryptoRequest},
     },
@@ -23,12 +23,6 @@ impl Client {
     /// Does not initialize any crypto state.
     pub fn new_test(settings: Option<ClientSettings>) -> Self {
         let client = Client::new(settings);
-        client
-            .platform()
-            .state()
-            .register_client_managed(std::sync::Arc::new(
-                MemoryRepository::<UserKeyState>::default(),
-            ));
         client
             .platform()
             .state()

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

@@ -10,17 +10,14 @@
 //!   [CompositeEncryptable](bitwarden_crypto::CompositeEncryptable), and
 //!   [Decryptable](bitwarden_crypto::Decryptable).
 
-use bitwarden_crypto::{
-    EncString, KeyStore, SymmetricCryptoKey, key_slot_ids, safe::PasswordProtectedKeyEnvelope,
-};
+use bitwarden_crypto::{EncString, KeyStore, SymmetricCryptoKey, key_slot_ids};
 
 #[cfg(feature = "internal")]
 pub mod account_cryptographic_state;
 #[cfg(feature = "internal")]
 pub mod crypto;
 #[cfg(feature = "internal")]
 mod crypto_client;
-use bitwarden_encoding::B64;
 #[cfg(feature = "internal")]
 pub use crypto_client::CryptoClient;
 
@@ -68,18 +65,6 @@ pub mod state_bridge;
 
 use crate::{OrganizationId, UserId};
 
-/// Represents the decrypted symmetric user-key of a user. This is held in ephemeral state of the
-/// client.
-#[derive(Serialize, Deserialize, Debug, Clone)]
-#[repr(transparent)]
-#[cfg_attr(feature = "wasm", derive(Tsify), tsify(into_wasm_abi, from_wasm_abi))]
-#[cfg_attr(feature = "uniffi", derive(uniffi::Record))]
-pub struct UserKeyState {
-    decrypted_user_key: B64,
-}
-
-bitwarden_state::register_repository_item!(String => UserKeyState, "UserKey");
-
 /// Represents the local user data key, wrapped by user key.
 /// This key is used to encrypt local user data (e.g., password generator history).
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -91,16 +76,6 @@ pub struct LocalUserDataKeyState {
 
 bitwarden_state::register_repository_item!(UserId => LocalUserDataKeyState, "LocalUserDataKey");
 
-/// Represents the PIN envelope in memory, when ephemeral PIN unlock is used.
-#[derive(Serialize, Deserialize, Debug, Clone)]
-#[cfg_attr(feature = "wasm", derive(Tsify), tsify(into_wasm_abi, from_wasm_abi))]
-#[cfg_attr(feature = "uniffi", derive(uniffi::Record))]
-pub struct EphemeralPinEnvelopeState {
-    pin_envelope: PasswordProtectedKeyEnvelope,
-}
-
-bitwarden_state::register_repository_item!(String => EphemeralPinEnvelopeState, "EphemeralPinEnvelope");
-
 key_slot_ids! {
     #[symmetric]
     pub enum SymmetricKeySlotId {

crates/bitwarden-core/src/key_management/wasm_unlock_state.rs

@@ -1,26 +1,19 @@
 //! The WASM sdk currently does not hold persistent SDK instances and instead re-createds SDK
 //! instances frequently. The unlock-state is lost, since the user-key is only held in the SDK. This
-//! file implements setting the user-key to WASM client-managed ephemeral state, so that
-//! SDK-re-creations have access to the user-key.
+//! file implements setting the user-key into the KM state bridge, so that SDK-re-creations have
+//! access to the user-key.
 //!
 //! This is not required on UNIFFI since there SDK instances live as long as the client is unlocked.
 //! Eventually, the WASM sdk will also hold SDK instances like described above.
 
 use bitwarden_crypto::SymmetricCryptoKey;
 use tracing::info;
 
-use crate::{
-    Client,
-    key_management::{self, SymmetricKeySlotId},
-};
-
-// The repository pattern requires us to specify a key. Here we use an empty string as the only
-// key for this repository map.
-const USER_KEY_REPOSITORY_KEY: &str = "";
+use crate::{Client, key_management::SymmetricKeySlotId};
 
 /// Error indicating inability to set the user key into state
 pub(crate) struct UnableToSetError;
-/// Sets the decrypted user key into the client-managed state, so that it survives re-creation of
+/// Sets the decrypted user key into the KM state bridge, so that it survives re-creation of
 /// the SDK
 pub(crate) async fn copy_user_key_to_client_managed_state(
     client: &Client,
@@ -36,67 +29,38 @@ pub(crate) async fn copy_user_key_to_client_managed_state(
             .clone()
     };
 
-    if let Ok(user_key_repository) = client
-        .platform()
-        .state()
-        .get::<key_management::UserKeyState>()
-    {
-        // We do not want to set the user-key if it is already set as that may trigger an observable
-        // loop in the client side which subscribes to the state
-        if let Ok(Some(existing_key)) = user_key_repository
-            .get(USER_KEY_REPOSITORY_KEY.to_string())
-            .await
-        {
-            if SymmetricCryptoKey::try_from(existing_key.decrypted_user_key)
-                .map_err(|_| UnableToSetError)?
-                == user_key
-            {
-                info!("User-key in client managed state is already up to date, skipping set");
-                return Ok(());
-            } else {
-                info!("User-key in client managed state is outdated, updating it");
-            }
-        } else {
-            info!("No user-key in client managed state, setting it");
+    let bridge = client.km_state_bridge();
+    if !bridge.is_bridge_registered() {
+        // No state bridge registered, older clients should just return gracefully.
+        info!("No state bridge registered, exiting gracefully");
+        return Ok(());
+    }
+
+    // We do not want to set the user-key if it is already set as that may trigger an observable
+    // loop in the client side which subscribes to the state
+    if let Some(existing_key) = bridge.get_user_key().await {
+        if existing_key == user_key {
+            info!("User-key in state bridge is already up to date, skipping set");
+            return Ok(());
         }
+        info!("User-key in state bridge is outdated, updating it");
     } else {
-        // UserKeyState repository does not exist, older clients should
-        // just return gracefully.
-        info!("No UserKeyState repository exists in client managed state, exiting gracefully");
-        return Ok(());
+        info!("No user-key in state bridge, setting it");
     }
 
-    info!("Setting the user-key to client managed-state from SDK");
-    // Set the user-key into the state repository.
-    client
-        .platform()
-        .state()
-        .get::<key_management::UserKeyState>()
-        .map_err(|_| UnableToSetError)?
-        .set(
-            USER_KEY_REPOSITORY_KEY.to_string(),
-            key_management::UserKeyState {
-                decrypted_user_key: user_key.to_base64(),
-            },
-        )
-        .await
-        .map_err(|_| UnableToSetError)
+    info!("Setting the user-key to the state bridge from SDK");
+    bridge.set_user_key(&user_key).await;
+    Ok(())
 }
 
 pub(crate) struct UnableToGetError;
 pub(crate) async fn get_user_key_from_client_managed_state(
     client: &Client,
 ) -> Result<SymmetricCryptoKey, UnableToGetError> {
-    info!("Getting the user-key from client managed-state in SDK");
-    // Get the user-key from the state repository.
-    let user_key_state = client
-        .platform()
-        .state()
-        .get::<key_management::UserKeyState>()
-        .map_err(|_| UnableToGetError)?
-        .get(USER_KEY_REPOSITORY_KEY.to_string())
+    info!("Getting the user-key from the state bridge in SDK");
+    client
+        .km_state_bridge()
+        .get_user_key()
         .await
-        .map_err(|_| UnableToGetError)?
-        .ok_or(UnableToGetError)?;
-    SymmetricCryptoKey::try_from(user_key_state.decrypted_user_key).map_err(|_| UnableToGetError)
+        .ok_or(UnableToGetError)
 }

crates/bitwarden-core/tests/register.rs

@@ -12,7 +12,7 @@ async fn test_register_initialize_crypto() {
     use bitwarden_core::{
         Client, UserId,
         key_management::{
-            EphemeralPinEnvelopeState, MasterPasswordUnlockData, UserKeyState,
+            MasterPasswordUnlockData,
             account_cryptographic_state::WrappedAccountCryptographicState,
             crypto::{InitUserCryptoMethod, InitUserCryptoRequest},
         },
@@ -21,26 +21,13 @@ async fn test_register_initialize_crypto() {
 
     let client = Client::new(None);
 
-    let user_key_repository = MemoryRepository::<UserKeyState>::default();
-    client
-        .platform()
-        .state()
-        .register_client_managed(std::sync::Arc::new(user_key_repository));
-
     client
         .platform()
         .state()
         .register_client_managed(std::sync::Arc::new(
             MemoryRepository::<LocalUserDataKeyState>::default(),
         ));
 
-    client
-        .platform()
-        .state()
-        .register_client_managed(std::sync::Arc::new(MemoryRepository::<
-            EphemeralPinEnvelopeState,
-        >::default()));
-
     let email = "test@bitwarden.com";
     let password = "test123";
     let kdf = Kdf::PBKDF2 {

crates/bitwarden-pm/src/migrations.rs

@@ -1,8 +1,6 @@
 //! Manages repository migrations for the Bitwarden SDK.
 
-use bitwarden_core::{
-    client::persisted_state::OrganizationSharedKey, key_management::UserKeyState,
-};
+use bitwarden_core::client::persisted_state::OrganizationSharedKey;
 use bitwarden_state::{
     SettingItem,
     repository::{RepositoryItem, RepositoryMigrationStep, RepositoryMigrations},
@@ -17,7 +15,6 @@ pub fn get_sdk_managed_migrations() -> RepositoryMigrations {
         // requires a separate migration step using `Remove(...)`.
         Add(Cipher::data()),
         Add(Folder::data()),
-        Add(UserKeyState::data()),
         Add(SettingItem::data()),
         Add(OrganizationSharedKey::data()),
     ])
@@ -36,9 +33,7 @@ macro_rules! create_client_managed_repositories {
             // <fully qualified path to the item>, <item type idenfier>, <field name>, <name of the repository implementation>
             ::bitwarden_vault::Cipher, Cipher, cipher, CipherRepository;
             ::bitwarden_vault::Folder, Folder, folder, FolderRepository;
-            ::bitwarden_core::key_management::UserKeyState, UserKeyState, user_key_state, UserKeyStateRepository;
             ::bitwarden_core::key_management::LocalUserDataKeyState, LocalUserDataKeyState, local_user_data_key_state, LocalUserDataKeyStateRepository;
-            ::bitwarden_core::key_management::EphemeralPinEnvelopeState, EphemeralPinEnvelopeState, ephemeral_pin_envelope_state, EphemeralPinEnvelopeStateRepository;
             ::bitwarden_core::client::persisted_state::OrganizationSharedKey, OrganizationSharedKey, organization_shared_key, OrganizationSharedKeyRepository;
         }
     };