[PM-29309] [BWA-209] bug: Fix TOTP countdown freeze when returning to Authenticator app (change Flow to StateFlow)

authenticator/src/main/kotlin/com/bitwarden/authenticator/data/authenticator/manager/TotpCodeManager.kt

@@ -3,19 +3,20 @@ package com.bitwarden.authenticator.data.authenticator.manager
 import com.bitwarden.authenticator.data.authenticator.datasource.disk.entity.AuthenticatorItemAlgorithm
 import com.bitwarden.authenticator.data.authenticator.manager.model.VerificationCodeItem
 import com.bitwarden.authenticator.data.authenticator.repository.model.AuthenticatorItem
-import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.StateFlow
 
 /**
  * Manages the flows for getting verification codes.
  */
 interface TotpCodeManager {
 
     /**
-     * Flow for getting a DataState with multiple verification code items.
+     * StateFlow for getting multiple verification code items. Returns a StateFlow that emits
+     * updated verification codes every second.
      */
     fun getTotpCodesFlow(
         itemList: List<AuthenticatorItem>,
-    ): Flow<List<VerificationCodeItem>>
+    ): StateFlow<List<VerificationCodeItem>>
 
     @Suppress("UndocumentedPublicClass")
     companion object {

authenticator/src/main/kotlin/com/bitwarden/authenticator/data/authenticator/manager/TotpCodeManagerImpl.kt

@@ -3,84 +3,137 @@ package com.bitwarden.authenticator.data.authenticator.manager
 import com.bitwarden.authenticator.data.authenticator.datasource.sdk.AuthenticatorSdkSource
 import com.bitwarden.authenticator.data.authenticator.manager.model.VerificationCodeItem
 import com.bitwarden.authenticator.data.authenticator.repository.model.AuthenticatorItem
+import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.currentCoroutineContext
 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.combine
 import kotlinx.coroutines.flow.flow
-import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.onCompletion
+import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.isActive
 import java.time.Clock
-import java.util.UUID
 import javax.inject.Inject
 
 private const val ONE_SECOND_MILLISECOND = 1000L
 
 /**
  * Primary implementation of [TotpCodeManager].
+ *
+ * This implementation uses per-item [StateFlow] caching to prevent flow recreation on each
+ * subscribe, ensuring smooth UI updates when returning from background. The pattern mirrors
+ * the Password Manager's [TotpCodeManagerImpl].
  */
 class TotpCodeManagerImpl @Inject constructor(
     private val authenticatorSdkSource: AuthenticatorSdkSource,
     private val clock: Clock,
+    private val dispatcherManager: DispatcherManager,
 ) : TotpCodeManager {
 
+    private val unconfinedScope = CoroutineScope(dispatcherManager.unconfined)
+
+    /**
+     * Cache of per-item StateFlows to prevent recreation on each subscribe.
+     * Key is the [AuthenticatorItem], value is the cached [StateFlow] for that item.
+     */
+    private val mutableItemVerificationCodeStateFlowMap =
+        mutableMapOf<AuthenticatorItem, StateFlow<VerificationCodeItem?>>()
+
     override fun getTotpCodesFlow(
         itemList: List<AuthenticatorItem>,
-    ): Flow<List<VerificationCodeItem>> {
+    ): StateFlow<List<VerificationCodeItem>> {
         if (itemList.isEmpty()) {
-            return flowOf(emptyList())
+            return MutableStateFlow(emptyList())
+        }
+
+        val stateFlows = itemList.map { getOrCreateItemStateFlow(it) }
+
+        return combine(stateFlows) { results ->
+            results.filterNotNull()
         }
-        val flows = itemList.map { it.toFlowOfVerificationCodes() }
-        return combine(flows) { it.toList() }
+            .stateIn(
+                scope = unconfinedScope,
+                started = SharingStarted.WhileSubscribed(),
+                initialValue = emptyList(),
+            )
     }
 
-    private fun AuthenticatorItem.toFlowOfVerificationCodes(): Flow<VerificationCodeItem> {
-        val otpUri = this.otpUri
-        return flow {
-            var item: VerificationCodeItem? = null
-            while (currentCoroutineContext().isActive) {
-                val time = (clock.millis() / ONE_SECOND_MILLISECOND).toInt()
-                if (item == null || item.isExpired(clock)) {
-                    // If the item is expired or we haven't generated our first item,
-                    // generate a new code using the SDK:
-                    item = authenticatorSdkSource
-                        .generateTotp(otpUri, clock.instant())
-                        .getOrNull()
-                        ?.let { response ->
-                            VerificationCodeItem(
-                                code = response.code,
-                                periodSeconds = response.period.toInt(),
-                                timeLeftSeconds = response.period.toInt() -
-                                    time % response.period.toInt(),
-                                issueTime = clock.millis(),
-                                id = when (source) {
-                                    is AuthenticatorItem.Source.Local -> source.cipherId
-                                    is AuthenticatorItem.Source.Shared -> UUID.randomUUID()
-                                        .toString()
-                                },
-                                issuer = issuer,
-                                label = label,
-                                source = source,
-                            )
-                        }
-                        ?: run {
-                            // We are assuming that our otp URIs can generate a valid code.
-                            // If they can't, we'll just silently omit that code from the list.
-                            currentCoroutineContext().cancel()
-                            return@flow
-                        }
-                } else {
-                    // Item is not expired, just update time left:
-                    item = item.copy(
-                        timeLeftSeconds = item.periodSeconds - (time % item.periodSeconds),
-                    )
+    /**
+     * Gets an existing [StateFlow] for the given [item] or creates a new one if it doesn't exist.
+     * Each item gets its own [CoroutineScope] to manage its lifecycle independently.
+     */
+    private fun getOrCreateItemStateFlow(
+        item: AuthenticatorItem,
+    ): StateFlow<VerificationCodeItem?> {
+        return mutableItemVerificationCodeStateFlowMap.getOrPut(item) {
+            // Define a per-item scope so that we can clear the Flow from the map when it is
+            // no longer needed.
+            val itemScope = CoroutineScope(dispatcherManager.unconfined)
+
+            createVerificationCodeFlow(item)
+                .onCompletion {
+                    mutableItemVerificationCodeStateFlowMap.remove(item)
+                    itemScope.cancel()
                 }
-                // Emit item
-                emit(item)
-                // Wait one second before heading to the top of the loop:
-                delay(ONE_SECOND_MILLISECOND)
+                .stateIn(
+                    scope = itemScope,
+                    started = SharingStarted.WhileSubscribed(),
+                    initialValue = null,
+                )
+        }
+    }
+
+    /**
+     * Creates a flow that emits [VerificationCodeItem] updates every second for the given [item].
+     */
+    private fun createVerificationCodeFlow(
+        item: AuthenticatorItem,
+    ): Flow<VerificationCodeItem?> = flow {
+        var verificationCodeItem: VerificationCodeItem? = null
+
+        while (currentCoroutineContext().isActive) {
+            val dateTime = clock.instant()
+            val time = dateTime.epochSecond.toInt()
+
+            if (verificationCodeItem == null || verificationCodeItem.isExpired(clock)) {
+                // If the item is expired, or we haven't generated our first item,
+                // generate a new code using the SDK:
+                authenticatorSdkSource
+                    .generateTotp(item.otpUri, dateTime)
+                    .onSuccess { response ->
+                        verificationCodeItem = VerificationCodeItem(
+                            code = response.code,
+                            periodSeconds = response.period.toInt(),
+                            timeLeftSeconds = response.period.toInt() -
+                                (time % response.period.toInt()),
+                            issueTime = clock.millis(),
+                            id = item.cipherId,
+                            issuer = item.issuer,
+                            label = item.label,
+                            source = item.source,
+                        )
+                    }
+                    .onFailure {
+                        // We are assuming that our otp URIs can generate a valid code.
+                        // If they can't, we'll just silently omit that code from the list.
+                        emit(null)
+                        return@flow
+                    }
+            } else {
+                // Item is not expired, just update time left:
+                verificationCodeItem = verificationCodeItem.copy(
+                    timeLeftSeconds = verificationCodeItem.periodSeconds -
+                        (time % verificationCodeItem.periodSeconds),
+                )
             }
+
+            emit(verificationCodeItem)
+            delay(ONE_SECOND_MILLISECOND)
         }
     }
 }

authenticator/src/main/kotlin/com/bitwarden/authenticator/data/authenticator/manager/di/AuthenticatorManagerModule.kt

@@ -3,6 +3,7 @@ package com.bitwarden.authenticator.data.authenticator.manager.di
 import com.bitwarden.authenticator.data.authenticator.datasource.sdk.AuthenticatorSdkSource
 import com.bitwarden.authenticator.data.authenticator.manager.TotpCodeManager
 import com.bitwarden.authenticator.data.authenticator.manager.TotpCodeManagerImpl
+import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
 import dagger.Module
 import dagger.Provides
 import dagger.hilt.InstallIn
@@ -22,8 +23,10 @@ object AuthenticatorManagerModule {
     fun provideTotpCodeManager(
         authenticatorSdkSource: AuthenticatorSdkSource,
         clock: Clock,
+        dispatcherManager: DispatcherManager,
     ): TotpCodeManager = TotpCodeManagerImpl(
         authenticatorSdkSource = authenticatorSdkSource,
         clock = clock,
+        dispatcherManager = dispatcherManager,
     )
 }

authenticator/src/main/kotlin/com/bitwarden/authenticator/data/authenticator/repository/AuthenticatorRepositoryImpl.kt

@@ -157,7 +157,7 @@ class AuthenticatorRepositoryImpl @Inject constructor(
             .flatMapLatest { it.toSharedVerificationCodesStateFlow() }
             .stateIn(
                 scope = unconfinedScope,
-                started = SharingStarted.WhileSubscribed(STOP_TIMEOUT_DELAY_MS),
+                started = SharingStarted.WhileSubscribed(),
                 initialValue = SharedVerificationCodesState.Loading,
             )
     }
@@ -171,8 +171,8 @@ class AuthenticatorRepositoryImpl @Inject constructor(
                         authenticatorData.items
                             .map { entity ->
                                 AuthenticatorItem(
+                                    cipherId = entity.id,
                                     source = AuthenticatorItem.Source.Local(
-                                        cipherId = entity.id,
                                         isFavorite = entity.favorite,
                                     ),
                                     otpUri = entity.toOtpAuthUriString(),
@@ -197,7 +197,7 @@ class AuthenticatorRepositoryImpl @Inject constructor(
             }
             .stateIn(
                 scope = unconfinedScope,
-                started = SharingStarted.WhileSubscribed(STOP_TIMEOUT_DELAY_MS),
+                started = SharingStarted.WhileSubscribed(),
                 initialValue = DataState.Loading,
             )
     }

authenticator/src/main/kotlin/com/bitwarden/authenticator/data/authenticator/repository/model/AuthenticatorItem.kt

@@ -4,12 +4,14 @@ package com.bitwarden.authenticator.data.authenticator.repository.model
  * Represents all the information required to generate TOTP verification codes, including both
  * local codes and codes shared from the main Bitwarden app.
  *
- * @param source Distinguishes between local and shared items.
- * @param otpUri OTP URI.
- * @param issuer The issuer of the codes.
- * @param label The label of the item.
+ * @property cipherId The cipher ID.
+ * @property source Distinguishes between local and shared items.
+ * @property otpUri OTP URI.
+ * @property issuer The issuer of the codes.
+ * @property label The label of the item.
  */
 data class AuthenticatorItem(
+    val cipherId: String,
     val source: Source,
     val otpUri: String,
     val issuer: String?,
@@ -24,22 +26,20 @@ data class AuthenticatorItem(
         /**
          * The item is from the local Authenticator app database.
          *
-         * @param cipherId Local cipher ID.
-         * @param isFavorite Whether the user has marked the item as a favorite.
+         * @property isFavorite Whether the user has marked the item as a favorite.
          */
         data class Local(
-            val cipherId: String,
             val isFavorite: Boolean,
         ) : Source()
 
         /**
          * The item is shared from the main Bitwarden app.
          *
-         * @param userId User ID from the main Bitwarden app. Used to group authenticator items
+         * @property userId User ID from the main Bitwarden app. Used to group authenticator items
          * by account.
-         * @param nameOfUser Username from the main Bitwarden app.
-         * @param email Email of the user.
-         * @param environmentLabel Label for the Bitwarden environment, like "bitwaren.com"
+         * @property nameOfUser Username from the main Bitwarden app.
+         * @property email Email of the user.
+         * @property environmentLabel Label for the Bitwarden environment, like "bitwaren.com"
          */
         data class Shared(
             val userId: String,

authenticator/src/main/kotlin/com/bitwarden/authenticator/data/authenticator/repository/util/SharedAccountDataExtensions.kt

@@ -29,6 +29,7 @@ fun List<SharedAccountData.Account>.toAuthenticatorItems(): List<AuthenticatorIt
                     ?: cipherData.username
 
                 AuthenticatorItem(
+                    cipherId = cipherData.id,
                     source = AuthenticatorItem.Source.Shared(
                         userId = sharedAccount.userId,
                         nameOfUser = sharedAccount.name,

authenticator/src/test/kotlin/com/bitwarden/authenticator/data/authenticator/manager/util/TotpCodeManagerTest.kt

@@ -4,6 +4,11 @@ import app.cash.turbine.test
 import com.bitwarden.authenticator.data.authenticator.datasource.sdk.AuthenticatorSdkSource
 import com.bitwarden.authenticator.data.authenticator.manager.TotpCodeManagerImpl
 import com.bitwarden.authenticator.data.authenticator.manager.model.VerificationCodeItem
+import com.bitwarden.core.data.manager.dispatcher.FakeDispatcherManager
+import com.bitwarden.core.data.util.asFailure
+import com.bitwarden.core.data.util.asSuccess
+import com.bitwarden.vault.TotpResponse
+import io.mockk.coEvery
 import io.mockk.mockk
 import kotlinx.coroutines.test.runTest
 import org.junit.jupiter.api.Assertions.assertEquals
@@ -19,18 +24,63 @@ class TotpCodeManagerTest {
         ZoneOffset.UTC,
     )
     private val authenticatorSdkSource: AuthenticatorSdkSource = mockk()
+    private val dispatcherManager = FakeDispatcherManager()
 
     private val manager = TotpCodeManagerImpl(
         authenticatorSdkSource = authenticatorSdkSource,
         clock = clock,
+        dispatcherManager = dispatcherManager,
     )
 
     @Test
-    fun `getTotpCodesFlow should return flow that emits empty list when input list is empty`() =
+    fun `getTotpCodesFlow should emit empty list when input list is empty`() =
         runTest {
             manager.getTotpCodesFlow(emptyList()).test {
                 assertEquals(emptyList<VerificationCodeItem>(), awaitItem())
-                awaitComplete()
+            }
+        }
+
+    @Test
+    fun `getTotpCodesFlow should emit data if a valid value is passed in`() =
+        runTest {
+            val totp = "otpUri"
+            val authenticatorItems = listOf(
+                createMockAuthenticatorItem(number = 1, otpUri = totp),
+            )
+            val code = "123456"
+            val totpResponse = TotpResponse(code = code, period = 30u)
+            coEvery {
+                authenticatorSdkSource.generateTotp(totp = totp, time = clock.instant())
+            } returns totpResponse.asSuccess()
+
+            val expected = createMockVerificationCodeItem(
+                number = 1,
+                code = code,
+                issueTime = clock.instant().toEpochMilli(),
+                timeLeftSeconds = 30,
+            )
+
+            manager.getTotpCodesFlow(authenticatorItems).test {
+                assertEquals(listOf(expected), awaitItem())
+            }
+        }
+
+    @Test
+    fun `getTotpCodesFlow should emit empty list if unable to generate auth code`() =
+        runTest {
+            val totp = "otpUri"
+            val authenticatorItems = listOf(
+                createMockAuthenticatorItem(number = 1, otpUri = totp),
+            )
+            coEvery {
+                authenticatorSdkSource.generateTotp(totp = totp, time = clock.instant())
+            } returns Exception().asFailure()
+
+            manager.getTotpCodesFlow(authenticatorItems).test {
+                assertEquals(
+                    emptyList<VerificationCodeItem>(),
+                    awaitItem(),
+                )
             }
         }
 }