From b81b79e21564d9ac19c60dca06d3e6fc63b82703 Mon Sep 17 00:00:00 2001
From: David Perez <david@livefront.com>
Date: Wed, 10 Apr 2024 18:10:54 -0500
Subject: [PATCH] BIT-2236: Trusted Device screen logs out when unauthenticated
 (#1253)

---
 .../trusteddevice/TrustedDeviceViewModel.kt   |  4 ++-
 .../TrustedDeviceViewModelTest.kt             | 25 +++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/trusteddevice/TrustedDeviceViewModel.kt b/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/trusteddevice/TrustedDeviceViewModel.kt
index f5b5bc6d7..231f710a2 100644
--- a/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/trusteddevice/TrustedDeviceViewModel.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/trusteddevice/TrustedDeviceViewModel.kt
@@ -5,6 +5,7 @@ import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.viewModelScope
 import com.x8bit.bitwarden.R
 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
+import com.x8bit.bitwarden.data.auth.repository.model.AuthState
 import com.x8bit.bitwarden.data.auth.repository.model.NewSsoUserResult
 import com.x8bit.bitwarden.data.platform.repository.EnvironmentRepository
 import com.x8bit.bitwarden.ui.platform.base.BaseViewModel
@@ -30,9 +31,10 @@ class TrustedDeviceViewModel @Inject constructor(
 ) : BaseViewModel<TrustedDeviceState, TrustedDeviceEvent, TrustedDeviceAction>(
     initialState = savedStateHandle[KEY_STATE]
         ?: run {
+            val isAuthenticated = authRepository.authStateFlow.value is AuthState.Authenticated
             val account = authRepository.userStateFlow.value?.activeAccount
             val trustedDevice = account?.trustedDevice
-            if (trustedDevice == null) authRepository.logout()
+            if (trustedDevice == null || !isAuthenticated) authRepository.logout()
             TrustedDeviceState(
                 dialogState = null,
                 emailAddress = account?.email.orEmpty(),
diff --git a/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/trusteddevice/TrustedDeviceViewModelTest.kt b/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/trusteddevice/TrustedDeviceViewModelTest.kt
index 4288d0ee6..2c60f9ae5 100644
--- a/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/trusteddevice/TrustedDeviceViewModelTest.kt
+++ b/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/trusteddevice/TrustedDeviceViewModelTest.kt
@@ -4,6 +4,7 @@ import androidx.lifecycle.SavedStateHandle
 import app.cash.turbine.test
 import com.x8bit.bitwarden.R
 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
+import com.x8bit.bitwarden.data.auth.repository.model.AuthState
 import com.x8bit.bitwarden.data.auth.repository.model.NewSsoUserResult
 import com.x8bit.bitwarden.data.auth.repository.model.UserState
 import com.x8bit.bitwarden.data.platform.repository.EnvironmentRepository
@@ -25,13 +26,35 @@ import org.junit.jupiter.api.Test
 
 class TrustedDeviceViewModelTest : BaseViewModelTest() {
 
+    private val mutableAuthStateFlow = MutableStateFlow(DEFAULT_AUTH_STATE)
     private val mutableUserStateFlow = MutableStateFlow<UserState?>(DEFAULT_USER_STATE)
     private val authRepository: AuthRepository = mockk {
+        every { authStateFlow } returns mutableAuthStateFlow
         every { userStateFlow } returns mutableUserStateFlow
         every { logout() } just runs
     }
     private val environmentRepo: FakeEnvironmentRepository = FakeEnvironmentRepository()
 
+    @Test
+    fun `on init should logout when Uninitialized`() {
+        mutableAuthStateFlow.value = AuthState.Uninitialized
+        createViewModel()
+
+        verify(exactly = 1) {
+            authRepository.logout()
+        }
+    }
+
+    @Test
+    fun `on init should logout when unauthenticated`() {
+        mutableAuthStateFlow.value = AuthState.Unauthenticated
+        createViewModel()
+
+        verify(exactly = 1) {
+            authRepository.logout()
+        }
+    }
+
     @Test
     fun `on init should logout when trusted device is not present`() {
         mutableUserStateFlow.value = DEFAULT_USER_STATE.copy(
@@ -222,6 +245,8 @@ private val DEFAULT_STATE: TrustedDeviceState = TrustedDeviceState(
     showMasterPasswordButton = false,
 )
 
+private val DEFAULT_AUTH_STATE: AuthState = AuthState.Authenticated(accessToken = "accessToken")
+
 private val TRUSTED_DEVICE = UserState.TrustedDevice(
     isDeviceTrusted = false,
     hasMasterPassword = false,