Restart activity to clear out in-memory data when locking or changing user (#1388)

app/src/main/java/com/x8bit/bitwarden/MainActivity.kt

@@ -94,6 +94,8 @@ class MainActivity : AppCompatActivity() {
                         handleCompleteAutofill(event)
                     }
 
+                    MainEvent.Recreate -> handleRecreate()
+
                     is MainEvent.ScreenCaptureSettingChange -> {
                         handleScreenCaptureSettingChange(event)
                     }
@@ -116,4 +118,8 @@ class MainActivity : AppCompatActivity() {
             window.addFlags(WindowManager.LayoutParams.FLAG_SECURE)
         }
     }
+
+    private fun handleRecreate() {
+        recreate()
+    }
 }

app/src/main/java/com/x8bit/bitwarden/MainViewModel.kt

@@ -5,20 +5,28 @@ import android.os.Parcelable
 import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.viewModelScope
 import com.bitwarden.core.CipherView
+import com.x8bit.bitwarden.data.auth.repository.AuthRepository
 import com.x8bit.bitwarden.data.auth.util.getPasswordlessRequestDataIntentOrNull
 import com.x8bit.bitwarden.data.autofill.manager.AutofillSelectionManager
 import com.x8bit.bitwarden.data.autofill.util.getAutofillSaveItemOrNull
 import com.x8bit.bitwarden.data.autofill.util.getAutofillSelectionDataOrNull
 import com.x8bit.bitwarden.data.platform.manager.SpecialCircumstanceManager
+import com.x8bit.bitwarden.data.platform.manager.garbage.GarbageCollectionManager
 import com.x8bit.bitwarden.data.platform.manager.model.SpecialCircumstance
 import com.x8bit.bitwarden.data.platform.repository.SettingsRepository
+import com.x8bit.bitwarden.data.vault.manager.model.VaultStateEvent
+import com.x8bit.bitwarden.data.vault.repository.VaultRepository
 import com.x8bit.bitwarden.ui.platform.base.BaseViewModel
 import com.x8bit.bitwarden.ui.platform.feature.settings.appearance.model.AppTheme
 import com.x8bit.bitwarden.ui.platform.manager.intent.IntentManager
 import com.x8bit.bitwarden.ui.platform.util.isMyVaultShortcut
 import com.x8bit.bitwarden.ui.platform.util.isPasswordGeneratorShortcut
 import dagger.hilt.android.lifecycle.HiltViewModel
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.drop
 import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.update
 import kotlinx.parcelize.Parcelize
@@ -29,12 +37,16 @@ private const val SPECIAL_CIRCUMSTANCE_KEY = "special-circumstance"
 /**
  * A view model that helps launch actions for the [MainActivity].
  */
