Merge pull request #7206 from vector-im/feature/mna/mutualize-pending-auth

Mutualize the pending auth handling (PSG-742)

changelog.d/7193.misc

@@ -0,0 +1 @@
+Mutualize the pending auth handling

vector/src/main/java/im/vector/app/features/auth/PendingAuthHandler.kt

@@ -0,0 +1,70 @@
+/*
+ * Copyright (c) 2022 New Vector Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package im.vector.app.features.auth
+
+import im.vector.app.core.di.ActiveSessionHolder
+import org.matrix.android.sdk.api.Matrix
+import org.matrix.android.sdk.api.auth.UIABaseAuth
+import org.matrix.android.sdk.api.auth.UserPasswordAuth
+import org.matrix.android.sdk.api.session.uia.exceptions.UiaCancelledException
+import org.matrix.android.sdk.api.util.fromBase64
+import timber.log.Timber
+import javax.inject.Inject
+import kotlin.coroutines.Continuation
+import kotlin.coroutines.resume
+import kotlin.coroutines.resumeWithException
+
+class PendingAuthHandler @Inject constructor(
+        private val matrix: Matrix,
+        private val activeSessionHolder: ActiveSessionHolder,
+) {
+    var uiaContinuation: Continuation<UIABaseAuth>? = null
+    var pendingAuth: UIABaseAuth? = null
+
+    fun ssoAuthDone() {
+        pendingAuth?.let {
+            Timber.d("ssoAuthDone, resuming action")
+            uiaContinuation?.resume(it)
+        } ?: run {
+            Timber.d("ssoAuthDone, cannot resume: no pendingAuth")
+            uiaContinuation?.resumeWithException(IllegalArgumentException())
+        }
+    }
+
+    fun passwordAuthDone(password: String) {
+        Timber.d("passwordAuthDone")
+        val decryptedPass = matrix.secureStorageService()
+                .loadSecureSecret<String>(
+                        inputStream = password.fromBase64().inputStream(),
+                        keyAlias = ReAuthActivity.DEFAULT_RESULT_KEYSTORE_ALIAS
+                )
+        uiaContinuation?.resume(
+                UserPasswordAuth(
+                        session = pendingAuth?.session,
+                        password = decryptedPass,
+                        user = activeSessionHolder.getActiveSession().myUserId
+                )
+        )
+    }
+
+    fun reAuthCancelled() {
+        Timber.d("reAuthCancelled")
+        uiaContinuation?.resumeWithException(UiaCancelledException())
+        uiaContinuation = null
+        pendingAuth = null
+    }
+}

vector/src/main/java/im/vector/app/features/crypto/recover/BootstrapSharedViewModel.kt

@@ -32,14 +32,13 @@ import im.vector.app.core.error.ErrorFormatter
 import im.vector.app.core.platform.VectorViewModel
 import im.vector.app.core.platform.WaitingViewData
 import im.vector.app.core.resources.StringProvider
-import im.vector.app.features.auth.ReAuthActivity
+import im.vector.app.features.auth.PendingAuthHandler
 import im.vector.app.features.raw.wellknown.SecureBackupMethod
 import im.vector.app.features.raw.wellknown.getElementWellknown
 import im.vector.app.features.raw.wellknown.isSecureBackupRequired
 import im.vector.app.features.raw.wellknown.secureBackupMethod
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
-import org.matrix.android.sdk.api.Matrix
 import org.matrix.android.sdk.api.auth.UIABaseAuth
 import org.matrix.android.sdk.api.auth.UserInteractiveAuthInterceptor
 import org.matrix.android.sdk.api.auth.UserPasswordAuth
@@ -57,10 +56,8 @@ import org.matrix.android.sdk.api.session.crypto.keysbackup.toKeysVersionResult
 import org.matrix.android.sdk.api.session.securestorage.RawBytesKeySpec
 import org.matrix.android.sdk.api.session.uia.DefaultBaseAuth
 import org.matrix.android.sdk.api.util.awaitCallback
-import org.matrix.android.sdk.api.util.fromBase64
 import java.io.OutputStream
 import kotlin.coroutines.Continuation
-import kotlin.coroutines.resume
 import kotlin.coroutines.resumeWithException
 
 class BootstrapSharedViewModel @AssistedInject constructor(
@@ -71,7 +68,7 @@ class BootstrapSharedViewModel @AssistedInject constructor(
         private val rawService: RawService,
         private val bootstrapTask: BootstrapCrossSigningTask,
         private val migrationTask: BackupToQuadSMigrationTask,
-        private val matrix: Matrix,
+        private val pendingAuthHandler: PendingAuthHandler,
 ) : VectorViewModel<BootstrapViewState, BootstrapActions, BootstrapViewEvents>(initialState) {
 
     private var doesKeyBackupExist: Boolean = false
@@ -85,11 +82,6 @@ class BootstrapSharedViewModel @AssistedInject constructor(
 
     companion object : MavericksViewModelFactory<BootstrapSharedViewModel, BootstrapViewState> by hiltMavericksViewModelFactory()
 
-//    private var _pendingSession: String? = null
-
-    var uiaContinuation: Continuation<UIABaseAuth>? = null
-    var pendingAuth: UIABaseAuth? = null
-
     init {
 
         setState {
@@ -272,21 +264,10 @@ class BootstrapSharedViewModel @AssistedInject constructor(
             is BootstrapActions.DoMigrateWithRecoveryKey -> {
                 startMigrationFlow(state.step, null, action.recoveryKey)
             }
-            BootstrapActions.SsoAuthDone -> {
-                uiaContinuation?.resume(DefaultBaseAuth(session = pendingAuth?.session ?: ""))
-            }
-            is BootstrapActions.PasswordAuthDone -> {
-                val decryptedPass = matrix.secureStorageService()
-                        .loadSecureSecret<String>(action.password.fromBase64().inputStream(), ReAuthActivity.DEFAULT_RESULT_KEYSTORE_ALIAS)
-                uiaContinuation?.resume(
-                        UserPasswordAuth(
-                                session = pendingAuth?.session,
-                                password = decryptedPass,
-                                user = session.myUserId
-                        )
-                )
-            }
+            BootstrapActions.SsoAuthDone -> pendingAuthHandler.ssoAuthDone()
+            is BootstrapActions.PasswordAuthDone -> pendingAuthHandler.passwordAuthDone(action.password)
             BootstrapActions.ReAuthCancelled -> {
+                pendingAuthHandler.reAuthCancelled()
                 setState {
                     copy(step = BootstrapStep.AccountReAuth(stringProvider.getString(R.string.authentication_error)))
                 }
@@ -402,13 +383,13 @@ class BootstrapSharedViewModel @AssistedInject constructor(
             override fun performStage(flowResponse: RegistrationFlowResponse, errCode: String?, promise: Continuation<UIABaseAuth>) {
                 when (flowResponse.nextUncompletedStage()) {
                     LoginFlowTypes.PASSWORD -> {
-                        pendingAuth = UserPasswordAuth(
+                        pendingAuthHandler.pendingAuth = UserPasswordAuth(
                                 // Note that _pendingSession may or may not be null, this is OK, it will be managed by the task
                                 session = flowResponse.session,
                                 user = session.myUserId,
                                 password = null
                         )
-                        uiaContinuation = promise
+                        pendingAuthHandler.uiaContinuation = promise
                         setState {
                             copy(
                                     step = BootstrapStep.AccountReAuth()
@@ -417,8 +398,8 @@ class BootstrapSharedViewModel @AssistedInject constructor(
                         _viewEvents.post(BootstrapViewEvents.RequestReAuth(flowResponse, errCode))
                     }
                     LoginFlowTypes.SSO -> {
-                        pendingAuth = DefaultBaseAuth(flowResponse.session)
-                        uiaContinuation = promise
+                        pendingAuthHandler.pendingAuth = DefaultBaseAuth(flowResponse.session)
+                        pendingAuthHandler.uiaContinuation = promise
                         setState {
                             copy(
                                     step = BootstrapStep.AccountReAuth()

vector/src/main/java/im/vector/app/features/settings/account/deactivation/DeactivateAccountViewModel.kt

@@ -23,22 +23,15 @@ import dagger.assisted.AssistedInject
 import im.vector.app.core.di.MavericksAssistedViewModelFactory
 import im.vector.app.core.di.hiltMavericksViewModelFactory
 import im.vector.app.core.platform.VectorViewModel
-import im.vector.app.features.auth.ReAuthActivity
+import im.vector.app.features.auth.PendingAuthHandler
 import kotlinx.coroutines.launch
-import org.matrix.android.sdk.api.Matrix
 import org.matrix.android.sdk.api.auth.UIABaseAuth
 import org.matrix.android.sdk.api.auth.UserInteractiveAuthInterceptor
-import org.matrix.android.sdk.api.auth.UserPasswordAuth
 import org.matrix.android.sdk.api.auth.registration.RegistrationFlowResponse
 import org.matrix.android.sdk.api.failure.isInvalidUIAAuth
 import org.matrix.android.sdk.api.session.Session
 import org.matrix.android.sdk.api.session.uia.DefaultBaseAuth
-import org.matrix.android.sdk.api.session.uia.exceptions.UiaCancelledException
-import org.matrix.android.sdk.api.util.fromBase64
-import timber.log.Timber
 import kotlin.coroutines.Continuation
-import kotlin.coroutines.resume
-import kotlin.coroutines.resumeWithException
 
 data class DeactivateAccountViewState(
         val dummy: Boolean = false
@@ -47,7 +40,7 @@ data class DeactivateAccountViewState(
 class DeactivateAccountViewModel @AssistedInject constructor(
         @Assisted private val initialState: DeactivateAccountViewState,
         private val session: Session,
-        private val matrix: Matrix,
+        private val pendingAuthHandler: PendingAuthHandler,
 ) :
         VectorViewModel<DeactivateAccountViewState, DeactivateAccountAction, DeactivateAccountViewEvents>(initialState) {
 
@@ -56,39 +49,18 @@ class DeactivateAccountViewModel @AssistedInject constructor(
         override fun create(initialState: DeactivateAccountViewState): DeactivateAccountViewModel
     }
 
-    var uiaContinuation: Continuation<UIABaseAuth>? = null
-    var pendingAuth: UIABaseAuth? = null
-
     override fun handle(action: DeactivateAccountAction) {
         when (action) {
             is DeactivateAccountAction.DeactivateAccount -> handleDeactivateAccount(action)
             DeactivateAccountAction.SsoAuthDone -> {
-                Timber.d("## UIA - FallBack success")
                 _viewEvents.post(DeactivateAccountViewEvents.Loading())
-                if (pendingAuth != null) {
-                    uiaContinuation?.resume(pendingAuth!!)
-                } else {
-                    uiaContinuation?.resumeWithException(IllegalArgumentException())
-                }
+                pendingAuthHandler.ssoAuthDone()
             }
             is DeactivateAccountAction.PasswordAuthDone -> {
                 _viewEvents.post(DeactivateAccountViewEvents.Loading())
-                val decryptedPass = matrix.secureStorageService()
-                        .loadSecureSecret<String>(action.password.fromBase64().inputStream(), ReAuthActivity.DEFAULT_RESULT_KEYSTORE_ALIAS)
-                uiaContinuation?.resume(
-                        UserPasswordAuth(
-                                session = pendingAuth?.session,
-                                password = decryptedPass,
-                                user = session.myUserId
-                        )
-                )
-            }
-            DeactivateAccountAction.ReAuthCancelled -> {
-                Timber.d("## UIA - Reauth cancelled")
-                uiaContinuation?.resumeWithException(UiaCancelledException())
-                uiaContinuation = null
-                pendingAuth = null
+                pendingAuthHandler.passwordAuthDone(action.password)
             }
+            DeactivateAccountAction.ReAuthCancelled -> pendingAuthHandler.reAuthCancelled()
         }
     }
 
@@ -102,8 +74,8 @@ class DeactivateAccountViewModel @AssistedInject constructor(
                         object : UserInteractiveAuthInterceptor {
                             override fun performStage(flowResponse: RegistrationFlowResponse, errCode: String?, promise: Continuation<UIABaseAuth>) {
                                 _viewEvents.post(DeactivateAccountViewEvents.RequestReAuth(flowResponse, errCode))
-                                pendingAuth = DefaultBaseAuth(session = flowResponse.session)
-                                uiaContinuation = promise
+                                pendingAuthHandler.pendingAuth = DefaultBaseAuth(session = flowResponse.session)
+                                pendingAuthHandler.uiaContinuation = promise
                             }
                         }
                 )

vector/src/main/java/im/vector/app/features/settings/crosssigning/CrossSigningSettingsViewModel.kt

@@ -24,12 +24,11 @@ import im.vector.app.core.di.MavericksAssistedViewModelFactory
 import im.vector.app.core.di.hiltMavericksViewModelFactory
 import im.vector.app.core.platform.VectorViewModel
 import im.vector.app.core.resources.StringProvider
-import im.vector.app.features.auth.ReAuthActivity
+import im.vector.app.features.auth.PendingAuthHandler
 import im.vector.app.features.login.ReAuthHelper
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.launch
-import org.matrix.android.sdk.api.Matrix
 import org.matrix.android.sdk.api.auth.UIABaseAuth
 import org.matrix.android.sdk.api.auth.UserInteractiveAuthInterceptor
 import org.matrix.android.sdk.api.auth.UserPasswordAuth
@@ -40,19 +39,17 @@ import org.matrix.android.sdk.api.session.Session
 import org.matrix.android.sdk.api.session.crypto.crosssigning.isVerified
 import org.matrix.android.sdk.api.session.uia.DefaultBaseAuth
 import org.matrix.android.sdk.api.util.awaitCallback
-import org.matrix.android.sdk.api.util.fromBase64
 import org.matrix.android.sdk.flow.flow
 import timber.log.Timber
 import kotlin.coroutines.Continuation
 import kotlin.coroutines.resume
-import kotlin.coroutines.resumeWithException
 
 class CrossSigningSettingsViewModel @AssistedInject constructor(
         @Assisted private val initialState: CrossSigningSettingsViewState,
         private val session: Session,
         private val reAuthHelper: ReAuthHelper,
         private val stringProvider: StringProvider,
-        private val matrix: Matrix,
+        private val pendingAuthHandler: PendingAuthHandler,
 ) : VectorViewModel<CrossSigningSettingsViewState, CrossSigningSettingsAction, CrossSigningSettingsViewEvents>(initialState) {
 
     init {
@@ -77,9 +74,6 @@ class CrossSigningSettingsViewModel @AssistedInject constructor(
                 }
     }
 
-    var uiaContinuation: Continuation<UIABaseAuth>? = null
-    var pendingAuth: UIABaseAuth? = null
-
     @AssistedFactory
     interface Factory : MavericksAssistedViewModelFactory<CrossSigningSettingsViewModel, CrossSigningSettingsViewState> {
         override fun create(initialState: CrossSigningSettingsViewState): CrossSigningSettingsViewModel
@@ -110,8 +104,8 @@ class CrossSigningSettingsViewModel @AssistedInject constructor(
                                             } else {
                                                 Timber.d("## UIA : initializeCrossSigning UIA > start reauth activity")
                                                 _viewEvents.post(CrossSigningSettingsViewEvents.RequestReAuth(flowResponse, errCode))
-                                                pendingAuth = DefaultBaseAuth(session = flowResponse.session)
-                                                uiaContinuation = promise
+                                                pendingAuthHandler.pendingAuth = DefaultBaseAuth(session = flowResponse.session)
+                                                pendingAuthHandler.uiaContinuation = promise
                                             }
                                         }
                                     }, it
@@ -125,31 +119,11 @@ class CrossSigningSettingsViewModel @AssistedInject constructor(
                 }
                 Unit
             }
-            is CrossSigningSettingsAction.SsoAuthDone -> {
-                Timber.d("## UIA - FallBack success")
-                if (pendingAuth != null) {
-                    uiaContinuation?.resume(pendingAuth!!)
-                } else {
-                    uiaContinuation?.resumeWithException(IllegalArgumentException())
-                }
-            }
-            is CrossSigningSettingsAction.PasswordAuthDone -> {
-                val decryptedPass = matrix.secureStorageService()
-                        .loadSecureSecret<String>(action.password.fromBase64().inputStream(), ReAuthActivity.DEFAULT_RESULT_KEYSTORE_ALIAS)
-                uiaContinuation?.resume(
-                        UserPasswordAuth(
-                                session = pendingAuth?.session,
-                                password = decryptedPass,
-                                user = session.myUserId
-                        )
-                )
-            }
+            is CrossSigningSettingsAction.SsoAuthDone -> pendingAuthHandler.ssoAuthDone()
+            is CrossSigningSettingsAction.PasswordAuthDone -> pendingAuthHandler.passwordAuthDone(action.password)
             CrossSigningSettingsAction.ReAuthCancelled -> {
-                Timber.d("## UIA - Reauth cancelled")
                 _viewEvents.post(CrossSigningSettingsViewEvents.HideModalWaitingView)
-                uiaContinuation?.resumeWithException(Exception())
-                uiaContinuation = null
-                pendingAuth = null
+                pendingAuthHandler.reAuthCancelled()
             }
         }
     }

vector/src/main/java/im/vector/app/features/settings/devices/DevicesViewModel.kt

@@ -32,7 +32,7 @@ import im.vector.app.core.di.hiltMavericksViewModelFactory
 import im.vector.app.core.platform.VectorViewModel
 import im.vector.app.core.resources.StringProvider
 import im.vector.app.core.utils.PublishDataSource
-import im.vector.app.features.auth.ReAuthActivity
+import im.vector.app.features.auth.PendingAuthHandler
 import im.vector.app.features.login.ReAuthHelper
 import im.vector.app.features.settings.devices.v2.list.CheckIfSessionIsInactiveUseCase
 import im.vector.app.features.settings.devices.v2.verification.GetEncryptionTrustLevelForDeviceUseCase
@@ -45,7 +45,6 @@ import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.sample
 import kotlinx.coroutines.launch
-import org.matrix.android.sdk.api.Matrix
 import org.matrix.android.sdk.api.MatrixCallback
 import org.matrix.android.sdk.api.NoOpMatrixCallback
 import org.matrix.android.sdk.api.auth.UIABaseAuth
@@ -67,13 +66,11 @@ import org.matrix.android.sdk.api.session.crypto.verification.VerificationTransa
 import org.matrix.android.sdk.api.session.crypto.verification.VerificationTxState
 import org.matrix.android.sdk.api.session.uia.DefaultBaseAuth
 import org.matrix.android.sdk.api.util.awaitCallback
-import org.matrix.android.sdk.api.util.fromBase64
 import org.matrix.android.sdk.flow.flow
 import timber.log.Timber
 import javax.net.ssl.HttpsURLConnection
 import kotlin.coroutines.Continuation
 import kotlin.coroutines.resume
-import kotlin.coroutines.resumeWithException
 
 data class DevicesViewState(
         val myDeviceId: String = "",
@@ -100,15 +97,12 @@ class DevicesViewModel @AssistedInject constructor(
         private val session: Session,
         private val reAuthHelper: ReAuthHelper,
         private val stringProvider: StringProvider,
-        private val matrix: Matrix,
+        private val pendingAuthHandler: PendingAuthHandler,
         private val checkIfSessionIsInactiveUseCase: CheckIfSessionIsInactiveUseCase,
         getCurrentSessionCrossSigningInfoUseCase: GetCurrentSessionCrossSigningInfoUseCase,
         private val getEncryptionTrustLevelForDeviceUseCase: GetEncryptionTrustLevelForDeviceUseCase,
 ) : VectorViewModel<DevicesViewState, DevicesAction, DevicesViewEvents>(initialState), VerificationService.Listener {
 
-    var uiaContinuation: Continuation<UIABaseAuth>? = null
-    var pendingAuth: UIABaseAuth? = null
-
     @AssistedFactory
     interface Factory : MavericksAssistedViewModelFactory<DevicesViewModel, DevicesViewState> {
         override fun create(initialState: DevicesViewState): DevicesViewModel
@@ -232,37 +226,9 @@ class DevicesViewModel @AssistedInject constructor(
             is DevicesAction.CompleteSecurity -> handleCompleteSecurity()
             is DevicesAction.MarkAsManuallyVerified -> handleVerifyManually(action)
             is DevicesAction.VerifyMyDeviceManually -> handleShowDeviceCryptoInfo(action)
-            is DevicesAction.SsoAuthDone -> {
-                // we should use token based auth
-                // _viewEvents.post(CrossSigningSettingsViewEvents.ShowModalWaitingView(null))
-                // will release the interactive auth interceptor
-                Timber.d("## UIA - FallBack success $pendingAuth , continuation: $uiaContinuation")
-                if (pendingAuth != null) {
-                    uiaContinuation?.resume(pendingAuth!!)
-                } else {
-                    uiaContinuation?.resumeWithException(IllegalArgumentException())
-                }
-                Unit
-            }
-            is DevicesAction.PasswordAuthDone -> {
-                val decryptedPass = matrix.secureStorageService()
-                        .loadSecureSecret<String>(action.password.fromBase64().inputStream(), ReAuthActivity.DEFAULT_RESULT_KEYSTORE_ALIAS)
-                uiaContinuation?.resume(
-                        UserPasswordAuth(
-                                session = pendingAuth?.session,
-                                password = decryptedPass,
-                                user = session.myUserId
-                        )
-                )
-                Unit
-            }
-            DevicesAction.ReAuthCancelled -> {
-                Timber.d("## UIA - Reauth cancelled")
-//                _viewEvents.post(DevicesViewEvents.Loading)
-                uiaContinuation?.resumeWithException(Exception())
-                uiaContinuation = null
-                pendingAuth = null
-            }
+            is DevicesAction.SsoAuthDone -> pendingAuthHandler.ssoAuthDone()
+            is DevicesAction.PasswordAuthDone -> pendingAuthHandler.passwordAuthDone(action.password)
+            DevicesAction.ReAuthCancelled -> pendingAuthHandler.reAuthCancelled()
             DevicesAction.ResetSecurity -> _viewEvents.post(DevicesViewEvents.PromptResetSecrets)
         }
     }
@@ -371,8 +337,8 @@ class DevicesViewModel @AssistedInject constructor(
                             } else {
                                 Timber.d("## UIA : deleteDevice UIA > start reauth activity")
                                 _viewEvents.post(DevicesViewEvents.RequestReAuth(flowResponse, errCode))
-                                pendingAuth = DefaultBaseAuth(session = flowResponse.session)
-                                uiaContinuation = promise
+                                pendingAuthHandler.pendingAuth = DefaultBaseAuth(session = flowResponse.session)
+                                pendingAuthHandler.uiaContinuation = promise
                             }
                         }
                     }, it)

vector/src/main/java/im/vector/app/features/settings/threepids/ThreePidsSettingsViewModel.kt

@@ -28,33 +28,26 @@ import im.vector.app.core.di.hiltMavericksViewModelFactory
 import im.vector.app.core.platform.VectorViewModel
 import im.vector.app.core.resources.StringProvider
 import im.vector.app.core.utils.ReadOnceTrue
-import im.vector.app.features.auth.ReAuthActivity
+import im.vector.app.features.auth.PendingAuthHandler
 import kotlinx.coroutines.launch
-import org.matrix.android.sdk.api.Matrix
 import org.matrix.android.sdk.api.auth.UIABaseAuth
 import org.matrix.android.sdk.api.auth.UserInteractiveAuthInterceptor
-import org.matrix.android.sdk.api.auth.UserPasswordAuth
 import org.matrix.android.sdk.api.auth.registration.RegistrationFlowResponse
 import org.matrix.android.sdk.api.session.Session
 import org.matrix.android.sdk.api.session.identity.ThreePid
 import org.matrix.android.sdk.api.session.uia.DefaultBaseAuth
-import org.matrix.android.sdk.api.util.fromBase64
 import org.matrix.android.sdk.flow.flow
-import timber.log.Timber
 import kotlin.coroutines.Continuation
-import kotlin.coroutines.resume
-import kotlin.coroutines.resumeWithException
 
 class ThreePidsSettingsViewModel @AssistedInject constructor(
         @Assisted initialState: ThreePidsSettingsViewState,
         private val session: Session,
         private val stringProvider: StringProvider,
-        private val matrix: Matrix,
+        private val pendingAuthHandler: PendingAuthHandler,
 ) : VectorViewModel<ThreePidsSettingsViewState, ThreePidsSettingsAction, ThreePidsSettingsViewEvents>(initialState) {
 
     // UIA session
     private var pendingThreePid: ThreePid? = null
-//    private var pendingSession: String? = null
 
     private suspend fun loadingSuspendable(block: suspend () -> Unit) {
         runCatching { block() }
@@ -126,42 +119,17 @@ class ThreePidsSettingsViewModel @AssistedInject constructor(
             is ThreePidsSettingsAction.CancelThreePid -> handleCancelThreePid(action)
             is ThreePidsSettingsAction.DeleteThreePid -> handleDeleteThreePid(action)
             is ThreePidsSettingsAction.ChangeUiState -> handleChangeUiState(action)
-            ThreePidsSettingsAction.SsoAuthDone -> {
-                Timber.d("## UIA - FallBack success")
-                if (pendingAuth != null) {
-                    uiaContinuation?.resume(pendingAuth!!)
-                } else {
-                    uiaContinuation?.resumeWithException(IllegalArgumentException())
-                }
-            }
-            is ThreePidsSettingsAction.PasswordAuthDone -> {
-                val decryptedPass = matrix.secureStorageService()
-                        .loadSecureSecret<String>(action.password.fromBase64().inputStream(), ReAuthActivity.DEFAULT_RESULT_KEYSTORE_ALIAS)
-                uiaContinuation?.resume(
-                        UserPasswordAuth(
-                                session = pendingAuth?.session,
-                                password = decryptedPass,
-                                user = session.myUserId
-                        )
-                )
-            }
-            ThreePidsSettingsAction.ReAuthCancelled -> {
-                Timber.d("## UIA - Reauth cancelled")
-                uiaContinuation?.resumeWithException(Exception())
-                uiaContinuation = null
-                pendingAuth = null
-            }
+            ThreePidsSettingsAction.SsoAuthDone -> pendingAuthHandler.ssoAuthDone()
+            is ThreePidsSettingsAction.PasswordAuthDone -> pendingAuthHandler.passwordAuthDone(action.password)
+            ThreePidsSettingsAction.ReAuthCancelled -> pendingAuthHandler.reAuthCancelled()
         }
     }
 
-    var uiaContinuation: Continuation<UIABaseAuth>? = null
-    var pendingAuth: UIABaseAuth? = null
-
     private val uiaInterceptor = object : UserInteractiveAuthInterceptor {
         override fun performStage(flowResponse: RegistrationFlowResponse, errCode: String?, promise: Continuation<UIABaseAuth>) {
             _viewEvents.post(ThreePidsSettingsViewEvents.RequestReAuth(flowResponse, errCode))
-            pendingAuth = DefaultBaseAuth(session = flowResponse.session)
-            uiaContinuation = promise
+            pendingAuthHandler.pendingAuth = DefaultBaseAuth(session = flowResponse.session)
+            pendingAuthHandler.uiaContinuation = promise
         }
     }
 

vector/src/test/java/im/vector/app/features/auth/PendingAuthHandlerTest.kt

@@ -0,0 +1,144 @@
+/*
+ * Copyright (c) 2022 New Vector Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package im.vector.app.features.auth
+
+import android.util.Base64
+import im.vector.app.test.fakes.FakeActiveSessionHolder
+import im.vector.app.test.fakes.FakeMatrix
+import io.mockk.every
+import io.mockk.just
+import io.mockk.mockk
+import io.mockk.mockkStatic
+import io.mockk.runs
+import io.mockk.unmockkAll
+import io.mockk.verify
+import org.amshove.kluent.shouldBe
+import org.junit.After
+import org.junit.Before
+import org.junit.Ignore
+import org.junit.Test
+import org.matrix.android.sdk.api.auth.UIABaseAuth
+import org.matrix.android.sdk.api.auth.UserPasswordAuth
+import org.matrix.android.sdk.api.session.uia.exceptions.UiaCancelledException
+import kotlin.coroutines.Continuation
+import kotlin.coroutines.resume
+import kotlin.coroutines.resumeWithException
+
+private const val A_PASSWORD = "a-password"
+private const val A_SESSION_ID = "session-id"
+
+class PendingAuthHandlerTest {
+
+    private val fakeMatrix = FakeMatrix()
+    private val fakeActiveSessionHolder = FakeActiveSessionHolder()
+
+    private val pendingAuthHandler = PendingAuthHandler(
+            matrix = fakeMatrix.instance,
+            activeSessionHolder = fakeActiveSessionHolder.instance,
+    )
+
+    @Before
+    fun setUp() {
+        mockkStatic(Base64::class)
+    }
+
+    @After
+    fun tearDown() {
+        unmockkAll()
+    }
+
+    @Test
+    fun `given a pending auth and continuation when SSO auth is done then continuation is resumed with pending auth`() {
+        // Given
+        val pendingAuth = mockk<UIABaseAuth>()
+        val continuation = mockk<Continuation<UIABaseAuth>>()
+        every { continuation.resume(any()) } just runs
+        pendingAuthHandler.pendingAuth = pendingAuth
+        pendingAuthHandler.uiaContinuation = continuation
+
+        // When
+        pendingAuthHandler.ssoAuthDone()
+
+        // Then
+        verify { continuation.resume(pendingAuth) }
+    }
+
+    @Test
+    @Ignore("Ignored due because of problem to mock the inline method continuation.resumeWithException")
+    fun `given missing pending auth and continuation when SSO auth is done then continuation is resumed with error`() {
+        // Given
+        val pendingAuth = null
+        val continuation = mockk<Continuation<UIABaseAuth>>()
+        every { continuation.resumeWithException(any()) } just runs
+        pendingAuthHandler.pendingAuth = pendingAuth
+        pendingAuthHandler.uiaContinuation = continuation
+
+        // When
+        pendingAuthHandler.ssoAuthDone()
+
+        // Then
+        verify { continuation.resumeWithException(match { it is IllegalArgumentException }) }
+    }
+
+    @Test
+    fun `given a password, pending auth and continuation when password auth is done then continuation is resumed with correct auth`() {
+        // Given
+        val pendingAuth = mockk<UIABaseAuth>()
+        every { pendingAuth.session } returns A_SESSION_ID
+        val continuation = mockk<Continuation<UIABaseAuth>>()
+        every { continuation.resume(any()) } just runs
+        pendingAuthHandler.pendingAuth = pendingAuth
+        pendingAuthHandler.uiaContinuation = continuation
+        val decryptedPwd = "decrypted-pwd"
+        val decodedPwd = byteArrayOf()
+        every { Base64.decode(A_PASSWORD, any()) } returns decodedPwd
+        fakeMatrix.fakeSecureStorageService.givenLoadSecureSecretReturns(decryptedPwd)
+        val expectedAuth = UserPasswordAuth(
+                session = A_SESSION_ID,
+                password = decryptedPwd,
+                user = fakeActiveSessionHolder.fakeSession.myUserId
+        )
+
+        // When
+        pendingAuthHandler.passwordAuthDone(A_PASSWORD)
+
+        // Then
+        verify {
+            fakeMatrix.fakeSecureStorageService.loadSecureSecret<String>(any(), ReAuthActivity.DEFAULT_RESULT_KEYSTORE_ALIAS)
+            continuation.resume(expectedAuth)
+        }
+    }
+
+    @Test
+    @Ignore("Ignored because of problem to mock the inline method continuation.resumeWithException")
+    fun `given pending auth and continuation when reAuth is cancelled then pending auth and continuation are reset`() {
+        // Given
+        val pendingAuth = mockk<UIABaseAuth>()
+        val continuation = mockk<Continuation<UIABaseAuth>>()
+        every { continuation.resumeWithException(any()) } just runs
+        pendingAuthHandler.pendingAuth = pendingAuth
+        pendingAuthHandler.uiaContinuation = continuation
+
+        // When
+        pendingAuthHandler.reAuthCancelled()
+
+        // Then
+        pendingAuthHandler.pendingAuth shouldBe null
+        pendingAuthHandler.uiaContinuation shouldBe null
+        verify { continuation.resumeWithException(match { it is UiaCancelledException }) }
+    }
+}

vector/src/test/java/im/vector/app/test/fakes/FakeMatrix.kt

@@ -0,0 +1,32 @@
+/*
+ * Copyright (c) 2022 New Vector Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package im.vector.app.test.fakes
+
+import io.mockk.every
+import io.mockk.mockk
+import org.matrix.android.sdk.api.Matrix
+
+class FakeMatrix(
+        val fakeSecureStorageService: FakeSecureStorageService = FakeSecureStorageService(),
+) {
+
+    val instance = mockk<Matrix>()
+
+    init {
+        every { instance.secureStorageService() } returns fakeSecureStorageService
+    }
+}

vector/src/test/java/im/vector/app/test/fakes/FakeSecureStorageService.kt

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) 2022 New Vector Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package im.vector.app.test.fakes
+
+import io.mockk.every
+import io.mockk.mockk
+import org.matrix.android.sdk.api.securestorage.SecureStorageService
+
+class FakeSecureStorageService : SecureStorageService by mockk() {
+
+    fun <T> givenLoadSecureSecretReturns(value: T?) {
+        every { loadSecureSecret<T>(any(), any()) } returns value
+    }
+}