Update login logic to handle TDE authentication (#1234)

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

@@ -123,6 +123,15 @@ interface AuthRepository : AuthenticatorProvider, AuthRequestManager {
      */
     suspend fun deleteAccount(password: String): DeleteAccountResult
 
+    /**
+     * Attempt to complete the trusted device login with the given [requestPrivateKey] and
+     * [asymmetricalKey]. This will unlock the vault and finish trusting the device.
+     */
+    suspend fun completeTdeLogin(
+        requestPrivateKey: String,
+        asymmetricalKey: String,
+    ): LoginResult
+
     /**
      * Attempt to login with the given email and password. Updated access token will be reflected
      * in [authStateFlow].

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

@@ -8,6 +8,7 @@ import com.bitwarden.crypto.Kdf
 import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountTokensJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.ForcePasswordResetReason
+import com.x8bit.bitwarden.data.auth.datasource.disk.model.UserStateJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.DeviceDataModel
 import com.x8bit.bitwarden.data.auth.datasource.network.model.GetTokenResponseJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.IdentityTokenAuthModel
@@ -18,6 +19,7 @@ import com.x8bit.bitwarden.data.auth.datasource.network.model.RegisterResponseJs
 import com.x8bit.bitwarden.data.auth.datasource.network.model.ResendEmailRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.ResetPasswordRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.SetPasswordRequestJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.TrustedDeviceUserDecryptionOptionsJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.TwoFactorAuthMethod
 import com.x8bit.bitwarden.data.auth.datasource.network.model.TwoFactorDataModel
 import com.x8bit.bitwarden.data.auth.datasource.network.service.AccountsService
@@ -29,6 +31,7 @@ import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSource
 import com.x8bit.bitwarden.data.auth.datasource.sdk.util.toInt
 import com.x8bit.bitwarden.data.auth.datasource.sdk.util.toKdfTypeJson
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestManager
+import com.x8bit.bitwarden.data.auth.manager.TrustedDeviceManager
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
 import com.x8bit.bitwarden.data.auth.repository.model.AuthState
 import com.x8bit.bitwarden.data.auth.repository.model.BreachCountResult
@@ -115,6 +118,7 @@ class AuthRepositoryImpl(
     private val settingsRepository: SettingsRepository,
     private val vaultRepository: VaultRepository,
     private val authRequestManager: AuthRequestManager,
+    private val trustedDeviceManager: TrustedDeviceManager,
     private val userLogoutManager: UserLogoutManager,
     private val policyManager: PolicyManager,
     pushManager: PushManager,
@@ -328,6 +332,36 @@ class AuthRepositoryImpl(
             )
     }
 
+    @Suppress("ReturnCount")
+    override suspend fun completeTdeLogin(
+        requestPrivateKey: String,
+        asymmetricalKey: String,
+    ): LoginResult {
+        val profile = authDiskSource.userState?.activeAccount?.profile
+            ?: return LoginResult.Error(errorMessage = null)
+        val userId = profile.userId
+        val privateKey = authDiskSource.getPrivateKey(userId = userId)
+            ?: return LoginResult.Error(errorMessage = null)
+        vaultRepository.unlockVault(
+            userId = userId,
+            email = profile.email,
+            kdf = profile.toSdkParams(),
+            privateKey = privateKey,
+            initUserCryptoMethod = InitUserCryptoMethod.AuthRequest(
+                requestPrivateKey = requestPrivateKey,
+                method = AuthRequestMethod.UserKey(protectedUserKey = asymmetricalKey),
+            ),
+            // We can separately unlock vault for organization data after
+            // receiving the sync response if this data is currently absent.
+            organizationKeys = null,
+        )
+
+        authDiskSource.storeUserKey(userId = userId, userKey = asymmetricalKey)
+        trustedDeviceManager.trustThisDeviceIfNecessary(userId = userId)
+        vaultRepository.syncIfNecessary()
+        return LoginResult.Success
+    }
+
     override suspend fun login(
         email: String,
         password: String,
@@ -1101,6 +1135,15 @@ class AuthRepositoryImpl(
             organizationIdentifier = orgIdentifier
         }
 
+        // Handle the Trusted Device Encryption flow
+        loginResponse.userDecryptionOptions?.trustedDeviceUserDecryptionOptions?.let { options ->
+            handleLoginCommonSuccessTrustedDeviceUserDecryptionOptions(
+                trustedDeviceDecryptionOptions = options,
+                userStateJson = userStateJson,
+                privateKey = requireNotNull(loginResponse.privateKey),
+            )
+        }
+
         // Remove any cached data after successfully logging in.
         identityTokenAuthModel = null
         twoFactorResponse = null
@@ -1138,8 +1181,7 @@ class AuthRepositoryImpl(
                     )
                 }
 
-            // Cache the password to verify against any password policies
-            // after the sync completes.
+            // Cache the password to verify against any password policies after the sync completes.
             passwordToCheck = it
         }
 
@@ -1182,7 +1224,11 @@ class AuthRepositoryImpl(
             ),
         )
         authDiskSource.userState = userStateJson
-        authDiskSource.storeUserKey(userId = userId, userKey = loginResponse.key)
+        loginResponse.key?.let {
+            // Only set the value if it's present, since we may have set it already
+            // when we completed the pending admin auth request.
+            authDiskSource.storeUserKey(userId = userId, userKey = it)
+        }
         authDiskSource.storePrivateKey(userId = userId, privateKey = loginResponse.privateKey)
         settingsRepository.setDefaultsIfNecessary(userId = userId)
         vaultRepository.syncIfNecessary()
@@ -1190,6 +1236,74 @@ class AuthRepositoryImpl(
         return LoginResult.Success
     }
 
+    /**
+     * A helper method to handle the [TrustedDeviceUserDecryptionOptionsJson] specific to TDE.
+     */
+    @Suppress("ReturnCount")
+    private suspend fun handleLoginCommonSuccessTrustedDeviceUserDecryptionOptions(
+        trustedDeviceDecryptionOptions: TrustedDeviceUserDecryptionOptionsJson,
+        userStateJson: UserStateJson,
+        privateKey: String,
+    ) {
+        val userId = userStateJson.activeUserId
+        val deviceKey = authDiskSource.getDeviceKey(userId = userId)
+        if (deviceKey == null) {
+            // A null device key means this device is not trusted.
+            val pendingRequest = authDiskSource.getPendingAuthRequest(userId = userId) ?: return
+            authRequestManager
+                .getAuthRequestIfApproved(pendingRequest.requestId)
+                .getOrNull()
+                ?.let { request ->
+                    // For approved requests the key will always be present.
+                    val userKey = requireNotNull(request.key)
+                    vaultRepository.unlockVault(
+                        userId = userId,
+                        email = userStateJson.activeAccount.profile.email,
+                        kdf = userStateJson.activeAccount.profile.toSdkParams(),
+                        privateKey = privateKey,
+                        initUserCryptoMethod = InitUserCryptoMethod.AuthRequest(
+                            requestPrivateKey = pendingRequest.requestPrivateKey,
+                            method = AuthRequestMethod.UserKey(protectedUserKey = userKey),
+                        ),
+                        // We can separately unlock vault for organization data after
+                        // receiving the sync response if this data is currently absent.
+                        organizationKeys = null,
+                    )
+                    authDiskSource.storeUserKey(userId = userId, userKey = userKey)
+                    trustedDeviceManager.trustThisDeviceIfNecessary(userId = userId)
+                }
+            authDiskSource.storePendingAuthRequest(
+                userId = userId,
+                pendingAuthRequest = null,
+            )
+            return
+        }
+
+        val encryptedPrivateKey = trustedDeviceDecryptionOptions.encryptedPrivateKey
+        val encryptedUserKey = trustedDeviceDecryptionOptions.encryptedUserKey
+        if (encryptedPrivateKey == null || encryptedUserKey == null) {
+            // If we have a device key but server is missing private key and user key, we
+            // need to clear the device key and let the user go through the TDE flow again.
+            authDiskSource.storeDeviceKey(userId = userId, deviceKey = null)
+            return
+        }
+        vaultRepository.unlockVault(
+            userId = userId,
+            email = userStateJson.activeAccount.profile.email,
+            kdf = userStateJson.activeAccount.profile.toSdkParams(),
+            privateKey = privateKey,
+            initUserCryptoMethod = InitUserCryptoMethod.DeviceKey(
+                deviceKey = deviceKey,
+                protectedDevicePrivateKey = encryptedPrivateKey,
+                deviceProtectedUserKey = encryptedUserKey,
+            ),
+            // We can separately unlock vault for organization data after
+            // receiving the sync response if this data is currently absent.
+            organizationKeys = null,
+        )
+        authDiskSource.storeUserKey(userId = userId, userKey = encryptedUserKey)
+    }
+
     /**
      * A helper method that processes the [GetTokenResponseJson.TwoFactorRequired] when logging in.
      */

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

@@ -8,6 +8,7 @@ import com.x8bit.bitwarden.data.auth.datasource.network.service.IdentityService
 import com.x8bit.bitwarden.data.auth.datasource.network.service.OrganizationService
 import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSource
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestManager
+import com.x8bit.bitwarden.data.auth.manager.TrustedDeviceManager
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
 import com.x8bit.bitwarden.data.auth.repository.AuthRepositoryImpl
@@ -48,6 +49,7 @@ object AuthRepositoryModule {
         settingsRepository: SettingsRepository,
         vaultRepository: VaultRepository,
         authRequestManager: AuthRequestManager,
+        trustedDeviceManager: TrustedDeviceManager,
         userLogoutManager: UserLogoutManager,
         pushManager: PushManager,
         policyManager: PolicyManager,
@@ -65,6 +67,7 @@ object AuthRepositoryModule {
         settingsRepository = settingsRepository,
         vaultRepository = vaultRepository,
         authRequestManager = authRequestManager,
+        trustedDeviceManager = trustedDeviceManager,
         userLogoutManager = userLogoutManager,
         pushManager = pushManager,
         policyManager = policyManager,

app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/loginwithdevice/LoginWithDeviceViewModel.kt

@@ -267,9 +267,9 @@ class LoginWithDeviceViewModel @Inject constructor(
             )
         }
         viewModelScope.launch {
-            when (state.loginWithDeviceType) {
+            val result = when (state.loginWithDeviceType) {
                 LoginWithDeviceType.OTHER_DEVICE -> {
-                    val result = authRepository.login(
+                    authRepository.login(
                         email = state.emailAddress,
                         requestId = loginData.requestId,
                         accessCode = loginData.accessCode,
@@ -278,15 +278,18 @@ class LoginWithDeviceViewModel @Inject constructor(
                         masterPasswordHash = loginData.masterPasswordHash,
                         captchaToken = loginData.captchaToken,
                     )
-                    sendAction(LoginWithDeviceAction.Internal.ReceiveLoginResult(result))
                 }
 
                 LoginWithDeviceType.SSO_ADMIN_APPROVAL,
                 LoginWithDeviceType.SSO_OTHER_DEVICE,
                 -> {
-                    sendEvent(LoginWithDeviceEvent.ShowToast("Not yet implemented!"))
+                    authRepository.completeTdeLogin(
+                        requestPrivateKey = loginData.privateKey,
+                        asymmetricalKey = loginData.asymmetricalKey,
+                    )
                 }
             }
+            sendAction(LoginWithDeviceAction.Internal.ReceiveLoginResult(result))
         }
     }
 

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

@@ -13,6 +13,7 @@ import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountTokensJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.EnvironmentUrlDataJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.ForcePasswordResetReason
+import com.x8bit.bitwarden.data.auth.datasource.disk.model.PendingAuthRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.UserStateJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.util.FakeAuthDiskSource
 import com.x8bit.bitwarden.data.auth.datasource.network.model.GetTokenResponseJson
@@ -30,6 +31,7 @@ import com.x8bit.bitwarden.data.auth.datasource.network.model.RegisterResponseJs
 import com.x8bit.bitwarden.data.auth.datasource.network.model.ResendEmailRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.ResetPasswordRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.SetPasswordRequestJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.TrustedDeviceUserDecryptionOptionsJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.TwoFactorAuthMethod
 import com.x8bit.bitwarden.data.auth.datasource.network.model.TwoFactorDataModel
 import com.x8bit.bitwarden.data.auth.datasource.network.model.UserDecryptionOptionsJson
@@ -45,7 +47,9 @@ import com.x8bit.bitwarden.data.auth.datasource.sdk.model.PasswordStrength.LEVEL
 import com.x8bit.bitwarden.data.auth.datasource.sdk.model.PasswordStrength.LEVEL_3
 import com.x8bit.bitwarden.data.auth.datasource.sdk.model.PasswordStrength.LEVEL_4
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestManager
+import com.x8bit.bitwarden.data.auth.manager.TrustedDeviceManager
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
+import com.x8bit.bitwarden.data.auth.manager.model.AuthRequest
 import com.x8bit.bitwarden.data.auth.repository.model.AuthState
 import com.x8bit.bitwarden.data.auth.repository.model.BreachCountResult
 import com.x8bit.bitwarden.data.auth.repository.model.DeleteAccountResult
@@ -187,6 +191,7 @@ class AuthRepositoryTest {
         } returns "AsymmetricEncString".asSuccess()
     }
     private val authRequestManager: AuthRequestManager = mockk()
+    private val trustedDeviceManager: TrustedDeviceManager = mockk()
     private val userLogoutManager: UserLogoutManager = mockk {
         every { logout(any(), any()) } just runs
     }
@@ -219,6 +224,7 @@ class AuthRepositoryTest {
         settingsRepository = settingsRepository,
         vaultRepository = vaultRepository,
         authRequestManager = authRequestManager,
+        trustedDeviceManager = trustedDeviceManager,
         userLogoutManager = userLogoutManager,
         dispatcherManager = dispatcherManager,
         pushManager = pushManager,
@@ -727,6 +733,78 @@ class AuthRepositoryTest {
         }
     }
 
+    @Test
+    fun `completeTdeLogin without active user fails`() = runTest {
+        val requestPrivateKey = "requestPrivateKey"
+        val asymmetricalKey = "asymmetricalKey"
+        val result = repository.completeTdeLogin(
+            requestPrivateKey = requestPrivateKey,
+            asymmetricalKey = asymmetricalKey,
+        )
+        assertEquals(LoginResult.Error(errorMessage = null), result)
+    }
+
+    @Test
+    fun `completeTdeLogin without private key fails`() = runTest {
+        fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+        val requestPrivateKey = "requestPrivateKey"
+        val asymmetricalKey = "asymmetricalKey"
+        val result = repository.completeTdeLogin(
+            requestPrivateKey = requestPrivateKey,
+            asymmetricalKey = asymmetricalKey,
+        )
+        assertEquals(LoginResult.Error(errorMessage = null), result)
+    }
+
+    @Test
+    fun `completeTdeLogin should unlock the vault and return success`() = runTest {
+        val requestPrivateKey = "requestPrivateKey"
+        val asymmetricalKey = "asymmetricalKey"
+        val privateKey = "privateKey"
+        fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+        fakeAuthDiskSource.storePrivateKey(userId = USER_ID_1, privateKey = privateKey)
+        coEvery {
+            vaultRepository.unlockVault(
+                userId = USER_ID_1,
+                email = SINGLE_USER_STATE_1.activeAccount.profile.email,
+                kdf = SINGLE_USER_STATE_1.activeAccount.profile.toSdkParams(),
+                privateKey = privateKey,
+                initUserCryptoMethod = InitUserCryptoMethod.AuthRequest(
+                    requestPrivateKey = requestPrivateKey,
+                    method = AuthRequestMethod.UserKey(protectedUserKey = asymmetricalKey),
+                ),
+                organizationKeys = null,
+            )
+        } returns VaultUnlockResult.Success
+        coEvery {
+            trustedDeviceManager.trustThisDeviceIfNecessary(userId = USER_ID_1)
+        } returns true.asSuccess()
+        coEvery { vaultRepository.syncIfNecessary() } just runs
+
+        val result = repository.completeTdeLogin(
+            requestPrivateKey = requestPrivateKey,
+            asymmetricalKey = asymmetricalKey,
+        )
+
+        coVerify(exactly = 1) {
+            vaultRepository.unlockVault(
+                userId = USER_ID_1,
+                email = SINGLE_USER_STATE_1.activeAccount.profile.email,
+                kdf = SINGLE_USER_STATE_1.activeAccount.profile.toSdkParams(),
+                privateKey = privateKey,
+                initUserCryptoMethod = InitUserCryptoMethod.AuthRequest(
+                    requestPrivateKey = requestPrivateKey,
+                    method = AuthRequestMethod.UserKey(protectedUserKey = asymmetricalKey),
+                ),
+                organizationKeys = null,
+            )
+            trustedDeviceManager.trustThisDeviceIfNecessary(userId = USER_ID_1)
+            vaultRepository.syncIfNecessary()
+        }
+        fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = asymmetricalKey)
+        assertEquals(LoginResult.Success, result)
+    }
+
     @Test
     fun `login when pre login fails should return Error with no message`() = runTest {
         coEvery {
@@ -1809,6 +1887,265 @@ class AuthRepositoryTest {
             verify { settingsRepository.setDefaultsIfNecessary(userId = USER_ID_1) }
         }
 
+    @Test
+    @Suppress("MaxLineLength")
+    fun `SSO login get token succeeds with trusted device key and no keys should return Success, clear device key, update AuthState, update stored keys, and sync`() =
+        runTest {
+            val deviceKey = "deviceKey"
+            fakeAuthDiskSource.storeDeviceKey(USER_ID_1, deviceKey)
+            val successResponse = GET_TOKEN_RESPONSE_SUCCESS.copy(
+                key = null,
+                userDecryptionOptions = USER_DECRYPTION_OPTIONS,
+            )
+            coEvery {
+                identityService.getToken(
+                    email = EMAIL,
+                    authModel = IdentityTokenAuthModel.SingleSignOn(
+                        ssoCode = SSO_CODE,
+                        ssoCodeVerifier = SSO_CODE_VERIFIER,
+                        ssoRedirectUri = SSO_REDIRECT_URI,
+                    ),
+                    captchaToken = null,
+                    uniqueAppId = UNIQUE_APP_ID,
+                )
+            } returns successResponse.asSuccess()
+            coEvery { vaultRepository.syncIfNecessary() } just runs
+            every {
+                successResponse.toUserState(
+                    previousUserState = null,
+                    environmentUrlData = EnvironmentUrlDataJson.DEFAULT_US,
+                )
+            } returns SINGLE_USER_STATE_1
+
+            val result = repository.login(
+                email = EMAIL,
+                ssoCode = SSO_CODE,
+                ssoCodeVerifier = SSO_CODE_VERIFIER,
+                ssoRedirectUri = SSO_REDIRECT_URI,
+                captchaToken = null,
+                organizationIdentifier = ORGANIZATION_IDENTIFIER,
+            )
+
+            assertEquals(LoginResult.Success, result)
+            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
+            fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = "privateKey")
+            fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = null)
+            fakeAuthDiskSource.assertDeviceKey(userId = USER_ID_1, deviceKey = null)
+            assertEquals(SINGLE_USER_STATE_1, fakeAuthDiskSource.userState)
+            coVerify(exactly = 1) {
+                identityService.getToken(
+                    email = EMAIL,
+                    authModel = IdentityTokenAuthModel.SingleSignOn(
+                        ssoCode = SSO_CODE,
+                        ssoCodeVerifier = SSO_CODE_VERIFIER,
+                        ssoRedirectUri = SSO_REDIRECT_URI,
+                    ),
+                    captchaToken = null,
+                    uniqueAppId = UNIQUE_APP_ID,
+                )
+                vaultRepository.syncIfNecessary()
+            }
+            verify(exactly = 1) {
+                settingsRepository.setDefaultsIfNecessary(userId = USER_ID_1)
+            }
+        }
+
+    @Test
+    @Suppress("MaxLineLength")
+    fun `SSO login get token succeeds with trusted device key should return Success, clear device key, update AuthState, update stored keys, and sync`() =
+        runTest {
+            val deviceKey = "deviceKey"
+            val encryptedUserKey = "encryptedUserKey"
+            val encryptedPrivateKey = "encryptedPrivateKey"
+            fakeAuthDiskSource.storeDeviceKey(USER_ID_1, deviceKey)
+            val successResponse = GET_TOKEN_RESPONSE_SUCCESS.copy(
+                key = null,
+                userDecryptionOptions = USER_DECRYPTION_OPTIONS.copy(
+                    trustedDeviceUserDecryptionOptions = TRUSTED_DEVICE_DECRYPTION_OPTIONS.copy(
+                        encryptedUserKey = encryptedUserKey,
+                        encryptedPrivateKey = encryptedPrivateKey,
+                    ),
+                ),
+            )
+            coEvery {
+                vaultRepository.unlockVault(
+                    userId = USER_ID_1,
+                    email = SINGLE_USER_STATE_1.activeAccount.profile.email,
+                    kdf = SINGLE_USER_STATE_1.activeAccount.profile.toSdkParams(),
+                    privateKey = requireNotNull(successResponse.privateKey),
+                    initUserCryptoMethod = InitUserCryptoMethod.DeviceKey(
+                        deviceKey = deviceKey,
+                        protectedDevicePrivateKey = encryptedPrivateKey,
+                        deviceProtectedUserKey = encryptedUserKey,
+                    ),
+                    organizationKeys = null,
+                )
+            } returns VaultUnlockResult.Success
+            coEvery {
+                identityService.getToken(
+                    email = EMAIL,
+                    authModel = IdentityTokenAuthModel.SingleSignOn(
+                        ssoCode = SSO_CODE,
+                        ssoCodeVerifier = SSO_CODE_VERIFIER,
+                        ssoRedirectUri = SSO_REDIRECT_URI,
+                    ),
+                    captchaToken = null,
+                    uniqueAppId = UNIQUE_APP_ID,
+                )
+            } returns successResponse.asSuccess()
+            coEvery { vaultRepository.syncIfNecessary() } just runs
+            every {
+                successResponse.toUserState(
+                    previousUserState = null,
+                    environmentUrlData = EnvironmentUrlDataJson.DEFAULT_US,
+                )
+            } returns SINGLE_USER_STATE_1
+
+            val result = repository.login(
+                email = EMAIL,
+                ssoCode = SSO_CODE,
+                ssoCodeVerifier = SSO_CODE_VERIFIER,
+                ssoRedirectUri = SSO_REDIRECT_URI,
+                captchaToken = null,
+                organizationIdentifier = ORGANIZATION_IDENTIFIER,
+            )
+
+            assertEquals(LoginResult.Success, result)
+            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
+            fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = "privateKey")
+            fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = encryptedUserKey)
+            fakeAuthDiskSource.assertDeviceKey(userId = USER_ID_1, deviceKey = deviceKey)
+            assertEquals(SINGLE_USER_STATE_1, fakeAuthDiskSource.userState)
+            coVerify(exactly = 1) {
+                identityService.getToken(
+                    email = EMAIL,
+                    authModel = IdentityTokenAuthModel.SingleSignOn(
+                        ssoCode = SSO_CODE,
+                        ssoCodeVerifier = SSO_CODE_VERIFIER,
+                        ssoRedirectUri = SSO_REDIRECT_URI,
+                    ),
+                    captchaToken = null,
+                    uniqueAppId = UNIQUE_APP_ID,
+                )
+                vaultRepository.unlockVault(
+                    userId = USER_ID_1,
+                    email = SINGLE_USER_STATE_1.activeAccount.profile.email,
+                    kdf = SINGLE_USER_STATE_1.activeAccount.profile.toSdkParams(),
+                    privateKey = requireNotNull(successResponse.privateKey),
+                    initUserCryptoMethod = InitUserCryptoMethod.DeviceKey(
+                        deviceKey = deviceKey,
+                        protectedDevicePrivateKey = encryptedPrivateKey,
+                        deviceProtectedUserKey = encryptedUserKey,
+                    ),
+                    organizationKeys = null,
+                )
+                vaultRepository.syncIfNecessary()
+            }
+            verify(exactly = 1) {
+                settingsRepository.setDefaultsIfNecessary(userId = USER_ID_1)
+            }
+        }
+
+    @Test
+    @Suppress("MaxLineLength")
+    fun `SSO login get token succeeds without trusted device key should return Success, unlock the vault with pending request, update AuthState, update stored keys, and sync`() =
+        runTest {
+            val pendingAuthRequest = PendingAuthRequestJson(
+                requestId = "requestId",
+                requestPrivateKey = "requestPrivateKey",
+            )
+            val successResponse = GET_TOKEN_RESPONSE_SUCCESS.copy(
+                key = null,
+                userDecryptionOptions = USER_DECRYPTION_OPTIONS,
+            )
+            val authRequestKey = "key"
+            val authRequest = mockk<AuthRequest> {
+                every { this@mockk.key } returns authRequestKey
+            }
+            coEvery {
+                authRequestManager.getAuthRequestIfApproved(pendingAuthRequest.requestId)
+            } returns authRequest.asSuccess()
+            fakeAuthDiskSource.storePendingAuthRequest(USER_ID_1, pendingAuthRequest)
+            coEvery {
+                vaultRepository.unlockVault(
+                    userId = USER_ID_1,
+                    email = SINGLE_USER_STATE_1.activeAccount.profile.email,
+                    kdf = SINGLE_USER_STATE_1.activeAccount.profile.toSdkParams(),
+                    privateKey = requireNotNull(successResponse.privateKey),
+                    initUserCryptoMethod = InitUserCryptoMethod.AuthRequest(
+                        requestPrivateKey = pendingAuthRequest.requestPrivateKey,
+                        method = AuthRequestMethod.UserKey(protectedUserKey = authRequestKey),
+                    ),
+                    organizationKeys = null,
+                )
+            } returns VaultUnlockResult.Success
+            coEvery {
+                identityService.getToken(
+                    email = EMAIL,
+                    authModel = IdentityTokenAuthModel.SingleSignOn(
+                        ssoCode = SSO_CODE,
+                        ssoCodeVerifier = SSO_CODE_VERIFIER,
+                        ssoRedirectUri = SSO_REDIRECT_URI,
+                    ),
+                    captchaToken = null,
+                    uniqueAppId = UNIQUE_APP_ID,
+                )
+            } returns successResponse.asSuccess()
+            coEvery { vaultRepository.syncIfNecessary() } just runs
+            coEvery {
+                trustedDeviceManager.trustThisDeviceIfNecessary(userId = USER_ID_1)
+            } returns true.asSuccess()
+            every {
+                successResponse.toUserState(
+                    previousUserState = null,
+                    environmentUrlData = EnvironmentUrlDataJson.DEFAULT_US,
+                )
+            } returns SINGLE_USER_STATE_1
+
+            val result = repository.login(
+                email = EMAIL,
+                ssoCode = SSO_CODE,
+                ssoCodeVerifier = SSO_CODE_VERIFIER,
+                ssoRedirectUri = SSO_REDIRECT_URI,
+                captchaToken = null,
+                organizationIdentifier = ORGANIZATION_IDENTIFIER,
+            )
+
+            assertEquals(LoginResult.Success, result)
+            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
+            fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = "privateKey")
+            fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = authRequestKey)
+            assertEquals(SINGLE_USER_STATE_1, fakeAuthDiskSource.userState)
+            coVerify(exactly = 1) {
+                identityService.getToken(
+                    email = EMAIL,
+                    authModel = IdentityTokenAuthModel.SingleSignOn(
+                        ssoCode = SSO_CODE,
+                        ssoCodeVerifier = SSO_CODE_VERIFIER,
+                        ssoRedirectUri = SSO_REDIRECT_URI,
+                    ),
+                    captchaToken = null,
+                    uniqueAppId = UNIQUE_APP_ID,
+                )
+                vaultRepository.unlockVault(
+                    userId = USER_ID_1,
+                    email = SINGLE_USER_STATE_1.activeAccount.profile.email,
+                    kdf = SINGLE_USER_STATE_1.activeAccount.profile.toSdkParams(),
+                    privateKey = requireNotNull(successResponse.privateKey),
+                    initUserCryptoMethod = InitUserCryptoMethod.AuthRequest(
+                        requestPrivateKey = pendingAuthRequest.requestPrivateKey,
+                        method = AuthRequestMethod.UserKey(protectedUserKey = authRequestKey),
+                    ),
+                    organizationKeys = null,
+                )
+                trustedDeviceManager.trustThisDeviceIfNecessary(userId = USER_ID_1)
+                vaultRepository.syncIfNecessary()
+            }
+            verify(exactly = 1) {
+                settingsRepository.setDefaultsIfNecessary(userId = USER_ID_1)
+            }
+        }
+
     @Suppress("MaxLineLength")
     @Test
     fun `SSO login get token succeeds when there is an existing user should switch to the new logged in user`() =