+@Suppress("LongParameterList")
 @HiltViewModel
 class MainViewModel @Inject constructor(
     private val autofillSelectionManager: AutofillSelectionManager,
     private val specialCircumstanceManager: SpecialCircumstanceManager,
+    private val garbageCollectionManager: GarbageCollectionManager,
     private val intentManager: IntentManager,
+    authRepository: AuthRepository,
     settingsRepository: SettingsRepository,
+    vaultRepository: VaultRepository,
     private val savedStateHandle: SavedStateHandle,
 ) : BaseViewModel<MainState, MainEvent, MainAction>(
     MainState(
@@ -72,6 +84,36 @@ class MainViewModel @Inject constructor(
                 sendEvent(MainEvent.ScreenCaptureSettingChange(isAllowed))
             }
             .launchIn(viewModelScope)
+
+        authRepository
+            .userStateFlow
+            .drop(count = 1)
+            // Trigger an action whenever the current user changes or we go into/out of a pending
+            // account state (which acts like switching to a temporary user).
+            .map { it?.activeUserId to it?.hasPendingAccountAddition }
+            .distinctUntilChanged()
+            .onEach {
+                // Switching between account states often involves some kind of animation (ex:
+                // account switcher) that we might want to give time to finish before triggering
+                // a refresh.
+                @Suppress("MagicNumber")
+                delay(500)
+                trySendAction(MainAction.Internal.CurrentUserStateChange)
+            }
+            .launchIn(viewModelScope)
+
+        vaultRepository
+            .vaultStateEventFlow
+            .onEach {
+                when (it) {
+                    is VaultStateEvent.Locked -> {
+                        trySendAction(MainAction.Internal.VaultUnlockStateChange)
+                    }
+
+                    is VaultStateEvent.Unlocked -> Unit
+                }
+            }
+            .launchIn(viewModelScope)
     }
 
     override fun handleAction(action: MainAction) {
@@ -80,7 +122,9 @@ class MainViewModel @Inject constructor(
                 handleAutofillSelectionReceive(action)
             }
 
+            is MainAction.Internal.CurrentUserStateChange -> handleCurrentUserStateChange()
             is MainAction.Internal.ThemeUpdate -> handleAppThemeUpdated(action)
+            is MainAction.Internal.VaultUnlockStateChange -> handleVaultUnlockStateChange()
             is MainAction.ReceiveFirstIntent -> handleFirstIntentReceived(action)
             is MainAction.ReceiveNewIntent -> handleNewIntentReceived(action)
         }
@@ -92,10 +136,18 @@ class MainViewModel @Inject constructor(
         sendEvent(MainEvent.CompleteAutofill(cipherView = action.cipherView))
     }
 
+    private fun handleCurrentUserStateChange() {
+        recreateUiAndGarbageCollect()
+    }
+
     private fun handleAppThemeUpdated(action: MainAction.Internal.ThemeUpdate) {
         mutableStateFlow.update { it.copy(theme = action.theme) }
     }
 
+    private fun handleVaultUnlockStateChange() {
+        recreateUiAndGarbageCollect()
+    }
+
     private fun handleFirstIntentReceived(action: MainAction.ReceiveFirstIntent) {
         handleIntent(
             intent = action.intent,
@@ -168,6 +220,11 @@ class MainViewModel @Inject constructor(
             }
         }
     }
+
+    private fun recreateUiAndGarbageCollect() {
+        sendEvent(MainEvent.Recreate)
+        garbageCollectionManager.tryCollect()
+    }
 }
 
 /**
@@ -203,12 +260,22 @@ sealed class MainAction {
             val cipherView: CipherView,
         ) : Internal()
 
+        /**
+         * Indicates a relevant change in the current user state.
+         */
+        data object CurrentUserStateChange : Internal()
+
         /**
          * Indicates that the app theme has changed.
          */
         data class ThemeUpdate(
             val theme: AppTheme,
         ) : Internal()
+
+        /**
+         * Indicates a relevant change in the current vault lock state.
+         */
+        data object VaultUnlockStateChange : Internal()
     }
 }
 
@@ -226,4 +293,9 @@ sealed class MainEvent {
      * Event indicating a change in the screen capture setting.
      */
     data class ScreenCaptureSettingChange(val isAllowed: Boolean) : MainEvent()
+
+    /**
+     * Event indicating that the UI should recreate itself.
+     */
+    data object Recreate : MainEvent()
 }

app/src/main/java/com/x8bit/bitwarden/data/vault/manager/VaultLockManager.kt

@@ -3,8 +3,10 @@ package com.x8bit.bitwarden.data.vault.manager
 import com.bitwarden.core.InitUserCryptoMethod
 import com.bitwarden.crypto.Kdf
 import com.bitwarden.sdk.ClientAuth
+import com.x8bit.bitwarden.data.vault.manager.model.VaultStateEvent
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockData
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockResult
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
 
 /**
@@ -16,6 +18,11 @@ interface VaultLockManager {
      */
     val vaultUnlockDataStateFlow: StateFlow<List<VaultUnlockData>>
 
+    /**
+     * Flow that emits whenever any vault is locked or unlocked.
+     */
+    val vaultStateEventFlow: Flow<VaultStateEvent>
+
     /**
      * Whether or not the vault is currently locked for the given [userId].
      */

app/src/main/java/com/x8bit/bitwarden/data/vault/manager/VaultLockManagerImpl.kt

@@ -18,10 +18,12 @@ import com.x8bit.bitwarden.data.platform.manager.model.AppForegroundState
 import com.x8bit.bitwarden.data.platform.repository.SettingsRepository
 import com.x8bit.bitwarden.data.platform.repository.model.VaultTimeout
 import com.x8bit.bitwarden.data.platform.repository.model.VaultTimeoutAction
+import com.x8bit.bitwarden.data.platform.repository.util.bufferedMutableSharedFlow
 import com.x8bit.bitwarden.data.platform.util.asSuccess
 import com.x8bit.bitwarden.data.platform.util.flatMap
 import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
 import com.x8bit.bitwarden.data.vault.datasource.sdk.model.InitializeCryptoResult
