Refactor PasswordEntry TOTP calculation into a cold flow (#1702)

2 files changed, 47 insertions, 69 deletions

diff --git a/format-common/src/main/kotlin/dev/msfjarvis/aps/data/passfile/PasswordEntry.kt b/format-common/src/main/kotlin/dev/msfjarvis/aps/data/passfile/PasswordEntry.kt
index 2bccd2cc..c36cb0ed 100644
--- a/format-common/src/main/kotlin/dev/msfjarvis/aps/data/passfile/PasswordEntry.kt
+++ b/format-common/src/main/kotlin/dev/msfjarvis/aps/data/passfile/PasswordEntry.kt
@@ -9,20 +9,15 @@ import com.github.michaelbull.result.mapBoth
 import dagger.assisted.Assisted
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
-import dev.msfjarvis.aps.util.coroutines.DispatcherProvider
 import dev.msfjarvis.aps.util.time.UserClock
 import dev.msfjarvis.aps.util.totp.Otp
 import dev.msfjarvis.aps.util.totp.TotpFinder
 import kotlin.collections.set
-import kotlin.time.Duration
 import kotlin.time.ExperimentalTime
-import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.awaitCancellation
 import kotlinx.coroutines.delay
-import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.asStateFlow
-import kotlinx.coroutines.flow.update
-import kotlinx.coroutines.launch
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.flow
 
 /** Represents a single entry in the password store. */
 @OptIn(ExperimentalTime::class)
@@ -30,20 +25,13 @@ public class PasswordEntry
 @AssistedInject
 constructor(
   /** A time source used to calculate the TOTP */
-  clock: UserClock,
+  private val clock: UserClock,
   /** [TotpFinder] implementation to extract data from a TOTP URI */
-  totpFinder: TotpFinder,
-  /** Instance of [DispatcherProvider] to select an IO dispatcher for emitting TOTP values. */
-  dispatcherProvider: DispatcherProvider,
-  /**
-   * A cancellable [CoroutineScope] inside which we constantly emit new TOTP values as time elapses
-   */
-  @Assisted scope: CoroutineScope,
+  private val totpFinder: TotpFinder,
   /** The content of this entry, as an array of bytes. */
   @Assisted bytes: ByteArray,
 ) {
 
-  private val _totp = MutableStateFlow("")
   private val content = bytes.decodeToString()
 
   /** The password text for this entry. Can be null. */
@@ -62,11 +50,21 @@ constructor(
   public val extraContentString: String
 
   /**
-   * A [StateFlow] providing the current TOTP. It will emit a single empty string on initialization
-   * which is replaced with a real TOTP if applicable. Call [hasTotp] to verify whether or not you
-   * need to observe this value.
+   * A [Flow] providing the current TOTP. It will start emitting only when collected. If this entry
+   * does not have a TOTP secret, the flow will never emit. Users should call [hasTotp] before
+   * collection to check if it is valid to collect this [Flow].
    */
-  public val totp: StateFlow<String> = _totp.asStateFlow()
+  public val totp: Flow<String> = flow {
+    if (totpSecret != null) {
+      repeat(Int.MAX_VALUE) {
+        val (otp, remainingTime) = calculateTotp()
+        emit(otp)
+        delay(remainingTime)
+      }
+    } else {
+      awaitCancellation()
+    }
+  }
 
   /**
    * String representation of [extraContent] but with authentication related data such as TOTP URIs
@@ -83,21 +81,6 @@ constructor(
     extraContent = generateExtraContentPairs()
     username = findUsername()
     totpSecret = totpFinder.findSecret(content)
-    if (totpSecret != null) {
-      scope.launch(dispatcherProvider.io()) {
-        val digits = totpFinder.findDigits(content)
-        val totpPeriod = totpFinder.findPeriod(content)
-        val totpAlgorithm = totpFinder.findAlgorithm(content)
-        val issuer = totpFinder.findIssuer(content)
-        val remainingTime = totpPeriod - (clock.millis() % totpPeriod)
-        updateTotp(clock.millis(), totpPeriod, totpAlgorithm, digits, issuer)
-        delay(Duration.seconds(remainingTime))
-        repeat(Int.MAX_VALUE) {
-          updateTotp(clock.millis(), totpPeriod, totpAlgorithm, digits, issuer)
-          delay(Duration.seconds(totpPeriod))
-        }
-      }
-    }
   }
 
   public fun hasTotp(): Boolean {
@@ -188,22 +171,25 @@ constructor(
     return null
   }
 
-  private fun updateTotp(
-    millis: Long,
-    totpPeriod: Long,
-    totpAlgorithm: String,
-    digits: String,
-    issuer: String?,
-  ) {
-    if (totpSecret != null) {
-      Otp.calculateCode(totpSecret, millis / (1000 * totpPeriod), totpAlgorithm, digits, issuer)
-        .mapBoth({ code -> _totp.update { code } }, { throwable -> throw throwable })
-    }
+  private fun calculateTotp(): Pair<String, Long> {
+    val digits = totpFinder.findDigits(content)
+    val totpPeriod = totpFinder.findPeriod(content)
+    val totpAlgorithm = totpFinder.findAlgorithm(content)
+    val issuer = totpFinder.findIssuer(content)
+    val millis = clock.millis()
+    val remainingTime = totpPeriod - (millis % totpPeriod)
+    Otp.calculateCode(totpSecret!!, millis / (1000 * totpPeriod), totpAlgorithm, digits, issuer)
+      .mapBoth(
+        { code ->
+          return code to remainingTime
+        },
+        { throwable -> throw throwable }
+      )
   }
 
   @AssistedFactory
   public interface Factory {
-    public fun create(scope: CoroutineScope, bytes: ByteArray): PasswordEntry
+    public fun create(bytes: ByteArray): PasswordEntry
   }
 
   internal companion object {
diff --git a/format-common/src/test/kotlin/dev/msfjarvis/aps/data/passfile/PasswordEntryTest.kt b/format-common/src/test/kotlin/dev/msfjarvis/aps/data/passfile/PasswordEntryTest.kt
index 0f63a74a..d8a1fdc2 100644
--- a/format-common/src/test/kotlin/dev/msfjarvis/aps/data/passfile/PasswordEntryTest.kt
+++ b/format-common/src/test/kotlin/dev/msfjarvis/aps/data/passfile/PasswordEntryTest.kt
@@ -6,10 +6,10 @@
 package dev.msfjarvis.aps.data.passfile
 
 import dev.msfjarvis.aps.test.CoroutineTestRule
+import dev.msfjarvis.aps.test.test2
 import dev.msfjarvis.aps.util.time.TestUserClock
 import dev.msfjarvis.aps.util.totp.TotpFinder
 import java.util.Locale
-import kotlin.test.Ignore
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertNotNull
@@ -17,7 +17,6 @@ import kotlin.test.assertNull
 import kotlin.test.assertTrue
 import kotlin.time.ExperimentalTime
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
 import org.junit.Rule
 
@@ -29,8 +28,6 @@ class PasswordEntryTest {
     PasswordEntry(
       fakeClock,
       testFinder,
-      coroutineTestRule.testDispatcherProvider,
-      TestScope(coroutineTestRule.testDispatcher),
       content.encodeToByteArray(),
     )
 
@@ -133,27 +130,22 @@ class PasswordEntryTest {
   }
 
   @Test
-  @Ignore("Timing with runTest seems hard to implement right now")
-  fun testGeneratesOtpFromTotpUri() {
-    runTest {
-      val entry = makeEntry("secret\nextra\n$TOTP_URI")
-      assertTrue(entry.hasTotp())
-      val code = entry.totp.value
-      assertNotNull(code) { "Generated OTP cannot be null" }
-      assertEquals("818800", code)
+  fun testGeneratesOtpFromTotpUri() = runTest {
+    val entry = makeEntry("secret\nextra\n$TOTP_URI")
+    assertTrue(entry.hasTotp())
+    entry.totp.test2 {
+      assertEquals("818800", expectMostRecentItem())
+      cancelAndIgnoreRemainingEvents()
     }
   }
 
   @Test
-  @Ignore("Timing with runTest seems hard to implement right now")
-  fun testGeneratesOtpWithOnlyUriInFile() {
-    runTest {
-      val entry = makeEntry(TOTP_URI)
-      assertNull(entry.password)
-      assertTrue(entry.hasTotp())
-      val code = entry.totp.value
-      assertNotNull(code) { "Generated OTP cannot be null" }
-      assertEquals("818800", code)
+  fun testGeneratesOtpWithOnlyUriInFile() = runTest {
+    val entry = makeEntry(TOTP_URI)
+    assertNull(entry.password)
+    entry.totp.test2 {
+      assertEquals("818800", expectMostRecentItem())
+      cancelAndIgnoreRemainingEvents()
     }
   }