feat(auth): prepare sessions after authentication PR 4

app/src/main/kotlin/com/wire/android/ui/authentication/create/code/CreateAccountCodeViewModel.kt

@@ -27,22 +27,25 @@ import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.ramcosta.composedestinations.generated.app.navArgs
 import com.wire.android.BuildConfig
-import com.wire.android.di.ClientScopeProvider
 import com.wire.android.di.DefaultWebSocketEnabledByDefault
 import com.wire.android.di.KaliumCoreLogic
+import com.wire.android.session.AppUserSessionPreparationResult
+import com.wire.android.session.UserSessionPreparationGate
 import com.wire.android.ui.authentication.create.common.CreateAccountFlowType
 import com.wire.android.ui.authentication.create.common.CreateAccountNavArgs
 import com.wire.android.ui.authentication.login.email.LoginEmailViewModel.Companion.RESEND_TIMER_DELAY
 import com.wire.android.ui.common.textfield.textAsFlow
 import com.wire.android.ui.registration.code.CreateAccountCodeResult
 import com.wire.android.util.WillNeverOccurError
 import com.wire.android.util.ui.CountdownTimer
+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.session.StoreSessionParam
 import com.wire.kalium.logic.data.user.UserId
 import com.wire.kalium.logic.feature.auth.AddAuthenticatedUserUseCase
 import com.wire.kalium.logic.feature.auth.autoVersioningAuth.AutoVersionAuthScopeUseCase
+import com.wire.kalium.logic.feature.UserSessionScope
 import com.wire.kalium.logic.feature.client.RegisterClientParam
 import com.wire.kalium.logic.feature.client.RegisterClientResult
 import com.wire.kalium.logic.feature.register.RegisterParam
@@ -59,7 +62,6 @@ class CreateAccountCodeViewModel @AssistedInject constructor(
     @Assisted savedStateHandle: SavedStateHandle,
     @KaliumCoreLogic private val coreLogic: CoreLogic,
     private val addAuthenticatedUser: AddAuthenticatedUserUseCase,
-    private val clientScopeProviderFactory: ClientScopeProvider.Factory,
     defaultServerConfig: ServerConfig.Links,
     @DefaultWebSocketEnabledByDefault private val defaultWebSocketEnabledByDefault: Boolean
 ) : ViewModel() {
@@ -68,6 +70,8 @@ class CreateAccountCodeViewModel @AssistedInject constructor(
         fun create(savedStateHandle: SavedStateHandle): CreateAccountCodeViewModel
     }
 
+    private val userSessionPreparationGate by lazy { UserSessionPreparationGate(coreLogic) }
+
     val createAccountNavArgs: CreateAccountNavArgs = savedStateHandle.navArgs()
 
     val serverConfig: ServerConfig.Links = createAccountNavArgs.customServerConfig ?: defaultServerConfig
@@ -205,7 +209,18 @@ class CreateAccountCodeViewModel @AssistedInject constructor(
                     is AddAuthenticatedUserUseCase.Result.Success -> it.userId
                 }
             }
