PM-20552: Ensure userState does not emit while the active user is unlocking (#5112)

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

@@ -283,6 +283,7 @@ class AuthRepositoryImpl(
         merge(
             mutableHasPendingAccountDeletionStateFlow,
             mutableUserStateTransactionCountStateFlow,
+            vaultRepository.isActiveUserUnlockingFlow,
         ),
     ) { array ->
         val userStateJson = array[0] as UserStateJson?
@@ -306,8 +307,11 @@ class AuthRepositoryImpl(
             firstTimeState = firstTimeState,
         )
     }
-        .filterNot { mutableHasPendingAccountDeletionStateFlow.value }
-        .filterNot { mutableUserStateTransactionCountStateFlow.value > 0 }
+        .filterNot {
+            mutableHasPendingAccountDeletionStateFlow.value ||
+                mutableUserStateTransactionCountStateFlow.value > 0 ||
+                vaultRepository.isActiveUserUnlockingFlow.value
+        }
         .stateIn(
             scope = unconfinedScope,
             started = SharingStarted.Eagerly,

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

@@ -2,6 +2,7 @@ package com.x8bit.bitwarden.data.vault.manager
 
 import com.bitwarden.core.InitUserCryptoMethod
 import com.bitwarden.crypto.Kdf
+import com.bitwarden.sdk.AuthClient
 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
@@ -17,6 +18,11 @@ interface VaultLockManager {
      */
     val vaultUnlockDataStateFlow: StateFlow<List<VaultUnlockData>>
 
+    /**
+     * Flow that indicates whether the active user is actively unlocking the vault.
+     */
+    val isActiveUserUnlockingFlow: StateFlow<Boolean>
+
     /**
      * Flow that emits whenever any vault is locked or unlocked.
      */

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

@@ -19,6 +19,7 @@ import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSource
 import com.x8bit.bitwarden.data.auth.manager.TrustedDeviceManager
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
 import com.x8bit.bitwarden.data.auth.repository.model.LogoutReason
+import com.x8bit.bitwarden.data.auth.repository.util.activeUserIdChangesFlow
 import com.x8bit.bitwarden.data.auth.repository.util.toSdkParams
 import com.x8bit.bitwarden.data.auth.repository.util.userAccountTokens
 import com.x8bit.bitwarden.data.auth.repository.util.userSwitchingChangesFlow
@@ -41,9 +42,11 @@ import com.x8bit.bitwarden.data.vault.repository.util.update
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.Job
+import kotlinx.coroutines.NonCancellable
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asSharedFlow
 import kotlinx.coroutines.flow.asStateFlow
@@ -56,8 +59,10 @@ import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.merge
 import kotlinx.coroutines.flow.onCompletion
 import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 import java.time.Clock
 import kotlin.time.Duration.Companion.minutes
 
@@ -99,6 +104,26 @@ class VaultLockManagerImpl(
     override val vaultUnlockDataStateFlow: StateFlow<List<VaultUnlockData>>
         get() = mutableVaultUnlockDataStateFlow.asStateFlow()
 
+    @OptIn(ExperimentalCoroutinesApi::class)
+    override val isActiveUserUnlockingFlow: StateFlow<Boolean>
+        get() = authDiskSource
+            .activeUserIdChangesFlow
+            .flatMapLatest { activeUserId ->
+                vaultUnlockDataStateFlow.map { vaultUnlockData ->
+                    vaultUnlockData.any {
+                        it.userId == activeUserId && it.status == VaultUnlockData.Status.UNLOCKING
+                    }
+                }
+            }
+            .distinctUntilChanged()
+            .stateIn(
+                scope = unconfinedScope,
+                started = SharingStarted.Lazily,
+                initialValue = mutableVaultUnlockDataStateFlow.value.any {
+                    it.userId == activeUserId && it.status == VaultUnlockData.Status.UNLOCKING
+                },
+            )
+
     override val vaultStateEventFlow: Flow<VaultStateEvent>
         get() = mutableVaultStateEventSharedFlow.asSharedFlow()
 
@@ -144,7 +169,7 @@ class VaultLockManagerImpl(
         privateKey: String,
         initUserCryptoMethod: InitUserCryptoMethod,
         organizationKeys: Map<String, String>?,
-    ): VaultUnlockResult =
+    ): VaultUnlockResult = withContext(context = NonCancellable) {
         flow {
             setVaultToUnlocking(userId = userId)
             emit(
@@ -202,10 +227,9 @@ class VaultLockManagerImpl(
                                 .also {
                                     if (it is VaultUnlockResult.Success) {
                                         clearInvalidUnlockCount(userId = userId)
-                                        setVaultToUnlocked(userId = userId)
-                                        trustedDeviceManager.trustThisDeviceIfNecessary(
-                                            userId = userId,
-                                        )
+                                        trustedDeviceManager
+                                            .trustThisDeviceIfNecessary(userId = userId)
+                                            .also { setVaultToUnlocked(userId = userId) }
                                     } else {
                                         incrementInvalidUnlockCount(userId = userId)
                                     }
@@ -216,6 +240,7 @@ class VaultLockManagerImpl(
         }
             .onCompletion { setVaultToNotUnlocking(userId = userId) }
             .first()
+    }
 
     override suspend fun waitUntilUnlocked(userId: String) {
         vaultUnlockDataStateFlow

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

@@ -170,9 +170,11 @@ class AuthRepositoryTest {
     private val haveIBeenPwnedService: HaveIBeenPwnedService = mockk()
     private val organizationService: OrganizationService = mockk()
     private val mutableVaultUnlockDataStateFlow = MutableStateFlow(VAULT_UNLOCK_DATA)
+    private val mutableIsActiveUserUnlockingFlow = MutableStateFlow(false)
     private val vaultRepository: VaultRepository = mockk {
         every { vaultUnlockDataStateFlow } returns mutableVaultUnlockDataStateFlow
         every { deleteVaultData(any()) } just runs
+        every { isActiveUserUnlockingFlow } returns mutableIsActiveUserUnlockingFlow
     }
     private val fakeAuthDiskSource = FakeAuthDiskSource()
     private val fakeEnvironmentRepository =

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

@@ -1,5 +1,6 @@
 package com.x8bit.bitwarden.data.vault.manager
 
+import android.annotation.SuppressLint
 import android.content.BroadcastReceiver
 import android.content.Context
 import android.content.Intent
@@ -32,7 +33,6 @@ import com.x8bit.bitwarden.data.vault.datasource.sdk.model.InitializeCryptoResul
 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
 import io.mockk.clearMocks
 import io.mockk.coEvery
 import io.mockk.coVerify
@@ -44,6 +44,7 @@ import io.mockk.slot
 import io.mockk.verify
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.async
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.asSharedFlow
@@ -64,6 +65,7 @@ import java.time.ZonedDateTime
 class VaultLockManagerTest {
     private val broadcastReceiver = slot<BroadcastReceiver>()
     private val context: Context = mockk {
+        @SuppressLint("UnspecifiedRegisterReceiverFlag")
         every { registerReceiver(capture(broadcastReceiver), any()) } returns null
     }
     private val fakeAuthDiskSource = FakeAuthDiskSource()
@@ -115,6 +117,7 @@ class VaultLockManagerTest {
     @Test
     fun `broadcast receiver should be registered on initialization`() {
         verify(exactly = 1) {
+            @SuppressLint("UnspecifiedRegisterReceiverFlag")
             context.registerReceiver(any(), any())
         }
     }
@@ -170,7 +173,7 @@ class VaultLockManagerTest {
                 assertEquals(VaultStateEvent.Locked(userId = USER_ID), awaitItem())
                 fakeAuthDiskSource.assertLastLockTimestamp(
                     userId = USER_ID,
-                    FIXED_CLOCK.instant(),
+                    expectedValue = FIXED_CLOCK.instant(),
                 )
             }
         }
@@ -213,6 +216,20 @@ class VaultLockManagerTest {
             }
         }
 
+    @Test
+    fun `isActiveUserUnlockingFlow should emit according to the current lock state`() = runTest {
+        // Ensure the vault is unlocked
+        fakeAuthDiskSource.userState = MOCK_USER_STATE
+        vaultLockManager.isActiveUserUnlockingFlow.test {
+            assertFalse(awaitItem())
+            verifyUnlockedVault(userId = USER_ID)
+            assertTrue(awaitItem())
+            assertFalse(awaitItem())
+            vaultLockManager.lockVault(userId = USER_ID, isUserInitiated = false)
+            expectNoEvents()
+        }
+    }
+
     @Test
     fun `app coming into background subsequent times should perform timeout action if necessary`() {
         setAccountTokens()
@@ -774,16 +791,15 @@ class VaultLockManagerTest {
         runTest {
             assertFalse(vaultLockManager.isVaultUnlocking(userId = USER_ID))
 
-            val unlockingJob = async {
-                verifyUnlockingVault(userId = USER_ID)
-            }
-            this.testScheduler.advanceUntilIdle()
+            // The async call will hang for 500ms
+            async { verifyUnlockingVault(userId = USER_ID) }
 
+            // We fast-forward 300ms, enough that the vault should be unlocking
+            this.testScheduler.advanceTimeBy(delayTimeMillis = 300L)
             assertTrue(vaultLockManager.isVaultUnlocking(userId = USER_ID))
 
-            unlockingJob.cancel()
-            this.testScheduler.advanceUntilIdle()
-
+            // We fast-forward another 300ms, enough that the vault should be done unlocking
+            this.testScheduler.advanceTimeBy(delayTimeMillis = 300L)
             assertFalse(vaultLockManager.isVaultUnlocking(userId = USER_ID))
         }
 
@@ -1560,7 +1576,7 @@ class VaultLockManagerTest {
 
     /**
      * Helper to ensures that the vault for the user with the given [userId] is actively unlocking.
-     * Note that this call will actively hang.
+     * Note that this call will delay for 500 ms.
      */
     private suspend fun verifyUnlockingVault(userId: String) {
         val kdf = MOCK_PROFILE.toSdkParams()
@@ -1582,7 +1598,10 @@ class VaultLockManagerTest {
                     ),
                 ),
             )
-        } just awaits
+        } coAnswers {
+            delay(timeMillis = 500L)
+            InitializeCryptoResult.AuthenticationError(error = Throwable()).asSuccess()
+        }
 
         vaultLockManager.unlockVault(
             userId = userId,