Fix verification failed issue, simplify verification logic (#3830)

* Simplify session verification:

- Reuse Rust `Client` instances created on the login process so we don't need to restore one right before the session verification.
- Remove unnecessary sources of verification state updates.
- Add an intermediate FTUE flow step which will display an indeterminate progress indicator instead of a blank screen.

* Remove unnecessary workaround: the SDK should already handle this

* Add regression tests for noop analytics service usage.

* Add `services.analytics.noop` module to the test dependencies

---------

Co-authored-by: Benoit Marty <benoit@matrix.org>

appnav/src/main/kotlin/io/element/android/appnav/di/MatrixClientsHolder.kt

@@ -27,10 +27,18 @@ private const val SAVE_INSTANCE_KEY = "io.element.android.x.di.MatrixClientsHold
 
 @SingleIn(AppScope::class)
 @ContributesBinding(AppScope::class)
-class MatrixClientsHolder @Inject constructor(private val authenticationService: MatrixAuthenticationService) : MatrixClientProvider {
+class MatrixClientsHolder @Inject constructor(
+    private val authenticationService: MatrixAuthenticationService,
+) : MatrixClientProvider {
     private val sessionIdsToMatrixClient = ConcurrentHashMap<SessionId, MatrixClient>()
     private val restoreMutex = Mutex()
 
+    init {
+        authenticationService.listenToNewMatrixClients { matrixClient ->
+            sessionIdsToMatrixClient[matrixClient.sessionId] = matrixClient
+        }
+    }
+
     fun removeAll() {
         sessionIdsToMatrixClient.clear()
     }

appnav/src/test/kotlin/io/element/android/appnav/di/MatrixClientsHolderTest.kt

@@ -81,4 +81,17 @@ class MatrixClientsHolderTest {
         matrixClientsHolder.restoreWithSavedState(savedStateMap)
         assertThat(matrixClientsHolder.getOrNull(A_SESSION_ID)).isEqualTo(fakeMatrixClient)
     }
+
+    @Test
+    fun `test AuthenticationService listenToNewMatrixClients emits a Client value and we save it`() = runTest {
+        val fakeAuthenticationService = FakeMatrixAuthenticationService()
+        val matrixClientsHolder = MatrixClientsHolder(fakeAuthenticationService)
+        assertThat(matrixClientsHolder.getOrNull(A_SESSION_ID)).isNull()
+
+        fakeAuthenticationService.givenMatrixClient(FakeMatrixClient(sessionId = A_SESSION_ID))
+        val loginSucceeded = fakeAuthenticationService.login("user", "pass")
+
+        assertThat(loginSucceeded.isSuccess).isTrue()
+        assertThat(matrixClientsHolder.getOrNull(A_SESSION_ID)).isNotNull()
+    }
 }

features/ftue/impl/build.gradle.kts

@@ -45,6 +45,7 @@ dependencies {
     testImplementation(libs.test.turbine)
     testImplementation(projects.libraries.matrix.test)
     testImplementation(projects.services.analytics.test)
+    testImplementation(projects.services.analytics.noop)
     testImplementation(projects.libraries.permissions.impl)
     testImplementation(projects.libraries.permissions.test)
     testImplementation(projects.libraries.preferences.test)

features/ftue/impl/src/main/kotlin/io/element/android/features/ftue/impl/FtueFlowNode.kt

@@ -8,7 +8,10 @@
 package io.element.android.features.ftue.impl
 
 import android.os.Parcelable
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.runtime.Composable
+import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.lifecycle.lifecycleScope
 import com.bumble.appyx.core.lifecycle.subscribe
@@ -31,12 +34,14 @@ import io.element.android.features.lockscreen.api.LockScreenEntryPoint
 import io.element.android.libraries.architecture.BackstackView
 import io.element.android.libraries.architecture.BaseFlowNode
 import io.element.android.libraries.architecture.createNode
+import io.element.android.libraries.designsystem.theme.components.CircularProgressIndicator
 import io.element.android.libraries.di.AppScope
 import io.element.android.libraries.di.SessionScope
 import io.element.android.services.analytics.api.AnalyticsService
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.launch
@@ -80,14 +85,17 @@ class FtueFlowNode @AssistedInject constructor(
         super.onBuilt()
 
         lifecycle.subscribe(onCreate = {
-            lifecycleScope.launch { moveToNextStep() }
+            moveToNextStepIfNeeded()
         })
 
         analyticsService.didAskUserConsent()
             .distinctUntilChanged()
-            .onEach {
+            .onEach { moveToNextStepIfNeeded() }
-                lifecycleScope.launch { moveToNextStep() }
+            .launchIn(lifecycleScope)
-            }
+
+        ftueState.isVerificationStatusKnown
+            .filter { it }
+            .onEach { moveToNextStepIfNeeded() }
             .launchIn(lifecycleScope)
     }
 
@@ -99,7 +107,7 @@ class FtueFlowNode @AssistedInject constructor(
             NavTarget.SessionVerification -> {
                 val callback = object : FtueSessionVerificationFlowNode.Callback {
                     override fun onDone() {
-                        lifecycleScope.launch { moveToNextStep() }
+                        moveToNextStepIfNeeded()
                     }
                 }
                 createNode<FtueSessionVerificationFlowNode>(buildContext, listOf(callback))
@@ -107,7 +115,7 @@ class FtueFlowNode @AssistedInject constructor(
             NavTarget.NotificationsOptIn -> {
                 val callback = object : NotificationsOptInNode.Callback {
                     override fun onNotificationsOptInFinished() {
-                        lifecycleScope.launch { moveToNextStep() }
+                        moveToNextStepIfNeeded()
                     }
                 }
                 createNode<NotificationsOptInNode>(buildContext, listOf(callback))
@@ -118,7 +126,7 @@ class FtueFlowNode @AssistedInject constructor(
             NavTarget.LockScreenSetup -> {
                 val callback = object : LockScreenEntryPoint.Callback {
                     override fun onSetupDone() {
-                        lifecycleScope.launch { moveToNextStep() }
+                        moveToNextStepIfNeeded()
                     }
                 }
                 lockScreenEntryPoint.nodeBuilder(this, buildContext, LockScreenEntryPoint.Target.Setup)
@@ -128,8 +136,11 @@ class FtueFlowNode @AssistedInject constructor(
         }
     }
 
-    private fun moveToNextStep() = lifecycleScope.launch {
+    private fun moveToNextStepIfNeeded() = lifecycleScope.launch {
         when (ftueState.getNextStep()) {
+            FtueStep.WaitingForInitialState -> {
+                backstack.newRoot(NavTarget.Placeholder)
+            }
             FtueStep.SessionVerification -> {
                 backstack.newRoot(NavTarget.SessionVerification)
             }
@@ -155,7 +166,14 @@ class FtueFlowNode @AssistedInject constructor(
     class PlaceholderNode @AssistedInject constructor(
         @Assisted buildContext: BuildContext,
         @Assisted plugins: List<Plugin>,
-    ) : Node(buildContext, plugins = plugins)
+    ) : Node(buildContext, plugins = plugins) {
+        @Composable
+        override fun View(modifier: Modifier) {
+            Box(modifier = modifier.fillMaxSize(), contentAlignment = Alignment.Center) {
+                CircularProgressIndicator()
+            }
+        }
+    }
 }
 
 private class NoOpBackstackHandlerStrategy<NavTarget : Any> : BaseBackPressHandlerStrategy<NavTarget, BackStack.State>() {

features/ftue/impl/src/main/kotlin/io/element/android/features/ftue/impl/state/DefaultFtueService.kt

@@ -24,18 +24,14 @@ import io.element.android.libraries.preferences.api.store.SessionPreferencesStor
 import io.element.android.services.analytics.api.AnalyticsService
 import io.element.android.services.toolbox.api.sdk.BuildVersionSdkIntProvider
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.FlowPreview
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.catch
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.flow.timeout
-import timber.log.Timber
 import javax.inject.Inject
-import kotlin.time.Duration.Companion.seconds
 
 @ContributesBinding(SessionScope::class)
 @SingleIn(SessionScope::class)
@@ -50,6 +46,14 @@ class DefaultFtueService @Inject constructor(
 ) : FtueService {
     override val state = MutableStateFlow<FtueState>(FtueState.Unknown)
 
+    /**
+     * This flow emits true when the FTUE flow is ready to be displayed.
+     * In this case, the FTUE flow is ready when the session verification status is known.
+     */
+    val isVerificationStatusKnown = sessionVerificationService.sessionVerifiedStatus
+        .map { it != SessionVerifiedStatus.Unknown }
+        .distinctUntilChanged()
+
     override suspend fun reset() {
         analyticsService.reset()
         if (sdkVersionProvider.isAtLeast(Build.VERSION_CODES.TIRAMISU)) {
@@ -70,7 +74,12 @@ class DefaultFtueService @Inject constructor(
 
     suspend fun getNextStep(currentStep: FtueStep? = null): FtueStep? =
         when (currentStep) {
-            null -> if (isSessionNotVerified()) {
+            null -> if (!isSessionVerificationStateReady()) {
+                FtueStep.WaitingForInitialState
+            } else {
+                getNextStep(FtueStep.WaitingForInitialState)
+            }
+            FtueStep.WaitingForInitialState -> if (isSessionNotVerified()) {
                 FtueStep.SessionVerification
             } else {
                 getNextStep(FtueStep.SessionVerification)
@@ -90,34 +99,18 @@ class DefaultFtueService @Inject constructor(
             } else {
                 getNextStep(FtueStep.AnalyticsOptIn)
             }
-            FtueStep.AnalyticsOptIn -> {
+            FtueStep.AnalyticsOptIn -> null
-                updateState()
-                null
-            }
         }
 
-    private suspend fun isAnyStepIncomplete(): Boolean {
+    private fun isSessionVerificationStateReady(): Boolean {
-        return listOf<suspend () -> Boolean>(
+        return sessionVerificationService.sessionVerifiedStatus.value != SessionVerifiedStatus.Unknown
-            { isSessionNotVerified() },
-            { shouldAskNotificationPermissions() },
-            { needsAnalyticsOptIn() },
-            { shouldDisplayLockscreenSetup() },
-        ).any { it() }
     }
 
-    @OptIn(FlowPreview::class)
     private suspend fun isSessionNotVerified(): Boolean {
-        // Wait for the first known (or ready) verification status
+        // Wait until the session verification status is known
-        val readyVerifiedSessionStatus = sessionVerificationService.sessionVerifiedStatus
+        isVerificationStatusKnown.filter { it }.first()
-            .filter { it != SessionVerifiedStatus.Unknown }
+
-            // This is not ideal, but there are some very rare cases when reading the flow seems to get stuck
+        return sessionVerificationService.sessionVerifiedStatus.value == SessionVerifiedStatus.NotVerified && !canSkipVerification()
-            .timeout(5.seconds)
-            .catch {
-                Timber.e(it, "Failed to get session verification status, assume it's not verified")
-                emit(SessionVerifiedStatus.NotVerified)
-            }
-            .first()
-        return readyVerifiedSessionStatus == SessionVerifiedStatus.NotVerified && !canSkipVerification()
     }
 
     private suspend fun canSkipVerification(): Boolean {
@@ -145,14 +138,17 @@ class DefaultFtueService @Inject constructor(
 
     @VisibleForTesting(otherwise = VisibleForTesting.PRIVATE)
     internal suspend fun updateState() {
+        val nextStep = getNextStep()
         state.value = when {
-            isAnyStepIncomplete() -> FtueState.Incomplete
+            // Final state, there aren't any more next steps
-            else -> FtueState.Complete
+            nextStep == null -> FtueState.Complete
+            else -> FtueState.Incomplete
         }
     }
 }
 
 sealed interface FtueStep {
+    data object WaitingForInitialState : FtueStep
     data object SessionVerification : FtueStep
     data object NotificationsOptIn : FtueStep
     data object AnalyticsOptIn : FtueStep

features/ftue/impl/src/test/kotlin/io/element/android/features/ftue/impl/DefaultFtueServiceTest.kt

@@ -23,6 +23,7 @@ import io.element.android.libraries.permissions.impl.FakePermissionStateProvider
 import io.element.android.libraries.preferences.api.store.SessionPreferencesStore
 import io.element.android.libraries.preferences.test.InMemorySessionPreferencesStore
 import io.element.android.services.analytics.api.AnalyticsService
+import io.element.android.services.analytics.noop.NoopAnalyticsService
 import io.element.android.services.analytics.test.FakeAnalyticsService
 import io.element.android.services.toolbox.test.sdk.FakeBuildVersionSdkIntProvider
 import io.element.android.tests.testutils.lambda.lambdaRecorder
@@ -73,6 +74,27 @@ class DefaultFtueServiceTest {
         assertThat(service.state.value).isEqualTo(FtueState.Complete)
     }
 
+    @Test
+    fun `given all checks being true with no analytics, FtueState is Complete`() = runTest {
+        val analyticsService = NoopAnalyticsService()
+        val sessionVerificationService = FakeSessionVerificationService()
+        val permissionStateProvider = FakePermissionStateProvider(permissionGranted = true)
+        val lockScreenService = FakeLockScreenService()
+        val service = createDefaultFtueService(
+            sessionVerificationService = sessionVerificationService,
+            analyticsService = analyticsService,
+            permissionStateProvider = permissionStateProvider,
+            lockScreenService = lockScreenService,
+        )
+
+        sessionVerificationService.emitVerifiedStatus(SessionVerifiedStatus.Verified)
+        permissionStateProvider.setPermissionGranted()
+        lockScreenService.setIsPinSetup(true)
+        service.updateState()
+
+        assertThat(service.state.value).isEqualTo(FtueState.Complete)
+    }
+
     @Test
     fun `traverse flow`() = runTest {
         val sessionVerificationService = FakeSessionVerificationService().apply {

libraries/matrix/api/src/main/kotlin/io/element/android/libraries/matrix/api/auth/MatrixAuthenticationService.kt

@@ -56,4 +56,7 @@ interface MatrixAuthenticationService {
     suspend fun loginWithOidc(callbackUrl: String): Result<SessionId>
 
     suspend fun loginWithQrCode(qrCodeData: MatrixQrCodeLoginData, progress: (QrCodeLoginStep) -> Unit): Result<SessionId>
+
+    /** Listen to new Matrix clients being created on authentication. */
+    fun listenToNewMatrixClients(lambda: (MatrixClient) -> Unit)
 }

libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/RustMatrixClientFactory.kt

@@ -25,6 +25,7 @@ import io.element.android.services.analytics.api.AnalyticsService
 import io.element.android.services.toolbox.api.systemclock.SystemClock
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.withContext
+import org.matrix.rustcomponents.sdk.Client
 import org.matrix.rustcomponents.sdk.ClientBuilder
 import org.matrix.rustcomponents.sdk.Session
 import org.matrix.rustcomponents.sdk.SlidingSyncVersion
@@ -51,8 +52,9 @@ class RustMatrixClientFactory @Inject constructor(
     private val timelineEventTypeFilterFactory: TimelineEventTypeFilterFactory,
     private val clientBuilderProvider: ClientBuilderProvider,
 ) {
+    private val sessionDelegate = RustClientSessionDelegate(sessionStore, appCoroutineScope, coroutineDispatchers)
+
     suspend fun create(sessionData: SessionData): RustMatrixClient = withContext(coroutineDispatchers.io) {
-        val sessionDelegate = RustClientSessionDelegate(sessionStore, appCoroutineScope, coroutineDispatchers)
         val client = getBaseClientBuilder(
             sessionPaths = sessionData.getSessionPaths(),
             passphrase = sessionData.passphrase,
@@ -60,18 +62,21 @@ class RustMatrixClientFactory @Inject constructor(
         )
             .homeserverUrl(sessionData.homeserverUrl)
             .username(sessionData.userId)
-            .setSessionDelegate(sessionDelegate)
             .use { it.build() }
 
         client.restoreSession(sessionData.toSession())
 
+        create(client)
+    }
+
+    suspend fun create(client: Client): RustMatrixClient {
+        val (anonymizedAccessToken, anonymizedRefreshToken) = client.session().anonymizedTokens()
+
         val syncService = client.syncService()
             .withUtdHook(UtdTracker(analyticsService))
             .finish()
 
-        val (anonymizedAccessToken, anonymizedRefreshToken) = sessionData.anonymizedTokens()
+        return RustMatrixClient(
-
-        RustMatrixClient(
             client = client,
             baseDirectory = baseDirectory,
             sessionStore = sessionStore,
@@ -98,6 +103,7 @@ class RustMatrixClientFactory @Inject constructor(
                 dataPath = sessionPaths.fileDirectory.absolutePath,
                 cachePath = sessionPaths.cacheDirectory.absolutePath,
             )
+            .setSessionDelegate(sessionDelegate)
             .passphrase(passphrase)
             .userAgent(userAgentProvider.provide())
             .addRootCertificates(userCertificatesProvider.provides())

libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/auth/RustMatrixAuthenticationService.kt

@@ -51,7 +51,6 @@ import org.matrix.rustcomponents.sdk.QrCodeData
 import org.matrix.rustcomponents.sdk.QrCodeDecodeException
 import org.matrix.rustcomponents.sdk.QrLoginProgress
 import org.matrix.rustcomponents.sdk.QrLoginProgressListener
-import org.matrix.rustcomponents.sdk.use
 import timber.log.Timber
 import uniffi.matrix_sdk.OidcAuthorizationData
 import javax.inject.Inject
@@ -77,6 +76,11 @@ class RustMatrixAuthenticationService @Inject constructor(
     private var currentClient: Client? = null
     private var currentHomeserver = MutableStateFlow<MatrixHomeServerDetails?>(null)
 
+    private var newMatrixClientObserver: ((MatrixClient) -> Unit)? = null
+    override fun listenToNewMatrixClients(lambda: (MatrixClient) -> Unit) {
+        newMatrixClientObserver = lambda
+    }
+
     private fun rotateSessionPath(): SessionPaths {
         sessionPaths?.deleteRecursively()
         return sessionPathsFactory.create()
@@ -155,7 +159,7 @@ class RustMatrixAuthenticationService @Inject constructor(
                         passphrase = pendingPassphrase,
                         sessionPaths = currentSessionPaths,
                     )
-                clear()
+                newMatrixClientObserver?.invoke(rustMatrixClientFactory.create(client))
                 sessionStore.storeData(sessionData)
                 SessionId(sessionData.userId)
             }.mapFailure { failure ->
@@ -226,9 +230,9 @@ class RustMatrixAuthenticationService @Inject constructor(
                     passphrase = pendingPassphrase,
                     sessionPaths = currentSessionPaths,
                 )
-                clear()
                 pendingOidcAuthorizationData?.close()
                 pendingOidcAuthorizationData = null
+                newMatrixClientObserver?.invoke(rustMatrixClientFactory.create(client))
                 sessionStore.storeData(sessionData)
                 SessionId(sessionData.userId)
             }.mapFailure { failure ->
@@ -256,15 +260,14 @@ class RustMatrixAuthenticationService @Inject constructor(
                     oidcConfiguration = oidcConfiguration,
                     progressListener = progressListener,
                 )
-                val sessionData = client.use { rustClient ->
+                val sessionData = client.session()
-                    rustClient.session()
+                    .toSessionData(
-                        .toSessionData(
+                        isTokenValid = true,
-                            isTokenValid = true,
+                        loginType = LoginType.QR,
-                            loginType = LoginType.QR,
+                        passphrase = pendingPassphrase,
-                            passphrase = pendingPassphrase,
+                        sessionPaths = emptySessionPaths,
-                            sessionPaths = emptySessionPaths,
+                    )
-                        )
+                newMatrixClientObserver?.invoke(rustMatrixClientFactory.create(client))
-                }
                 sessionStore.storeData(sessionData)
                 SessionId(sessionData.userId)
             }.mapFailure {

libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/verification/RustSessionVerificationService.kt

@@ -18,15 +18,13 @@ import io.element.android.libraries.matrix.api.verification.VerificationFlowStat
 import io.element.android.libraries.matrix.impl.util.cancelAndDestroy
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.NonCancellable
-import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
-import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.sync.Mutex
@@ -61,11 +59,13 @@ class RustSessionVerificationService(
     private val _sessionVerifiedStatus = MutableStateFlow<SessionVerifiedStatus>(SessionVerifiedStatus.Unknown)
     override val sessionVerifiedStatus: StateFlow<SessionVerifiedStatus> = _sessionVerifiedStatus.asStateFlow()
 
+    private val recoveryState = MutableStateFlow(RecoveryState.UNKNOWN)
+
     // Listen for changes in verification status and update accordingly
     private val verificationStateListenerTaskHandle = encryptionService.verificationStateListener(object : VerificationStateListener {
         override fun onUpdate(status: VerificationState) {
             Timber.d("New verification state: $status")
-            sessionCoroutineScope.launch { updateVerificationStatus() }
+            _sessionVerifiedStatus.value = status.map()
         }
     })
 
@@ -74,7 +74,7 @@ class RustSessionVerificationService(
         override fun onUpdate(status: RecoveryState) {
             Timber.d("New recovery state: $status")
             // We could check the `RecoveryState`, but it's easier to just use the verification state directly
-            sessionCoroutineScope.launch { updateVerificationStatus() }
+            recoveryState.value = status
         }
     })
 
@@ -88,22 +88,7 @@ class RustSessionVerificationService(
         verificationStatus == SessionVerifiedStatus.NotVerified
     }
 
-    init {
+    private var isOwnVerification = true
-        // Update initial state in case sliding sync isn't ready
-        sessionCoroutineScope.launch { updateVerificationStatus() }
-
-        isReady.onEach { isReady ->
-            if (isReady) {
-                Timber.d("Starting verification service")
-                // Immediate status update
-                updateVerificationStatus()
-            } else {
-                Timber.d("Stopping verification service")
-                updateVerificationStatus()
-            }
-        }
-            .launchIn(sessionCoroutineScope)
-    }
 
     override fun didReceiveVerificationRequest(details: RustSessionVerificationRequestDetails) {
         listener?.onIncomingSessionRequest(details.map())
@@ -135,6 +120,8 @@ class RustSessionVerificationService(
     }
 
     override suspend fun acknowledgeVerificationRequest(details: SessionVerificationRequestDetails) = tryOrFail {
+        isOwnVerification = false
+        initVerificationControllerIfNeeded()
         verificationController.acknowledgeVerificationRequest(
             senderId = details.senderId.value,
             flowId = details.flowId.value,
@@ -179,14 +166,22 @@ class RustSessionVerificationService(
             // Ideally this should be `verificationController?.isVerified().orFalse()` but for some reason it returns false if run immediately
             // It also sometimes unexpectedly fails to report the session as verified, so we have to handle that possibility and fail if needed
             runCatching {
-                withTimeout(30.seconds) {
+                withTimeout(20.seconds) {
-                    while (encryptionService.verificationState() != VerificationState.VERIFIED) {
+                    // Wait until the SDK reports the state as verified
-                        delay(100)
+                    sessionVerifiedStatus.first { it == SessionVerifiedStatus.Verified }
-                    }
                 }
             }
                 .onSuccess {
-                    // Order here is important, first set the flow state as finished, then update the verification status
+                    if (isOwnVerification) {
+                        // Try waiting for the final recovery state for better UX, but don't block the verification state on it
+                        tryOrNull {
+                            withTimeout(10.seconds) {
+                                // Wait until the recovery state is either fully loaded or we check it's explicitly disabled
+                                recoveryState.first { it == RecoveryState.ENABLED || it == RecoveryState.DISABLED }
+                            }
+                        }
+                    }
+
                     _verificationFlowState.value = VerificationFlowState.DidFinish
                     updateVerificationStatus()
                 }
@@ -209,6 +204,7 @@ class RustSessionVerificationService(
     // end-region
 
     override suspend fun reset(cancelAnyPendingVerificationAttempt: Boolean) {
+        isOwnVerification = true
         if (isReady.value && cancelAnyPendingVerificationAttempt) {
             // Cancel any pending verification attempt
             tryOrNull { verificationController.cancelVerification() }
@@ -237,37 +233,20 @@ class RustSessionVerificationService(
         }
     }
 
-    private suspend fun updateVerificationStatus() {
+    private fun updateVerificationStatus() {
-        if (verificationFlowState.value == VerificationFlowState.DidFinish) {
+        runCatching {
-            // Calling `encryptionService.verificationState()` performs a network call and it will deadlock if there is no network
+            _sessionVerifiedStatus.value = encryptionService.verificationState().map()
-            // So we need to check that *only* if we know there is network connection, which is the case when the verification flow just finished
+            Timber.d("New verification status: ${_sessionVerifiedStatus.value}")
-            Timber.d("Updating verification status: flow just finished")
-            runCatching {
-                encryptionService.waitForE2eeInitializationTasks()
-            }.onSuccess {
-                _sessionVerifiedStatus.value = when (encryptionService.verificationState()) {
-                    VerificationState.UNKNOWN -> SessionVerifiedStatus.Unknown
-                    VerificationState.VERIFIED -> SessionVerifiedStatus.Verified
-                    VerificationState.UNVERIFIED -> SessionVerifiedStatus.NotVerified
-                }
-                Timber.d("New verification status: ${_sessionVerifiedStatus.value}")
-            }
-        } else {
-            // Otherwise, just check the current verification status from the session verification controller instead
-            Timber.d("Updating verification status: flow is pending or was finished some time ago")
-            runCatching {
-                initVerificationControllerIfNeeded()
-                _sessionVerifiedStatus.value = if (encryptionService.verificationState() == VerificationState.VERIFIED) {
-                    SessionVerifiedStatus.Verified
-                } else {
-                    SessionVerifiedStatus.NotVerified
-                }
-                Timber.d("New verification status: ${_sessionVerifiedStatus.value}")
-            }
         }
     }
 }
 
+private fun VerificationState.map() = when (this) {
+    VerificationState.UNKNOWN -> SessionVerifiedStatus.Unknown
+    VerificationState.VERIFIED -> SessionVerifiedStatus.Verified
+    VerificationState.UNVERIFIED -> SessionVerifiedStatus.NotVerified
+}
+
 private fun RustSessionVerificationData.map(): SessionVerificationData {
     return use { sessionVerificationData ->
         when (sessionVerificationData) {

libraries/matrix/test/src/main/kotlin/io/element/android/libraries/matrix/test/auth/FakeMatrixAuthenticationService.kt

@@ -18,6 +18,7 @@ import io.element.android.libraries.matrix.api.auth.qrlogin.QrCodeLoginStep
 import io.element.android.libraries.matrix.api.core.SessionId
 import io.element.android.libraries.matrix.test.A_SESSION_ID
 import io.element.android.libraries.matrix.test.A_USER_ID
+import io.element.android.libraries.matrix.test.FakeMatrixClient
 import io.element.android.libraries.sessionstorage.api.LoggedInState
 import io.element.android.tests.testutils.lambda.lambdaError
 import io.element.android.tests.testutils.lambda.lambdaRecorder
@@ -41,6 +42,7 @@ class FakeMatrixAuthenticationService(
     private var loginError: Throwable? = null
     private var changeServerError: Throwable? = null
     private var matrixClient: MatrixClient? = null
+    private var onAuthenticationListener: ((MatrixClient) -> Unit)? = null
 
     var getLatestSessionIdLambda: (() -> SessionId?) = { null }
 
@@ -55,6 +57,7 @@ class FakeMatrixAuthenticationService(
             return it.invoke(sessionId)
         }
         return if (matrixClient != null) {
+            onAuthenticationListener?.invoke(matrixClient!!)
             Result.success(matrixClient!!)
         } else {
             Result.failure(IllegalStateException())
@@ -74,7 +77,10 @@ class FakeMatrixAuthenticationService(
     }
 
     override suspend fun login(username: String, password: String): Result<SessionId> = simulateLongTask {
-        loginError?.let { Result.failure(it) } ?: Result.success(A_USER_ID)
+        loginError?.let { Result.failure(it) } ?: run {
+            onAuthenticationListener?.invoke(matrixClient ?: FakeMatrixClient())
+            Result.success(A_USER_ID)
+        }
     }
 
     override suspend fun importCreatedSession(externalSession: ExternalSession): Result<SessionId> = simulateLongTask {
@@ -90,13 +96,21 @@ class FakeMatrixAuthenticationService(
     }
 
     override suspend fun loginWithOidc(callbackUrl: String): Result<SessionId> = simulateLongTask {
-        loginError?.let { Result.failure(it) } ?: Result.success(A_USER_ID)
+        loginError?.let { Result.failure(it) } ?: run {
+            onAuthenticationListener?.invoke(matrixClient ?: FakeMatrixClient())
+            Result.success(A_USER_ID)
+        }
     }
 
     override suspend fun loginWithQrCode(qrCodeData: MatrixQrCodeLoginData, progress: (QrCodeLoginStep) -> Unit): Result<SessionId> = simulateLongTask {
+        onAuthenticationListener?.invoke(matrixClient ?: FakeMatrixClient())
         loginWithQrCodeResult(qrCodeData, progress)
     }
 
+    override fun listenToNewMatrixClients(lambda: (MatrixClient) -> Unit) {
+        onAuthenticationListener = lambda
+    }
+
     fun givenOidcError(throwable: Throwable?) {
         oidcError = throwable
     }