-            registerClient(storedUserId, registerParam.password).let {
+            val sessionScope = when (val preparation = userSessionPreparationGate.prepare(storedUserId)) {
+                is AppUserSessionPreparationResult.Ready -> preparation.sessionScope
+                is AppUserSessionPreparationResult.Failed -> {
+                    updateCodeErrorState(
+                        CreateAccountCodeResult.Error.DialogError.GenericError(
+                            CoreFailure.Unknown(IllegalStateException("User session preparation failed: ${preparation.reason}"))
+                        )
+                    )
+                    return@launch
+                }
+            }
+            registerClient(sessionScope, registerParam.password).let {
                 when (it) {
                     is RegisterClientResult.Failure -> {
                         updateCodeErrorState(it.toCodeError(storedUserId))
@@ -229,8 +244,8 @@ class CreateAccountCodeViewModel @AssistedInject constructor(
         codeState = codeState.copy(loading = false, result = codeError)
     }
 
-    private suspend fun registerClient(userId: UserId, password: String) =
-        clientScopeProviderFactory.create(userId).clientScope.getOrRegister(
+    private suspend fun registerClient(sessionScope: UserSessionScope, password: String) =
+        sessionScope.client.getOrRegister(
             RegisterClientParam(
                 password = password,
                 capabilities = null,

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

@@ -32,6 +32,7 @@ import com.wire.kalium.logic.data.user.UserId
 import com.wire.kalium.logic.feature.auth.AddAuthenticatedUserUseCase
 import com.wire.kalium.logic.feature.auth.AuthenticationResult
 import com.wire.kalium.logic.feature.auth.DomainLookupUseCase
+import com.wire.kalium.logic.feature.UserSessionScope
 import com.wire.kalium.logic.feature.client.RegisterClientResult
 
 @Suppress("TooManyFunctions")
@@ -60,6 +61,13 @@ open class LoginViewModel(
         capabilities: List<ClientCapability>? = null,
     ): RegisterClientResult = loginExtension.registerClient(userId, password, secondFactorVerificationCode, capabilities)
 
+    suspend fun registerClient(
+        sessionScope: UserSessionScope,
+        password: String?,
+        secondFactorVerificationCode: String? = null,
+        capabilities: List<ClientCapability>? = null,
+    ): RegisterClientResult = loginExtension.registerClient(sessionScope, password, secondFactorVerificationCode, capabilities)
+
     internal suspend fun isInitialSyncCompleted(userId: UserId): Boolean = loginExtension.isInitialSyncCompleted(userId)
 }
 

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

@@ -22,6 +22,8 @@ import com.wire.android.datastore.UserDataStoreProvider
 import com.wire.android.di.ClientScopeProvider
 import com.wire.kalium.logic.data.client.ClientCapability
 import com.wire.kalium.logic.data.user.UserId
+import com.wire.kalium.logic.feature.UserSessionScope
+import com.wire.kalium.logic.feature.client.GetOrRegisterClientUseCase
 import com.wire.kalium.logic.feature.client.RegisterClientParam
 import com.wire.kalium.logic.feature.client.RegisterClientResult
 import kotlinx.coroutines.flow.first
@@ -38,16 +40,35 @@ class LoginViewModelExtension(
         capabilities: List<ClientCapability>? = null,
     ): RegisterClientResult {
         val clientScope = clientScopeProviderFactory.create(userId).clientScope
-        return clientScope.getOrRegister(
-            RegisterClientParam(
-                password = password,
-                capabilities = capabilities,
-                secondFactorVerificationCode = secondFactorVerificationCode,
-                modelPostfix = if (BuildConfig.PRIVATE_BUILD) " [${BuildConfig.FLAVOR}_${BuildConfig.BUILD_TYPE}]" else null
-            )
-        )
+        return registerClient(clientScope.getOrRegister, password, secondFactorVerificationCode, capabilities)
     }
 
+    suspend fun registerClient(
+        sessionScope: UserSessionScope,
+        password: String?,
+        secondFactorVerificationCode: String? = null,
+        capabilities: List<ClientCapability>? = null,
+    ): RegisterClientResult = registerClient(
+        sessionScope.client.getOrRegister,
+        password,
+        secondFactorVerificationCode,
+        capabilities,
+    )
+
+    private suspend fun registerClient(
+        getOrRegister: GetOrRegisterClientUseCase,
+        password: String?,
+        secondFactorVerificationCode: String?,
+        capabilities: List<ClientCapability>?,
+    ): RegisterClientResult = getOrRegister(
+        RegisterClientParam(
+            password = password,
+            capabilities = capabilities,
+            secondFactorVerificationCode = secondFactorVerificationCode,
+            modelPostfix = if (BuildConfig.PRIVATE_BUILD) " [${BuildConfig.FLAVOR}_${BuildConfig.BUILD_TYPE}]" else null
+        )
+    )
+
     internal suspend fun isInitialSyncCompleted(userId: UserId): Boolean =
         userDataStoreProvider.getOrCreate(userId).initialSyncCompleted.first()
 }

app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailViewModel.kt

@@ -27,11 +27,14 @@ import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
 import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.viewModelScope
+import com.wire.android.appLogger
 import com.wire.android.datastore.UserDataStoreProvider
 import com.wire.android.datastore.GlobalDataStore
 import com.wire.android.di.ClientScopeProvider
 import com.wire.android.di.DefaultWebSocketEnabledByDefault
 import com.wire.android.di.KaliumCoreLogic
+import com.wire.android.session.AppUserSessionPreparationResult
+import com.wire.android.session.UserSessionPreparationGate
 import com.wire.android.ui.authentication.toBackendConfigUrl
 import com.wire.android.ui.authentication.login.LoginNavArgs
 import com.wire.android.ui.authentication.login.LoginSavedInputStore
@@ -137,6 +140,7 @@ class LoginEmailViewModel(
     )
 
     private val preFilledUserIdentifier: PreFilledUserIdentifierType = loginNavArgs.userHandle ?: PreFilledUserIdentifierType.None
+    private val userSessionPreparationGate by lazy { UserSessionPreparationGate(coreLogic) }
 
     val userIdentifierTextState: TextFieldState = TextFieldState()
     val passwordTextState: TextFieldState = TextFieldState()
@@ -327,9 +331,24 @@ class LoginEmailViewModel(
                 }
             }
 
+            val preparedSessionScope = when (val preparation = userSessionPreparationGate.prepare(storedUserId)) {
+                is AppUserSessionPreparationResult.Ready -> preparation.sessionScope
+                is AppUserSessionPreparationResult.Failed -> {
+                    restorePreviousSession()
+                    updateEmailFlowState(
+                        LoginState.Error.DialogError.GenericError(
+                            com.wire.kalium.common.error.CoreFailure.Unknown(
+                                IllegalStateException("User session preparation failed: ${preparation.reason}")
+                            )
+                        )
+                    )
+                    return@launch
+                }
+            }
+
             withContext(dispatchers.io()) {
                 if (coreLogic.getGlobalScope().validateEmailUseCase(userIdentifierTextState.text.toString())) {
-                    coreLogic.getSessionScope(storedUserId).users.persistSelfUserEmail(userIdentifierTextState.text.toString())
+                    preparedSessionScope.users.persistSelfUserEmail(userIdentifierTextState.text.toString())
                 } else {
                     null
                 }
@@ -343,7 +362,7 @@ class LoginEmailViewModel(
 
             withContext(dispatchers.io()) {
                 registerClient(
-                    userId = storedUserId,
+                    sessionScope = preparedSessionScope,
                     password = passwordTextState.text.toString(),
                 )
             }.let {
@@ -370,14 +389,28 @@ class LoginEmailViewModel(
     }
 
     private suspend fun revertNewSession() {
-        loginJobData.value?.newSessionUserId?.let { newSessionUserId ->
-            // logout to cancel all session-related actions, remove all sensitive data and free up resources
-            coreLogic.getSessionScope(newSessionUserId).logout(reason = LogoutReason.SELF_HARD_LOGOUT, waitUntilCompletes = true)
-            // delete the session to make it seem like the session was never logged in
-            coreLogic.getGlobalScope().deleteSession(newSessionUserId)
+        val jobData = loginJobData.value
+        jobData?.newSessionUserId?.let { newSessionUserId ->
+            when (val preparation = userSessionPreparationGate.prepare(newSessionUserId)) {
+                is AppUserSessionPreparationResult.Ready -> {
+                    // logout to cancel session actions and remove sensitive data before deleting the session
+                    preparation.sessionScope.logout(
+                        reason = LogoutReason.SELF_HARD_LOGOUT,
+                        waitUntilCompletes = true,
+                    )
+                    coreLogic.getGlobalScope().deleteSession(newSessionUserId)
+                }
+
+                is AppUserSessionPreparationResult.Failed -> appLogger.w(
+                    "Login rollback skipped database deletion because session preparation failed: ${preparation.reason}"
+                )
+            }
         }
-        // set the previous session back
-        coreLogic.getGlobalScope().session.updateCurrentSession(loginJobData.value?.previousSessionUserId)
+        restorePreviousSession(jobData)
+    }
+
+    private suspend fun restorePreviousSession(jobData: LoginJobData? = loginJobData.value) {
+        coreLogic.getGlobalScope().session.updateCurrentSession(jobData?.previousSessionUserId)
     }
 
     private suspend fun revertLogin() {

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

@@ -18,6 +18,8 @@
 package com.wire.android.ui.authentication.login.sso
 
 import com.wire.android.appLogger
+import com.wire.android.session.AppUserSessionPreparationResult
+import com.wire.android.session.UserSessionPreparationGate
 import com.wire.kalium.common.error.CoreFailure
 import com.wire.kalium.logic.CoreLogic
 import com.wire.kalium.logic.configuration.server.ServerConfig
@@ -48,6 +50,7 @@ class LoginSSOViewModelExtension(
     private val coreLogic: CoreLogic,
     private val defaultWebSocketEnabledByDefault: Boolean,
 ) {
+    private val userSessionPreparationGate by lazy { UserSessionPreparationGate(coreLogic) }
     suspend fun withAuthenticationScope(
         serverConfig: ServerConfig.Links,
         onAuthScopeFailure: (AutoVersionAuthScopeUseCase.Result.Failure) -> Unit,
@@ -149,15 +152,27 @@ class LoginSSOViewModelExtension(
         when (authenticatedUserResult) {
             AddAuthenticatedUserUseCase.Result.Failure.SsoIdentityChanged -> onSsoIdentityChanged(session)
             is AddAuthenticatedUserUseCase.Result.Failure -> onAddAuthenticatedUserFailure(authenticatedUserResult)
-            is AddAuthenticatedUserUseCase.Result.Success -> onSuccess(authenticatedUserResult.userId)
+            is AddAuthenticatedUserUseCase.Result.Success ->
+                when (val preparation = userSessionPreparationGate.prepare(authenticatedUserResult.userId)) {
+                    is AppUserSessionPreparationResult.Ready -> onSuccess(authenticatedUserResult.userId)
+                    is AppUserSessionPreparationResult.Failed -> onAddAuthenticatedUserFailure(
+                        AddAuthenticatedUserUseCase.Result.Failure.Generic(preparation.toCoreFailure())
+                    )
+                }
         }
     }
 
-    @Suppress("TooGenericExceptionCaught")
+    @Suppress("TooGenericExceptionCaught", "NestedBlockDepth")
     suspend fun replaceRetainedSsoSession(session: StoreSessionParam): ReplaceRetainedSsoSessionResult =
         try {
             val userId = session.accountTokens.userId
-            coreLogic.getSessionScope(userId).logout(
+            val retainedScope = when (val preparation = userSessionPreparationGate.prepare(userId)) {
+                is AppUserSessionPreparationResult.Ready -> preparation.sessionScope
+                is AppUserSessionPreparationResult.Failed -> return ReplaceRetainedSsoSessionResult.Failure(
+                    AddAuthenticatedUserUseCase.Result.Failure.Generic(preparation.toCoreFailure())
+                )
+            }
+            retainedScope.logout(
                 reason = LogoutReason.SELF_HARD_LOGOUT,
                 waitUntilCompletes = true
             )
@@ -167,8 +182,16 @@ class LoginSSOViewModelExtension(
                     ReplaceRetainedSsoSessionResult.Failure(
                         AddAuthenticatedUserUseCase.Result.Failure.Generic(deleteResult.cause)
                     )
-                DeleteSessionUseCase.Result.Success ->
-                    addAuthenticatedUser(session, replace = false).toReplaceRetainedSsoSessionResult()
+                DeleteSessionUseCase.Result.Success -> when (val addResult = addAuthenticatedUser(session, replace = false)) {
+                    is AddAuthenticatedUserUseCase.Result.Failure -> ReplaceRetainedSsoSessionResult.Failure(addResult)
+                    is AddAuthenticatedUserUseCase.Result.Success ->
+                        when (val preparation = userSessionPreparationGate.prepare(addResult.userId)) {
+                            is AppUserSessionPreparationResult.Ready -> ReplaceRetainedSsoSessionResult.Success(addResult.userId)
+                            is AppUserSessionPreparationResult.Failed -> ReplaceRetainedSsoSessionResult.Failure(
+                                AddAuthenticatedUserUseCase.Result.Failure.Generic(preparation.toCoreFailure())
+                            )
+                        }
+                }
             }
         } catch (exception: CancellationException) {
             throw exception
@@ -179,6 +202,9 @@ class LoginSSOViewModelExtension(
         }
 }
 
+private fun AppUserSessionPreparationResult.Failed.toCoreFailure(): CoreFailure =
+    CoreFailure.Unknown(IllegalStateException("User session preparation failed: $reason"))
+
 private fun AddAuthenticatedUserUseCase.Result.toReplaceRetainedSsoSessionResult(): ReplaceRetainedSsoSessionResult =
     when (this) {
         is AddAuthenticatedUserUseCase.Result.Failure -> ReplaceRetainedSsoSessionResult.Failure(this)