Migrate SDK to CryptoService

crates/bitwarden-core/src/auth/auth_request.rs

@@ -8,7 +8,7 @@ use bitwarden_crypto::{EncString, KeyDecryptable, SymmetricCryptoKey};
 
 #[cfg(feature = "internal")]
 use crate::client::encryption_settings::EncryptionSettingsError;
-use crate::{error::Error, Client};
+use crate::{error::Error, key_management::SymmetricKeyRef, Client};
 
 #[cfg_attr(feature = "uniffi", derive(uniffi::Record))]
 pub struct AuthRequestResponse {
@@ -83,8 +83,11 @@ pub(crate) fn approve_auth_request(
 ) -> Result<AsymmetricEncString, Error> {
     let public_key = AsymmetricPublicCryptoKey::from_der(&STANDARD.decode(public_key)?)?;
 
-    let enc = client.internal.get_encryption_settings()?;
-    let key = enc.get_key(&None)?;
+    let crypto_service = client.internal.get_crypto_service();
+    let ctx = crypto_service.context();
+
+    #[allow(deprecated)]
+    let key = ctx.dangerous_get_symmetric_key(SymmetricKeyRef::User)?;
 
     Ok(AsymmetricEncString::encrypt_rsa2048_oaep_sha1(
         &key.to_vec(),
@@ -235,21 +238,22 @@ mod tests {
 
         // We can validate that the vault is unlocked correctly by confirming the user key is the
         // same
-        assert_eq!(
-            existing_device
-                .internal
-                .get_encryption_settings()
-                .unwrap()
-                .get_key(&None)
+        let existing_key_b64 = {
+            let ctx = existing_device.internal.get_crypto_service().context();
+            #[allow(deprecated)]
+            ctx.dangerous_get_symmetric_key(SymmetricKeyRef::User)
                 .unwrap()
-                .to_base64(),
-            new_device
-                .internal
-                .get_encryption_settings()
-                .unwrap()
-                .get_key(&None)
+                .to_base64()
+        };
+
+        let new_key_b64 = {
+            let ctx = new_device.internal.get_crypto_service().context();
+            #[allow(deprecated)]
+            ctx.dangerous_get_symmetric_key(SymmetricKeyRef::User)
                 .unwrap()
                 .to_base64()
-        );
+        };
+
+        assert_eq!(existing_key_b64, new_key_b64,);
     }
 }

crates/bitwarden-core/src/auth/client_auth.rs

@@ -162,9 +162,11 @@ impl<'a> ClientAuth<'a> {
 
 #[cfg(feature = "internal")]
 fn trust_device(client: &Client) -> Result<TrustDeviceResponse> {
-    let enc = client.internal.get_encryption_settings()?;
+    use crate::key_management::SymmetricKeyRef;
 
-    let user_key = enc.get_key(&None)?;
+    let ctx = client.internal.get_crypto_service().context();
+    #[allow(deprecated)]
+    let user_key = ctx.dangerous_get_symmetric_key(SymmetricKeyRef::User)?;
 
     Ok(DeviceKey::trust_device(user_key)?)
 }

crates/bitwarden-core/src/auth/password/validate.rs

@@ -44,6 +44,8 @@ pub(crate) fn validate_password_user_key(
     password: String,
     encrypted_user_key: String,
 ) -> Result<String> {
+    use crate::key_management::SymmetricKeyRef;
+
     let login_method = client
         .internal
         .get_login_method()
@@ -59,9 +61,10 @@ pub(crate) fn validate_password_user_key(
                     .decrypt_user_key(encrypted_user_key.parse()?)
                     .map_err(|_| "wrong password")?;
 
-                let enc = client.internal.get_encryption_settings()?;
-
-                let existing_key = enc.get_key(&None)?;
+                let crypto_service = client.internal.get_crypto_service();
+                let ctx = crypto_service.context();
+                #[allow(deprecated)]
+                let existing_key = ctx.dangerous_get_symmetric_key(SymmetricKeyRef::User)?;
 
                 if user_key.to_vec() != existing_key.to_vec() {
                     return Err("wrong user key".into());

crates/bitwarden-core/src/auth/pin.rs

@@ -3,6 +3,7 @@ use bitwarden_crypto::{EncString, PinKey};
 use crate::{
     client::{LoginMethod, UserLoginMethod},
     error::{Error, Result},
+    key_management::SymmetricKeyRef,
     Client,
 };
 
@@ -24,8 +25,9 @@ pub(crate) fn validate_pin(
     match login_method {
         UserLoginMethod::Username { email, kdf, .. }
         | UserLoginMethod::ApiKey { email, kdf, .. } => {
-            let enc = client.internal.get_encryption_settings()?;
-            let user_key = enc.get_key(&None)?;
+            let ctx = client.internal.get_crypto_service().context();
+            #[allow(deprecated)]
+            let user_key = ctx.dangerous_get_symmetric_key(SymmetricKeyRef::User)?;
 
             let pin_key = PinKey::derive(pin.as_bytes(), email.as_bytes(), kdf)?;
 

crates/bitwarden-core/src/auth/renew.rs

@@ -10,6 +10,7 @@ use crate::{
     auth::api::{request::ApiTokenRequest, response::IdentityTokenResponse},
     client::{internal::InternalClient, LoginMethod, UserLoginMethod},
     error::{Error, Result},
+    key_management::SymmetricKeyRef,
 };
 
 pub(crate) async fn renew_token(client: &InternalClient) -> Result<()> {
@@ -70,10 +71,14 @@ pub(crate) async fn renew_token(client: &InternalClient) -> Result<()> {
                     .send(&config)
                     .await?;
 
-                    if let (IdentityTokenResponse::Payload(r), Some(state_file), Ok(enc_settings)) =
-                        (&result, state_file, client.get_encryption_settings())
+                    if let (IdentityTokenResponse::Payload(r), Some(state_file)) =
+                        (&result, state_file)
                     {
-                        if let Ok(enc_key) = enc_settings.get_key(&None) {
+                        let ctx = client.get_crypto_service().context();
+
+                        #[allow(deprecated)]
+                        if let Ok(enc_key) = ctx.dangerous_get_symmetric_key(SymmetricKeyRef::User)
+                        {
                             let state =
                                 ClientState::new(r.access_token.clone(), enc_key.to_base64());
                             _ = state::set(state_file, access_token, state);

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

@@ -1,5 +1,6 @@
 use std::sync::{Arc, RwLock};
 
+use bitwarden_crypto::service::CryptoService;
 use reqwest::header::{self, HeaderValue};
 
 use super::internal::InternalClient;
@@ -79,6 +80,7 @@ impl Client {
                 })),
                 external_client,
                 encryption_settings: RwLock::new(None),
+                crypto_service: CryptoService::new(),
             },
         }
     }

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

@@ -1,14 +1,15 @@
-use std::collections::HashMap;
-
-use bitwarden_crypto::{AsymmetricCryptoKey, CryptoError, KeyContainer, SymmetricCryptoKey};
+use bitwarden_crypto::{service::CryptoService, AsymmetricCryptoKey, SymmetricCryptoKey};
 #[cfg(feature = "internal")]
 use bitwarden_crypto::{AsymmetricEncString, EncString, MasterKey};
 use thiserror::Error;
 use uuid::Uuid;
 
 #[cfg(feature = "internal")]
 use crate::error::Result;
-use crate::VaultLocked;
+use crate::{
+    key_management::{AsymmetricKeyRef, SymmetricKeyRef},
+    VaultLocked,
+};
 
 #[derive(Debug, Error)]
 pub enum EncryptionSettingsError {
@@ -30,9 +31,7 @@ pub enum EncryptionSettingsError {
 
 #[derive(Clone)]
 pub struct EncryptionSettings {
-    user_key: SymmetricCryptoKey,
-    pub(crate) private_key: Option<AsymmetricCryptoKey>,
-    org_keys: HashMap<Uuid, SymmetricCryptoKey>,
+    crypto_service: CryptoService<SymmetricKeyRef, AsymmetricKeyRef>,
 }
 
 impl std::fmt::Debug for EncryptionSettings {
@@ -48,10 +47,12 @@ impl EncryptionSettings {
         master_key: MasterKey,
         user_key: EncString,
         private_key: EncString,
+        crypto_service: &CryptoService<SymmetricKeyRef, AsymmetricKeyRef>,
     ) -> Result<Self, EncryptionSettingsError> {
         // Decrypt the user key
+
         let user_key = master_key.decrypt_user_key(user_key)?;
-        Self::new_decrypted_key(user_key, private_key)
+        Self::new_decrypted_key(user_key, private_key, crypto_service)
     }
 
     /// Initialize the encryption settings with the decrypted user key and the encrypted user
@@ -62,6 +63,7 @@ impl EncryptionSettings {
     pub(crate) fn new_decrypted_key(
         user_key: SymmetricCryptoKey,
         private_key: EncString,
+        crypto_service: &CryptoService<SymmetricKeyRef, AsymmetricKeyRef>,
     ) -> Result<Self, EncryptionSettingsError> {
         use bitwarden_crypto::KeyDecryptable;
         use log::warn;
@@ -83,21 +85,35 @@ impl EncryptionSettings {
             // )
         };
 
+        #[allow(deprecated)]
+        {
+            let mut ctx = crypto_service.context_mut();
+            ctx.set_symmetric_key(SymmetricKeyRef::User, user_key)?;
+            if let Some(private_key) = private_key {
+                ctx.set_asymmetric_key(AsymmetricKeyRef::UserPrivateKey, private_key)?;
+            }
+        }
+
         Ok(EncryptionSettings {
-            user_key,
-            private_key,
-            org_keys: HashMap::new(),
+            crypto_service: crypto_service.clone(),
         })
     }
 
     /// Initialize the encryption settings with only a single decrypted key.
     /// This is used only for logging in Secrets Manager with an access token
     #[cfg(feature = "secrets")]
-    pub(crate) fn new_single_key(key: SymmetricCryptoKey) -> Self {
+    pub(crate) fn new_single_key(
+        key: SymmetricCryptoKey,
+        crypto_service: &CryptoService<SymmetricKeyRef, AsymmetricKeyRef>,
+    ) -> Self {
+        #[allow(deprecated)]
+        crypto_service
+            .context_mut()
+            .set_symmetric_key(SymmetricKeyRef::User, key)
+            .expect("Mutable context");
+
         EncryptionSettings {
-            user_key: key,
-            private_key: None,
-            org_keys: HashMap::new(),
+            crypto_service: crypto_service.clone(),
         }
     }
 
@@ -106,53 +122,26 @@ impl EncryptionSettings {
         &mut self,
         org_enc_keys: Vec<(Uuid, AsymmetricEncString)>,
     ) -> Result<&Self, EncryptionSettingsError> {
-        use bitwarden_crypto::KeyDecryptable;
+        let mut ctx = self.crypto_service.context_mut();
 
         // Make sure we only keep the keys given in the arguments and not any of the previous
         // ones, which might be from organizations that the user is no longer a part of anymore
-        self.org_keys.clear();
+        ctx.retain_symmetric_keys(|key_ref| !matches!(key_ref, SymmetricKeyRef::Organization(_)));
 
         // FIXME: [PM-11690] - Early abort to handle private key being corrupt
         if org_enc_keys.is_empty() {
             return Ok(self);
         }
 
-        let private_key = self
-            .private_key
-            .as_ref()
-            .ok_or(EncryptionSettingsError::MissingPrivateKey)?;
-
         // Decrypt the org keys with the private key
         for (org_id, org_enc_key) in org_enc_keys {
-            let mut dec: Vec<u8> = org_enc_key.decrypt_with_key(private_key)?;
-
-            let org_key = SymmetricCryptoKey::try_from(dec.as_mut_slice())?;
-
-            self.org_keys.insert(org_id, org_key);
+            ctx.decrypt_symmetric_key_with_asymmetric_key(
+                AsymmetricKeyRef::UserPrivateKey,
+                SymmetricKeyRef::Organization(org_id),
+                &org_enc_key,
+            )?;
         }
 
         Ok(self)
     }
-
-    pub fn get_key(&self, org_id: &Option<Uuid>) -> Result<&SymmetricCryptoKey, CryptoError> {
-        // If we don't have a private key set (to decode multiple org keys), we just use the main
-        // user key
-        if self.private_key.is_none() {
-            return Ok(&self.user_key);
-        }
-
-        match org_id {
-            Some(org_id) => self
-                .org_keys
-                .get(org_id)
-                .ok_or(CryptoError::MissingKey(*org_id)),
-            None => Ok(&self.user_key),
-        }
-    }
-}
-
-impl KeyContainer for EncryptionSettings {
-    fn get_key(&self, org_id: &Option<Uuid>) -> Result<&SymmetricCryptoKey, CryptoError> {
-        EncryptionSettings::get_key(self, org_id)
-    }
 }

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

@@ -1,5 +1,6 @@
 use std::sync::{Arc, RwLock};
 
+use bitwarden_crypto::service::CryptoService;
 #[cfg(any(feature = "internal", feature = "secrets"))]
 use bitwarden_crypto::SymmetricCryptoKey;
 #[cfg(feature = "internal")]
@@ -13,6 +14,7 @@ use crate::{
     auth::renew::renew_token,
     client::{encryption_settings::EncryptionSettings, login_method::LoginMethod},
     error::{Result, VaultLocked},
+    key_management::{AsymmetricKeyRef, SymmetricKeyRef},
     DeviceType,
 };
 #[cfg(feature = "internal")]
@@ -59,6 +61,8 @@ pub struct InternalClient {
     pub(crate) external_client: reqwest::Client,
 
     pub(super) encryption_settings: RwLock<Option<Arc<EncryptionSettings>>>,
+
+    pub(super) crypto_service: CryptoService<SymmetricKeyRef, AsymmetricKeyRef>,
 }
 
 impl InternalClient {
@@ -175,6 +179,10 @@ impl InternalClient {
             .ok_or(VaultLocked)
     }
 
+    pub fn get_crypto_service(&self) -> &CryptoService<SymmetricKeyRef, AsymmetricKeyRef> {
+        &self.crypto_service
+    }
+
     #[cfg(feature = "internal")]
     pub(crate) fn initialize_user_crypto_master_key(
         &self,
@@ -189,6 +197,7 @@ impl InternalClient {
             master_key,
             user_key,
             private_key,
+            &self.crypto_service,
         )?));
 
         Ok(())
@@ -204,7 +213,7 @@ impl InternalClient {
             .encryption_settings
             .write()
             .expect("RwLock is not poisoned") = Some(Arc::new(
-            EncryptionSettings::new_decrypted_key(user_key, private_key)?,
+            EncryptionSettings::new_decrypted_key(user_key, private_key, &self.crypto_service)?,
         ));
 
         Ok(())
@@ -226,8 +235,10 @@ impl InternalClient {
         *self
             .encryption_settings
             .write()
-            .expect("RwLock is not poisoned") =
-            Some(Arc::new(EncryptionSettings::new_single_key(key)));
+            .expect("RwLock is not poisoned") = Some(Arc::new(EncryptionSettings::new_single_key(
+            key,
+            &self.crypto_service,
+        )));
     }
 
     #[cfg(feature = "internal")]