Force the UserState to re-evaluate on authentication change (#1291)

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

@@ -50,6 +50,8 @@ import com.x8bit.bitwarden.data.auth.repository.model.ResendEmailResult
 import com.x8bit.bitwarden.data.auth.repository.model.ResetPasswordResult
 import com.x8bit.bitwarden.data.auth.repository.model.SetPasswordResult
 import com.x8bit.bitwarden.data.auth.repository.model.SwitchAccountResult
+import com.x8bit.bitwarden.data.auth.repository.model.UserAccountTokens
+import com.x8bit.bitwarden.data.auth.repository.model.UserOrganizations
 import com.x8bit.bitwarden.data.auth.repository.model.UserState
 import com.x8bit.bitwarden.data.auth.repository.model.ValidatePasswordResult
 import com.x8bit.bitwarden.data.auth.repository.model.VaultUnlockType
@@ -63,6 +65,8 @@ 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.toUserStateJson
 import com.x8bit.bitwarden.data.auth.repository.util.toUserStateJsonWithPassword
+import com.x8bit.bitwarden.data.auth.repository.util.userAccountTokens
+import com.x8bit.bitwarden.data.auth.repository.util.userAccountTokensFlow
 import com.x8bit.bitwarden.data.auth.repository.util.userOrganizationsList
 import com.x8bit.bitwarden.data.auth.repository.util.userOrganizationsListFlow
 import com.x8bit.bitwarden.data.auth.util.KdfParamsConstants.DEFAULT_PBKDF2_ITERATIONS
@@ -82,6 +86,7 @@ import com.x8bit.bitwarden.data.vault.datasource.network.model.PolicyTypeJson
 import com.x8bit.bitwarden.data.vault.datasource.network.model.SyncResponseJson
 import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
 import com.x8bit.bitwarden.data.vault.repository.VaultRepository
+import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockData
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockResult
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -214,8 +219,10 @@ class AuthRepositoryImpl(
             initialValue = AuthState.Uninitialized,
         )
 
+    @Suppress("UNCHECKED_CAST", "MagicNumber")
     override val userStateFlow: StateFlow<UserState?> = combine(
         authDiskSource.userStateFlow,
+        authDiskSource.userAccountTokensFlow,
         authDiskSource.userOrganizationsListFlow,
         vaultRepository.vaultUnlockDataStateFlow,
         mutableHasPendingAccountAdditionStateFlow,
@@ -224,20 +231,19 @@ class AuthRepositoryImpl(
             mutableHasPendingAccountDeletionStateFlow,
             mutableUserStateTransactionCountStateFlow,
         ),
-    ) {
-            userStateJson,
-            userOrganizationsList,
-            vaultState,
-            hasPendingAccountAddition,
-            _,
-        ->
+    ) { array ->
+        val userStateJson = array[0] as UserStateJson?
+        val userAccountTokens = array[1] as List<UserAccountTokens>
+        val userOrganizationsList = array[2] as List<UserOrganizations>
+        val vaultState = array[3] as List<VaultUnlockData>
+        val hasPendingAccountAddition = array[4] as Boolean
         userStateJson?.toUserState(
             vaultState = vaultState,
+            userAccountTokens = userAccountTokens,
             userOrganizationsList = userOrganizationsList,
             hasPendingAccountAddition = hasPendingAccountAddition,
             isBiometricsEnabledProvider = ::isBiometricsEnabled,
             vaultUnlockTypeProvider = ::getVaultUnlockType,
-            isLoggedInProvider = ::isUserLoggedIn,
             isDeviceTrustedProvider = ::isDeviceTrusted,
         )
     }
@@ -250,11 +256,11 @@ class AuthRepositoryImpl(
                 .userState
                 ?.toUserState(
                     vaultState = vaultRepository.vaultUnlockDataStateFlow.value,
+                    userAccountTokens = authDiskSource.userAccountTokens,
                     userOrganizationsList = authDiskSource.userOrganizationsList,
                     hasPendingAccountAddition = mutableHasPendingAccountAdditionStateFlow.value,
                     isBiometricsEnabledProvider = ::isBiometricsEnabled,
                     vaultUnlockTypeProvider = ::getVaultUnlockType,
-                    isLoggedInProvider = ::isUserLoggedIn,
                     isDeviceTrustedProvider = ::isDeviceTrusted,
                 ),
         )
@@ -1136,10 +1142,6 @@ class AuthRepositoryImpl(
         userId: String,
     ): Boolean = authDiskSource.getDeviceKey(userId = userId) != null
 
-    private fun isUserLoggedIn(
-        userId: String,
-    ): Boolean = authDiskSource.getAccountTokens(userId = userId)?.isLoggedIn == true
-
     private fun getVaultUnlockType(
         userId: String,
     ): VaultUnlockType =

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

@@ -0,0 +1,15 @@
+package com.x8bit.bitwarden.data.auth.repository.model
+
+/**
+ * Associates the [accessToken] and [refreshToken] with the given [userId].
+ */
+data class UserAccountTokens(
+    val userId: String,
+    val accessToken: String?,
+    val refreshToken: String?,
+) {
+    /**
+     * Returns `true` if the user is logged in, `false otherwise.
+     */
+    val isLoggedIn: Boolean get() = accessToken != null
+}

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

@@ -1,6 +1,7 @@
 package com.x8bit.bitwarden.data.auth.repository.util
 
 import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
+import com.x8bit.bitwarden.data.auth.repository.model.UserAccountTokens
 import com.x8bit.bitwarden.data.auth.repository.model.UserOrganizations
 import com.x8bit.bitwarden.data.auth.repository.model.UserSwitchingData
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -55,6 +56,50 @@ val AuthDiskSource.userOrganizationsListFlow: Flow<List<UserOrganizations>>
             }
             .distinctUntilChanged()
 
