Skip to content

Commit 8585b9f

Browse files
david-livefronttajchertclaude
authored
[PM-29309] [BWA-209] bug: Fix TOTP countdown freeze when returning to Authenticator app (change Flow to StateFlow) (#6764)
Co-authored-by: Michal Tajchert <michal.tajchert@lite.tech> Co-authored-by: Claude <noreply@anthropic.com>
1 parent 09a2abf commit 8585b9f

11 files changed

Lines changed: 205 additions & 78 deletions

File tree

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,19 +3,20 @@ package com.bitwarden.authenticator.data.authenticator.manager
33
import com.bitwarden.authenticator.data.authenticator.datasource.disk.entity.AuthenticatorItemAlgorithm
44
import com.bitwarden.authenticator.data.authenticator.manager.model.VerificationCodeItem
55
import com.bitwarden.authenticator.data.authenticator.repository.model.AuthenticatorItem
6-
import kotlinx.coroutines.flow.Flow
6+
import kotlinx.coroutines.flow.StateFlow
77

88
/**
99
* Manages the flows for getting verification codes.
1010
*/
1111
interface TotpCodeManager {
1212

1313
/**
14-
* Flow for getting a DataState with multiple verification code items.
14+
* StateFlow for getting multiple verification code items. Returns a StateFlow that emits
15+
* updated verification codes every second.
1516
*/
1617
fun getTotpCodesFlow(
1718
itemList: List<AuthenticatorItem>,
18-
): Flow<List<VerificationCodeItem>>
19+
): StateFlow<List<VerificationCodeItem>>
1920

2021
@Suppress("UndocumentedPublicClass")
2122
companion object {

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

Lines changed: 103 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -3,84 +3,137 @@ package com.bitwarden.authenticator.data.authenticator.manager
33
import com.bitwarden.authenticator.data.authenticator.datasource.sdk.AuthenticatorSdkSource
44
import com.bitwarden.authenticator.data.authenticator.manager.model.VerificationCodeItem
55
import com.bitwarden.authenticator.data.authenticator.repository.model.AuthenticatorItem
6+
import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
7+
import kotlinx.coroutines.CoroutineScope
68
import kotlinx.coroutines.cancel
79
import kotlinx.coroutines.currentCoroutineContext
810
import kotlinx.coroutines.delay
911
import kotlinx.coroutines.flow.Flow
12+
import kotlinx.coroutines.flow.MutableStateFlow
13+
import kotlinx.coroutines.flow.SharingStarted
14+
import kotlinx.coroutines.flow.StateFlow
1015
import kotlinx.coroutines.flow.combine
1116
import kotlinx.coroutines.flow.flow
12-
import kotlinx.coroutines.flow.flowOf
17+
import kotlinx.coroutines.flow.onCompletion
18+
import kotlinx.coroutines.flow.stateIn
1319
import kotlinx.coroutines.isActive
1420
import java.time.Clock
15-
import java.util.UUID
1621
import javax.inject.Inject
1722

1823
private const val ONE_SECOND_MILLISECOND = 1000L
1924

2025
/**
2126
* Primary implementation of [TotpCodeManager].
27+
*
28+
* This implementation uses per-item [StateFlow] caching to prevent flow recreation on each
29+
* subscribe, ensuring smooth UI updates when returning from background. The pattern mirrors
30+
* the Password Manager's [TotpCodeManagerImpl].
2231
*/
2332
class TotpCodeManagerImpl @Inject constructor(
2433
private val authenticatorSdkSource: AuthenticatorSdkSource,
2534
private val clock: Clock,
35+
private val dispatcherManager: DispatcherManager,
2636
) : TotpCodeManager {
2737

38+
private val unconfinedScope = CoroutineScope(dispatcherManager.unconfined)
39+
40+
/**
41+
* Cache of per-item StateFlows to prevent recreation on each subscribe.
42+
* Key is the [AuthenticatorItem], value is the cached [StateFlow] for that item.
43+
*/
44+
private val mutableItemVerificationCodeStateFlowMap =
45+
mutableMapOf<AuthenticatorItem, StateFlow<VerificationCodeItem?>>()
46+
2847
override fun getTotpCodesFlow(
2948
itemList: List<AuthenticatorItem>,
30-
): Flow<List<VerificationCodeItem>> {
49+
): StateFlow<List<VerificationCodeItem>> {
3150
if (itemList.isEmpty()) {
32-
return flowOf(emptyList())
51+
return MutableStateFlow(emptyList())
52+
}
53+
54+
val stateFlows = itemList.map { getOrCreateItemStateFlow(it) }
55+
56+
return combine(stateFlows) { results ->
57+
results.filterNotNull()
3358
}
34-
val flows = itemList.map { it.toFlowOfVerificationCodes() }
35-
return combine(flows) { it.toList() }
59+
.stateIn(
60+
scope = unconfinedScope,
61+
started = SharingStarted.WhileSubscribed(),
62+
initialValue = emptyList(),
63+
)
3664
}
3765

38-
private fun AuthenticatorItem.toFlowOfVerificationCodes(): Flow<VerificationCodeItem> {
39-
val otpUri = this.otpUri
40-
return flow {
41-
var item: VerificationCodeItem? = null
42-
while (currentCoroutineContext().isActive) {
43-
val time = (clock.millis() / ONE_SECOND_MILLISECOND).toInt()
44-
if (item == null || item.isExpired(clock)) {
45-
// If the item is expired or we haven't generated our first item,
46-
// generate a new code using the SDK:
47-
item = authenticatorSdkSource
48-
.generateTotp(otpUri, clock.instant())
49-
.getOrNull()
50-
?.let { response ->
51-
VerificationCodeItem(
52-
code = response.code,
53-
periodSeconds = response.period.toInt(),
54-
timeLeftSeconds = response.period.toInt() -
55-
time % response.period.toInt(),
56-
issueTime = clock.millis(),
57-
id = when (source) {
58-
is AuthenticatorItem.Source.Local -> source.cipherId
59-
is AuthenticatorItem.Source.Shared -> UUID.randomUUID()
60-
.toString()
61-
},
62-
issuer = issuer,
63-
label = label,
64-
source = source,
65-
)
66-
}
67-
?: run {
68-
// We are assuming that our otp URIs can generate a valid code.
69-
// If they can't, we'll just silently omit that code from the list.
70-
currentCoroutineContext().cancel()
71-
return@flow
72-
}
73-
} else {
74-
// Item is not expired, just update time left:
75-
item = item.copy(
76-
timeLeftSeconds = item.periodSeconds - (time % item.periodSeconds),
77-
)
66+
/**
67+
* Gets an existing [StateFlow] for the given [item] or creates a new one if it doesn't exist.
68+
* Each item gets its own [CoroutineScope] to manage its lifecycle independently.
69+
*/
70+
private fun getOrCreateItemStateFlow(
71+
item: AuthenticatorItem,
72+
): StateFlow<VerificationCodeItem?> {
73+
return mutableItemVerificationCodeStateFlowMap.getOrPut(item) {
74+
// Define a per-item scope so that we can clear the Flow from the map when it is
75+
// no longer needed.
76+
val itemScope = CoroutineScope(dispatcherManager.unconfined)
77+
78+
createVerificationCodeFlow(item)
79+
.onCompletion {
80+
mutableItemVerificationCodeStateFlowMap.remove(item)
81+
itemScope.cancel()
7882
}
79-
// Emit item
80-
emit(item)
81-
// Wait one second before heading to the top of the loop:
82-
delay(ONE_SECOND_MILLISECOND)
83+
.stateIn(
84+
scope = itemScope,
85+
started = SharingStarted.WhileSubscribed(),
86+
initialValue = null,
87+
)
88+
}
89+
}
90+
91+
/**
92+
* Creates a flow that emits [VerificationCodeItem] updates every second for the given [item].
93+
*/
94+
private fun createVerificationCodeFlow(
95+
item: AuthenticatorItem,
96+
): Flow<VerificationCodeItem?> = flow {
97+
var verificationCodeItem: VerificationCodeItem? = null
98+
99+
while (currentCoroutineContext().isActive) {
100+
val dateTime = clock.instant()
101+
val time = dateTime.epochSecond.toInt()
102+
103+
if (verificationCodeItem == null || verificationCodeItem.isExpired(clock)) {
104+
// If the item is expired, or we haven't generated our first item,
105+
// generate a new code using the SDK:
106+
authenticatorSdkSource
107+
.generateTotp(item.otpUri, dateTime)
108+
.onSuccess { response ->
109+
verificationCodeItem = VerificationCodeItem(
110+
code = response.code,
111+
periodSeconds = response.period.toInt(),
112+
timeLeftSeconds = response.period.toInt() -
113+
(time % response.period.toInt()),
114+
issueTime = clock.millis(),
115+
id = item.cipherId,
116+
issuer = item.issuer,
117+
label = item.label,
118+
source = item.source,
119+
)
120+
}
121+
.onFailure {
122+
// We are assuming that our otp URIs can generate a valid code.
123+
// If they can't, we'll just silently omit that code from the list.
124+
emit(null)
125+
return@flow
126+
}
127+
} else {
128+
// Item is not expired, just update time left:
129+
verificationCodeItem = verificationCodeItem.copy(
130+
timeLeftSeconds = verificationCodeItem.periodSeconds -
131+
(time % verificationCodeItem.periodSeconds),
132+
)
83133
}
134+
135+
emit(verificationCodeItem)
136+
delay(ONE_SECOND_MILLISECOND)
84137
}
85138
}
86139
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package com.bitwarden.authenticator.data.authenticator.manager.di
33
import com.bitwarden.authenticator.data.authenticator.datasource.sdk.AuthenticatorSdkSource
44
import com.bitwarden.authenticator.data.authenticator.manager.TotpCodeManager
55
import com.bitwarden.authenticator.data.authenticator.manager.TotpCodeManagerImpl
6+
import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
67
import dagger.Module
78
import dagger.Provides
89
import dagger.hilt.InstallIn
@@ -22,8 +23,10 @@ object AuthenticatorManagerModule {
2223
fun provideTotpCodeManager(
2324
authenticatorSdkSource: AuthenticatorSdkSource,
2425
clock: Clock,
26+
dispatcherManager: DispatcherManager,
2527
): TotpCodeManager = TotpCodeManagerImpl(
2628
authenticatorSdkSource = authenticatorSdkSource,
2729
clock = clock,
30+
dispatcherManager = dispatcherManager,
2831
)
2932
}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,7 @@ class AuthenticatorRepositoryImpl @Inject constructor(
157157
.flatMapLatest { it.toSharedVerificationCodesStateFlow() }
158158
.stateIn(
159159
scope = unconfinedScope,
160-
started = SharingStarted.WhileSubscribed(STOP_TIMEOUT_DELAY_MS),
160+
started = SharingStarted.WhileSubscribed(),
161161
initialValue = SharedVerificationCodesState.Loading,
162162
)
163163
}
@@ -171,8 +171,8 @@ class AuthenticatorRepositoryImpl @Inject constructor(
171171
authenticatorData.items
172172
.map { entity ->
173173
AuthenticatorItem(
174+
cipherId = entity.id,
174175
source = AuthenticatorItem.Source.Local(
175-
cipherId = entity.id,
176176
isFavorite = entity.favorite,
177177
),
178178
otpUri = entity.toOtpAuthUriString(),
@@ -197,7 +197,7 @@ class AuthenticatorRepositoryImpl @Inject constructor(
197197
}
198198
.stateIn(
199199
scope = unconfinedScope,
200-
started = SharingStarted.WhileSubscribed(STOP_TIMEOUT_DELAY_MS),
200+
started = SharingStarted.WhileSubscribed(),
201201
initialValue = DataState.Loading,
202202
)
203203
}

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

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,14 @@ package com.bitwarden.authenticator.data.authenticator.repository.model
44
* Represents all the information required to generate TOTP verification codes, including both
55
* local codes and codes shared from the main Bitwarden app.
66
*
7-
* @param source Distinguishes between local and shared items.
8-
* @param otpUri OTP URI.
9-
* @param issuer The issuer of the codes.
10-
* @param label The label of the item.
7+
* @property cipherId The cipher ID.
8+
* @property source Distinguishes between local and shared items.
9+
* @property otpUri OTP URI.
10+
* @property issuer The issuer of the codes.
11+
* @property label The label of the item.
1112
*/
1213
data class AuthenticatorItem(
14+
val cipherId: String,
1315
val source: Source,
1416
val otpUri: String,
1517
val issuer: String?,
@@ -24,22 +26,20 @@ data class AuthenticatorItem(
2426
/**
2527
* The item is from the local Authenticator app database.
2628
*
27-
* @param cipherId Local cipher ID.
28-
* @param isFavorite Whether the user has marked the item as a favorite.
29+
* @property isFavorite Whether the user has marked the item as a favorite.
2930
*/
3031
data class Local(
31-
val cipherId: String,
3232
val isFavorite: Boolean,
3333
) : Source()
3434

3535
/**
3636
* The item is shared from the main Bitwarden app.
3737
*
38-
* @param userId User ID from the main Bitwarden app. Used to group authenticator items
38+
* @property userId User ID from the main Bitwarden app. Used to group authenticator items
3939
* by account.
40-
* @param nameOfUser Username from the main Bitwarden app.
41-
* @param email Email of the user.
42-
* @param environmentLabel Label for the Bitwarden environment, like "bitwaren.com"
40+
* @property nameOfUser Username from the main Bitwarden app.
41+
* @property email Email of the user.
42+
* @property environmentLabel Label for the Bitwarden environment, like "bitwaren.com"
4343
*/
4444
data class Shared(
4545
val userId: String,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ fun List<SharedAccountData.Account>.toAuthenticatorItems(): List<AuthenticatorIt
2929
?: cipherData.username
3030

3131
AuthenticatorItem(
32+
cipherId = cipherData.id,
3233
source = AuthenticatorItem.Source.Shared(
3334
userId = sharedAccount.userId,
3435
nameOfUser = sharedAccount.name,

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

Lines changed: 52 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,11 @@ import app.cash.turbine.test
44
import com.bitwarden.authenticator.data.authenticator.datasource.sdk.AuthenticatorSdkSource
55
import com.bitwarden.authenticator.data.authenticator.manager.TotpCodeManagerImpl
66
import com.bitwarden.authenticator.data.authenticator.manager.model.VerificationCodeItem
7+
import com.bitwarden.core.data.manager.dispatcher.FakeDispatcherManager
8+
import com.bitwarden.core.data.util.asFailure
9+
import com.bitwarden.core.data.util.asSuccess
10+
import com.bitwarden.vault.TotpResponse
11+
import io.mockk.coEvery
712
import io.mockk.mockk
813
import kotlinx.coroutines.test.runTest
914
import org.junit.jupiter.api.Assertions.assertEquals
@@ -19,18 +24,63 @@ class TotpCodeManagerTest {
1924
ZoneOffset.UTC,
2025
)
2126
private val authenticatorSdkSource: AuthenticatorSdkSource = mockk()
27+
private val dispatcherManager = FakeDispatcherManager()
2228

2329
private val manager = TotpCodeManagerImpl(
2430
authenticatorSdkSource = authenticatorSdkSource,
2531
clock = clock,
32+
dispatcherManager = dispatcherManager,
2633
)
2734

2835
@Test
29-
fun `getTotpCodesFlow should return flow that emits empty list when input list is empty`() =
36+
fun `getTotpCodesFlow should emit empty list when input list is empty`() =
3037
runTest {
3138
manager.getTotpCodesFlow(emptyList()).test {
3239
assertEquals(emptyList<VerificationCodeItem>(), awaitItem())
33-
awaitComplete()
40+
}
41+
}
42+
43+
@Test
44+
fun `getTotpCodesFlow should emit data if a valid value is passed in`() =
45+
runTest {
46+
val totp = "otpUri"
47+
val authenticatorItems = listOf(
48+
createMockAuthenticatorItem(number = 1, otpUri = totp),
49+
)
50+
val code = "123456"
51+
val totpResponse = TotpResponse(code = code, period = 30u)
52+
coEvery {
53+
authenticatorSdkSource.generateTotp(totp = totp, time = clock.instant())
54+
} returns totpResponse.asSuccess()
55+
56+
val expected = createMockVerificationCodeItem(
57+
number = 1,
58+
code = code,
59+
issueTime = clock.instant().toEpochMilli(),
60+
timeLeftSeconds = 30,
61+
)
62+
63+
manager.getTotpCodesFlow(authenticatorItems).test {
64+
assertEquals(listOf(expected), awaitItem())
65+
}
66+
}
67+
68+
@Test
69+
fun `getTotpCodesFlow should emit empty list if unable to generate auth code`() =
70+
runTest {
71+
val totp = "otpUri"
72+
val authenticatorItems = listOf(
73+
createMockAuthenticatorItem(number = 1, otpUri = totp),
74+
)
75+
coEvery {
76+
authenticatorSdkSource.generateTotp(totp = totp, time = clock.instant())
77+
} returns Exception().asFailure()
78+
79+
manager.getTotpCodesFlow(authenticatorItems).test {
80+
assertEquals(
81+
emptyList<VerificationCodeItem>(),
82+
awaitItem(),
83+
)
3484
}
3585
}
3686
}

0 commit comments

Comments
 (0)