+import com.x8bit.bitwarden.data.vault.manager.model.VaultStateEvent
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockData
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockResult
 import com.x8bit.bitwarden.data.vault.repository.util.statusFor
@@ -29,8 +31,10 @@ import com.x8bit.bitwarden.data.vault.repository.util.toVaultUnlockResult
 import com.x8bit.bitwarden.data.vault.repository.util.update
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asSharedFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.first
@@ -74,10 +78,14 @@ class VaultLockManagerImpl(
 
     private val mutableVaultUnlockDataStateFlow =
         MutableStateFlow<List<VaultUnlockData>>(emptyList())
+    private val mutableVaultStateEventSharedFlow = bufferedMutableSharedFlow<VaultStateEvent>()
 
     override val vaultUnlockDataStateFlow: StateFlow<List<VaultUnlockData>>
         get() = mutableVaultUnlockDataStateFlow.asStateFlow()
 
+    override val vaultStateEventFlow: Flow<VaultStateEvent>
+        get() = mutableVaultStateEventSharedFlow.asSharedFlow()
+
     init {
         observeAppForegroundChanges()
         observeUserSwitchingChanges()
@@ -230,15 +238,20 @@ class VaultLockManagerImpl(
     }
 
     private fun setVaultToUnlocked(userId: String) {
+        val wasVaultUnlocked = isVaultUnlocked(userId = userId)
         mutableVaultUnlockDataStateFlow.update {
             it.update(userId, VaultUnlockData.Status.UNLOCKED)
         }
         // If we are unlocking an account with a timeout of Never, we should make sure to store the
         // auto-unlock key.
         storeUserAutoUnlockKeyIfNecessary(userId = userId)
+        if (!wasVaultUnlocked) {
+            mutableVaultStateEventSharedFlow.tryEmit(VaultStateEvent.Unlocked(userId = userId))
+        }
     }
 
     private fun setVaultToLocked(userId: String) {
+        val wasVaultLocked = !isVaultUnlocked(userId = userId) && !isVaultUnlocking(userId = userId)
         vaultSdkSource.clearCrypto(userId = userId)
         mutableVaultUnlockDataStateFlow.update {
             it.update(userId, null)
@@ -247,6 +260,9 @@ class VaultLockManagerImpl(
             userId = userId,
             userAutoUnlockKey = null,
         )
+        if (!wasVaultLocked) {
+            mutableVaultStateEventSharedFlow.tryEmit(VaultStateEvent.Locked(userId = userId))
+        }
     }
 
     private fun setVaultToUnlocking(userId: String) {

app/src/main/java/com/x8bit/bitwarden/data/vault/manager/model/VaultStateEvent.kt

@@ -0,0 +1,16 @@
+package com.x8bit.bitwarden.data.vault.manager.model
+
+/**
+ * Represents an event that indicates if the vault for a particular user is not locked or unlocked.
+ */
+sealed class VaultStateEvent {
+    /**
+     * Indicates that the vault has been locked for the given [userId].
+     */
+    data class Locked(val userId: String) : VaultStateEvent()
+
+    /**
+     * Indicates that the vault has been unlocked for the given [userId].
+     */
+    data class Unlocked(val userId: String) : VaultStateEvent()
+}

app/src/test/java/com/x8bit/bitwarden/MainViewModelTest.kt

@@ -4,6 +4,8 @@ import android.content.Intent
 import androidx.lifecycle.SavedStateHandle
 import app.cash.turbine.test
 import com.bitwarden.core.CipherView
+import com.x8bit.bitwarden.data.auth.repository.AuthRepository
+import com.x8bit.bitwarden.data.auth.repository.model.UserState
 import com.x8bit.bitwarden.data.auth.util.getPasswordlessRequestDataIntentOrNull
 import com.x8bit.bitwarden.data.autofill.manager.AutofillSelectionManager
 import com.x8bit.bitwarden.data.autofill.manager.AutofillSelectionManagerImpl
@@ -12,17 +14,23 @@ import com.x8bit.bitwarden.data.autofill.model.AutofillSelectionData
 import com.x8bit.bitwarden.data.autofill.util.getAutofillSaveItemOrNull
 import com.x8bit.bitwarden.data.autofill.util.getAutofillSelectionDataOrNull
 import com.x8bit.bitwarden.data.platform.manager.SpecialCircumstanceManagerImpl
+import com.x8bit.bitwarden.data.platform.manager.garbage.GarbageCollectionManager
 import com.x8bit.bitwarden.data.platform.manager.model.PasswordlessRequestData
 import com.x8bit.bitwarden.data.platform.manager.model.SpecialCircumstance
 import com.x8bit.bitwarden.data.platform.repository.SettingsRepository
+import com.x8bit.bitwarden.data.platform.repository.util.bufferedMutableSharedFlow
+import com.x8bit.bitwarden.data.vault.manager.model.VaultStateEvent
+import com.x8bit.bitwarden.data.vault.repository.VaultRepository
 import com.x8bit.bitwarden.ui.platform.base.BaseViewModelTest
 import com.x8bit.bitwarden.ui.platform.feature.settings.appearance.model.AppTheme
 import com.x8bit.bitwarden.ui.platform.manager.intent.IntentManager
 import com.x8bit.bitwarden.ui.platform.util.isMyVaultShortcut
 import com.x8bit.bitwarden.ui.platform.util.isPasswordGeneratorShortcut
 import io.mockk.every
+import io.mockk.just
 import io.mockk.mockk
 import io.mockk.mockkStatic
+import io.mockk.runs
 import io.mockk.unmockkStatic
 import io.mockk.verify
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -36,6 +44,10 @@ import org.junit.jupiter.api.Test
 class MainViewModelTest : BaseViewModelTest() {
 
     private val autofillSelectionManager: AutofillSelectionManager = AutofillSelectionManagerImpl()
+    private val mutableUserStateFlow = MutableStateFlow<UserState?>(null)
+    private val authRepository = mockk<AuthRepository> {
+        every { userStateFlow } returns mutableUserStateFlow
+    }
     private val mutableAppThemeFlow = MutableStateFlow(AppTheme.DEFAULT)
     private val mutableScreenCaptureAllowedFlow = MutableStateFlow(true)
     private val settingsRepository = mockk<SettingsRepository> {
@@ -43,6 +55,13 @@ class MainViewModelTest : BaseViewModelTest() {
         every { appThemeStateFlow } returns mutableAppThemeFlow
         every { isScreenCaptureAllowedStateFlow } returns mutableScreenCaptureAllowedFlow
     }
+    private val mutableVaultStateEventFlow = bufferedMutableSharedFlow<VaultStateEvent>()
+    private val vaultRepository = mockk<VaultRepository> {
+        every { vaultStateEventFlow } returns mutableVaultStateEventFlow
+    }
+    private val garbageCollectionManager = mockk<GarbageCollectionManager> {
+        every { tryCollect() } just runs
+    }
     private val specialCircumstanceManager = SpecialCircumstanceManagerImpl()
     private val intentManager: IntentManager = mockk {
         every { getShareDataFromIntent(any()) } returns null
@@ -91,6 +110,77 @@ class MainViewModelTest : BaseViewModelTest() {
         )
     }
 
+    @Test
+    fun `user state updates should emit Recreate event and trigger garbage collection`() = runTest {
+        val userId1 = "userId1"
+        val userId2 = "userId12"
+        val viewModel = createViewModel()
+
+        viewModel.eventFlow.test {
+            // Ignore initial screen capture event
+            awaitItem()
+
+            mutableUserStateFlow.value = UserState(
+                activeUserId = userId1,
+                accounts = listOf(
+                    mockk<UserState.Account> {
+                        every { userId } returns userId1
+                    },
+                ),
+                hasPendingAccountAddition = false,
+            )
+            assertEquals(MainEvent.Recreate, awaitItem())
+
+            mutableUserStateFlow.value = UserState(
+                activeUserId = userId1,
+                accounts = listOf(
+                    mockk<UserState.Account> {
+                        every { userId } returns userId1
+                    },
+                ),
+                hasPendingAccountAddition = true,
+            )
+            assertEquals(MainEvent.Recreate, awaitItem())
+
+            mutableUserStateFlow.value = UserState(
+                activeUserId = userId2,
+                accounts = listOf(
+                    mockk<UserState.Account> {
+                        every { userId } returns userId1
+                    },
+                    mockk<UserState.Account> {
+                        every { userId } returns userId2
+                    },
+                ),
+                hasPendingAccountAddition = true,
+            )
+            assertEquals(MainEvent.Recreate, awaitItem())
+        }
+        verify(exactly = 3) {
+            garbageCollectionManager.tryCollect()
+        }
+    }
+
+    @Test
+    fun `vault state lock events should emit Recreate event and trigger garbage collection`() =
+        runTest {
+            val viewModel = createViewModel()
+
+            viewModel.eventFlow.test {
+                // Ignore initial screen capture event
+                awaitItem()
+
+                mutableVaultStateEventFlow.tryEmit(VaultStateEvent.Unlocked(userId = "userId"))
+                expectNoEvents()
+
+                mutableVaultStateEventFlow.tryEmit(VaultStateEvent.Locked(userId = "userId"))
+                assertEquals(MainEvent.Recreate, awaitItem())
+            }
+            verify(exactly = 1) {
+                garbageCollectionManager.tryCollect()
+            }
+        }
+
     @Test
     fun `autofill selection updates should emit CompleteAutofill events`() = runTest {
         val viewModel = createViewModel()
@@ -441,7 +531,10 @@ class MainViewModelTest : BaseViewModelTest() {
     ) = MainViewModel(
         autofillSelectionManager = autofillSelectionManager,
         specialCircumstanceManager = specialCircumstanceManager,
+        garbageCollectionManager = garbageCollectionManager,
+        authRepository = authRepository,
         settingsRepository = settingsRepository,
+        vaultRepository = vaultRepository,
         intentManager = intentManager,
         savedStateHandle = savedStateHandle.apply {
             set(SPECIAL_CIRCUMSTANCE_KEY, initialSpecialCircumstance)

app/src/test/java/com/x8bit/bitwarden/data/vault/manager/VaultLockManagerTest.kt

@@ -1,5 +1,6 @@
 package com.x8bit.bitwarden.data.vault.manager
 
+import app.cash.turbine.test
 import com.bitwarden.core.InitOrgCryptoRequest
 import com.bitwarden.core.InitUserCryptoMethod
 import com.bitwarden.core.InitUserCryptoRequest
@@ -22,6 +23,7 @@ import com.x8bit.bitwarden.data.platform.util.asFailure
 import com.x8bit.bitwarden.data.platform.util.asSuccess
 import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
 import com.x8bit.bitwarden.data.vault.datasource.sdk.model.InitializeCryptoResult
+import com.x8bit.bitwarden.data.vault.manager.model.VaultStateEvent
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockData
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockResult
 import io.mockk.awaits
@@ -86,6 +88,56 @@ class VaultLockManagerTest {
         elapsedRealtimeMillisProvider = { elapsedRealtimeMillis },
     )
 
+    @Test
+    fun `vaultStateEventFlow should emit Locked event when vault state changes to locked`() =
+        runTest {
+            // Ensure the vault is unlocked
+            verifyUnlockedVault(userId = USER_ID)
+
+            vaultLockManager.vaultStateEventFlow.test {
+                vaultLockManager.lockVault(userId = USER_ID)
+                assertEquals(VaultStateEvent.Locked(userId = USER_ID), awaitItem())
+            }
+        }
+
+    @Test
+    fun `vaultStateEventFlow should not emit Locked event when vault state remains locked`() =
+        runTest {
+            // Ensure the vault is locked
+            vaultLockManager.lockVault(userId = USER_ID)
+
+            vaultLockManager.vaultStateEventFlow.test {
+                vaultLockManager.lockVault(userId = USER_ID)
+                expectNoEvents()
+            }
+        }
+
+    @Test
+    fun `vaultStateEventFlow should emit Unlocked event when vault state changes to unlocked`() =
+        runTest {
+            // Ensure the vault is locked
+            vaultLockManager.lockVault(userId = USER_ID)
+
+            vaultLockManager.vaultStateEventFlow.test {
+                verifyUnlockedVault(userId = USER_ID)
+                assertEquals(VaultStateEvent.Unlocked(userId = USER_ID), awaitItem())
+            }
+        }
+
+    @Test
+    fun `vaultStateEventFlow should not emit Unlocked event when vault state remains unlocked`() =
+        runTest {
+            // Ensure the vault is unlocked
+            verifyUnlockedVault(userId = USER_ID)
+
+            vaultLockManager.vaultStateEventFlow.test {
+                // There is no great way to directly call the internal setVaultToUnlocked
+                // but that will be called internally again when syncing.
+                vaultLockManager.syncVaultState(userId = USER_ID)
+                expectNoEvents()
+            }
+        }
+
     @Test
     fun `app going into background should update the current user's last active time`() {
         fakeAuthDiskSource.userState = MOCK_USER_STATE