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

12 files changed, 103 insertions, 95 deletions

diff --git a/app/src/main/java/dev/msfjarvis/aps/ui/autofill/AutofillDecryptActivity.kt b/app/src/main/java/dev/msfjarvis/aps/ui/autofill/AutofillDecryptActivity.kt
index ccd4f649..b6b8a78a 100644
--- a/app/src/main/java/dev/msfjarvis/aps/ui/autofill/AutofillDecryptActivity.kt
+++ b/app/src/main/java/dev/msfjarvis/aps/ui/autofill/AutofillDecryptActivity.kt
@@ -43,7 +43,6 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
 import logcat.LogPriority.ERROR
-import logcat.asLog
 import logcat.logcat
 import me.msfjarvis.openpgpktx.util.OpenPgpApi
 import me.msfjarvis.openpgpktx.util.OpenPgpServiceConnection
@@ -204,7 +203,7 @@ class AutofillDecryptActivity : AppCompatActivity() {
                   val entry =
                     withContext(Dispatchers.IO) {
                       @Suppress("BlockingMethodInNonBlockingContext")
-                      passwordEntryFactory.create(lifecycleScope, decryptedOutput.toByteArray())
+                      passwordEntryFactory.create(decryptedOutput.toByteArray())
                     }
                   AutofillPreferences.credentialsFromStoreEntry(
                     this,
diff --git a/app/src/main/java/dev/msfjarvis/aps/ui/autofill/AutofillDecryptActivityV2.kt b/app/src/main/java/dev/msfjarvis/aps/ui/autofill/AutofillDecryptActivityV2.kt
index 555366c4..a784f64b 100644
--- a/app/src/main/java/dev/msfjarvis/aps/ui/autofill/AutofillDecryptActivityV2.kt
+++ b/app/src/main/java/dev/msfjarvis/aps/ui/autofill/AutofillDecryptActivityV2.kt
@@ -162,7 +162,7 @@ class AutofillDecryptActivityV2 : AppCompatActivity() {
           }
           .onSuccess { result ->
             return runCatching {
-              val entry = passwordEntryFactory.create(lifecycleScope, result.toByteArray())
+              val entry = passwordEntryFactory.create(result.toByteArray())
               AutofillPreferences.credentialsFromStoreEntry(this, file, entry, directoryStructure)
             }
               .getOrElse { e ->
diff --git a/app/src/main/java/dev/msfjarvis/aps/ui/crypto/DecryptActivity.kt b/app/src/main/java/dev/msfjarvis/aps/ui/crypto/DecryptActivity.kt
index 1fd4ee60..413e6f6a 100644
--- a/app/src/main/java/dev/msfjarvis/aps/ui/crypto/DecryptActivity.kt
+++ b/app/src/main/java/dev/msfjarvis/aps/ui/crypto/DecryptActivity.kt
@@ -31,6 +31,7 @@ import kotlin.time.ExperimentalTime
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.collect
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
@@ -178,7 +179,7 @@ class DecryptActivity : BasePgpActivity(), OpenPgpServiceConnection.OnBound {
           startAutoDismissTimer()
           runCatching {
             val showPassword = settings.getBoolean(PreferenceKeys.SHOW_PASSWORD, true)
-            val entry = passwordEntryFactory.create(lifecycleScope, outputStream.toByteArray())
+            val entry = passwordEntryFactory.create(outputStream.toByteArray())
 
             if (settings.getBoolean(PreferenceKeys.COPY_ON_DECRYPT, false)) {
               copyPasswordToClipboard(entry.password)
@@ -193,7 +194,7 @@ class DecryptActivity : BasePgpActivity(), OpenPgpServiceConnection.OnBound {
             }
 
             if (entry.hasTotp()) {
-              items.add(FieldItem.createOtpField(entry.totp.value))
+              items.add(FieldItem.createOtpField(entry.totp.first()))
             }
 
             if (!entry.username.isNullOrBlank()) {
diff --git a/app/src/main/java/dev/msfjarvis/aps/ui/crypto/DecryptActivityV2.kt b/app/src/main/java/dev/msfjarvis/aps/ui/crypto/DecryptActivityV2.kt
index 5b84478e..f5c02e49 100644
--- a/app/src/main/java/dev/msfjarvis/aps/ui/crypto/DecryptActivityV2.kt
+++ b/app/src/main/java/dev/msfjarvis/aps/ui/crypto/DecryptActivityV2.kt
@@ -31,6 +31,7 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.collectLatest
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
@@ -166,7 +167,7 @@ class DecryptActivityV2 : BasePgpActivity() {
     require(result.size() != 0) { "Incorrect password" }
     startAutoDismissTimer()
 
-    val entry = passwordEntryFactory.create(lifecycleScope, result.toByteArray())
+    val entry = passwordEntryFactory.create(result.toByteArray())
     passwordEntry = entry
     createPasswordUi(entry)
   }
@@ -182,7 +183,7 @@ class DecryptActivityV2 : BasePgpActivity() {
       }
 
       if (entry.hasTotp()) {
-        items.add(FieldItem.createOtpField(entry.totp.value))
+        items.add(FieldItem.createOtpField(entry.totp.first()))
       }
 
       if (!entry.username.isNullOrBlank()) {
diff --git a/app/src/main/java/dev/msfjarvis/aps/ui/crypto/PasswordCreationActivity.kt b/app/src/main/java/dev/msfjarvis/aps/ui/crypto/PasswordCreationActivity.kt
index 1b23dc0a..37a16b16 100644
--- a/app/src/main/java/dev/msfjarvis/aps/ui/crypto/PasswordCreationActivity.kt
+++ b/app/src/main/java/dev/msfjarvis/aps/ui/crypto/PasswordCreationActivity.kt
@@ -293,10 +293,7 @@ class PasswordCreationActivity : BasePgpActivity(), OpenPgpServiceConnection.OnB
               // User wants to disable username encryption, so we extract the
               // username from the encrypted extras and use it as the filename.
               val entry =
-                passwordEntryFactory.create(
-                  lifecycleScope,
-                  "PASSWORD\n${extraContent.text}".encodeToByteArray()
-                )
+                passwordEntryFactory.create("PASSWORD\n${extraContent.text}".encodeToByteArray())
               val username = entry.username
 
               // username should not be null here by the logic in
@@ -370,10 +367,7 @@ class PasswordCreationActivity : BasePgpActivity(), OpenPgpServiceConnection.OnB
     with(binding) {
       // Use PasswordEntry to parse extras for username
       val entry =
-        passwordEntryFactory.create(
-          lifecycleScope,
-          "PLACEHOLDER\n${extraContent.text}".encodeToByteArray()
-        )
+        passwordEntryFactory.create("PLACEHOLDER\n${extraContent.text}".encodeToByteArray())
       encryptUsername.apply {
         if (visibility != View.VISIBLE) return@apply
         val hasUsernameInFileName = filename.text.toString().isNotBlank()
@@ -529,7 +523,7 @@ class PasswordCreationActivity : BasePgpActivity(), OpenPgpServiceConnection.OnB
 
               if (shouldGeneratePassword) {
                 val directoryStructure = AutofillPreferences.directoryStructure(applicationContext)
-                val entry = passwordEntryFactory.create(lifecycleScope, content.encodeToByteArray())
+                val entry = passwordEntryFactory.create(content.encodeToByteArray())
                 returnIntent.putExtra(RETURN_EXTRA_PASSWORD, entry.password)
                 val username = entry.username ?: directoryStructure.getUsernameFor(file)
                 returnIntent.putExtra(RETURN_EXTRA_USERNAME, username)
diff --git a/app/src/main/java/dev/msfjarvis/aps/ui/crypto/PasswordCreationActivityV2.kt b/app/src/main/java/dev/msfjarvis/aps/ui/crypto/PasswordCreationActivityV2.kt
index a191ba61..e6ee7a09 100644
--- a/app/src/main/java/dev/msfjarvis/aps/ui/crypto/PasswordCreationActivityV2.kt
+++ b/app/src/main/java/dev/msfjarvis/aps/ui/crypto/PasswordCreationActivityV2.kt
@@ -223,10 +223,7 @@ class PasswordCreationActivityV2 : BasePgpActivity() {
               // User wants to disable username encryption, so we extract the
               // username from the encrypted extras and use it as the filename.
               val entry =
-                passwordEntryFactory.create(
-                  lifecycleScope,
-                  "PASSWORD\n${extraContent.text}".encodeToByteArray()
-                )
+                passwordEntryFactory.create("PASSWORD\n${extraContent.text}".encodeToByteArray())
               val username = entry.username
 
               // username should not be null here by the logic in
@@ -300,10 +297,7 @@ class PasswordCreationActivityV2 : BasePgpActivity() {
     with(binding) {
       // Use PasswordEntry to parse extras for username
       val entry =
-        passwordEntryFactory.create(
-          lifecycleScope,
-          "PLACEHOLDER\n${extraContent.text}".encodeToByteArray()
-        )
+        passwordEntryFactory.create("PLACEHOLDER\n${extraContent.text}".encodeToByteArray())
       encryptUsername.apply {
         if (visibility != View.VISIBLE) return@apply
         val hasUsernameInFileName = filename.text.toString().isNotBlank()
@@ -406,7 +400,7 @@ class PasswordCreationActivityV2 : BasePgpActivity() {
 
           if (shouldGeneratePassword) {
             val directoryStructure = AutofillPreferences.directoryStructure(applicationContext)
-            val entry = passwordEntryFactory.create(lifecycleScope, content.encodeToByteArray())
+            val entry = passwordEntryFactory.create(content.encodeToByteArray())
             returnIntent.putExtra(RETURN_EXTRA_PASSWORD, entry.password)
             val username = entry.username ?: directoryStructure.getUsernameFor(file)
             returnIntent.putExtra(RETURN_EXTRA_USERNAME, username)
diff --git a/app/src/main/java/dev/msfjarvis/aps/util/autofill/AutofillPreferences.kt b/app/src/main/java/dev/msfjarvis/aps/util/autofill/AutofillPreferences.kt
index fdbd15af..2028fedb 100644
--- a/app/src/main/java/dev/msfjarvis/aps/util/autofill/AutofillPreferences.kt
+++ b/app/src/main/java/dev/msfjarvis/aps/util/autofill/AutofillPreferences.kt
@@ -15,6 +15,8 @@ import dev.msfjarvis.aps.util.services.getDefaultUsername
 import dev.msfjarvis.aps.util.settings.PreferenceKeys
 import java.io.File
 import java.nio.file.Paths
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.runBlocking
 
 enum class DirectoryStructure(val value: String) {
   EncryptedUsername("encrypted_username"),
@@ -141,6 +143,6 @@ object AutofillPreferences {
     // Always give priority to a username stored in the encrypted extras
     val username =
       entry.username ?: directoryStructure.getUsernameFor(file) ?: context.getDefaultUsername()
-    return Credentials(username, entry.password, entry.totp.value)
+    return Credentials(username, entry.password, runBlocking { entry.totp.first() })
   }
 }
diff --git a/coroutine-utils-testing/build.gradle.kts b/coroutine-utils-testing/build.gradle.kts
index 96d87dc1..c4f7400b 100644
--- a/coroutine-utils-testing/build.gradle.kts
+++ b/coroutine-utils-testing/build.gradle.kts
@@ -11,4 +11,5 @@ dependencies {
   implementation(projects.coroutineUtils)
   implementation(libs.testing.junit)
   implementation(libs.kotlin.coroutines.test)
+  api(libs.testing.turbine)
 }
diff --git a/coroutine-utils-testing/src/main/kotlin/dev/msfjarvis/aps/test/turbine.ext.kt b/coroutine-utils-testing/src/main/kotlin/dev/msfjarvis/aps/test/turbine.ext.kt
new file mode 100644
index 00000000..596be1a1
--- /dev/null
+++ b/coroutine-utils-testing/src/main/kotlin/dev/msfjarvis/aps/test/turbine.ext.kt
@@ -0,0 +1,37 @@
+/*
+ * Copyright © 2014-2022 The Android Password Store Authors. All Rights Reserved.
+ * SPDX-License-Identifier: GPL-3.0-only
+ */
+
+package dev.msfjarvis.aps.test
+
+import app.cash.turbine.FlowTurbine
+import app.cash.turbine.test
+import kotlin.coroutines.coroutineContext
+import kotlin.time.Duration
+import kotlin.time.Duration.Companion.seconds
+import kotlin.time.ExperimentalTime
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.flowOn
+import kotlinx.coroutines.test.TestCoroutineScheduler
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+
+/**
+ * Wrapper for [test] that implements compatibility with kotlinx.coroutines 1.6.0
+ *
+ * @see "https://github.com/cashapp/turbine/issues/42#issuecomment-1000317026"
+ */
+@ExperimentalTime
+@ExperimentalCoroutinesApi
+public suspend fun <T> Flow<T>.test2(
+  timeout: Duration = 1.seconds,
+  validate: suspend FlowTurbine<T>.() -> Unit,
+) {
+  val testScheduler = coroutineContext[TestCoroutineScheduler]
+  return if (testScheduler == null) {
+    test(timeout, validate)
+  } else {
+    flowOn(UnconfinedTestDispatcher(testScheduler)).test(timeout, validate)
+  }
+}
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()
     }
   }
 
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 6ce369c8..e2375fa4 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -93,6 +93,7 @@ testing-junit = "junit:junit:4.13.2"
 testing-kotlintest-junit = { module = "org.jetbrains.kotlin:kotlin-test-junit", version.ref = "kotlin" }
 testing-robolectric = "org.robolectric:robolectric:4.7.3"
 testing-sharedPrefsMock = "com.github.android-password-store:shared-preferences-fake:2.0.0"
+testing-turbine = "app.cash.turbine:turbine:0.7.0"
 
 [bundles]
 androidxLifecycle = ["androidx-lifecycle-common", "androidx-lifecycle-livedataKtx", "androidx-lifecycle-viewmodelKtx"]