@@ -3665,6 +4002,18 @@ class AuthRepositoryTest {
             refreshToken = REFRESH_TOKEN_2,
             tokenType = "Bearer",
         )
+        private val TRUSTED_DEVICE_DECRYPTION_OPTIONS = TrustedDeviceUserDecryptionOptionsJson(
+            encryptedPrivateKey = null,
+            encryptedUserKey = null,
+            hasAdminApproval = false,
+            hasLoginApprovingDevice = false,
+            hasManageResetPasswordPermission = false,
+        )
+        private val USER_DECRYPTION_OPTIONS = UserDecryptionOptionsJson(
+            hasMasterPassword = false,
+            trustedDeviceUserDecryptionOptions = TRUSTED_DEVICE_DECRYPTION_OPTIONS,
+            keyConnectorUserDecryptionOptions = null,
+        )
         private val GET_TOKEN_RESPONSE_SUCCESS = GetTokenResponseJson.Success(
             accessToken = ACCESS_TOKEN,
             refreshToken = "refreshToken",

app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/loginwithdevice/LoginWithDeviceViewModelTest.kt

@@ -208,6 +208,13 @@ class LoginWithDeviceViewModelTest : BaseViewModelTest() {
     @Test
     fun `on createAuthRequestWithUpdates Success with SSO_ADMIN_APPROVAL should emit toast`() =
         runTest {
+            coEvery {
+                authRepository.completeTdeLogin(
+                    asymmetricalKey = DEFAULT_LOGIN_DATA.asymmetricalKey,
+                    requestPrivateKey = DEFAULT_LOGIN_DATA.privateKey,
+                )
+            } returns LoginResult.Success
+
             val initialViewState = DEFAULT_CONTENT_VIEW_STATE.copy(
                 loginWithDeviceType = LoginWithDeviceType.SSO_ADMIN_APPROVAL,
             )
@@ -217,8 +224,8 @@ class LoginWithDeviceViewModelTest : BaseViewModelTest() {
             )
             val viewModel = createViewModel(initialState)
 
-            viewModel.stateEventFlow(backgroundScope) { stateFlow, eventFlow ->
-                assertEquals(initialState, stateFlow.awaitItem())
+            viewModel.stateFlow.test {
+                assertEquals(initialState, awaitItem())
                 mutableCreateAuthRequestWithUpdatesFlow.tryEmit(
                     CreateAuthRequestResult.Success(AUTH_REQUEST, AUTH_REQUEST_RESPONSE),
                 )
@@ -232,12 +239,24 @@ class LoginWithDeviceViewModelTest : BaseViewModelTest() {
                         ),
                         loginData = DEFAULT_LOGIN_DATA,
                     ),
-                    stateFlow.awaitItem(),
+                    awaitItem(),
                 )
-
                 assertEquals(
-                    LoginWithDeviceEvent.ShowToast("Not yet implemented!"),
-                    eventFlow.awaitItem(),
+                    initialState.copy(
+                        viewState = initialViewState.copy(
+                            fingerprintPhrase = "",
+                        ),
+                        dialogState = null,
+                        loginData = DEFAULT_LOGIN_DATA,
+                    ),
+                    awaitItem(),
+                )
+            }
+
+            coVerify(exactly = 1) {
+                authRepository.completeTdeLogin(
+                    asymmetricalKey = DEFAULT_LOGIN_DATA.asymmetricalKey,
+                    requestPrivateKey = DEFAULT_LOGIN_DATA.privateKey,
                 )
             }
         }