Chore: Consolidate functions for setting the MasterPasswordUnlock data (#7203)

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

@@ -109,10 +109,9 @@ import com.x8bit.bitwarden.data.auth.repository.util.policyInformation
 import com.x8bit.bitwarden.data.auth.repository.util.toAccountCryptographicState
 import com.x8bit.bitwarden.data.auth.repository.util.toDeviceInfo
 import com.x8bit.bitwarden.data.auth.repository.util.toOrganizations
-import com.x8bit.bitwarden.data.auth.repository.util.toRemovedPasswordUserStateJson
 import com.x8bit.bitwarden.data.auth.repository.util.toSdkParams
 import com.x8bit.bitwarden.data.auth.repository.util.toUserState
-import com.x8bit.bitwarden.data.auth.repository.util.toUserStateJsonWithPassword
+import com.x8bit.bitwarden.data.auth.repository.util.updateMasterPasswordUnlock
 import com.x8bit.bitwarden.data.auth.repository.util.userSwitchingChangesFlow
 import com.x8bit.bitwarden.data.auth.util.KdfParamsConstants.DEFAULT_PBKDF2_ITERATIONS
 import com.x8bit.bitwarden.data.auth.util.YubiKeyResult
@@ -1058,7 +1057,10 @@ class AuthRepositoryImpl(
                     MigrateExistingUserToKeyConnectorResult.Success -> {
                         authDiskSource.userState = authDiskSource
                             .userState
-                            ?.toRemovedPasswordUserStateJson(userId = userId)
+                            ?.updateMasterPasswordUnlock(
+                                userId = userId,
+                                masterPasswordUnlock = null,
+                            )
                         vaultRepository.sync()
                         settingsRepository.setDefaultsIfNecessary(userId = userId)
                         RemovePasswordResult.Success
@@ -1180,7 +1182,8 @@ class AuthRepositoryImpl(
                     .map { response }
             }
             .onSuccess { response ->
-                authDiskSource.userState = authDiskSource.userState?.toUserStateJsonWithPassword(
+                authDiskSource.userState = authDiskSource.userState?.updateMasterPasswordUnlock(
+                    userId = userId,
                     masterPasswordUnlock = MasterPasswordUnlockData(
                         kdf = profile.toSdkParams(),
                         masterKeyWrappedUserKey = response.newKey,
@@ -1241,7 +1244,8 @@ class AuthRepositoryImpl(
                     userId = userId,
                     accountCryptographicState = response.accountCryptographicState,
                 )
-                authDiskSource.userState = authDiskSource.userState?.toUserStateJsonWithPassword(
+                authDiskSource.userState = authDiskSource.userState?.updateMasterPasswordUnlock(
+                    userId = userId,
                     masterPasswordUnlock = response.masterPasswordUnlock,
                 )
                 this.organizationIdentifier = null
@@ -1304,7 +1308,8 @@ class AuthRepositoryImpl(
                         )
                         authDiskSource.userState = authDiskSource
                             .userState
-                            ?.toUserStateJsonWithPassword(
+                            ?.updateMasterPasswordUnlock(
+                                userId = userId,
                                 masterPasswordUnlock = MasterPasswordUnlockData(
                                     kdf = profile.toSdkParams(),
                                     masterKeyWrappedUserKey = response.encryptedUserKey,

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

@@ -24,36 +24,6 @@ import com.x8bit.bitwarden.data.platform.manager.model.FirstTimeState
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockData
 import com.x8bit.bitwarden.data.vault.repository.util.statusFor
 
-/**
- * Updates the given [UserStateJson] with the data to indicate that the password has been removed.
- * The original will be returned if the [userId] does not match any accounts in the [UserStateJson].
- */
-fun UserStateJson.toRemovedPasswordUserStateJson(
-    userId: String,
-): UserStateJson {
-    val account = this.accounts[userId] ?: return this
-    val profile = account.profile
-    val updatedUserDecryptionOptions = profile
-        .userDecryptionOptions
-        ?.copy(
-            hasMasterPassword = false,
-            masterPasswordUnlock = null,
-        )
-        ?: UserDecryptionOptionsJson(
-            hasMasterPassword = false,
-            trustedDeviceUserDecryptionOptions = null,
-            keyConnectorUserDecryptionOptions = null,
-            masterPasswordUnlock = null,
-        )
-    val updatedProfile = profile.copy(userDecryptionOptions = updatedUserDecryptionOptions)
-    val updatedAccount = account.copy(profile = updatedProfile)
-    return this.copy(
-        accounts = accounts
-            .toMutableMap()
-            .apply { replace(userId, updatedAccount) },
-    )
-}
-
 /**
  * Updates the given [UserStateJson] with the data from the [syncResponse] to return a new
  * [UserStateJson]. The original will be returned if the sync response does not match any accounts
@@ -134,44 +104,43 @@ private fun SyncResponseJson.Profile.getForcePasswordResetReason(
 }
 
 /**
- * Updates the [UserStateJson] to set the `hasMasterPassword` value to `true` after a user sets
- * their password.
+ * Updates the [UserStateJson] by setting the `hasMasterPassword` and `masterPasswordUnlock` values
+ * for the given [userId]. If the user is not present in the `UserStateJson`, nothing is updated.
  */
-fun UserStateJson.toUserStateJsonWithPassword(
-    masterPasswordUnlock: MasterPasswordUnlockData,
+fun UserStateJson.updateMasterPasswordUnlock(
+    userId: String,
+    masterPasswordUnlock: MasterPasswordUnlockData?,
 ): UserStateJson {
-    val account = this.activeAccount
+    val account = accounts[userId] ?: return this
     val profile = account.profile
     val userDecryptionOptions = profile.userDecryptionOptions
-    val masterPasswordUnlockJson = MasterPasswordUnlockDataJson(
-        salt = masterPasswordUnlock.salt,
-        kdf = masterPasswordUnlock.kdf.toKdfRequestModel(),
-        masterKeyWrappedUserKey = masterPasswordUnlock.masterKeyWrappedUserKey,
+    val masterPasswordUnlockJson = masterPasswordUnlock?.let {
+        MasterPasswordUnlockDataJson(
+            salt = it.salt,
+            kdf = it.kdf.toKdfRequestModel(),
+            masterKeyWrappedUserKey = it.masterKeyWrappedUserKey,
+        )
+    }
+    val updatedProfile = profile.copy(
+        forcePasswordResetReason = null,
+        userDecryptionOptions = userDecryptionOptions
+            ?.copy(
+                hasMasterPassword = masterPasswordUnlockJson != null,
+                masterPasswordUnlock = masterPasswordUnlockJson,
+            )
+            ?: UserDecryptionOptionsJson(
+                hasMasterPassword = masterPasswordUnlockJson != null,
+                keyConnectorUserDecryptionOptions = null,
+                trustedDeviceUserDecryptionOptions = null,
+                masterPasswordUnlock = masterPasswordUnlockJson,
+            ),
     )
-    val updatedProfile = profile
-        .copy(
-            forcePasswordResetReason = null,
-            userDecryptionOptions = userDecryptionOptions
-                ?.copy(
-                    hasMasterPassword = true,
-                    masterPasswordUnlock = masterPasswordUnlockJson,
-                )
-                ?: UserDecryptionOptionsJson(
-                    hasMasterPassword = true,
-                    keyConnectorUserDecryptionOptions = null,
-                    trustedDeviceUserDecryptionOptions = null,
-                    masterPasswordUnlock = masterPasswordUnlockJson,
-                ),
-        )
     val updatedAccount = account.copy(profile = updatedProfile)
-    return this
-        .copy(
-            accounts = accounts
-                .toMutableMap()
-                .apply {
-                    replace(activeUserId, updatedAccount)
-                },
-        )
+    return this.copy(
+        accounts = accounts
+            .toMutableMap()
+            .apply { replace(userId, updatedAccount) },
+    )
 }
 
 /**

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

@@ -131,9 +131,9 @@ import com.x8bit.bitwarden.data.auth.repository.util.CookieCallbackResult
 import com.x8bit.bitwarden.data.auth.repository.util.DuoCallbackTokenResult
 import com.x8bit.bitwarden.data.auth.repository.util.SsoCallbackResult
 import com.x8bit.bitwarden.data.auth.repository.util.WebAuthResult
-import com.x8bit.bitwarden.data.auth.repository.util.toRemovedPasswordUserStateJson
 import com.x8bit.bitwarden.data.auth.repository.util.toSdkParams
 import com.x8bit.bitwarden.data.auth.repository.util.toUserState
+import com.x8bit.bitwarden.data.auth.repository.util.updateMasterPasswordUnlock
 import com.x8bit.bitwarden.data.auth.util.YubiKeyResult
 import com.x8bit.bitwarden.data.auth.util.toSdkParams
 import com.x8bit.bitwarden.data.platform.datasource.disk.util.FakeSettingsDiskSource
@@ -330,7 +330,7 @@ class AuthRepositoryTest {
     fun beforeEach() {
         mockkStatic(
             GetTokenResponseJson.Success::toUserState,
-            UserStateJson::toRemovedPasswordUserStateJson,
+            UserStateJson::updateMasterPasswordUnlock,
         )
         mockkConstructor(
             NoActiveUserException::class,
@@ -348,7 +348,7 @@ class AuthRepositoryTest {
     fun tearDown() {
         unmockkStatic(
             GetTokenResponseJson.Success::toUserState,
-            UserStateJson::toRemovedPasswordUserStateJson,
+            UserStateJson::updateMasterPasswordUnlock,
         )
         mockkConstructor(
             NoActiveUserException::class,
@@ -5036,7 +5036,10 @@ class AuthRepositoryTest {
                 )
             } returns MigrateExistingUserToKeyConnectorResult.Success.asSuccess()
             every {
-                SINGLE_USER_STATE_1.toRemovedPasswordUserStateJson(userId = USER_ID_1)
+                SINGLE_USER_STATE_1.updateMasterPasswordUnlock(
+                    userId = USER_ID_1,
+                    masterPasswordUnlock = null,
+                )
             } returns SINGLE_USER_STATE_1
             every { vaultRepository.sync() } just runs
             every { settingsRepository.setDefaultsIfNecessary(userId = USER_ID_1) } just runs
@@ -5045,7 +5048,10 @@ class AuthRepositoryTest {
 
             assertEquals(RemovePasswordResult.Success, result)
             verify(exactly = 1) {
-                SINGLE_USER_STATE_1.toRemovedPasswordUserStateJson(userId = USER_ID_1)
+                SINGLE_USER_STATE_1.updateMasterPasswordUnlock(
+                    userId = USER_ID_1,
+                    masterPasswordUnlock = null,
+                )
                 vaultRepository.sync()
                 settingsRepository.setDefaultsIfNecessary(userId = USER_ID_1)
             }

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

@@ -59,20 +59,23 @@ class UserStateJsonExtensionsTest {
     }
 
     @Test
-    fun `toRemovedPasswordUserStateJson should do nothing for a non-matching account`() {
+    fun `updateMasterPasswordUnlock should do nothing for a non-matching account`() {
         val originalUserState = UserStateJson(
             activeUserId = "activeUserId",
             accounts = mapOf("activeUserId" to mockk()),
         )
         assertEquals(
             originalUserState,
-            originalUserState.toRemovedPasswordUserStateJson(userId = "nonActiveUserId"),
+            originalUserState.updateMasterPasswordUnlock(
+                userId = "nonActiveUserId",
+                masterPasswordUnlock = null,
+            ),
         )
     }
 
     @Suppress("MaxLineLength")
     @Test
-    fun `toRemovedPasswordUserStateJson should create user decryption options without a password if not present`() {
+    fun `updateMasterPasswordUnlock should create user decryption options without a password if not present`() {
         val originalProfile = AccountJson.Profile(
             userId = "activeUserId",
             email = "email",
@@ -118,13 +121,16 @@ class UserStateJsonExtensionsTest {
                     ),
                 ),
             ),
-            originalUserState.toRemovedPasswordUserStateJson(userId = "activeUserId"),
+            originalUserState.updateMasterPasswordUnlock(
+                userId = "activeUserId",
+                masterPasswordUnlock = null,
+            ),
         )
     }
 
     @Suppress("MaxLineLength")
     @Test
-    fun `toRemovedPasswordUserStateJson should update user decryption options to not have a password`() {
+    fun `updateMasterPasswordUnlock should update user decryption options to not have a password`() {
         val originalProfile = AccountJson.Profile(
             userId = "activeUserId",
             email = "email",
@@ -179,7 +185,10 @@ class UserStateJsonExtensionsTest {
                     ),
                 ),
             ),
-            originalUserState.toRemovedPasswordUserStateJson(userId = "activeUserId"),
+            originalUserState.updateMasterPasswordUnlock(
+                userId = "activeUserId",
+                masterPasswordUnlock = null,
+            ),
         )
     }
 
@@ -420,7 +429,7 @@ class UserStateJsonExtensionsTest {
 
     @Suppress("MaxLineLength")
     @Test
-    fun `toUserStateJsonWithPassword with masterPasswordUnlock should update active account to set hasMasterPassword and masterPasswordUnlock`() {
+    fun `updateMasterPasswordUnlock with masterPasswordUnlock should update account to set hasMasterPassword and masterPasswordUnlock`() {
         val originalProfile = AccountJson.Profile(
             userId = "activeUserId",
             email = "email",
@@ -481,7 +490,10 @@ class UserStateJsonExtensionsTest {
                     "activeUserId" to originalAccount,
                 ),
             )
-                .toUserStateJsonWithPassword(masterPasswordUnlock = masterPasswordUnlock),
+                .updateMasterPasswordUnlock(
+                    userId = "activeUserId",
+                    masterPasswordUnlock = masterPasswordUnlock,
+                ),
         )
     }