+/**
+ * Returns the current list of [UserAccountTokens].
+ */
+val AuthDiskSource.userAccountTokens: List<UserAccountTokens>
+    get() = this
+        .userState
+        ?.accounts
+        .orEmpty()
+        .map { (userId, _) ->
+            val accountTokens = this.getAccountTokens(userId = userId)
+            UserAccountTokens(
+                userId = userId,
+                accessToken = accountTokens?.accessToken,
+                refreshToken = accountTokens?.refreshToken,
+            )
+        }
+
+/**
+ * Returns a [Flow] that emits distinct updates to [UserAccountTokens].
+ */
+@OptIn(ExperimentalCoroutinesApi::class)
+val AuthDiskSource.userAccountTokensFlow: Flow<List<UserAccountTokens>>
+    get() = this
+        .userStateFlow
+        .flatMapLatest { userStateJson ->
+            combine(
+                userStateJson
+                    ?.accounts
+                    .orEmpty()
+                    .map { (userId, _) ->
+                        this
+                            .getAccountTokensFlow(userId = userId)
+                            .map {
+                                UserAccountTokens(
+                                    userId = userId,
+                                    accessToken = it?.accessToken,
+                                    refreshToken = it?.refreshToken,
+                                )
+                            }
+                    },
+            ) { it.toList() }
+        }
+        .distinctUntilChanged()
+
 /**
  * Returns a [Flow] that emits every time the active user is changed.
  */

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

@@ -2,6 +2,7 @@ package com.x8bit.bitwarden.data.auth.repository.util
 
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.UserStateJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.UserDecryptionOptionsJson
+import com.x8bit.bitwarden.data.auth.repository.model.UserAccountTokens
 import com.x8bit.bitwarden.data.auth.repository.model.UserOrganizations
 import com.x8bit.bitwarden.data.auth.repository.model.UserState
 import com.x8bit.bitwarden.data.auth.repository.model.VaultUnlockType
@@ -77,11 +78,11 @@ fun UserStateJson.toUserStateJsonWithPassword(): UserStateJson {
 @Suppress("LongParameterList")
 fun UserStateJson.toUserState(
     vaultState: List<VaultUnlockData>,
+    userAccountTokens: List<UserAccountTokens>,
     userOrganizationsList: List<UserOrganizations>,
     hasPendingAccountAddition: Boolean,
     isBiometricsEnabledProvider: (userId: String) -> Boolean,
     vaultUnlockTypeProvider: (userId: String) -> VaultUnlockType,
-    isLoggedInProvider: (userId: String) -> Boolean,
     isDeviceTrustedProvider: (userId: String) -> Boolean,
 ): UserState =
     UserState(
@@ -120,7 +121,9 @@ fun UserStateJson.toUserState(
                         .environmentUrlData
                         .toEnvironmentUrlsOrDefault(),
                     isPremium = profile.hasPremium == true,
-                    isLoggedIn = isLoggedInProvider(userId),
+                    isLoggedIn = userAccountTokens
+                        .find { it.userId == userId }
+                        ?.isLoggedIn == true,
                     isVaultUnlocked = vaultUnlocked,
                     needsPasswordReset = needsPasswordReset,
                     organizations = userOrganizationsList