Update logic for handling the pin protected user key (#6169)

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/datasource/disk/AuthDiskSourceImpl.kt

@@ -145,9 +145,6 @@ class AuthDiskSourceImpl(
         storeInvalidUnlockAttempts(userId = userId, invalidUnlockAttempts = null)
         storeUserKey(userId = userId, userKey = null)
         storeUserAutoUnlockKey(userId = userId, userAutoUnlockKey = null)
-        storePinProtectedUserKey(userId = userId, pinProtectedUserKey = null)
-        storePinProtectedUserKeyEnvelope(userId = userId, pinProtectedUserKeyEnvelope = null)
-        storeEncryptedPin(userId = userId, encryptedPin = null)
         storePrivateKey(userId = userId, privateKey = null)
         storeAccountKeys(userId = userId, accountKeys = null)
         storeOrganizationKeys(userId = userId, organizationKeys = null)
@@ -163,9 +160,13 @@ class AuthDiskSourceImpl(
         storeShowImportLogins(userId = userId, showImportLogins = null)
         storeLastLockTimestamp(userId = userId, lastLockTimestamp = null)
 
-        // Do not remove the DeviceKey or PendingAuthRequest on logout, these are persisted
+        // Certain values are never removed as required by the feature requirements:
-        // indefinitely unless the TDE flow explicitly removes them.
+        // * EncryptedPin
-        // Do not remove OnboardingStatus we want to keep track of this even after logout.
+        // * PinProtectedUserKey
+        // * PinProtectedUserKeyEnvelope
+        // * DeviceKey
+        // * PendingAuthRequest
+        // * OnboardingStatus
     }
 
     override fun getAuthenticatorSyncUnlockKey(userId: String): String? =

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/manager/UserLogoutManagerImpl.kt

@@ -77,16 +77,10 @@ class UserLogoutManagerImpl(
         if (isExpired) {
             showToast(message = BitwardenString.login_expired)
         }
-        authDiskSource.storeAccountTokens(
-            userId = userId,
-            accountTokens = null,
-        )
 
         // Save any data that will still need to be retained after otherwise clearing all dat
         val vaultTimeoutInMinutes = settingsDiskSource.getVaultTimeoutInMinutes(userId = userId)
         val vaultTimeoutAction = settingsDiskSource.getVaultTimeoutAction(userId = userId)
-        val pinProtectedUserKeyEnvelope = authDiskSource
-            .getPinProtectedUserKeyEnvelope(userId = userId)
 
         switchUserIfAvailable(
             currentUserId = userId,
@@ -108,10 +102,6 @@ class UserLogoutManagerImpl(
                 vaultTimeoutAction = vaultTimeoutAction,
             )
         }
-        authDiskSource.storePinProtectedUserKeyEnvelope(
-            userId = userId,
-            pinProtectedUserKeyEnvelope = pinProtectedUserKeyEnvelope,
-        )
     }
 
     private fun clearData(userId: String) {

app/src/main/kotlin/com/x8bit/bitwarden/data/vault/manager/PinProtectedUserKeyManager.kt

@@ -0,0 +1,30 @@
+package com.x8bit.bitwarden.data.vault.manager
+
+/**
+ * Manager class to manage the pin-protected user key.
+ */
+interface PinProtectedUserKeyManager {
+    /**
+     * Checks if the given [userId] has an associated encrypted PIN key but not a pin-protected
+     * user key. This indicates a scenario in which a user has requested PIN unlocking but requires
+     * master-password unlocking on app restart. This function may then be called after such an
+     * unlock to derive a pin-protected user key and store it in memory for use for any subsequent
+     * unlocks during this current app session.
+     *
+     * If the user's vault has not yet been unlocked, this call will do nothing.
+     *
+     * @param userId The ID of the user to check.
+     */
+    suspend fun deriveTemporaryPinProtectedUserKeyIfNecessary(userId: String)
+
+    /**
+     * Migrates the PIN-protected user key for the given user if needed.
+     *
+     * If an encrypted PIN exists and no PIN-protected user key envelope is present, enrolls the
+     * PIN with the encrypted PIN and stores the resulting envelope.
+     * Optionally marks the envelope as in-memory only if the PIN-protected user key is not present.
+     *
+     * @param userId The ID of the user for whom to migrate the PIN-protected user key.
+     */
+    suspend fun migratePinProtectedUserKeyIfNeeded(userId: String)
+}

app/src/main/kotlin/com/x8bit/bitwarden/data/vault/manager/PinProtectedUserKeyManagerImpl.kt

@@ -0,0 +1,95 @@
+package com.x8bit.bitwarden.data.vault.manager
+
+import com.bitwarden.core.EnrollPinResponse
+import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
+import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
+import timber.log.Timber
+
+/**
+ * The default implementation of the [PinProtectedUserKeyManager].
+ */
+internal class PinProtectedUserKeyManagerImpl(
+    private val authDiskSource: AuthDiskSource,
+    private val vaultSdkSource: VaultSdkSource,
+) : PinProtectedUserKeyManager {
+    override suspend fun deriveTemporaryPinProtectedUserKeyIfNecessary(userId: String) {
+        val encryptedPin = authDiskSource.getEncryptedPin(userId = userId) ?: return
+        if (authDiskSource.getPinProtectedUserKeyEnvelope(userId = userId) != null) return
+
+        this
+            .enrollPinWithEncryptedPin(
+                userId = userId,
+                encryptedPin = encryptedPin,
+                inMemoryOnly = true,
+            )
+            .onSuccess {
+                Timber.d("[Auth] Set PIN-protected user key in memory")
+            }
+    }
+
+    override suspend fun migratePinProtectedUserKeyIfNeeded(userId: String) {
+        val encryptedPin = authDiskSource.getEncryptedPin(userId = userId) ?: return
+        if (authDiskSource.getPinProtectedUserKeyEnvelope(userId = userId) != null) return
+
+        val inMemoryOnly = authDiskSource.getPinProtectedUserKey(userId = userId) == null
+        this
+            .enrollPinWithEncryptedPin(
+                userId = userId,
+                encryptedPin = encryptedPin,
+                inMemoryOnly = inMemoryOnly,
+            )
+            .onSuccess {
+                if (inMemoryOnly) {
+                    Timber.d("[Auth] Set PIN-protected user key in memory")
+                } else {
+                    Timber.d("[Auth] Migrated from legacy PIN to PIN-protected user key envelope")
+                }
+            }
+    }
+
+    private suspend fun enrollPinWithEncryptedPin(
+        userId: String,
+        encryptedPin: String,
+        inMemoryOnly: Boolean,
+    ): Result<EnrollPinResponse> =
+        vaultSdkSource
+            .enrollPinWithEncryptedPin(userId = userId, encryptedPin = encryptedPin)
+            .onSuccess { enrollPinResponse ->
+                storePinData(
+                    userId = userId,
+                    encryptedPin = enrollPinResponse.userKeyEncryptedPin,
+                    pinProtectedUserKeyEnvelope = enrollPinResponse.pinProtectedUserKeyEnvelope,
+                    inMemoryOnly = inMemoryOnly,
+                )
+            }
+            .onFailure {
+                storePinData(
+                    userId = userId,
+                    encryptedPin = null,
+                    pinProtectedUserKeyEnvelope = null,
+                    inMemoryOnly = false,
+                )
+            }
+
+    private fun storePinData(
+        userId: String,
+        encryptedPin: String?,
+        pinProtectedUserKeyEnvelope: String?,
+        inMemoryOnly: Boolean,
+    ) {
+        authDiskSource.storeEncryptedPin(userId = userId, encryptedPin = encryptedPin)
+        authDiskSource.storePinProtectedUserKeyEnvelope(
+            userId = userId,
+            pinProtectedUserKeyEnvelope = pinProtectedUserKeyEnvelope,
+            inMemoryOnly = inMemoryOnly,
+        )
+        // This property is deprecated and we should be migrated to the PinProtectedUserKeyEnvelope.
+        // Because of this, we always clear this value and it should always be cleared at the disk
+        // level, not the in-memory level.
+        authDiskSource.storePinProtectedUserKey(
+            userId = userId,
+            pinProtectedUserKey = null,
+            inMemoryOnly = false,
+        )
+    }
+}

app/src/main/kotlin/com/x8bit/bitwarden/data/vault/manager/VaultLockManagerImpl.kt

@@ -91,6 +91,7 @@ class VaultLockManagerImpl(
     private val userLogoutManager: UserLogoutManager,
     private val trustedDeviceManager: TrustedDeviceManager,
     private val kdfManager: KdfManager,
+    private val pinProtectedUserKeyManager: PinProtectedUserKeyManager,
     dispatcherManager: DispatcherManager,
     context: Context,
 ) : VaultLockManager {
@@ -234,7 +235,8 @@ class VaultLockManagerImpl(
                                         trustedDeviceManager
                                             .trustThisDeviceIfNecessary(userId = userId)
                                         updateKdfIfNeeded(initUserCryptoMethod)
-                                        migratePinProtectedUserKeyIfNeeded(userId = userId)
+                                        pinProtectedUserKeyManager
+                                            .migratePinProtectedUserKeyIfNeeded(userId = userId)
                                         setVaultToUnlocked(userId = userId)
                                     } else {
                                         incrementInvalidUnlockCount(userId = userId)
@@ -308,47 +310,6 @@ class VaultLockManagerImpl(
             )
     }
 
-    /**
-     * Migrates the PIN-protected user key for the given user if needed.
-     *
-     * If an encrypted PIN exists and no PIN-protected user key envelope is present,
-     * enrolls the PIN with the encrypted PIN and stores the resulting envelope.
-     * Optionally marks the envelope as in-memory only if the PIN-protected user key is not present.
-     *
-     * @param userId The ID of the user for whom to migrate the PIN-protected user key.
-     */
-    private suspend fun migratePinProtectedUserKeyIfNeeded(
-        userId: String,
-    ) {
-        val encryptedPin = authDiskSource.getEncryptedPin(userId) ?: return
-        if (authDiskSource.getPinProtectedUserKeyEnvelope(userId) != null) return
-
-        val inMemoryOnly = authDiskSource.getPinProtectedUserKey(userId) == null
-
-        vaultSdkSource.enrollPinWithEncryptedPin(userId, encryptedPin)
-            .onSuccess { enrollPinResponse ->
-                authDiskSource.storeEncryptedPin(
-                    userId = userId,
-                    encryptedPin = enrollPinResponse.userKeyEncryptedPin,
-                )
-                authDiskSource.storePinProtectedUserKeyEnvelope(
-                    userId = userId,
-                    pinProtectedUserKeyEnvelope = enrollPinResponse.pinProtectedUserKeyEnvelope,
-                    inMemoryOnly = inMemoryOnly,
-                )
-                authDiskSource.storePinProtectedUserKey(
-                    userId = userId,
-                    pinProtectedUserKey = null,
-                    inMemoryOnly = inMemoryOnly,
-                )
-                if (inMemoryOnly) {
-                    Timber.d("[Auth] Set PIN-protected user key in memory")
-                } else {
-                    Timber.d("[Auth] Migrated from legacy PIN to PIN-protected user key envelope")
-                }
-            }
-    }
-
     /**
      * Increments the stored invalid unlock count for the given [userId] and automatically logs out
      * if this new value is greater than [MAXIMUM_INVALID_UNLOCK_ATTEMPTS].

app/src/main/kotlin/com/x8bit/bitwarden/data/vault/manager/di/VaultManagerModule.kt

@@ -31,6 +31,8 @@ import com.x8bit.bitwarden.data.vault.manager.FileManager
 import com.x8bit.bitwarden.data.vault.manager.FileManagerImpl
 import com.x8bit.bitwarden.data.vault.manager.FolderManager
 import com.x8bit.bitwarden.data.vault.manager.FolderManagerImpl
+import com.x8bit.bitwarden.data.vault.manager.PinProtectedUserKeyManager
+import com.x8bit.bitwarden.data.vault.manager.PinProtectedUserKeyManagerImpl
 import com.x8bit.bitwarden.data.vault.manager.SendManager
 import com.x8bit.bitwarden.data.vault.manager.SendManagerImpl
 import com.x8bit.bitwarden.data.vault.manager.TotpCodeManager
@@ -146,6 +148,7 @@ object VaultManagerModule {
         dispatcherManager: DispatcherManager,
         trustedDeviceManager: TrustedDeviceManager,
         kdfManager: KdfManager,
+        pinProtectedUserKeyManager: PinProtectedUserKeyManager,
     ): VaultLockManager =
         VaultLockManagerImpl(
             context = context,
@@ -160,6 +163,18 @@ object VaultManagerModule {
             dispatcherManager = dispatcherManager,
             trustedDeviceManager = trustedDeviceManager,
             kdfManager = kdfManager,
+            pinProtectedUserKeyManager = pinProtectedUserKeyManager,
+        )
+
+    @Provides
+    @Singleton
+    fun providePinProtectedUserKeyManager(
+        authDiskSource: AuthDiskSource,
+        vaultSdkSource: VaultSdkSource,
+    ): PinProtectedUserKeyManager =
+        PinProtectedUserKeyManagerImpl(
+            authDiskSource = authDiskSource,
+            vaultSdkSource = vaultSdkSource,
         )
 
     @Provides

app/src/main/kotlin/com/x8bit/bitwarden/data/vault/repository/VaultRepositoryImpl.kt

@@ -28,6 +28,7 @@ import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
 import com.x8bit.bitwarden.data.vault.manager.CipherManager
 import com.x8bit.bitwarden.data.vault.manager.CredentialExchangeImportManager
 import com.x8bit.bitwarden.data.vault.manager.FolderManager
+import com.x8bit.bitwarden.data.vault.manager.PinProtectedUserKeyManager
 import com.x8bit.bitwarden.data.vault.manager.SendManager
 import com.x8bit.bitwarden.data.vault.manager.TotpCodeManager
 import com.x8bit.bitwarden.data.vault.manager.VaultLockManager
@@ -40,7 +41,6 @@ import com.x8bit.bitwarden.data.vault.repository.model.GenerateTotpResult
 import com.x8bit.bitwarden.data.vault.repository.model.ImportCredentialsResult
 import com.x8bit.bitwarden.data.vault.repository.model.TotpCodeResult
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockResult
-import com.x8bit.bitwarden.data.vault.repository.util.logTag
 import com.x8bit.bitwarden.data.vault.repository.util.toEncryptedSdkCipher
 import com.x8bit.bitwarden.data.vault.repository.util.toEncryptedSdkFolder
 import com.x8bit.bitwarden.data.vault.repository.util.toSdkAccount
@@ -79,6 +79,7 @@ class VaultRepositoryImpl(
     private val totpCodeManager: TotpCodeManager,
     private val vaultSyncManager: VaultSyncManager,
     private val credentialExchangeImportManager: CredentialExchangeImportManager,
+    private val pinProtectedUserKeyManager: PinProtectedUserKeyManager,
     dispatcherManager: DispatcherManager,
 ) : VaultRepository,
     CipherManager by cipherManager,
@@ -328,11 +329,8 @@ class VaultRepositoryImpl(
                         )
                         authDiskSource.storeUserBiometricInitVector(userId = userId, iv = cipher.iv)
                     }
-                    deriveTemporaryPinProtectedUserKeyIfNecessary(
+                    pinProtectedUserKeyManager.deriveTemporaryPinProtectedUserKeyIfNecessary(
                         userId = userId,
-                        initUserCryptoMethod = InitUserCryptoMethod.DecryptedKey(
-                            decryptedUserKey = decryptedUserKey,
-                        ),
                     )
                 }
             }
@@ -369,9 +367,8 @@ class VaultRepositoryImpl(
             )
             .also {
                 if (it is VaultUnlockResult.Success) {
-                    deriveTemporaryPinProtectedUserKeyIfNecessary(
+                    pinProtectedUserKeyManager.deriveTemporaryPinProtectedUserKeyIfNecessary(
                         userId = userId,
-                        initUserCryptoMethod = initUserCryptoMethod,
                     )
                 }
             }
@@ -530,54 +527,6 @@ class VaultRepositoryImpl(
             )
     }
 
-    /**
-     * Checks if the given [userId] has an associated encrypted PIN key but not a pin-protected user
-     * key. This indicates a scenario in which a user has requested PIN unlocking but requires
-     * master-password unlocking on app restart. This function may then be called after such an
-     * unlock to derive a pin-protected user key and store it in memory for use for any subsequent
-     * unlocks during this current app session.
-     *
-     * If the user's vault has not yet been unlocked, this call will do nothing.
-     *
-     * @param userId The ID of the user to check.
-     * @param initUserCryptoMethod The method used to initialize the user's crypto.
-     */
-    private suspend fun deriveTemporaryPinProtectedUserKeyIfNecessary(
-        userId: String,
-        initUserCryptoMethod: InitUserCryptoMethod,
-    ) {
-        Timber.d("[Auth] Vault unlocked, method: ${initUserCryptoMethod.logTag}")
-        val encryptedPin = authDiskSource.getEncryptedPin(userId = userId) ?: return
-        val existingPinProtectedUserKeyEnvelope = authDiskSource
-            .getPinProtectedUserKeyEnvelope(
-                userId = userId,
-            )
-        if (existingPinProtectedUserKeyEnvelope != null) return
-
-        vaultSdkSource
-            .enrollPinWithEncryptedPin(
-                userId = userId,
-                encryptedPin = encryptedPin,
-            )
-            .onSuccess { enrollPinResponse ->
-                authDiskSource.storeEncryptedPin(
-                    userId = userId,
-                    encryptedPin = enrollPinResponse.userKeyEncryptedPin,
-                )
-                authDiskSource.storePinProtectedUserKeyEnvelope(
-                    userId = userId,
-                    pinProtectedUserKeyEnvelope = enrollPinResponse.pinProtectedUserKeyEnvelope,
-                    inMemoryOnly = true,
-                )
-                authDiskSource.storePinProtectedUserKey(
-                    userId = userId,
-                    pinProtectedUserKey = null,
-                    inMemoryOnly = true,
-                )
-                Timber.d("[Auth] Set PIN-protected user key in memory")
-            }
-    }
-
     private suspend fun unlockVaultForUser(
         userId: String,
         initUserCryptoMethod: InitUserCryptoMethod,

app/src/main/kotlin/com/x8bit/bitwarden/data/vault/repository/di/VaultRepositoryModule.kt

@@ -7,6 +7,7 @@ import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
 import com.x8bit.bitwarden.data.vault.manager.CipherManager
 import com.x8bit.bitwarden.data.vault.manager.CredentialExchangeImportManager
 import com.x8bit.bitwarden.data.vault.manager.FolderManager
+import com.x8bit.bitwarden.data.vault.manager.PinProtectedUserKeyManager
 import com.x8bit.bitwarden.data.vault.manager.SendManager
 import com.x8bit.bitwarden.data.vault.manager.TotpCodeManager
 import com.x8bit.bitwarden.data.vault.manager.VaultLockManager
@@ -40,6 +41,7 @@ object VaultRepositoryModule {
         totpCodeManager: TotpCodeManager,
         vaultSyncManager: VaultSyncManager,
         credentialExchangeImportManager: CredentialExchangeImportManager,
+        pinProtectedUserKeyManager: PinProtectedUserKeyManager,
     ): VaultRepository = VaultRepositoryImpl(
         vaultDiskSource = vaultDiskSource,
         vaultSdkSource = vaultSdkSource,
@@ -52,5 +54,6 @@ object VaultRepositoryModule {
         totpCodeManager = totpCodeManager,
         vaultSyncManager = vaultSyncManager,
         credentialExchangeImportManager = credentialExchangeImportManager,
+        pinProtectedUserKeyManager = pinProtectedUserKeyManager,
     )
 }

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/datasource/disk/AuthDiskSourceTest.kt

@@ -267,9 +267,15 @@ class AuthDiskSourceTest {
             userId = userId,
             biometricsKey = "1234-9876-0192",
         )
+        val pinProtectedUserKey = "pinProtectedUserKey"
         authDiskSource.storePinProtectedUserKey(
             userId = userId,
-            pinProtectedUserKey = "pinProtectedUserKey",
+            pinProtectedUserKey = pinProtectedUserKey,
+        )
+        val pinProtectedUserKeyEnvelope = "pinProtectedUserKeyEnvelope"
+        authDiskSource.storePinProtectedUserKeyEnvelope(
+            userId = userId,
+            pinProtectedUserKeyEnvelope = pinProtectedUserKeyEnvelope,
         )
         authDiskSource.storeInvalidUnlockAttempts(
             userId = userId,
@@ -304,7 +310,8 @@ class AuthDiskSourceTest {
                 refreshToken = "refreshToken",
             ),
         )
-        authDiskSource.storeEncryptedPin(userId = userId, encryptedPin = "encryptedPin")
+        val encryptedPin = "encryptedPin"
+        authDiskSource.storeEncryptedPin(userId = userId, encryptedPin = encryptedPin)
         authDiskSource.storeMasterPasswordHash(userId = userId, passwordHash = "passwordHash")
         authDiskSource.storeAuthenticatorSyncUnlockKey(
             userId = userId,
@@ -326,12 +333,16 @@ class AuthDiskSourceTest {
             OnboardingStatus.AUTOFILL_SETUP,
             authDiskSource.getOnboardingStatus(userId = userId),
         )
+        assertEquals(encryptedPin, authDiskSource.getEncryptedPin(userId = userId))
+        assertEquals(pinProtectedUserKey, authDiskSource.getPinProtectedUserKey(userId = userId))
+        assertEquals(
+            pinProtectedUserKeyEnvelope,
+            authDiskSource.getPinProtectedUserKeyEnvelope(userId = userId),
+        )
 
         // These should be cleared
         assertNull(authDiskSource.getUserBiometricInitVector(userId = userId))
         assertNull(authDiskSource.getUserBiometricUnlockKey(userId = userId))
-        assertNull(authDiskSource.getPinProtectedUserKey(userId = userId))
-        assertNull(authDiskSource.getPinProtectedUserKeyEnvelope(userId = userId))
         assertNull(authDiskSource.getInvalidUnlockAttempts(userId = userId))
         assertNull(authDiskSource.getUserKey(userId = userId))
         assertNull(authDiskSource.getUserAutoUnlockKey(userId = userId))
@@ -341,7 +352,6 @@ class AuthDiskSourceTest {
         assertNull(authDiskSource.getOrganizations(userId = userId))
         assertNull(authDiskSource.getPolicies(userId = userId))
         assertNull(authDiskSource.getAccountTokens(userId = userId))
-        assertNull(authDiskSource.getEncryptedPin(userId = userId))
         assertNull(authDiskSource.getMasterPasswordHash(userId = userId))
         assertNull(authDiskSource.getShouldUseKeyConnector(userId = userId))
         assertNull(authDiskSource.getIsTdeLoginComplete(userId = userId))

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/datasource/disk/util/FakeAuthDiskSource.kt

@@ -84,8 +84,6 @@ class FakeAuthDiskSource : AuthDiskSource {
         storedPrivateKeys.remove(userId)
         storedTwoFactorTokens.clear()
         storedUserAutoUnlockKeys.remove(userId)
-        storedPinProtectedUserKeys.remove(userId)
-        storedEncryptedPins.remove(userId)
         storedOrganizations.remove(userId)
         storedPolicies.remove(userId)
         storedAccountTokens.remove(userId)
@@ -98,7 +96,6 @@ class FakeAuthDiskSource : AuthDiskSource {
         mutableOrganizationsFlowMap.remove(userId)
         mutablePoliciesFlowMap.remove(userId)
         mutableAccountTokensFlowMap.remove(userId)
-        mutablePinProtectedUserKeyEnvelopesFlowMap.remove(userId)
     }
 
     private fun getMutablePinProtectedUserKeyEnvelopeFlow(

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/manager/UserLogoutManagerTest.kt

@@ -33,13 +33,6 @@ import java.time.ZonedDateTime
 @ExtendWith(MainDispatcherExtension::class)
 class UserLogoutManagerTest {
     private val authDiskSource: AuthDiskSource = mockk {
-        every { storeAccountTokens(userId = any(), accountTokens = null) } just runs
-        every {
-            storePinProtectedUserKeyEnvelope(
-                userId = any(),
-                pinProtectedUserKeyEnvelope = any(),
-            )
-        } just runs
         every { userState = any() } just runs
         every { clearData(any()) } just runs
     }
@@ -149,7 +142,6 @@ class UserLogoutManagerTest {
 
         userLogoutManager.softLogout(userId = userId, reason = LogoutReason.Timeout)
 
-        verify { authDiskSource.storeAccountTokens(userId = userId, accountTokens = null) }
         assertDataCleared(userId = userId)
 
         verify(exactly = 1) {
@@ -170,10 +162,6 @@ class UserLogoutManagerTest {
                 userId = userId,
                 vaultTimeoutAction = vaultTimeoutAction,
             )
-            authDiskSource.storePinProtectedUserKeyEnvelope(
-                userId = userId,
-                pinProtectedUserKeyEnvelope = pinProtectedUserKey,
-            )
         }
     }
 
@@ -198,7 +186,6 @@ class UserLogoutManagerTest {
         userLogoutManager.softLogout(userId = userId, reason = LogoutReason.Timeout)
 
         verify(exactly = 1) {
-            authDiskSource.storeAccountTokens(userId = userId, accountTokens = null)
             authDiskSource.userState = UserStateJson(
                 activeUserId = USER_ID_2,
                 accounts = MULTI_USER_STATE.accounts,
@@ -212,10 +199,6 @@ class UserLogoutManagerTest {
                 userId = userId,
                 vaultTimeoutAction = vaultTimeoutAction,
             )
-            authDiskSource.storePinProtectedUserKeyEnvelope(
-                userId = userId,
-                pinProtectedUserKeyEnvelope = pinProtectedUserKey,
-            )
         }
     }
 
@@ -241,7 +224,6 @@ class UserLogoutManagerTest {
         userLogoutManager.softLogout(userId = userId, reason = LogoutReason.SecurityStamp)
 
         verify(exactly = 1) {
-            authDiskSource.storeAccountTokens(userId = userId, accountTokens = null)
             authDiskSource.userState = UserStateJson(
                 activeUserId = USER_ID_2,
                 accounts = MULTI_USER_STATE.accounts,
@@ -255,10 +237,6 @@ class UserLogoutManagerTest {
                 userId = userId,
                 vaultTimeoutAction = vaultTimeoutAction,
             )
-            authDiskSource.storePinProtectedUserKeyEnvelope(
-                userId = userId,
-                pinProtectedUserKeyEnvelope = pinProtectedUserKey,
-            )
         }
     }
 

app/src/test/kotlin/com/x8bit/bitwarden/data/vault/manager/PinProtectedUserKeyManagerTest.kt

@@ -0,0 +1,279 @@
+package com.x8bit.bitwarden.data.vault.manager
+
+import com.bitwarden.core.EnrollPinResponse
+import com.bitwarden.core.data.util.asFailure
+import com.bitwarden.core.data.util.asSuccess
+import com.x8bit.bitwarden.data.auth.datasource.disk.util.FakeAuthDiskSource
+import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
+import io.mockk.coEvery
+import io.mockk.coVerify
+import io.mockk.mockk
+import kotlinx.coroutines.test.runTest
+import org.junit.jupiter.api.Test
+
+class PinProtectedUserKeyManagerTest {
+
+    private val fakeAuthDiskSource: FakeAuthDiskSource = FakeAuthDiskSource()
+    private val vaultSdkSource: VaultSdkSource = mockk()
+
+    private val pinProtectedUserKeyManager: PinProtectedUserKeyManager =
+        PinProtectedUserKeyManagerImpl(
+            authDiskSource = fakeAuthDiskSource,
+            vaultSdkSource = vaultSdkSource,
+        )
+
+    @Test
+    fun `deriveTemporaryPinProtectedUserKeyIfNecessary without encryptedKey does nothing`() =
+        runTest {
+            fakeAuthDiskSource.storeEncryptedPin(userId = USER_ID, encryptedPin = null)
+
+            pinProtectedUserKeyManager.deriveTemporaryPinProtectedUserKeyIfNecessary(
+                userId = USER_ID,
+            )
+
+            coVerify(exactly = 0) {
+                vaultSdkSource.enrollPinWithEncryptedPin(userId = any(), encryptedPin = any())
+            }
+        }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `deriveTemporaryPinProtectedUserKeyIfNecessary with encrypted key and existing pinProtectedUserKeyEnvelope does nothing`() =
+        runTest {
+            fakeAuthDiskSource.storeEncryptedPin(userId = USER_ID, encryptedPin = "encryptedPin")
+            fakeAuthDiskSource.storePinProtectedUserKeyEnvelope(
+                userId = USER_ID,
+                pinProtectedUserKeyEnvelope = "pinProtectedUserKeyEnvelope",
+            )
+
+            pinProtectedUserKeyManager.deriveTemporaryPinProtectedUserKeyIfNecessary(
+                userId = USER_ID,
+            )
+
+            coVerify(exactly = 0) {
+                vaultSdkSource.enrollPinWithEncryptedPin(userId = any(), encryptedPin = any())
+            }
+        }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `deriveTemporaryPinProtectedUserKeyIfNecessary with enrollment success should store new pin data`() =
+        runTest {
+            val encryptedPin = "encryptedPin"
+            val pinProtectedUserKeyEnvelope = "pinProtectedUserKeyEnvelope"
+            val userKeyEncryptedPin = "userKeyEncryptedPin"
+            val enrollPinResponse = EnrollPinResponse(
+                pinProtectedUserKeyEnvelope = pinProtectedUserKeyEnvelope,
+                userKeyEncryptedPin = userKeyEncryptedPin,
+            )
+            fakeAuthDiskSource.storeEncryptedPin(userId = USER_ID, encryptedPin = encryptedPin)
+            fakeAuthDiskSource.storePinProtectedUserKeyEnvelope(
+                userId = USER_ID,
+                pinProtectedUserKeyEnvelope = null,
+            )
+            coEvery {
+                vaultSdkSource.enrollPinWithEncryptedPin(
+                    userId = USER_ID,
+                    encryptedPin = encryptedPin,
+                )
+            } returns enrollPinResponse.asSuccess()
+
+            pinProtectedUserKeyManager.deriveTemporaryPinProtectedUserKeyIfNecessary(
+                userId = USER_ID,
+            )
+
+            coVerify(exactly = 1) {
+                vaultSdkSource.enrollPinWithEncryptedPin(
+                    userId = USER_ID,
+                    encryptedPin = encryptedPin,
+                )
+            }
+            fakeAuthDiskSource.assertEncryptedPin(
+                userId = USER_ID,
+                encryptedPin = userKeyEncryptedPin,
+            )
+            fakeAuthDiskSource.assertPinProtectedUserKeyEnvelope(
+                userId = USER_ID,
+                pinProtectedUserKeyEnvelope = pinProtectedUserKeyEnvelope,
+                inMemoryOnly = true,
+            )
+            fakeAuthDiskSource.assertPinProtectedUserKey(
+                userId = USER_ID,
+                pinProtectedUserKey = null,
+                inMemoryOnly = false,
+            )
+        }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `deriveTemporaryPinProtectedUserKeyIfNecessary with enrollment failure should clear all pin data`() =
+        runTest {
+            val encryptedPin = "encryptedPin"
+            val error = Throwable("Fail!")
+            fakeAuthDiskSource.storeEncryptedPin(userId = USER_ID, encryptedPin = encryptedPin)
+            fakeAuthDiskSource.storePinProtectedUserKeyEnvelope(
+                userId = USER_ID,
+                pinProtectedUserKeyEnvelope = null,
+            )
+            coEvery {
+                vaultSdkSource.enrollPinWithEncryptedPin(
+                    userId = USER_ID,
+                    encryptedPin = encryptedPin,
+                )
+            } returns error.asFailure()
+
+            pinProtectedUserKeyManager.deriveTemporaryPinProtectedUserKeyIfNecessary(
+                userId = USER_ID,
+            )
+
+            coVerify(exactly = 1) {
+                vaultSdkSource.enrollPinWithEncryptedPin(
+                    userId = USER_ID,
+                    encryptedPin = encryptedPin,
+                )
+            }
+            fakeAuthDiskSource.assertEncryptedPin(
+                userId = USER_ID,
+                encryptedPin = null,
+            )
+            fakeAuthDiskSource.assertPinProtectedUserKeyEnvelope(
+                userId = USER_ID,
+                pinProtectedUserKeyEnvelope = null,
+                inMemoryOnly = false,
+            )
+            fakeAuthDiskSource.assertPinProtectedUserKey(
+                userId = USER_ID,
+                pinProtectedUserKey = null,
+                inMemoryOnly = false,
+            )
+        }
+
+    @Test
+    fun `migratePinProtectedUserKeyIfNeeded without encryptedKey does nothing`() =
+        runTest {
+            fakeAuthDiskSource.storeEncryptedPin(userId = USER_ID, encryptedPin = null)
+
+            pinProtectedUserKeyManager.migratePinProtectedUserKeyIfNeeded(userId = USER_ID)
+
+            coVerify(exactly = 0) {
+                vaultSdkSource.enrollPinWithEncryptedPin(userId = any(), encryptedPin = any())
+            }
+        }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `migratePinProtectedUserKeyIfNeeded with encrypted key and existing pinProtectedUserKeyEnvelope does nothing`() =
+        runTest {
+            fakeAuthDiskSource.storeEncryptedPin(userId = USER_ID, encryptedPin = "encryptedPin")
+            fakeAuthDiskSource.storePinProtectedUserKeyEnvelope(
+                userId = USER_ID,
+                pinProtectedUserKeyEnvelope = "pinProtectedUserKeyEnvelope",
+            )
+
+            pinProtectedUserKeyManager.migratePinProtectedUserKeyIfNeeded(userId = USER_ID)
+
+            coVerify(exactly = 0) {
+                vaultSdkSource.enrollPinWithEncryptedPin(userId = any(), encryptedPin = any())
+            }
+        }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `migratePinProtectedUserKeyIfNeeded with enrollment success should store new pin data in memory`() =
+        runTest {
+            val encryptedPin = "encryptedPin"
+            val pinProtectedUserKeyEnvelope = "pinProtectedUserKeyEnvelope"
+            val userKeyEncryptedPin = "userKeyEncryptedPin"
+            val enrollPinResponse = EnrollPinResponse(
+                pinProtectedUserKeyEnvelope = pinProtectedUserKeyEnvelope,
+                userKeyEncryptedPin = userKeyEncryptedPin,
+            )
+            fakeAuthDiskSource.storeEncryptedPin(userId = USER_ID, encryptedPin = encryptedPin)
+            fakeAuthDiskSource.storePinProtectedUserKeyEnvelope(
+                userId = USER_ID,
+                pinProtectedUserKeyEnvelope = null,
+            )
+            fakeAuthDiskSource.storePinProtectedUserKey(
+                userId = USER_ID,
+                pinProtectedUserKey = null,
+            )
+            coEvery {
+                vaultSdkSource.enrollPinWithEncryptedPin(
+                    userId = USER_ID,
+                    encryptedPin = encryptedPin,
+                )
+            } returns enrollPinResponse.asSuccess()
+
+            pinProtectedUserKeyManager.migratePinProtectedUserKeyIfNeeded(userId = USER_ID)
+
+            coVerify(exactly = 1) {
+                vaultSdkSource.enrollPinWithEncryptedPin(
+                    userId = USER_ID,
+                    encryptedPin = encryptedPin,
+                )
+            }
+            fakeAuthDiskSource.assertEncryptedPin(
+                userId = USER_ID,
+                encryptedPin = userKeyEncryptedPin,
+            )
+            fakeAuthDiskSource.assertPinProtectedUserKeyEnvelope(
+                userId = USER_ID,
+                pinProtectedUserKeyEnvelope = pinProtectedUserKeyEnvelope,
+                inMemoryOnly = true,
+            )
+            fakeAuthDiskSource.assertPinProtectedUserKey(
+                userId = USER_ID,
+                pinProtectedUserKey = null,
+                inMemoryOnly = false,
+            )
+        }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `migratePinProtectedUserKeyIfNeeded with enrollment failure should clear all pin data at disk level`() =
+        runTest {
+            val encryptedPin = "encryptedPin"
+            val pinProtectedUserKey = "pinProtectedUserKey"
+            val error = Throwable("Fail!")
+            fakeAuthDiskSource.storeEncryptedPin(userId = USER_ID, encryptedPin = encryptedPin)
+            fakeAuthDiskSource.storePinProtectedUserKeyEnvelope(
+                userId = USER_ID,
+                pinProtectedUserKeyEnvelope = null,
+            )
+            fakeAuthDiskSource.storePinProtectedUserKey(
+                userId = USER_ID,
+                pinProtectedUserKey = pinProtectedUserKey,
+            )
+            coEvery {
+                vaultSdkSource.enrollPinWithEncryptedPin(
+                    userId = USER_ID,
+                    encryptedPin = encryptedPin,
+                )
+            } returns error.asFailure()
+
+            pinProtectedUserKeyManager.migratePinProtectedUserKeyIfNeeded(userId = USER_ID)
+
+            coVerify(exactly = 1) {
+                vaultSdkSource.enrollPinWithEncryptedPin(
+                    userId = USER_ID,
+                    encryptedPin = encryptedPin,
+                )
+            }
+            fakeAuthDiskSource.assertEncryptedPin(
+                userId = USER_ID,
+                encryptedPin = null,
+            )
+            fakeAuthDiskSource.assertPinProtectedUserKeyEnvelope(
+                userId = USER_ID,
+                pinProtectedUserKeyEnvelope = null,
+                inMemoryOnly = false,
+            )
+            fakeAuthDiskSource.assertPinProtectedUserKey(
+                userId = USER_ID,
+                pinProtectedUserKey = null,
+                inMemoryOnly = false,
+            )
+        }
+}
+
+private const val USER_ID: String = "user_id"

app/src/test/kotlin/com/x8bit/bitwarden/data/vault/manager/VaultLockManagerTest.kt

@@ -115,6 +115,9 @@ class VaultLockManagerTest {
             updateKdfToMinimumsIfNeeded(password = any())
         } returns UpdateKdfMinimumsResult.Success
     }
+    private val pinProtectedUserKeyManager: PinProtectedUserKeyManager = mockk {
+        coEvery { migratePinProtectedUserKeyIfNeeded(userId = any()) } just runs
+    }
 
     private val vaultLockManager: VaultLockManager = VaultLockManagerImpl(
         context = context,
@@ -129,6 +132,7 @@ class VaultLockManagerTest {
         trustedDeviceManager = trustedDeviceManager,
         dispatcherManager = fakeDispatcherManager,
         kdfManager = kdfManager,
+        pinProtectedUserKeyManager = pinProtectedUserKeyManager,
     )
 
     @Test
@@ -1678,12 +1682,6 @@ class VaultLockManagerTest {
                     request = InitOrgCryptoRequest(organizationKeys = organizationKeys),
                 )
             } returns InitializeCryptoResult.Success.asSuccess()
-            coEvery {
-                vaultSdkSource.enrollPinWithEncryptedPin(
-                    userId = USER_ID,
-                    encryptedPin = userKeyEncryptedPin,
-                )
-            } returns enrollResponse.asSuccess()
             coEvery {
                 trustedDeviceManager.trustThisDeviceIfNecessary(userId = USER_ID)
             } returns false.asSuccess()
@@ -1692,18 +1690,6 @@ class VaultLockManagerTest {
                 vaultLockManager.vaultUnlockDataStateFlow.value,
             )
             mutableVaultTimeoutStateFlow.value = VaultTimeout.ThirtyMinutes
-            fakeAuthDiskSource.storeUserAutoUnlockKey(
-                userId = USER_ID,
-                userAutoUnlockKey = null,
-            )
-            fakeAuthDiskSource.storeEncryptedPin(
-                userId = USER_ID,
-                encryptedPin = userKeyEncryptedPin,
-            )
-            fakeAuthDiskSource.storePinProtectedUserKey(
-                userId = USER_ID,
-                pinProtectedUserKey = userKeyEncryptedPin,
-            )
 
             val result = vaultLockManager.unlockVault(
                 userId = USER_ID,
@@ -1730,19 +1716,6 @@ class VaultLockManagerTest {
                 vaultLockManager.vaultUnlockDataStateFlow.value,
             )
 
-            fakeAuthDiskSource.assertUserAutoUnlockKey(
-                userId = USER_ID,
-                userAutoUnlockKey = null,
-            )
-            fakeAuthDiskSource.assertMasterPasswordHash(
-                userId = USER_ID,
-                passwordHash = "hashedPassword",
-            )
-            fakeAuthDiskSource.assertPinProtectedUserKeyEnvelope(
-                userId = USER_ID,
-                pinProtectedUserKeyEnvelope = pinProtectedUserKeyEnvelope,
-                inMemoryOnly = false,
-            )
             coVerify(exactly = 1) {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
@@ -1764,6 +1737,7 @@ class VaultLockManagerTest {
                     request = InitOrgCryptoRequest(organizationKeys = organizationKeys),
                 )
                 trustedDeviceManager.trustThisDeviceIfNecessary(userId = USER_ID)
+                pinProtectedUserKeyManager.migratePinProtectedUserKeyIfNeeded(userId = USER_ID)
             }
         }
 

app/src/test/kotlin/com/x8bit/bitwarden/data/vault/repository/VaultRepositoryTest.kt

@@ -3,7 +3,6 @@ package com.x8bit.bitwarden.data.vault.repository
 import app.cash.turbine.test
 import com.bitwarden.collections.CollectionView
 import com.bitwarden.core.DateTime
-import com.bitwarden.core.EnrollPinResponse
 import com.bitwarden.core.InitUserCryptoMethod
 import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
 import com.bitwarden.core.data.manager.dispatcher.FakeDispatcherManager
@@ -42,6 +41,7 @@ import com.x8bit.bitwarden.data.vault.datasource.sdk.model.createMockSdkFolder
 import com.x8bit.bitwarden.data.vault.datasource.sdk.model.createMockSdkSend
 import com.x8bit.bitwarden.data.vault.datasource.sdk.model.createMockSendView
 import com.x8bit.bitwarden.data.vault.manager.CredentialExchangeImportManager
+import com.x8bit.bitwarden.data.vault.manager.PinProtectedUserKeyManager
 import com.x8bit.bitwarden.data.vault.manager.TotpCodeManager
 import com.x8bit.bitwarden.data.vault.manager.VaultLockManager
 import com.x8bit.bitwarden.data.vault.manager.VaultSyncManager
@@ -111,6 +111,9 @@ class VaultRepositoryTest {
         every { sendDataStateFlow } returns mutableSendDataStateFlow
     }
     private val credentialExchangeImportManager: CredentialExchangeImportManager = mockk()
+    private val pinProtectedUserKeyManager: PinProtectedUserKeyManager = mockk {
+        coEvery { deriveTemporaryPinProtectedUserKeyIfNecessary(userId = any()) } just runs
+    }
 
     private val vaultRepository = VaultRepositoryImpl(
         vaultDiskSource = vaultDiskSource,
@@ -124,6 +127,7 @@ class VaultRepositoryTest {
         sendManager = mockk(),
         vaultSyncManager = vaultSyncManager,
         credentialExchangeImportManager = credentialExchangeImportManager,
+        pinProtectedUserKeyManager = pinProtectedUserKeyManager,
     )
 
     @BeforeEach
@@ -186,69 +190,6 @@ class VaultRepositoryTest {
             )
         }
 
-    @Suppress("MaxLineLength")
-    @Test
-    fun `unlockVaultWithBiometrics with VaultLockManager Success and no encrypted PIN should unlock for the current user and return Success`() =
-        runTest {
-            val userId = MOCK_USER_STATE.activeUserId
-            val privateKey = "mockPrivateKey-1"
-            val biometricsKey = "asdf1234"
-            fakeAuthDiskSource.userState = MOCK_USER_STATE
-            val encryptedBytes = byteArrayOf(1, 1)
-            val initVector = byteArrayOf(2, 2)
-            val cipher = mockk<Cipher> {
-                every { doFinal(any()) } returns encryptedBytes
-                every { iv } returns initVector
-            }
-            coEvery {
-                vaultLockManager.unlockVault(
-                    userId = userId,
-                    email = "email",
-                    kdf = MOCK_PROFILE.toSdkParams(),
-                    privateKey = privateKey,
-                    signingKey = null,
-                    securityState = null,
-                    initUserCryptoMethod = InitUserCryptoMethod.DecryptedKey(
-                        decryptedUserKey = biometricsKey,
-                    ),
-                    organizationKeys = null,
-                )
-            } returns VaultUnlockResult.Success
-            fakeAuthDiskSource.apply {
-                storeUserBiometricInitVector(userId = userId, iv = null)
-                storeUserBiometricUnlockKey(userId = userId, biometricsKey = biometricsKey)
-                storePrivateKey(userId = userId, privateKey = privateKey)
-            }
-
-            val result = vaultRepository.unlockVaultWithBiometrics(cipher = cipher)
-
-            assertEquals(VaultUnlockResult.Success, result)
-            coVerify {
-                vaultLockManager.unlockVault(
-                    userId = userId,
-                    email = "email",
-                    kdf = MOCK_PROFILE.toSdkParams(),
-                    privateKey = privateKey,
-                    signingKey = null,
-                    securityState = null,
-                    initUserCryptoMethod = InitUserCryptoMethod.DecryptedKey(
-                        decryptedUserKey = biometricsKey,
-                    ),
-                    organizationKeys = null,
-                )
-            }
-            coVerify(exactly = 0) {
-                vaultSdkSource.enrollPinWithEncryptedPin(any(), any())
-            }
-            fakeAuthDiskSource.apply {
-                assertBiometricsKey(
-                    userId = userId,
-                    biometricsKey = encryptedBytes.toString(Charsets.ISO_8859_1),
-                )
-                assertBiometricInitVector(userId = userId, iv = initVector)
-            }
-        }
-
     @Suppress("MaxLineLength")
     @Test
     fun `unlockVaultWithBiometrics with failure to decode biometrics key should return BiometricDecodingError`() =
@@ -275,9 +216,6 @@ class VaultRepositoryTest {
                 ),
                 result,
             )
-            coVerify(exactly = 0) {
-                vaultSdkSource.enrollPinWithEncryptedPin(any(), any())
-            }
         }
 
     @Suppress("MaxLineLength")
@@ -302,9 +240,6 @@ class VaultRepositoryTest {
                 VaultUnlockResult.BiometricDecodingError(error = error),
                 result,
             )
-            coVerify(exactly = 0) {
-                vaultSdkSource.enrollPinWithEncryptedPin(any(), any())
-            }
         }
 
     @Suppress("MaxLineLength")
@@ -357,9 +292,6 @@ class VaultRepositoryTest {
                     organizationKeys = null,
                 )
             }
-            coVerify(exactly = 0) {
-                vaultSdkSource.enrollPinWithEncryptedPin(any(), any())
-            }
         }
 
     @Suppress("MaxLineLength")
@@ -367,13 +299,7 @@ class VaultRepositoryTest {
     fun `unlockVaultWithBiometrics with VaultLockManager Success and a stored encrypted pin should unlock for the current user, derive a new pin-protected key, and return Success`() =
         runTest {
             val userId = MOCK_USER_STATE.activeUserId
-            val encryptedPin = "encryptedPin"
             val privateKey = "mockPrivateKey-1"
-            val pinProtectedUserKeyEnvelope = "pinProtectedUserKeyEnvelope"
-            val enrollResponse = EnrollPinResponse(
-                pinProtectedUserKeyEnvelope = pinProtectedUserKeyEnvelope,
-                userKeyEncryptedPin = encryptedPin,
-            )
             val biometricsKey = "asdf1234"
             fakeAuthDiskSource.userState = MOCK_USER_STATE
             val encryptedBytes = byteArrayOf(1, 1)
@@ -382,12 +308,6 @@ class VaultRepositoryTest {
                 every { doFinal(any()) } returns encryptedBytes
                 every { iv } returns initVector
             }
-            coEvery {
-                vaultSdkSource.enrollPinWithEncryptedPin(
-                    userId = userId,
-                    encryptedPin = encryptedPin,
-                )
-            } returns enrollResponse.asSuccess()
             coEvery {
                 vaultLockManager.unlockVault(
                     userId = userId,
@@ -406,32 +326,11 @@ class VaultRepositoryTest {
                 storeUserBiometricInitVector(userId = userId, iv = null)
                 storeUserBiometricUnlockKey(userId = userId, biometricsKey = biometricsKey)
                 storePrivateKey(userId = userId, privateKey = privateKey)
-                storeEncryptedPin(userId = userId, encryptedPin = encryptedPin)
-                storePinProtectedUserKey(
-                    userId = userId,
-                    pinProtectedUserKey = null,
-                    inMemoryOnly = true,
-                )
-                storePinProtectedUserKeyEnvelope(
-                    userId = userId,
-                    pinProtectedUserKeyEnvelope = null,
-                    inMemoryOnly = true,
-                )
             }
 
             val result = vaultRepository.unlockVaultWithBiometrics(cipher = cipher)
 
             assertEquals(VaultUnlockResult.Success, result)
-            fakeAuthDiskSource.assertPinProtectedUserKey(
-                userId = userId,
-                pinProtectedUserKey = null,
-                inMemoryOnly = true,
-            )
-            fakeAuthDiskSource.assertPinProtectedUserKeyEnvelope(
-                userId = userId,
-                pinProtectedUserKeyEnvelope = pinProtectedUserKeyEnvelope,
-                inMemoryOnly = true,
-            )
             coVerify {
                 vaultLockManager.unlockVault(
                     userId = userId,
@@ -445,11 +344,8 @@ class VaultRepositoryTest {
                     ),
                     organizationKeys = null,
                 )
-            }
+                pinProtectedUserKeyManager.deriveTemporaryPinProtectedUserKeyIfNecessary(
-            coEvery {
-                vaultSdkSource.enrollPinWithEncryptedPin(
                     userId = userId,
-                    encryptedPin = encryptedPin,
                 )
             }
             fakeAuthDiskSource.apply {
@@ -616,100 +512,13 @@ class VaultRepositoryTest {
             )
         }
 
-    @Suppress("MaxLineLength")
-    @Test
-    fun `unlockVaultWithMasterPassword with VaultLockManager Success and no encrypted PIN should unlock for the current user and return Success`() =
-        runTest {
-            val userId = "mockId-1"
-            val mockVaultUnlockResult = VaultUnlockResult.Success
-            val userKeyEncryptedPin = "encryptedPin"
-            val pinProtectedUserKeyEnvelope = "pinProtectedUserKeyEnvelope"
-            val enrollResponse = EnrollPinResponse(
-                pinProtectedUserKeyEnvelope = pinProtectedUserKeyEnvelope,
-                userKeyEncryptedPin = userKeyEncryptedPin,
-            )
-            coEvery {
-                vaultSdkSource.enrollPinWithEncryptedPin(any(), any())
-            } returns enrollResponse.asSuccess()
-            prepareStateForUnlocking(unlockResult = mockVaultUnlockResult)
-            fakeAuthDiskSource.apply {
-                storeEncryptedPin(
-                    userId = userId,
-                    encryptedPin = null,
-                )
-                storePinProtectedUserKey(
-                    userId = userId,
-                    pinProtectedUserKey = null,
-                    inMemoryOnly = true,
-                )
-                storePinProtectedUserKeyEnvelope(
-                    userId = userId,
-                    pinProtectedUserKeyEnvelope = null,
-                    inMemoryOnly = true,
-                )
-            }
-
-            val result = vaultRepository.unlockVaultWithMasterPassword(
-                masterPassword = "mockPassword-1",
-            )
-
-            assertEquals(
-                mockVaultUnlockResult,
-                result,
-            )
-            coVerify {
-                vaultLockManager.unlockVault(
-                    userId = userId,
-                    email = "email",
-                    kdf = MOCK_PROFILE.toSdkParams(),
-                    privateKey = "mockPrivateKey-1",
-                    signingKey = null,
-                    securityState = null,
-                    initUserCryptoMethod = InitUserCryptoMethod.Password(
-                        password = "mockPassword-1",
-                        userKey = "mockKey-1",
-                    ),
-                    organizationKeys = createMockOrganizationKeys(number = 1),
-                )
-            }
-            coVerify(exactly = 0) { vaultSdkSource.enrollPinWithEncryptedPin(any(), any()) }
-        }
-
     @Suppress("MaxLineLength")
     @Test
     fun `unlockVaultWithMasterPassword with VaultLockManager Success and a stored encrypted pin should unlock for the current user, derive a new pin-protected key, and return Success`() =
         runTest {
             val userId = "mockId-1"
-            val pinProtectedUserKey = "pinProtectedUserkeyEnvelope"
-            val encryptedPin = "userKeyEncryptedPin"
-            val enrollResponse = EnrollPinResponse(
-                pinProtectedUserKeyEnvelope = pinProtectedUserKey,
-                userKeyEncryptedPin = encryptedPin,
-            )
             val mockVaultUnlockResult = VaultUnlockResult.Success
-            coEvery {
-                vaultSdkSource.enrollPinWithEncryptedPin(
-                    userId = userId,
-                    encryptedPin = encryptedPin,
-                )
-            } returns enrollResponse.asSuccess()
             prepareStateForUnlocking(unlockResult = mockVaultUnlockResult)
-            fakeAuthDiskSource.apply {
-                storeEncryptedPin(
-                    userId = userId,
-                    encryptedPin = encryptedPin,
-                )
-                storePinProtectedUserKey(
-                    userId = userId,
-                    pinProtectedUserKey = null,
-                    inMemoryOnly = true,
-                )
-                storePinProtectedUserKeyEnvelope(
-                    userId = userId,
-                    pinProtectedUserKeyEnvelope = null,
-                    inMemoryOnly = true,
-                )
-            }
 
             val result = vaultRepository.unlockVaultWithMasterPassword(
                 masterPassword = "mockPassword-1",
@@ -719,16 +528,6 @@ class VaultRepositoryTest {
                 mockVaultUnlockResult,
                 result,
             )
-            fakeAuthDiskSource.assertPinProtectedUserKey(
-                userId = userId,
-                pinProtectedUserKey = null,
-                inMemoryOnly = true,
-            )
-            fakeAuthDiskSource.assertPinProtectedUserKeyEnvelope(
-                userId = userId,
-                pinProtectedUserKeyEnvelope = pinProtectedUserKey,
-                inMemoryOnly = true,
-            )
             coVerify {
                 vaultLockManager.unlockVault(
                     userId = userId,
@@ -743,11 +542,8 @@ class VaultRepositoryTest {
                     ),
                     organizationKeys = createMockOrganizationKeys(number = 1),
                 )
-            }
+                pinProtectedUserKeyManager.deriveTemporaryPinProtectedUserKeyIfNecessary(
-            coEvery {
-                vaultSdkSource.enrollPinWithEncryptedPin(
                     userId = userId,
-                    encryptedPin = encryptedPin,
                 )
             }
         }