feat: Restore crypto state after login (WPB-23555)

app/src/main/kotlin/com/wire/android/di/accountScoped/BackupModule.kt

@@ -67,4 +67,14 @@ class BackupModule {
     @Provides
     fun provideOnboardingBackupUseCase(backupScope: BackupScope) =
         backupScope.createUnEncryptedCopy
+
+    @ViewModelScoped
+    @Provides
+    fun provideBackupAndUploadCryptoState(backupScope: BackupScope) =
+        backupScope.backupAndUploadCryptoState
+
+    @ViewModelScoped
+    @Provides
+    fun provideSetLastDeviceIdUseCase(backupScope: BackupScope) =
+        backupScope.setLastDeviceId
 }

app/src/main/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOViewModel.kt

@@ -27,6 +27,7 @@ import androidx.compose.runtime.setValue
 import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.viewModelScope
 import com.ramcosta.composedestinations.generated.app.navArgs
+import com.wire.android.appLogger
 import com.wire.android.config.DefaultServerConfig
 import com.wire.android.datastore.UserDataStoreProvider
 import com.wire.android.di.ClientScopeProvider
@@ -40,25 +41,31 @@ import com.wire.android.ui.common.dialogs.CustomServerDetailsDialogState
 import com.wire.android.ui.common.textfield.textAsFlow
 import com.wire.android.util.EMPTY
 import com.wire.android.util.deeplink.DeepLinkResult
+import com.wire.android.util.dispatchers.DispatcherProvider
 import com.wire.kalium.common.error.CoreFailure
 import com.wire.kalium.logic.CoreLogic
 import com.wire.kalium.logic.configuration.server.ServerConfig
+import com.wire.kalium.logic.data.logout.LogoutReason
+import com.wire.kalium.logic.data.user.UserId
 import com.wire.kalium.logic.feature.auth.AddAuthenticatedUserUseCase
 import com.wire.kalium.logic.feature.auth.AuthenticationScope
 import com.wire.kalium.logic.feature.auth.DomainLookupUseCase
+import com.wire.kalium.logic.feature.auth.IsNomadProfilesEnabledUseCase
 import com.wire.kalium.logic.feature.auth.ValidateEmailUseCase
 import com.wire.kalium.logic.feature.auth.autoVersioningAuth.AutoVersionAuthScopeUseCase
 import com.wire.kalium.logic.feature.auth.sso.SSOInitiateLoginResult
 import com.wire.kalium.logic.feature.auth.sso.SSOLoginSessionResult
+import com.wire.kalium.logic.feature.backup.RestoreCryptoStateResult
 import com.wire.kalium.logic.feature.client.RegisterClientResult
 import dagger.hilt.android.lifecycle.HiltViewModel
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 import javax.inject.Inject
 
