[PM-23280] Use masterPasswordUnlock KDF settings on vault unlock (#6026)

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/datasource/sdk/util/KdfExtensions.kt

@@ -5,6 +5,8 @@ import com.bitwarden.network.model.KdfJson
 import com.bitwarden.network.model.KdfTypeJson
 import com.bitwarden.network.model.KdfTypeJson.ARGON2_ID
 import com.bitwarden.network.model.KdfTypeJson.PBKDF2_SHA256
+import com.x8bit.bitwarden.data.auth.util.KdfParamsConstants.DEFAULT_ARGON2_MEMORY
+import com.x8bit.bitwarden.data.auth.util.KdfParamsConstants.DEFAULT_ARGON2_PARALLELISM
 
 /**
  * Convert a [Kdf] to a [KdfTypeJson].
@@ -34,3 +36,16 @@ fun Kdf.toKdfRequestModel(): KdfJson =
             parallelism = null,
         )
     }
+
+/**
+ * Convert a [KdfJson] to a [Kdf].
+ */
+fun KdfJson.toKdf(): Kdf =
+    when (this.kdfType) {
+        ARGON2_ID -> Kdf.Argon2id(
+            iterations = iterations.toUInt(),
+            memory = memory?.toUInt() ?: DEFAULT_ARGON2_MEMORY.toUInt(),
+            parallelism = parallelism?.toUInt() ?: DEFAULT_ARGON2_PARALLELISM.toUInt(),
+        )
+        PBKDF2_SHA256 -> Kdf.Pbkdf2(iterations = iterations.toUInt())
+    }

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryImpl.kt

@@ -51,6 +51,7 @@ import com.x8bit.bitwarden.data.auth.datasource.disk.model.OnboardingStatus
 import com.x8bit.bitwarden.data.auth.datasource.network.model.DeviceDataModel
 import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSource
 import com.x8bit.bitwarden.data.auth.datasource.sdk.util.toInt
+import com.x8bit.bitwarden.data.auth.datasource.sdk.util.toKdf
 import com.x8bit.bitwarden.data.auth.datasource.sdk.util.toKdfTypeJson
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestManager
 import com.x8bit.bitwarden.data.auth.manager.KdfManager
@@ -2046,10 +2047,20 @@ class AuthRepositoryImpl(
         initUserCryptoMethod: InitUserCryptoMethod,
     ): VaultUnlockResult {
         val userId = accountProfile.userId
+        val kdfParams = (initUserCryptoMethod as? InitUserCryptoMethod.Password)
+            ?.let {
+                accountProfile
+                    .userDecryptionOptions
+                    ?.masterPasswordUnlock
+                    ?.kdf
+                    ?.toKdf()
+            }
+            ?: accountProfile.toSdkParams()
+
         return vaultRepository.unlockVault(
             userId = userId,
             email = accountProfile.email,
-            kdf = accountProfile.toSdkParams(),
+            kdf = kdfParams,
             privateKey = privateKey,
             signingKey = signingKey,
             securityState = securityState,

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryTest.kt

@@ -30,6 +30,7 @@ import com.bitwarden.network.model.GetTokenResponseJson
 import com.bitwarden.network.model.IdentityTokenAuthModel
 import com.bitwarden.network.model.KdfTypeJson
 import com.bitwarden.network.model.KeyConnectorMasterKeyResponseJson
+import com.bitwarden.network.model.MasterPasswordUnlockDataJson
 import com.bitwarden.network.model.OrganizationAutoEnrollStatusResponseJson
 import com.bitwarden.network.model.OrganizationKeysResponseJson
 import com.bitwarden.network.model.OrganizationType
@@ -76,6 +77,7 @@ import com.x8bit.bitwarden.data.auth.datasource.sdk.model.PasswordStrength.LEVEL
 import com.x8bit.bitwarden.data.auth.datasource.sdk.model.PasswordStrength.LEVEL_2
 import com.x8bit.bitwarden.data.auth.datasource.sdk.model.PasswordStrength.LEVEL_3
 import com.x8bit.bitwarden.data.auth.datasource.sdk.model.PasswordStrength.LEVEL_4
+import com.x8bit.bitwarden.data.auth.datasource.sdk.util.toKdfRequestModel
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestManager
 import com.x8bit.bitwarden.data.auth.manager.KdfManager
 import com.x8bit.bitwarden.data.auth.manager.KeyConnectorManager
@@ -6985,6 +6987,79 @@ class AuthRepositoryTest {
             )
         }
 
+    @Test
+    fun `unlockVault uses user decryption options for KDF when init method is password`() =
+        runTest {
+            val successResponse = GET_TOKEN_WITH_ACCOUNT_KEYS_RESPONSE_SUCCESS
+            coEvery {
+                identityService.preLogin(email = EMAIL)
+            } returns PRE_LOGIN_SUCCESS.asSuccess()
+            coEvery {
+                identityService.getToken(
+                    email = EMAIL,
+                    authModel = IdentityTokenAuthModel.MasterPassword(
+                        username = EMAIL,
+                        password = PASSWORD_HASH,
+                    ),
+                    uniqueAppId = UNIQUE_APP_ID,
+                )
+            } returns successResponse.asSuccess()
+            coEvery {
+                vaultRepository.unlockVault(
+                    userId = USER_ID_1,
+                    email = EMAIL,
+                    kdf = ACCOUNT_2.profile.toSdkParams(),
+                    privateKey = successResponse.accountKeys!!
+                        .publicKeyEncryptionKeyPair
+                        .wrappedPrivateKey,
+                    signingKey = successResponse.accountKeys
+                        ?.signatureKeyPair
+                        ?.wrappedSigningKey,
+                    securityState = successResponse.accountKeys
+                        ?.securityState
+                        ?.securityState,
+                    initUserCryptoMethod = InitUserCryptoMethod.Password(
+                        password = PASSWORD,
+                        userKey = successResponse.key!!,
+                    ),
+                    organizationKeys = null,
+                )
+            } returns VaultUnlockResult.Success
+            coEvery { vaultRepository.syncIfNecessary() } just runs
+            every {
+                GET_TOKEN_WITH_ACCOUNT_KEYS_RESPONSE_SUCCESS.toUserState(
+                    previousUserState = null,
+                    environmentUrlData = EnvironmentUrlDataJson.DEFAULT_US,
+                )
+            } returns SINGLE_USER_STATE_1_WITH_DECRYPTION_OPTIONS
+
+            repository.login(email = EMAIL, password = PASSWORD)
+
+            coVerify {
+                vaultRepository.unlockVault(
+                    userId = USER_ID_1,
+                    email = EMAIL,
+                    kdf = ACCOUNT_2.profile.toSdkParams(),
+                    privateKey = successResponse.accountKeys!!
+                        .publicKeyEncryptionKeyPair
+                        .wrappedPrivateKey,
+                    signingKey = successResponse.accountKeys
+                        ?.signatureKeyPair
+                        ?.wrappedSigningKey,
+                    securityState = successResponse.accountKeys
+                        ?.securityState
+                        ?.securityState,
+                    initUserCryptoMethod = InitUserCryptoMethod.Password(
+                        password = PASSWORD,
+                        userKey = successResponse.key!!,
+                    ),
+                    organizationKeys = null,
+                )
+                vaultRepository.syncIfNecessary()
+                settingsRepository.storeUserHasLoggedInValue(userId = USER_ID_1)
+            }
+        }
+
     companion object {
         private val FIXED_CLOCK: Clock = Clock.fixed(
             Instant.parse("2023-10-27T12:00:00Z"),
@@ -7175,6 +7250,28 @@ class AuthRepositoryTest {
             ),
         )
 
+        private val MOCK_MASTER_PASSWORD_UNLOCK_DATA = MasterPasswordUnlockDataJson(
+            salt = "mockSalt",
+            kdf = ACCOUNT_2.profile.toSdkParams().toKdfRequestModel(),
+            masterKeyWrappedUserKey = "masterKeyWrappedUserKeyMock",
+        )
+
+        private val SINGLE_USER_STATE_1_WITH_DECRYPTION_OPTIONS = UserStateJson(
+            activeUserId = USER_ID_1,
+            accounts = mapOf(
+                USER_ID_1 to ACCOUNT_1.copy(
+                    profile = ACCOUNT_1.profile.copy(
+                        userDecryptionOptions = UserDecryptionOptionsJson(
+                            hasMasterPassword = true,
+                            keyConnectorUserDecryptionOptions = null,
+                            trustedDeviceUserDecryptionOptions = null,
+                            masterPasswordUnlock = MOCK_MASTER_PASSWORD_UNLOCK_DATA,
+                        ),
+                    ),
+                ),
+            ),
+        )
+
         private val SINGLE_USER_STATE_2 = UserStateJson(
             activeUserId = USER_ID_2,
             accounts = mapOf(