Fix: "Reset identity" flow leaves backup disabled #5075 (#6420)

* Don't cancel the resetOidc job in onStart or onDestroy of ResetIdentityFlowNode

* Add logging around the launch and completion of reserOidc

* Some improvements to make sure we always cancel the reset job.

Also, the flow can be considered done when the key backup is enabled, at that point we should already be verified.

* Don't cancel the `ResetIdentityFlowManager` when starting a reset

This also cancels the check that will call `onDone` when the flow finishes successfully.

It seems like it worked for me locally because of some race condition.

---------

Co-authored-by: Jorge Martín <jorgem@element.io>

features/securebackup/impl/src/main/kotlin/io/element/android/features/securebackup/impl/reset/ResetIdentityFlowManager.kt

@@ -11,15 +11,13 @@ package io.element.android.features.securebackup.impl.reset
 import dev.zacsweers.metro.Inject
 import io.element.android.libraries.architecture.AsyncData
 import io.element.android.libraries.di.annotations.SessionCoroutineScope
+import io.element.android.libraries.matrix.api.encryption.BackupState
 import io.element.android.libraries.matrix.api.encryption.EncryptionService
 import io.element.android.libraries.matrix.api.encryption.IdentityResetHandle
-import io.element.android.libraries.matrix.api.verification.SessionVerificationService
-import io.element.android.libraries.matrix.api.verification.SessionVerifiedStatus
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.filterIsInstance
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.launch
 
