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>
2024-11-08 16:42:27 +01:00
parent 979c4faafe
commit 49e1cfed42
11 changed files with 176 additions and 113 deletions
--- a/appnav/src/main/kotlin/io/element/android/appnav/di/MatrixClientsHolder.kt
+++ b/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()
    }
--- a/appnav/src/test/kotlin/io/element/android/appnav/di/MatrixClientsHolderTest.kt
+++ b/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()
+    }
 }
--- a/features/ftue/impl/build.gradle.kts
+++ b/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)
--- a/features/ftue/impl/src/main/kotlin/io/element/android/features/ftue/impl/FtueFlowNode.kt
+++ b/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 {
-                lifecycleScope.launch { moveToNextStep() }
-            }
+            .onEach { moveToNextStepIfNeeded() }
+            .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>() {
--- a/features/ftue/impl/src/main/kotlin/io/element/android/features/ftue/impl/state/DefaultFtueService.kt
+++ b/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 -> {
-                updateState()
-                null
-            }
+            FtueStep.AnalyticsOptIn -> null
        }

-    private suspend fun isAnyStepIncomplete(): Boolean {
-        return listOf<suspend () -> Boolean>(
-            { isSessionNotVerified() },
-            { shouldAskNotificationPermissions() },
-            { needsAnalyticsOptIn() },
-            { shouldDisplayLockscreenSetup() },
-        ).any { it() }
+    private fun isSessionVerificationStateReady(): Boolean {
+        return sessionVerificationService.sessionVerifiedStatus.value != SessionVerifiedStatus.Unknown
    }

-    @OptIn(FlowPreview::class)
    private suspend fun isSessionNotVerified(): Boolean {
-        // Wait for the first known (or ready) verification status
-        val readyVerifiedSessionStatus = sessionVerificationService.sessionVerifiedStatus
-            .filter { it != SessionVerifiedStatus.Unknown }
-            // This is not ideal, but there are some very rare cases when reading the flow seems to get stuck
-            .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()
+        // Wait until the session verification status is known
+        isVerificationStatusKnown.filter { it }.first()
+
+        return sessionVerificationService.sessionVerifiedStatus.value == 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
-            else -> FtueState.Complete
+            // Final state, there aren't any more next steps
+            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
--- a/features/ftue/impl/src/test/kotlin/io/element/android/features/ftue/impl/DefaultFtueServiceTest.kt
+++ b/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 {
--- a/libraries/matrix/api/src/main/kotlin/io/element/android/libraries/matrix/api/auth/MatrixAuthenticationService.kt
+++ b/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)
 }
--- a/libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/RustMatrixClientFactory.kt
+++ b/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()
-
-        RustMatrixClient(
+        return 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())
--- a/libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/auth/RustMatrixAuthenticationService.kt
+++ b/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 ->
-                    rustClient.session()
-                        .toSessionData(
-                            isTokenValid = true,
-                            loginType = LoginType.QR,
-                            passphrase = pendingPassphrase,
-                            sessionPaths = emptySessionPaths,
-                        )
-                }
+                val sessionData = client.session()
+                    .toSessionData(
+                        isTokenValid = true,
+                        loginType = LoginType.QR,
+                        passphrase = pendingPassphrase,
+                        sessionPaths = emptySessionPaths,
+                    )
+                newMatrixClientObserver?.invoke(rustMatrixClientFactory.create(client))
                sessionStore.storeData(sessionData)
                SessionId(sessionData.userId)
            }.mapFailure {
--- a/libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/verification/RustSessionVerificationService.kt
+++ b/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 {
-        // 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)
-    }
+    private var isOwnVerification = true

    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) {
-                    while (encryptionService.verificationState() != VerificationState.VERIFIED) {
-                        delay(100)
-                    }
+                withTimeout(20.seconds) {
+                    // Wait until the SDK reports the state as verified
+                    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() {
-        if (verificationFlowState.value == VerificationFlowState.DidFinish) {
-            // Calling `encryptionService.verificationState()` performs a network call and it will deadlock if there is no network
-            // 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("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 updateVerificationStatus() {
+        runCatching {
+            _sessionVerifiedStatus.value = encryptionService.verificationState().map()
+            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) {
--- a/libraries/matrix/test/src/main/kotlin/io/element/android/libraries/matrix/test/auth/FakeMatrixAuthenticationService.kt
+++ b/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
    }