Add the continue button flow for TDE (#1248)

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

@@ -9,6 +9,7 @@ import com.x8bit.bitwarden.data.auth.repository.model.BreachCountResult
 import com.x8bit.bitwarden.data.auth.repository.model.DeleteAccountResult
 import com.x8bit.bitwarden.data.auth.repository.model.KnownDeviceResult
 import com.x8bit.bitwarden.data.auth.repository.model.LoginResult
+import com.x8bit.bitwarden.data.auth.repository.model.NewSsoUserResult
 import com.x8bit.bitwarden.data.auth.repository.model.OrganizationDomainSsoDetailsResult
 import com.x8bit.bitwarden.data.auth.repository.model.PasswordHintResult
 import com.x8bit.bitwarden.data.auth.repository.model.PasswordStrengthResult
@@ -123,6 +124,12 @@ interface AuthRepository : AuthenticatorProvider, AuthRequestManager {
      */
     suspend fun deleteAccount(password: String): DeleteAccountResult
 
+    /**
+     * Attempt to create a new user via SSO and log them into their account. Upon success the new
+     * user will also have the vault automatically unlocked for them.
+     */
+    suspend fun createNewSsoUser(): NewSsoUserResult
+
     /**
      * Attempt to complete the trusted device login with the given [requestPrivateKey] and
      * [asymmetricalKey]. This will unlock the vault and finish trusting the device.

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

@@ -38,6 +38,7 @@ import com.x8bit.bitwarden.data.auth.repository.model.BreachCountResult
 import com.x8bit.bitwarden.data.auth.repository.model.DeleteAccountResult
 import com.x8bit.bitwarden.data.auth.repository.model.KnownDeviceResult
 import com.x8bit.bitwarden.data.auth.repository.model.LoginResult
+import com.x8bit.bitwarden.data.auth.repository.model.NewSsoUserResult
 import com.x8bit.bitwarden.data.auth.repository.model.OrganizationDomainSsoDetailsResult
 import com.x8bit.bitwarden.data.auth.repository.model.PasswordHintResult
 import com.x8bit.bitwarden.data.auth.repository.model.PasswordStrengthResult
@@ -332,6 +333,61 @@ class AuthRepositoryImpl(
             )
     }
 
+    @Suppress("ReturnCount")
+    override suspend fun createNewSsoUser(): NewSsoUserResult {
+        val account = authDiskSource.userState?.activeAccount ?: return NewSsoUserResult.Failure
+        val orgIdentifier = rememberedOrgIdentifier ?: return NewSsoUserResult.Failure
+        val userId = account.profile.userId
+        return organizationService
+            .getOrganizationAutoEnrollStatus(orgIdentifier)
+            .flatMap { orgAutoEnrollStatus ->
+                organizationService
+                    .getOrganizationKeys(orgAutoEnrollStatus.organizationId)
+                    .flatMap { organizationKeys ->
+                        authSdkSource.makeRegisterTdeKeysAndUnlockVault(
+                            userId = userId,
+                            orgPublicKey = organizationKeys.publicKey,
+                            rememberDevice = authDiskSource.shouldTrustDevice,
+                        )
+                    }
+                    .flatMap { keys ->
+                        accountsService
+                            .createAccountKeys(
+                                publicKey = keys.publicKey,
+                                encryptedPrivateKey = keys.privateKey,
+                            )
+                            .map { keys }
+                    }
+                    .flatMap { keys ->
+                        organizationService
+                            .organizationResetPasswordEnroll(
+                                organizationId = orgAutoEnrollStatus.organizationId,
+                                userId = userId,
+                                passwordHash = null,
+                                resetPasswordKey = keys.adminReset,
+                            )
+                            .map { keys }
+                    }
+                    .onSuccess { keys ->
+                        authDiskSource.storePrivateKey(
+                            userId = userId,
+                            privateKey = keys.privateKey,
+                        )
+                        keys.deviceKey?.let { trustDeviceResponse ->
+                            trustedDeviceManager.trustThisDevice(
+                                userId = userId,
+                                trustDeviceResponse = trustDeviceResponse,
+                            )
+                        }
+                        vaultRepository.syncVaultState(userId = userId)
+                    }
+            }
+            .fold(
+                onSuccess = { NewSsoUserResult.Success },
+                onFailure = { NewSsoUserResult.Failure },
+            )
+    }
+
     @Suppress("ReturnCount")
     override suspend fun completeTdeLogin(
         requestPrivateKey: String,

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

@@ -0,0 +1,16 @@
+package com.x8bit.bitwarden.data.auth.repository.model
+
+/**
+ * Models result of creating a new user via SSO.
+ */
+sealed class NewSsoUserResult {
+    /**
+     * A new user has successfully been created.
+     */
+    data object Success : NewSsoUserResult()
+
+    /**
+     * There was an error while truing to create the new user.
+     */
+    data object Failure : NewSsoUserResult()
+}

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.ClientAuth
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockData
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockResult
 import kotlinx.coroutines.flow.StateFlow
@@ -55,4 +56,14 @@ interface VaultLockManager {
      * Suspends until the vault for the given [userId] is unlocked.
      */
     suspend fun waitUntilUnlocked(userId: String)
+
+    /**
+     * This will check the vault lock state for the given user and ensure that the
+     * [vaultUnlockDataStateFlow] is up-to-date.
+     *
+     * This is only meant to be used when the SDK unlocks the vault as a side-effect of some other
+     * function, such as [ClientAuth.makeRegisterTdeKeys]. When using the regular [unlockVault]
+     * functions, this is not necessary.
+     */
+    suspend fun syncVaultState(userId: String)
 }

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

@@ -187,6 +187,18 @@ class VaultLockManagerImpl(
             .first { unlockedUserIds -> userId in unlockedUserIds }
     }
 
+    override suspend fun syncVaultState(userId: String) {
+        // There is no proper way to query if the vault is actually unlocked or not but we can
+        // attempt to retrieve the user encryption key. If it fails, then the vault is locked and
+        // if it succeeds, then the vault is unlocked.
+        vaultSdkSource
+            .getUserEncryptionKey(userId = userId)
+            .fold(
+                onFailure = { setVaultToLocked(userId = userId) },
+                onSuccess = { setVaultToUnlocked(userId = userId) },
+            )
+    }
+
     /**
      * Increments the stored invalid unlock count for the given [userId] and automatically logs out
      * if this new value is greater than [MAXIMUM_INVALID_UNLOCK_ATTEMPTS].

app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/trusteddevice/TrustedDeviceViewModel.kt

@@ -2,13 +2,17 @@ package com.x8bit.bitwarden.ui.auth.feature.trusteddevice
 
 import android.os.Parcelable
 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.NewSsoUserResult
 import com.x8bit.bitwarden.data.platform.repository.EnvironmentRepository
 import com.x8bit.bitwarden.ui.platform.base.BaseViewModel
 import com.x8bit.bitwarden.ui.platform.base.util.Text
 import com.x8bit.bitwarden.ui.platform.base.util.asText
 import dagger.hilt.android.lifecycle.HiltViewModel
 import kotlinx.coroutines.flow.update
+import kotlinx.coroutines.launch
 import kotlinx.parcelize.Parcelize
 import javax.inject.Inject
 
@@ -17,6 +21,7 @@ private const val KEY_STATE = "state"
 /**
  * Manages application state for the Trusted Device screen.
  */
+@Suppress("TooManyFunctions")
 @HiltViewModel
 class TrustedDeviceViewModel @Inject constructor(
     savedStateHandle: SavedStateHandle,
@@ -52,6 +57,37 @@ class TrustedDeviceViewModel @Inject constructor(
             TrustedDeviceAction.ApproveWithDeviceClick -> handleApproveWithDeviceClick()
             TrustedDeviceAction.ApproveWithPasswordClick -> handleApproveWithPasswordClick()
             TrustedDeviceAction.NotYouClick -> handleNotYouClick()
+            is TrustedDeviceAction.Internal -> handleInternalAction(action)
+        }
+    }
+
+    private fun handleInternalAction(action: TrustedDeviceAction.Internal) {
+        when (action) {
+            is TrustedDeviceAction.Internal.ReceiveNewSsoUserResult -> {
+                handleReceiveNewSsoUserResult(action)
+            }
+        }
+    }
+
+    private fun handleReceiveNewSsoUserResult(
+        action: TrustedDeviceAction.Internal.ReceiveNewSsoUserResult,
+    ) {
+        when (action.result) {
+            NewSsoUserResult.Failure -> {
+                mutableStateFlow.update {
+                    it.copy(
+                        dialogState = TrustedDeviceState.DialogState.Error(
+                            title = R.string.an_error_has_occurred.asText(),
+                            message = R.string.generic_error_message.asText(),
+                        ),
+                    )
+                }
+            }
+
+            NewSsoUserResult.Success -> {
+                mutableStateFlow.update { it.copy(dialogState = null) }
+                // Should automatically navigate to a logged in state.
+            }
         }
     }
 
@@ -68,7 +104,16 @@ class TrustedDeviceViewModel @Inject constructor(
     }
 
     private fun handleContinueClick() {
-        sendEvent(TrustedDeviceEvent.ShowToast("Not yet implemented".asText()))
+        mutableStateFlow.update {
+            it.copy(
+                dialogState = TrustedDeviceState.DialogState.Loading(R.string.loading.asText()),
+            )
+        }
+        authRepository.shouldTrustDevice = state.isRemembered
+        viewModelScope.launch {
+            val result = authRepository.createNewSsoUser()
+            sendAction(TrustedDeviceAction.Internal.ReceiveNewSsoUserResult(result))
+        }
     }
 
     private fun handleApproveWithAdminClick() {
@@ -196,4 +241,16 @@ sealed class TrustedDeviceAction {
      * Indicates that the "Not you?" text was clicked.
      */
     data object NotYouClick : TrustedDeviceAction()
+
+    /**
+     * Actions for internal use by the ViewModel.
+     */
+    sealed class Internal : TrustedDeviceAction() {
+        /**
+         * Indicates a new SSO user result has been received.
+         */
+        data class ReceiveNewSsoUserResult(
+            val result: NewSsoUserResult,
+        ) : Internal()
+    }
 }