-@Suppress("LongParameterList")
+@Suppress("LongParameterList", "TooManyFunctions")
 @HiltViewModel
 class LoginSSOViewModel(
     private val savedStateHandle: SavedStateHandle,
@@ -68,7 +75,8 @@ class LoginSSOViewModel(
     clientScopeProviderFactory: ClientScopeProvider.Factory,
     userDataStoreProvider: UserDataStoreProvider,
     private val ssoExtension: LoginSSOViewModelExtension,
-    serverConfig: ServerConfig.Links
+    serverConfig: ServerConfig.Links,
+    private val dispatchers: DispatcherProvider,
 ) : LoginViewModel(
     savedStateHandle,
     clientScopeProviderFactory,
@@ -90,6 +98,7 @@ class LoginSSOViewModel(
         userDataStoreProvider: UserDataStoreProvider,
         serverConfig: ServerConfig.Links,
         @DefaultWebSocketEnabledByDefault defaultWebSocketEnabledByDefault: Boolean,
+        dispatchers: DispatcherProvider,
     ) : this(
         savedStateHandle,
         addAuthenticatedUser,
@@ -98,7 +107,8 @@ class LoginSSOViewModel(
         clientScopeProviderFactory,
         userDataStoreProvider,
         LoginSSOViewModelExtension(addAuthenticatedUser, coreLogic, defaultWebSocketEnabledByDefault),
-        serverConfig
+        serverConfig,
+        dispatchers,
     )
 
     var openWebUrl = MutableSharedFlow<Pair<String, ServerConfig.Links>>()
@@ -237,16 +247,38 @@ class LoginSSOViewModel(
                 onSSOLoginFailure = { updateSSOFlowState(it.toLoginError()) },
                 onAddAuthenticatedUserFailure = { updateSSOFlowState(it.toLoginError()) },
                 onSuccess = { storedUserId ->
-                    registerClient(storedUserId, null).let {
-                        when (it) {
-                            is RegisterClientResult.Success ->
+                    appLogger.i("$TAG SSO session established successfully for userId: $storedUserId, checking Nomad status")
+                    val isNomadEnabled = withContext(dispatchers.io()) {
+                        val result = coreLogic.getSessionScope(storedUserId).authenticationScope.isNomadProfilesEnabled()
+                        result is IsNomadProfilesEnabledUseCase.Result.Success && result.isEnabled
+                    }
+                    if (!isNomadEnabled) {
+                        appLogger.i("$TAG Nomad not enabled, proceeding with regular login")
+                        registerClientAndUpdateState(storedUserId, setLastDeviceId = false)
+                    } else {
+                        appLogger.i("$TAG Nomad enabled, attempting crypto state restore")
+                        when (
+                            withContext(dispatchers.io()) {
+                                coreLogic.getSessionScope(storedUserId).backup.restoreCryptoState()
+                            }
+                        ) {
+                            is RestoreCryptoStateResult.Success -> {
                                 updateSSOFlowState(LoginState.Success(isInitialSyncCompleted(storedUserId), false))
+                            }
 
-                            is RegisterClientResult.Failure ->
-                                updateSSOFlowState(it.toLoginError())
+                            is RestoreCryptoStateResult.NoBackupAvailable -> {
+                                registerClientAndUpdateState(storedUserId, setLastDeviceId = true)
+                            }
 
-                            is RegisterClientResult.E2EICertificateRequired ->
-                                updateSSOFlowState(LoginState.Success(isInitialSyncCompleted(storedUserId), true))
+                            is RestoreCryptoStateResult.Failure -> {
+                                appLogger.e("$TAG Failed to restore crypto state during SSO login")
+                                revertSSOSession(storedUserId)
+                                updateSSOFlowState(
+                                    LoginState.Error.DialogError.GenericError(
+                                        CoreFailure.Unknown(Exception("Failed to restore crypto state"))
+                                    )
+                                )
+                            }
                         }
                     }
                 }
@@ -272,6 +304,37 @@ class LoginSSOViewModel(
         }
     }
 
+    private suspend fun registerClientAndUpdateState(userId: UserId, setLastDeviceId: Boolean = false) {
+        withContext(dispatchers.io()) {
+            registerClient(userId = userId, password = null)
+        }.let {
+            when (it) {
+                is RegisterClientResult.Success -> {
+                    if (setLastDeviceId) {
+                        coreLogic.getSessionScope(userId).backup.setLastDeviceId(it.client.id.value)
+                    }
+                    updateSSOFlowState(LoginState.Success(isInitialSyncCompleted(userId), false))
+                }
+
+                is RegisterClientResult.E2EICertificateRequired ->
+                    updateSSOFlowState(LoginState.Success(isInitialSyncCompleted(userId), true))
+
+                is RegisterClientResult.Failure.TooManyClients ->
+                    updateSSOFlowState(LoginState.Error.TooManyDevicesError)
+
+                is RegisterClientResult.Failure -> {
+                    revertSSOSession(userId)
+                    updateSSOFlowState(it.toLoginError())
+                }
+            }
+        }
+    }
+
+    private suspend fun revertSSOSession(userId: UserId) {
+        coreLogic.getSessionScope(userId).logout(reason = LogoutReason.SELF_HARD_LOGOUT, waitUntilCompletes = true)
+        coreLogic.getGlobalScope().deleteSession(userId)
+    }
+
     private fun openWebUrl(url: String, customServerConfig: ServerConfig.Links) {
         viewModelScope.launch {
             updateSSOFlowState(LoginState.Default)
@@ -281,6 +344,7 @@ class LoginSSOViewModel(
 
     companion object {
         const val SSO_CODE_SAVED_STATE_KEY = "sso_code"
+        private const val TAG = "[LoginSSOViewModel]"
     }
 
     private fun consumePendingNomadServiceUrl(): String? = pendingNomadServiceUrl.also {

app/src/main/kotlin/com/wire/android/ui/common/dialogs/NomadAccountBlocksLoginDialog.kt

@@ -25,8 +25,10 @@ import com.wire.android.ui.common.VisibilityState
 import com.wire.android.ui.common.WireDialog
 import com.wire.android.ui.common.WireDialogButtonProperties
 import com.wire.android.ui.common.WireDialogButtonType
+import com.wire.android.ui.common.preview.MultipleThemePreviews
 import com.wire.android.ui.common.visbility.VisibilityState
 import com.wire.android.ui.common.wireDialogPropertiesBuilder
+import com.wire.android.ui.theme.WireTheme
 
 @Composable
 fun NomadAccountBlocksLoginDialog(
@@ -50,3 +52,14 @@ fun NomadAccountBlocksLoginDialog(
         )
     }
 }
+
+@Composable
+@MultipleThemePreviews
+fun PreviewNomadAccountBlocksLoginDialog() {
+    WireTheme {
+        NomadAccountBlocksLoginDialog(
+            dialogState = VisibilityState(true),
+            onActionButtonClicked = {}
+        )
+    }
+}