Update set-password logic to accommodate TDE admin force password update (#1272)

app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/model/SetPasswordRequestJson.kt

@@ -35,7 +35,7 @@ data class SetPasswordRequestJson(
     val key: String,
 
     @SerialName("keys")
-    val keys: Keys,
+    val keys: Keys?,
 
     @SerialName("orgIdentifier")
     val organizationIdentifier: String,

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

@@ -780,6 +780,7 @@ class AuthRepositoryImpl(
             .userState
             ?.activeAccount
             ?: return SetPasswordResult.Error
+        val userId = activeAccount.profile.userId
 
         // Update the saved master password hash.
         val passwordHash = authSdkSource
@@ -791,13 +792,27 @@ class AuthRepositoryImpl(
             )
             .getOrElse { return@setPassword SetPasswordResult.Error }
 
-        return authSdkSource
-            .makeRegisterKeys(
-                email = activeAccount.profile.email,
-                password = password,
-                kdf = activeAccount.profile.toSdkParams(),
-            )
-            .flatMap { keyResponse ->
+        return when (activeAccount.profile.forcePasswordResetReason) {
+            ForcePasswordResetReason.TDE_USER_WITHOUT_PASSWORD_HAS_PASSWORD_RESET_PERMISSION -> {
+                vaultSdkSource
+                    .updatePassword(userId = userId, newPassword = password)
+                    .map { it.newKey to null }
+            }
+
+            ForcePasswordResetReason.ADMIN_FORCE_PASSWORD_RESET,
+            ForcePasswordResetReason.WEAK_MASTER_PASSWORD_ON_LOGIN,
+            null,
+            -> {
+                authSdkSource
+                    .makeRegisterKeys(
+                        email = activeAccount.profile.email,
+                        password = password,
+                        kdf = activeAccount.profile.toSdkParams(),
+                    )
+                    .map { it.encryptedUserKey to it.keys }
+            }
+        }
+            .flatMap { (encryptedUserKey, rsaKeys) ->
                 accountsService
                     .setPassword(
                         body = SetPasswordRequestJson(
@@ -808,28 +823,27 @@ class AuthRepositoryImpl(
                             kdfMemory = activeAccount.profile.kdfMemory,
                             kdfParallelism = activeAccount.profile.kdfParallelism,
                             kdfType = activeAccount.profile.kdfType,
-                            key = keyResponse.encryptedUserKey,
-                            keys = RegisterRequestJson.Keys(
-                                publicKey = keyResponse.keys.public,
-                                encryptedPrivateKey = keyResponse.keys.private,
-                            ),
+                            key = encryptedUserKey,
+                            keys = rsaKeys?.let {
+                                RegisterRequestJson.Keys(
+                                    publicKey = it.public,
+                                    encryptedPrivateKey = it.private,
+                                )
+                            },
                         ),
                     )
                     .onSuccess {
-                        authDiskSource.storePrivateKey(
-                            userId = activeAccount.profile.userId,
-                            privateKey = keyResponse.keys.private,
-                        )
-                        authDiskSource.storeUserKey(
-                            userId = activeAccount.profile.userId,
-                            userKey = keyResponse.encryptedUserKey,
-                        )
+                        rsaKeys?.private?.let {
+                            authDiskSource.storePrivateKey(userId = userId, privateKey = it)
+                        }
+                        authDiskSource.storeUserKey(userId = userId, userKey = encryptedUserKey)
                     }
             }
             .flatMap {
                 when (vaultRepository.unlockVaultWithMasterPassword(password)) {
                     is VaultUnlockResult.Success -> {
                         enrollUserInPasswordReset(
+                            userId = userId,
                             organizationIdentifier = organizationIdentifier,
                             passwordHash = passwordHash,
                         )
@@ -844,13 +858,8 @@ class AuthRepositoryImpl(
                 }
             }
             .onSuccess {
-                authDiskSource.storeMasterPasswordHash(
-                    userId = activeAccount.profile.userId,
-                    passwordHash = passwordHash,
-                )
-
+                authDiskSource.storeMasterPasswordHash(userId = userId, passwordHash = passwordHash)
                 authDiskSource.userState = authDiskSource.userState?.toUserStateJsonWithPassword()
-
                 this.organizationIdentifier = null
             }
             .fold(
@@ -1054,14 +1063,12 @@ class AuthRepositoryImpl(
      * Enrolls the active user in password reset if their organization requires it.
      */
     private suspend fun enrollUserInPasswordReset(
+        userId: String,
         organizationIdentifier: String,
         passwordHash: String,
-    ): Result<Unit> {
-        val userId = activeUserId ?: return IllegalStateException("No active user").asFailure()
-        return organizationService
-            .getOrganizationAutoEnrollStatus(
-                organizationIdentifier = organizationIdentifier,
-            )
+    ): Result<Unit> =
+        organizationService
+            .getOrganizationAutoEnrollStatus(organizationIdentifier = organizationIdentifier)
             .flatMap { statusResponse ->
                 if (statusResponse.isResetPasswordEnabled) {
                     organizationService
@@ -1084,7 +1091,6 @@ class AuthRepositoryImpl(
                     Unit.asSuccess()
                 }
             }
-    }
 
     /**
      * Get the remembered two-factor token associated with the user's email, if applicable.

app/src/main/java/com/x8bit/bitwarden/data/auth/repository/util/UserStateJsonExtensions.kt

@@ -46,10 +46,11 @@ fun UserStateJson.toUpdatedUserStateJson(
  * their password.
  */
 fun UserStateJson.toUserStateJsonWithPassword(): UserStateJson {
-    val account = this.accounts[activeUserId] ?: return this
+    val account = this.activeAccount
     val profile = account.profile
     val updatedProfile = profile
         .copy(
+            forcePasswordResetReason = null,
             userDecryptionOptions = profile
                 .userDecryptionOptions
                 ?.copy(hasMasterPassword = true)

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

@@ -3462,6 +3462,45 @@ class AuthRepositoryTest {
         fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = null)
     }
 
+    @Test
+    fun `setPassword with vaultSdkSource updatePassword failure should return Error`() = runTest {
+        val password = "password"
+        val passwordHash = "passwordHash"
+        val kdf = SINGLE_USER_STATE_1.activeAccount.profile.toSdkParams()
+        fakeAuthDiskSource.userState = SINGLE_USER_STATE_1.copy(
+            accounts = mapOf(
+                USER_ID_1 to ACCOUNT_1.copy(
+                    profile = PROFILE_1.copy(
+                        forcePasswordResetReason = ForcePasswordResetReason
+                            .TDE_USER_WITHOUT_PASSWORD_HAS_PASSWORD_RESET_PERMISSION,
+                    ),
+                ),
+            ),
+        )
+        coEvery {
+            authSdkSource.hashPassword(
+                email = EMAIL,
+                password = password,
+                kdf = kdf,
+                purpose = HashPurpose.SERVER_AUTHORIZATION,
+            )
+        } returns passwordHash.asSuccess()
+        coEvery {
+            vaultSdkSource.updatePassword(userId = USER_ID_1, newPassword = password)
+        } returns Throwable("Fail").asFailure()
+
+        val result = repository.setPassword(
+            organizationIdentifier = "organizationId",
+            password = password,
+            passwordHint = "passwordHint",
+        )
+
+        assertEquals(SetPasswordResult.Error, result)
+        fakeAuthDiskSource.assertMasterPasswordHash(userId = USER_ID_1, passwordHash = null)
+        fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = null)
+        fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = null)
+    }
+
     @Test
     fun `setPassword with accountsService setPassword failure should return Error`() = runTest {
         val password = "password"
@@ -3636,6 +3675,129 @@ class AuthRepositoryTest {
         }
     }
 
+    @Test
+    fun `setPassword with updatePassword success should return Success`() = runTest {
+        val password = "password"
+        val passwordHash = "passwordHash"
+        val passwordHint = "passwordHint"
+        val organizationIdentifier = ORGANIZATION_IDENTIFIER
+        val organizationId = "orgId"
+        val encryptedUserKey = "encryptedUserKey"
+        val publicOrgKey = "publicOrgKey"
+        val resetPasswordKey = "resetPasswordKey"
+        val userState = SINGLE_USER_STATE_1.copy(
+            accounts = mapOf(
+                USER_ID_1 to ACCOUNT_1.copy(
+                    profile = PROFILE_1.copy(
+                        forcePasswordResetReason = ForcePasswordResetReason
+                            .TDE_USER_WITHOUT_PASSWORD_HAS_PASSWORD_RESET_PERMISSION,
+                    ),
+                ),
+            ),
+        )
+        val profile = userState.activeAccount.profile
+        val kdf = profile.toSdkParams()
+        val updatePasswordResponse = UpdatePasswordResponse(
+            passwordHash = passwordHash,
+            newKey = encryptedUserKey,
+        )
+        val setPasswordRequestJson = SetPasswordRequestJson(
+            passwordHash = passwordHash,
+            passwordHint = passwordHint,
+            organizationIdentifier = organizationIdentifier,
+            kdfIterations = profile.kdfIterations,
+            kdfMemory = profile.kdfMemory,
+            kdfParallelism = profile.kdfParallelism,
+            kdfType = profile.kdfType,
+            key = encryptedUserKey,
+            keys = null,
+        )
+        fakeAuthDiskSource.userState = userState
+        coEvery {
+            authSdkSource.hashPassword(
+                email = EMAIL,
+                password = password,
+                kdf = kdf,
+                purpose = HashPurpose.SERVER_AUTHORIZATION,
+            )
+        } returns passwordHash.asSuccess()
+        coEvery {
+            vaultSdkSource.updatePassword(userId = USER_ID_1, newPassword = password)
+        } returns updatePasswordResponse.asSuccess()
+        coEvery {
+            accountsService.setPassword(body = setPasswordRequestJson)
+        } returns Unit.asSuccess()
+        coEvery {
+            organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
+        } returns OrganizationAutoEnrollStatusResponseJson(
+            organizationId = organizationId,
+            isResetPasswordEnabled = true,
+        )
+            .asSuccess()
+        coEvery {
+            organizationService.getOrganizationKeys(organizationId)
+        } returns OrganizationKeysResponseJson(
+            privateKey = "",
+            publicKey = publicOrgKey,
+        )
+            .asSuccess()
+        coEvery {
+            organizationService.organizationResetPasswordEnroll(
+                organizationId = organizationId,
+                userId = profile.userId,
+                passwordHash = passwordHash,
+                resetPasswordKey = resetPasswordKey,
+            )
+        } returns Unit.asSuccess()
+        coEvery {
+            vaultSdkSource.getResetPasswordKey(
+                orgPublicKey = publicOrgKey,
+                userId = profile.userId,
+            )
+        } returns resetPasswordKey.asSuccess()
+        coEvery {
+            vaultRepository.unlockVaultWithMasterPassword(password)
+        } returns VaultUnlockResult.Success
+
+        val result = repository.setPassword(
+            organizationIdentifier = organizationIdentifier,
+            password = password,
+            passwordHint = passwordHint,
+        )
+
+        assertEquals(SetPasswordResult.Success, result)
+        fakeAuthDiskSource.assertMasterPasswordHash(
+            userId = USER_ID_1,
+            passwordHash = passwordHash,
+        )
+        fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = null)
+        fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = encryptedUserKey)
+        fakeAuthDiskSource.assertUserState(SINGLE_USER_STATE_1_WITH_PASS)
+        coVerify(exactly = 1) {
+            authSdkSource.hashPassword(
+                email = EMAIL,
+                password = password,
+                kdf = kdf,
+                purpose = HashPurpose.SERVER_AUTHORIZATION,
+            )
+            vaultSdkSource.updatePassword(userId = USER_ID_1, newPassword = password)
+            accountsService.setPassword(body = setPasswordRequestJson)
+            organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
+            organizationService.getOrganizationKeys(organizationId)
+            organizationService.organizationResetPasswordEnroll(
+                organizationId = organizationId,
+                userId = profile.userId,
+                passwordHash = passwordHash,
+                resetPasswordKey = resetPasswordKey,
+            )
+            vaultRepository.unlockVaultWithMasterPassword(password)
+            vaultSdkSource.getResetPasswordKey(
+                orgPublicKey = publicOrgKey,
+                userId = profile.userId,
+            )
+        }
+    }
+
     @Test
     fun `setPassword with unlockVaultWithMasterPassword error should return Failure`() = runTest {
         val password = "password"
@@ -4515,23 +4677,24 @@ class AuthRepositoryTest {
             masterPasswordPolicyOptions = null,
             userDecryptionOptions = null,
         )
+        private val PROFILE_1 = AccountJson.Profile(
+            userId = USER_ID_1,
+            email = EMAIL,
+            isEmailVerified = true,
+            name = "Bitwarden Tester",
+            hasPremium = false,
+            stamp = null,
+            organizationId = null,
+            avatarColorHex = null,
+            forcePasswordResetReason = null,
+            kdfType = KdfTypeJson.ARGON2_ID,
+            kdfIterations = 600000,
+            kdfMemory = 16,
+            kdfParallelism = 4,
+            userDecryptionOptions = null,
+        )
         private val ACCOUNT_1 = AccountJson(
-            profile = AccountJson.Profile(
-                userId = USER_ID_1,
-                email = EMAIL,
-                isEmailVerified = true,
-                name = "Bitwarden Tester",
-                hasPremium = false,
-                stamp = null,
-                organizationId = null,
-                avatarColorHex = null,
-                forcePasswordResetReason = null,
-                kdfType = KdfTypeJson.ARGON2_ID,
-                kdfIterations = 600000,
-                kdfMemory = 16,
-                kdfParallelism = 4,
-                userDecryptionOptions = null,
-            ),
+            profile = PROFILE_1,
             settings = AccountJson.Settings(
                 environmentUrlData = null,
             ),

app/src/test/java/com/x8bit/bitwarden/data/auth/repository/util/UserStateJsonExtensionsTest.kt

@@ -3,6 +3,7 @@ package com.x8bit.bitwarden.data.auth.repository.util
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountTokensJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.EnvironmentUrlDataJson
+import com.x8bit.bitwarden.data.auth.datasource.disk.model.ForcePasswordResetReason
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.UserStateJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.KdfTypeJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.KeyConnectorUserDecryptionOptionsJson
@@ -96,8 +97,9 @@ class UserStateJsonExtensionsTest {
         )
     }
 
+    @Suppress("MaxLineLength")
     @Test
-    fun `toUserStateJsonWithPassword should update correct account to set needsMasterPassword`() {
+    fun `toUserStateJsonWithPassword should update active account to set hasMasterPassword and clear forcePasswordResetReason`() {
         val originalProfile = AccountJson.Profile(
             userId = "activeUserId",
             email = "email",
@@ -107,7 +109,8 @@ class UserStateJsonExtensionsTest {
             organizationId = null,
             avatarColorHex = null,
             hasPremium = true,
-            forcePasswordResetReason = null,
+            forcePasswordResetReason = ForcePasswordResetReason
+                .TDE_USER_WITHOUT_PASSWORD_HAS_PASSWORD_RESET_PERMISSION,
             kdfType = KdfTypeJson.ARGON2_ID,
             kdfIterations = 600000,
             kdfMemory = 16,
@@ -125,6 +128,7 @@ class UserStateJsonExtensionsTest {
                 accounts = mapOf(
                     "activeUserId" to originalAccount.copy(
                         profile = originalProfile.copy(
+                            forcePasswordResetReason = null,
                             userDecryptionOptions = UserDecryptionOptionsJson(
                                 hasMasterPassword = true,
                                 keyConnectorUserDecryptionOptions = null,