Save the current environment to UserStateJson when logging in (#260)

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

@@ -22,6 +22,7 @@ import com.x8bit.bitwarden.data.auth.repository.util.toUserState
 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
+import com.x8bit.bitwarden.data.platform.repository.EnvironmentRepository
 import com.x8bit.bitwarden.data.platform.util.asFailure
 import com.x8bit.bitwarden.data.platform.util.asSuccess
 import com.x8bit.bitwarden.data.platform.util.flatMap
@@ -47,6 +48,7 @@ class AuthRepositoryImpl constructor(
     private val identityService: IdentityService,
     private val authSdkSource: AuthSdkSource,
     private val authDiskSource: AuthDiskSource,
+    private val environmentRepository: EnvironmentRepository,
     private val vaultRepository: VaultRepository,
     dispatcherManager: DispatcherManager,
 ) : AuthRepository {
@@ -126,6 +128,9 @@ class AuthRepositoryImpl constructor(
                         authDiskSource.userState = it
                             .toUserState(
                                 previousUserState = authDiskSource.userState,
+                                environmentUrlData = environmentRepository
+                                    .environment
+                                    .environmentUrlData,
                             )
                             .also { userState ->
                                 authDiskSource.storeUserKey(

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

@@ -8,6 +8,7 @@ import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSource
 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
 import com.x8bit.bitwarden.data.auth.repository.AuthRepositoryImpl
 import com.x8bit.bitwarden.data.platform.manager.dispatcher.DispatcherManager
+import com.x8bit.bitwarden.data.platform.repository.EnvironmentRepository
 import com.x8bit.bitwarden.data.vault.repository.VaultRepository
 import dagger.Module
 import dagger.Provides
@@ -32,6 +33,7 @@ object AuthRepositoryModule {
         authSdkSource: AuthSdkSource,
         authDiskSource: AuthDiskSource,
         dispatchers: DispatcherManager,
+        environmentRepository: EnvironmentRepository,
         vaultRepository: VaultRepository,
     ): AuthRepository = AuthRepositoryImpl(
         accountsService = accountsService,
@@ -40,6 +42,7 @@ object AuthRepositoryModule {
         authDiskSource = authDiskSource,
         haveIBeenPwnedService = haveIBeenPwnedService,
         dispatcherManager = dispatchers,
+        environmentRepository = environmentRepository,
         vaultRepository = vaultRepository,
     )
 }

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

@@ -1,15 +1,20 @@
 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.EnvironmentUrlDataJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.UserStateJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.GetTokenResponseJson
 
 /**
- * Converts the given [GetTokenResponseJson.Success] to a [UserStateJson], given the current
- * [previousUserState].
+ * Converts the given [GetTokenResponseJson.Success] to a [UserStateJson], given the following
+ * additional information:
+ *
+ * - the [previousUserState]
+ * - the current [environmentUrlData]
  */
 fun GetTokenResponseJson.Success.toUserState(
     previousUserState: UserStateJson?,
+    environmentUrlData: EnvironmentUrlDataJson,
 ): UserStateJson {
     val accessToken = this.accessToken
 
@@ -40,7 +45,7 @@ fun GetTokenResponseJson.Success.toUserState(
             refreshToken = this.refreshToken,
         ),
         settings = AccountJson.Settings(
-            environmentUrlData = null,
+            environmentUrlData = environmentUrlData,
         ),
     )
 

app/src/main/java/com/x8bit/bitwarden/data/platform/repository/EnvironmentRepositoryImpl.kt

@@ -1,5 +1,6 @@
 package com.x8bit.bitwarden.data.platform.repository
 
+import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.EnvironmentUrlDataJson
 import com.x8bit.bitwarden.data.platform.datasource.disk.EnvironmentDiskSource
 import com.x8bit.bitwarden.data.platform.manager.dispatcher.DispatcherManager
@@ -8,7 +9,9 @@ import com.x8bit.bitwarden.data.platform.repository.util.toEnvironmentUrls
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.stateIn
 
 /**
@@ -16,6 +19,7 @@ import kotlinx.coroutines.flow.stateIn
  */
 class EnvironmentRepositoryImpl(
     private val environmentDiskSource: EnvironmentDiskSource,
+    private val authDiskSource: AuthDiskSource,
     dispatcherManager: DispatcherManager,
 ) : EnvironmentRepository {
 
@@ -38,6 +42,20 @@ class EnvironmentRepositoryImpl(
                 started = SharingStarted.Lazily,
                 initialValue = Environment.Us,
             )
+
+    init {
+        authDiskSource
+            .userStateFlow
+            .onEach { userState ->
+                // If the active account has environment data, set that as the current value.
+                userState
+                    ?.activeAccount
+                    ?.settings
+                    ?.environmentUrlData
+                    ?.let { environmentDiskSource.preAuthEnvironmentUrlData = it }
+            }
+            .launchIn(scope)
+    }
 }
 
 /**

app/src/main/java/com/x8bit/bitwarden/data/platform/repository/di/PlatformRepositoryModule.kt

@@ -1,5 +1,6 @@
 package com.x8bit.bitwarden.data.platform.repository.di
 
+import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
 import com.x8bit.bitwarden.data.platform.datasource.disk.EnvironmentDiskSource
 import com.x8bit.bitwarden.data.platform.manager.dispatcher.DispatcherManager
 import com.x8bit.bitwarden.data.platform.repository.EnvironmentRepository
@@ -21,10 +22,12 @@ object PlatformRepositoryModule {
     @Singleton
     fun provideEnvironmentRepository(
         environmentDiskSource: EnvironmentDiskSource,
+        authDiskSource: AuthDiskSource,
         dispatcherManager: DispatcherManager,
     ): EnvironmentRepository =
         EnvironmentRepositoryImpl(
             environmentDiskSource = environmentDiskSource,
+            authDiskSource = authDiskSource,
             dispatcherManager = dispatcherManager,
         )
 }