@@ -27,7 +25,6 @@ import kotlinx.coroutines.launch
 class ResetIdentityFlowManager(
     private val encryptionService: EncryptionService,
     @SessionCoroutineScope private val sessionCoroutineScope: CoroutineScope,
-    private val sessionVerificationService: SessionVerificationService,
 ) {
     private val resetHandleFlow: MutableStateFlow<AsyncData<IdentityResetHandle?>> = MutableStateFlow(AsyncData.Uninitialized)
     val currentHandleFlow: StateFlow<AsyncData<IdentityResetHandle?>> = resetHandleFlow
@@ -35,7 +32,7 @@ class ResetIdentityFlowManager(
 
     fun whenResetIsDone(block: () -> Unit) {
         whenResetIsDoneWaitingJob = sessionCoroutineScope.launch {
-            sessionVerificationService.sessionVerifiedStatus.filterIsInstance<SessionVerifiedStatus.Verified>().first()
+            encryptionService.backupStateStateFlow.first { it == BackupState.ENABLED }
             block()
         }
     }

features/securebackup/impl/src/main/kotlin/io/element/android/features/securebackup/impl/reset/ResetIdentityFlowNode.kt

@@ -16,8 +16,6 @@ import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.getValue
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.window.DialogProperties
-import androidx.lifecycle.DefaultLifecycleObserver
-import androidx.lifecycle.LifecycleOwner
 import com.bumble.appyx.core.modality.BuildContext
 import com.bumble.appyx.core.node.Node
 import com.bumble.appyx.core.plugin.Plugin
@@ -42,7 +40,7 @@ import io.element.android.libraries.matrix.api.encryption.IdentityOidcResetHandl
 import io.element.android.libraries.matrix.api.encryption.IdentityPasswordResetHandle
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
-import kotlinx.coroutines.flow.collectLatest
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.launch
 import kotlinx.parcelize.Parcelize
 import timber.log.Timber
@@ -81,24 +79,9 @@ class ResetIdentityFlowNode(
     override fun onBuilt() {
         super.onBuilt()
 
-        lifecycle.addObserver(object : DefaultLifecycleObserver {
-            override fun onStart(owner: LifecycleOwner) {
-                // If the custom tab / Web browser was opened, we need to cancel the reset job
-                // when we come back to the node if the reset wasn't successful
-                sessionCoroutineScope.launch {
-                    cancelResetJob()
-
-                    resetIdentityFlowManager.whenResetIsDone {
-                        callback.onDone()
-                    }
-                }
-            }
-
-            override fun onDestroy(owner: LifecycleOwner) {
-                // Make sure we cancel the reset job when the node is destroyed, just in case
-                sessionCoroutineScope.launch { cancelResetJob() }
-            }
-        })
+        resetIdentityFlowManager.whenResetIsDone {
+            callback.onDone()
+        }
     }
 
     override fun resolve(navTarget: NavTarget, buildContext: BuildContext): Node {
@@ -122,34 +105,53 @@ class ResetIdentityFlowNode(
     }
 
     private fun CoroutineScope.startReset() = launch {
-        resetIdentityFlowManager.getResetHandle()
-            .collectLatest { state ->
-                when (state) {
-                    is AsyncData.Failure -> {
-                        cancelResetJob()
-                        Timber.e(state.error, "Could not load the reset identity handle.")
+        // Instead of cancelling the reset job on every ON_START, we can do it before starting a new attempt
+        cancelResetJob()
+
+        val handleResult = resetIdentityFlowManager.getResetHandle()
+            // We're only interested in the success/failure case, and we need this flow to stop by itself
+            // since each call to `startReset` will create a new one
+            .first { it.isSuccess() || it.isFailure() }
+
+        when (handleResult) {
+            is AsyncData.Failure -> {
+                cancelResetJob()
+                Timber.e(handleResult.error, "Could not load the reset identity handle.")
+            }
+            is AsyncData.Success -> {
+                when (val handle = handleResult.data) {
+                    null -> {
+                        Timber.d("No reset handle return, the reset is done.")
                     }
-                    is AsyncData.Success -> {
-                        when (val handle = state.data) {
-                            null -> {
-                                Timber.d("No reset handle return, the reset is done.")
-                            }
-                            is IdentityOidcResetHandle -> {
-                                activity.openUrlInChromeCustomTab(null, darkTheme, handle.url)
-                                resetJob = launch { handle.resetOidc() }
-                            }
-                            is IdentityPasswordResetHandle -> backstack.push(NavTarget.ResetPassword)
-                        }
+                    is IdentityOidcResetHandle -> {
+                        Timber.d("Launching reset confirmation in MAS")
+                        activity.openUrlInChromeCustomTab(null, darkTheme, handle.url)
+                        Timber.d("Starting resetOidc")
+                        resetJob = launch { handle.resetOidc() }
+                        resetJob?.invokeOnCompletion { Timber.d("resetOidc ended") }
                     }
-                    else -> Unit
+                    is IdentityPasswordResetHandle -> backstack.push(NavTarget.ResetPassword)
                 }
             }
+            else -> Unit
+        }
     }
 
-    private suspend fun cancelResetJob() {
+    override fun performUpNavigation(): Boolean {
+        val navigatesUp = super.performUpNavigation()
+
+        // This intercepts the back navigation so we only cancel this job when the user actually navigates up
+        if (navigatesUp) {
+            sessionCoroutineScope.launch { resetIdentityFlowManager.cancel() }
+            cancelResetJob()
+        }
+
+        return navigatesUp
+    }
+
+    private fun cancelResetJob() {
         resetJob?.cancel()
         resetJob = null
-        resetIdentityFlowManager.cancel()
     }
 
     @Composable

features/securebackup/impl/src/test/kotlin/io/element/android/features/securebackup/impl/reset/ResetIdentityFlowManagerTest.kt

@@ -11,11 +11,10 @@ package io.element.android.features.securebackup.impl.reset
 import app.cash.turbine.test
 import com.google.common.truth.Truth.assertThat
 import io.element.android.libraries.architecture.AsyncData
+import io.element.android.libraries.matrix.api.encryption.BackupState
 import io.element.android.libraries.matrix.api.encryption.IdentityResetHandle
-import io.element.android.libraries.matrix.api.verification.SessionVerifiedStatus
 import io.element.android.libraries.matrix.test.encryption.FakeEncryptionService
 import io.element.android.libraries.matrix.test.encryption.FakeIdentityPasswordResetHandle
-import io.element.android.libraries.matrix.test.verification.FakeSessionVerificationService
 import io.element.android.tests.testutils.lambda.lambdaRecorder
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.TestScope
@@ -104,9 +103,9 @@ class ResetIdentityFlowManagerTest {
 
     @OptIn(ExperimentalCoroutinesApi::class)
     @Test
-    fun `whenResetIsDone - will trigger the lambda when verification status is verified`() = runTest {
-        val verificationService = FakeSessionVerificationService()
-        val flowManager = createFlowManager(sessionVerificationService = verificationService)
+    fun `whenResetIsDone - will trigger the lambda when key backup is enabled`() = runTest {
+        val encryptionService = FakeEncryptionService()
+        val flowManager = createFlowManager(encryptionService = encryptionService)
         var isDone = false
 
         flowManager.whenResetIsDone {
@@ -115,25 +114,31 @@ class ResetIdentityFlowManagerTest {
 
         assertThat(isDone).isFalse()
 
-        verificationService.emitVerifiedStatus(SessionVerifiedStatus.Unknown)
+        encryptionService.emitBackupState(BackupState.UNKNOWN)
         advanceUntilIdle()
         assertThat(isDone).isFalse()
 
-        verificationService.emitVerifiedStatus(SessionVerifiedStatus.NotVerified)
+        encryptionService.emitBackupState(BackupState.DISABLING)
         advanceUntilIdle()
         assertThat(isDone).isFalse()
 
-        verificationService.emitVerifiedStatus(SessionVerifiedStatus.Verified)
+        encryptionService.emitBackupState(BackupState.WAITING_FOR_SYNC)
+        advanceUntilIdle()
+        assertThat(isDone).isFalse()
+
+        encryptionService.emitBackupState(BackupState.ENABLING)
+        advanceUntilIdle()
+        assertThat(isDone).isFalse()
+
+        encryptionService.emitBackupState(BackupState.ENABLED)
         advanceUntilIdle()
         assertThat(isDone).isTrue()
     }
 
     private fun TestScope.createFlowManager(
         encryptionService: FakeEncryptionService = FakeEncryptionService(),
-        sessionVerificationService: FakeSessionVerificationService = FakeSessionVerificationService(),
     ) = ResetIdentityFlowManager(
         encryptionService = encryptionService,
         sessionCoroutineScope = this,
-        sessionVerificationService = sessionVerificationService,
     )
 }