[PM-31128] Add reinit_user_crypto for mobile

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

@@ -31,7 +31,8 @@ use crate::{
     },
     error::NotAuthenticatedError,
     key_management::{
-        MasterPasswordUnlockData, SecurityState, V2UpgradeToken,
+        MasterPasswordUnlockData, PrivateKeySlotId, SecurityState, SigningKeySlotId,
+        SymmetricKeySlotId, V2UpgradeToken,
         account_cryptographic_state::WrappedAccountCryptographicState, state_bridge::StateBridge,
     },
 };
@@ -292,6 +293,15 @@ impl InternalClient {
         // Note: The actual key does not get logged unless the crypto crate has the
         // dangerous-crypto-debug feature enabled, so this is safe
         info!("Setting user key with ID {:?}", user_key_id);
+
+        // The key store should not already have any keys initialized
+        if ctx.has_symmetric_key(SymmetricKeySlotId::User)
+            || ctx.has_private_key(PrivateKeySlotId::UserPrivateKey)
+            || ctx.has_signing_key(SigningKeySlotId::UserSigningKey)
+        {
+            return Err(EncryptionSettingsError::CryptoInitialization);
+        }
+
         // The user key gets set to the local context frame here; It then gets persisted to the
         // context when the cryptographic state was unwrapped correctly, so that there is no
         // risk of a partial / incorrect setup.

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

@@ -53,10 +53,6 @@ pub enum AccountCryptographyInitializationError {
     /// The decrypted data is corrupt.
     #[error("Signature or mac verification failed, the data may have been tampered with")]
     TamperedData,
-    /// The key store is already initialized with account keys. Currently, updating keys is not a
-    /// supported operation
-    #[error("Key store is already initialized")]
-    KeyStoreAlreadyInitialized,
     /// A generic cryptographic error occurred.
     #[error("A generic cryptographic error occurred: {0}")]
     GenericCrypto(CryptoError),
@@ -491,13 +487,6 @@ impl WrappedAccountCryptographicState {
         store: &KeyStore<KeySlotIds>,
         mut ctx: KeyStoreContext<KeySlotIds>,
     ) -> Result<(), AccountCryptographyInitializationError> {
-        if ctx.has_symmetric_key(SymmetricKeySlotId::User)
-            || ctx.has_private_key(PrivateKeySlotId::UserPrivateKey)
-            || ctx.has_signing_key(SigningKeySlotId::UserSigningKey)
-        {
-            return Err(AccountCryptographyInitializationError::KeyStoreAlreadyInitialized);
-        }
-
         match self {
             WrappedAccountCryptographicState::V1 { private_key } => {
                 info!(state = ?self, "Initializing V1 account cryptographic state");

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

@@ -4,6 +4,8 @@
 //! the SDK is fully implemented. When porting functionality from `client` the mobile clients should
 //! be updated to consume the regular code paths and in this module should eventually disappear.
 
+#[cfg(feature = "uniffi")]
+mod reinit_user_crypto;
 use std::collections::HashMap;
 
 use bitwarden_api_api::models::AccountKeysRequestModel;
@@ -19,6 +21,10 @@ use bitwarden_crypto::{
 };
 use bitwarden_encoding::B64;
 use bitwarden_error::bitwarden_error;
+#[cfg(feature = "uniffi")]
+pub(super) use reinit_user_crypto::reinit_user_crypto;
+#[cfg(feature = "uniffi")]
+pub use reinit_user_crypto::{ReinitUserCryptoError, ReinitUserCryptoRequest};
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use tracing::info;
@@ -392,8 +398,7 @@ pub(super) async fn initialize_user_crypto(
             .map_err(|_| EncryptionSettingsError::UserKeyStateUpdateFailed)?;
     }
 
-    initialize_user_local_data_key(client).await?;
-    PinLockSystem::on_unlock(&PinLockSystem::with_client(client)).await;
+    on_unlock_handler(client).await?;
 
     client
         .internal
@@ -1159,6 +1164,17 @@ async fn initialize_user_local_data_key(client: &Client) -> Result<(), Encryptio
         .map_err(|_| EncryptionSettingsError::LocalUserDataKeyLoadFailed)
 }
 
+/// Runs the code needed post unlock used by `initialize_user_crypto` and `reinit_user_crypto`.
+///
+/// Both code paths leave the SDK in an unlocked state with the active user key in the key store,
+/// and both need to ensure derived per-user state is consistent with that user key before clients
+/// can use the session.
+async fn on_unlock_handler(client: &Client) -> Result<(), EncryptionSettingsError> {
+    initialize_user_local_data_key(client).await?;
+    PinLockSystem::on_unlock(&PinLockSystem::with_client(client)).await;
+    Ok(())
+}
+
 /// Create the data needed to register for JIT master password
 pub(crate) fn make_user_jit_master_password_registration(
     client: &Client,

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

@@ -19,6 +19,10 @@ use super::crypto::{
     make_user_password_registration, verify_asymmetric_keys,
 };
 use crate::key_management::V2UpgradeToken;
+#[cfg(feature = "uniffi")]
+use crate::key_management::crypto::{
+    ReinitUserCryptoError, ReinitUserCryptoRequest, reinit_user_crypto,
+};
 #[cfg(feature = "internal")]
 use crate::key_management::{
     SymmetricKeySlotId,
@@ -343,6 +347,20 @@ impl CryptoClient {
     }
 }
 
+#[cfg(feature = "uniffi")]
+impl CryptoClient {
+    /// Re-initialize the user's cryptographic state during an unlock session.
+    ///
+    /// Requires the SDK to be unlocked. Replaces the in-memory account
+    /// cryptographic state with the provided one, and upgrades the active user key from V1 to V2.
+    pub async fn reinit_user_crypto(
+        &self,
+        req: ReinitUserCryptoRequest,
+    ) -> Result<(), ReinitUserCryptoError> {
+        reinit_user_crypto(&self.client, req).await
+    }
+}
+
 impl Client {
     /// Access to crypto functionality.
     pub fn crypto(&self) -> CryptoClient {

crates/bitwarden-uniffi/src/crypto.rs

@@ -2,7 +2,7 @@ use bitwarden_core::key_management::{
     V2UpgradeToken,
     crypto::{
         DeriveKeyConnectorRequest, DerivePinKeyResponse, EnrollPinResponse, InitOrgCryptoRequest,
-        InitUserCryptoRequest, UpdateKdfResponse, UpdatePasswordResponse,
+        InitUserCryptoRequest, ReinitUserCryptoRequest, UpdateKdfResponse, UpdatePasswordResponse,
     },
 };
 use bitwarden_crypto::{EncString, Kdf, RotateableKeySet, UnsignedSharedKey};
@@ -28,6 +28,13 @@ impl CryptoClient {
         Ok(self.0.initialize_org_crypto(req).await?)
     }
 
+    /// Re-initialize the user's cryptographic state during an unlock session for handling a synced
+    /// v2 upgrade token. Requires the SDK to be unlocked. See
+    /// [`bitwarden_core::key_management::CryptoClient::reinit_user_crypto`].
+    pub async fn reinit_user_crypto(&self, req: ReinitUserCryptoRequest) -> Result<()> {
+        Ok(self.0.reinit_user_crypto(req).await?)
+    }
+
     /// Get the uses's decrypted encryption key. Note: It's very important
     /// to keep this key safe, as it can be used to decrypt all of the user's data
     pub async fn get_user_encryption_key(&self) -> Result<B64> {

crates/bitwarden-uniffi/src/error.rs

@@ -23,6 +23,8 @@ pub enum BitwardenError {
     #[error(transparent)]
     MobileCrypto(#[from] bitwarden_core::key_management::crypto::CryptoClientError),
     #[error(transparent)]
+    ReinitUserCrypto(#[from] bitwarden_core::key_management::crypto::ReinitUserCryptoError),
+    #[error(transparent)]
     AuthValidate(#[from] bitwarden_core::auth::AuthValidateError),
     #[error(transparent)]
     ApproveAuthRequest(#[from] bitwarden_core::auth::ApproveAuthRequestError),

crates/bitwarden-uniffi/swift/integration-tests/Tests/IntegrationTests/ReinitUserCryptoTests.swift

@@ -0,0 +1,91 @@
+import BitwardenSdk
+import XCTest
+
+/// Swift integration tests for `CryptoClient.reinitUserCrypto`.
+final class ReinitUserCryptoTests: XCTestCase {
+    var stateBridge: InMemoryStateBridge!
+
+    override func setUp() async throws {
+        try await super.setUp()
+        stateBridge = InMemoryStateBridge()
+    }
+
+    func testReturnsNotUnlockedWhenLocked() async throws {
+        // No `initializeUserCrypto` call — the user-key slot is empty.
+        let client = Client(tokenProvider: MockTokenProvider(), settings: nil)
+        client.kmStateBridge().registerBridgeImpl(bridgeImpl: stateBridge)
+
+        let req = ReinitUserCryptoRequest(
+            accountCryptographicState: makeV2AccountCryptographicState(),
+            upgradeToken: makeMockUpgradeToken()
+        )
+
+        do {
+            try await client.crypto().reinitUserCrypto(req: req)
+            XCTFail("expected ReinitUserCryptoError.NotUnlocked")
+        } catch BitwardenError.ReinitUserCrypto(let inner) {
+            guard case .NotUnlocked = inner else {
+                XCTFail("expected .NotUnlocked, got \(inner)")
+                return
+            }
+        }
+    }
+
+    func testReturnsAlreadyV2WhenActiveUserIsV2() async throws {
+        let client = try await makeV2InitializedClient(stateBridge: stateBridge)
+
+        let req = ReinitUserCryptoRequest(
+            accountCryptographicState: makeV2AccountCryptographicState(),
+            upgradeToken: makeMockUpgradeToken()
+        )
+
+        do {
+            try await client.crypto().reinitUserCrypto(req: req)
+            XCTFail("expected ReinitUserCryptoError.AlreadyV2Encryption")
+        } catch BitwardenError.ReinitUserCrypto(let inner) {
+            guard case .AlreadyV2Encryption = inner else {
+                XCTFail("expected .AlreadyV2Encryption, got \(inner)")
+                return
+            }
+        }
+    }
+
+    func testUpgradesV1ToV2WithValidToken() async throws {
+        let client = try await makeV1InitializedClient(stateBridge: stateBridge)
+
+        let upgradeToken = makeValidUpgradeToken()
+        await stateBridge.setV2UpgradeToken(value: upgradeToken)
+
+        let req = ReinitUserCryptoRequest(
+            accountCryptographicState: makeV2AccountCryptographicState(),
+            upgradeToken: upgradeToken
+        )
+
+        try await client.crypto().reinitUserCrypto(req: req)
+
+        // After a successful V1→V2 reinit, the active user key in the slot
+        // must be the V2 test-vector key (returned base64-encoded by
+        // `getUserEncryptionKey` for V2 keys via COSE serialization).
+        let userKey = try await client.crypto().getUserEncryptionKey()
+        XCTAssertEqual(userKey, TEST_VECTOR_USER_KEY_V2_B64)
+    }
+
+    func testInvalidUpgradeTokenReturnsError() async throws {
+        let client = try await makeV1InitializedClient(stateBridge: stateBridge)
+
+        let req = ReinitUserCryptoRequest(
+            accountCryptographicState: makeV2AccountCryptographicState(),
+            upgradeToken: makeMockUpgradeToken()
+        )
+
+        do {
+            try await client.crypto().reinitUserCrypto(req: req)
+            XCTFail("expected ReinitUserCryptoError.InvalidUpgradeToken")
+        } catch BitwardenError.ReinitUserCrypto(let inner) {
+            guard case .InvalidUpgradeToken = inner else {
+                XCTFail("expected .InvalidUpgradeToken, got \(inner)")
+                return
+            }
+        }
+    }
+}