Add Organizations to UserState.Account (#432)

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

@@ -23,16 +23,14 @@ private const val ORGANIZATION_KEYS_KEY = "$BASE_KEY:encOrgKeys"
 /**
  * Primary implementation of [AuthDiskSource].
  */
+@Suppress("TooManyFunctions")
 class AuthDiskSourceImpl(
     sharedPreferences: SharedPreferences,
     private val json: Json,
 ) : BaseDiskSource(sharedPreferences = sharedPreferences),
     AuthDiskSource {
-    private val mutableOrganizationsFlow =
-        MutableSharedFlow<List<SyncResponseJson.Profile.Organization>?>(
-            replay = 1,
-            extraBufferCapacity = Int.MAX_VALUE,
-        )
+    private val mutableOrganizationsFlowMap =
+        mutableMapOf<String, MutableSharedFlow<List<SyncResponseJson.Profile.Organization>?>>()
 
     override val uniqueAppId: String
         get() = getString(key = UNIQUE_APP_ID_KEY) ?: generateAndStoreUniqueAppId()
@@ -108,7 +106,7 @@ class AuthDiskSourceImpl(
     override fun getOrganizationsFlow(
         userId: String,
     ): Flow<List<SyncResponseJson.Profile.Organization>?> =
-        mutableOrganizationsFlow
+        getMutableOrganizationsFlow(userId = userId)
             .onSubscription { emit(getOrganizations(userId = userId)) }
 
     override fun storeOrganizations(
@@ -119,7 +117,7 @@ class AuthDiskSourceImpl(
             key = "${ORGANIZATIONS_KEY}_$userId",
             value = organizations?.let { json.encodeToString(it) },
         )
-        mutableOrganizationsFlow.tryEmit(organizations)
+        getMutableOrganizationsFlow(userId = userId).tryEmit(organizations)
     }
 
     private fun generateAndStoreUniqueAppId(): String =
@@ -129,4 +127,14 @@ class AuthDiskSourceImpl(
             .also {
                 putString(key = UNIQUE_APP_ID_KEY, value = it)
             }
+
+    private fun getMutableOrganizationsFlow(
+        userId: String,
+    ): MutableSharedFlow<List<SyncResponseJson.Profile.Organization>?> =
+        mutableOrganizationsFlowMap.getOrPut(userId) {
+            MutableSharedFlow(
+                replay = 1,
+                extraBufferCapacity = Int.MAX_VALUE,
+            )
+        }
 }

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

@@ -27,6 +27,8 @@ import com.x8bit.bitwarden.data.auth.repository.util.CaptchaCallbackTokenResult
 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.userOrganizationsList
+import com.x8bit.bitwarden.data.auth.repository.util.userOrganizationsListFlow
 import com.x8bit.bitwarden.data.auth.util.KdfParamsConstants.DEFAULT_PBKDF2_ITERATIONS
 import com.x8bit.bitwarden.data.auth.util.toSdkParams
 import com.x8bit.bitwarden.data.platform.manager.dispatcher.DispatcherManager
@@ -36,6 +38,7 @@ import com.x8bit.bitwarden.data.platform.util.flatMap
 import com.x8bit.bitwarden.data.vault.repository.VaultRepository
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -94,14 +97,17 @@ class AuthRepositoryImpl constructor(
             initialValue = AuthState.Uninitialized,
         )
 
+    @OptIn(ExperimentalCoroutinesApi::class)
     override val userStateFlow: StateFlow<UserState?> = combine(
         authDiskSource.userStateFlow,
+        authDiskSource.userOrganizationsListFlow,
         vaultRepository.vaultStateFlow,
         mutableSpecialCircumstanceStateFlow,
-    ) { userStateJson, vaultState, specialCircumstance ->
+    ) { userStateJson, userOrganizationsList, vaultState, specialCircumstance ->
         userStateJson
             ?.toUserState(
                 vaultState = vaultState,
+                userOrganizationsList = userOrganizationsList,
                 specialCircumstance = specialCircumstance,
             )
     }
@@ -112,6 +118,7 @@ class AuthRepositoryImpl constructor(
                 .userState
                 ?.toUserState(
                     vaultState = vaultRepository.vaultStateFlow.value,
+                    userOrganizationsList = authDiskSource.userOrganizationsList,
                     specialCircumstance = mutableSpecialCircumstanceStateFlow.value,
                 ),
         )

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

@@ -0,0 +1,12 @@
+package com.x8bit.bitwarden.data.auth.repository.model
+
+/**
+ * Represents an organization a user may be a member of.
+ *
+ * @property id The ID of the organization.
+ * @property name The name of the organization (if applicable).
+ */
+data class Organization(
+    val id: String,
+    val name: String?,
+)

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

@@ -0,0 +1,9 @@
+package com.x8bit.bitwarden.data.auth.repository.model
+
+/**
+ * Associates a list of [organizations] with the given [userId].
+ */
+data class UserOrganizations(
+    val userId: String,
+    val organizations: List<Organization>,
+)

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

@@ -43,6 +43,7 @@ data class UserState(
      * @property environment The [Environment] associated with the user's account.
      * @property isPremium `true` if the account has a premium membership.
      * @property isVaultUnlocked Whether or not the user's vault is currently unlocked.
+     * @property organizations List of [Organization]s the user is associated with, if any.
      */
     data class Account(
         val userId: String,
@@ -52,6 +53,7 @@ data class UserState(
         val environment: Environment,
         val isPremium: Boolean,
         val isVaultUnlocked: Boolean,
+        val organizations: List<Organization>,
     )
 
     /**

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

@@ -0,0 +1,55 @@
+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.UserOrganizations
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.map
+
+/**
+ * Returns the current list of [UserOrganizations].
+ */
+val AuthDiskSource.userOrganizationsList: List<UserOrganizations>
+    get() = this
+        .userState
+        ?.accounts
+        .orEmpty()
+        .map { (userId, _) ->
+            UserOrganizations(
+                userId = userId,
+                organizations = this
+                    .getOrganizations(userId = userId)
+                    .orEmpty()
+                    .toOrganizations(),
+            )
+        }
+
+/**
+ * Returns a [Flow] that emits distinct updates to [UserOrganizations].
+ */
+@OptIn(ExperimentalCoroutinesApi::class)
+val AuthDiskSource.userOrganizationsListFlow: Flow<List<UserOrganizations>>
+    get() =
+        this
+            .userStateFlow
+            .flatMapLatest { userStateJson ->
+                combine(
+                    userStateJson
+                        ?.accounts
+                        .orEmpty()
+                        .map { (userId, _) ->
+                            this
+                                .getOrganizationsFlow(userId = userId)
+                                .map {
+                                    UserOrganizations(
+                                        userId = userId,
+                                        organizations = it.orEmpty().toOrganizations(),
+                                    )
+                                }
+                        },
+                ) { values -> values.toList() }
+            }
+            .distinctUntilChanged()

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

@@ -0,0 +1,20 @@
+package com.x8bit.bitwarden.data.auth.repository.util
+
+import com.x8bit.bitwarden.data.auth.repository.model.Organization
+import com.x8bit.bitwarden.data.vault.datasource.network.model.SyncResponseJson
+
+/**
+ * Maps the given [SyncResponseJson.Profile.Organization] to an [Organization].
+ */
+fun SyncResponseJson.Profile.Organization.toOrganization(): Organization =
+    Organization(
+        id = this.id,
+        name = this.name,
+    )
+
+/**
+ * Maps the given list of [SyncResponseJson.Profile.Organization] to a list of
+ * [Organization]s.
+ */
+fun List<SyncResponseJson.Profile.Organization>.toOrganizations(): List<Organization> =
+    this.map { it.toOrganization() }

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

@@ -1,6 +1,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.repository.model.UserOrganizations
 import com.x8bit.bitwarden.data.auth.repository.model.UserState
 import com.x8bit.bitwarden.data.platform.repository.util.toEnvironmentUrlsOrDefault
 import com.x8bit.bitwarden.data.vault.datasource.network.model.SyncResponseJson
@@ -42,6 +43,7 @@ fun UserStateJson.toUpdatedUserStateJson(
  */
 fun UserStateJson.toUserState(
     vaultState: VaultState,
+    userOrganizationsList: List<UserOrganizations>,
     specialCircumstance: UserState.SpecialCircumstance?,
 ): UserState =
     UserState(
@@ -63,6 +65,10 @@ fun UserStateJson.toUserState(
                         .toEnvironmentUrlsOrDefault(),
                     isPremium = accountJson.profile.hasPremium == true,
                     isVaultUnlocked = userId in vaultState.unlockedVaultUserIds,
+                    organizations = userOrganizationsList
+                        .find { it.userId == userId }
+                        ?.organizations
+                        .orEmpty(),
                 )
             },
         specialCircumstance = specialCircumstance,