Use libsignal for TOTP operations.

app/src/main/java/org/thoughtcrime/securesms/components/settings/app/account/authenticator/InMemoryTotpApi.kt

@@ -1,92 +0,0 @@
-/*
- * Copyright 2026 Signal Messenger, LLC
- * SPDX-License-Identifier: AGPL-3.0-only
- */
-
-package org.thoughtcrime.securesms.components.settings.app.account.authenticator
-
-import org.signal.libsignal.net.RequestResult
-import java.security.SecureRandom
-
-/**
- * Stands in for the service until the gRPC methods land. Nothing here is persisted or sent anywhere, so it lasts only
- * as long as the process does.
- *
- * It doesn't verify one-time passwords -- any code confirms the pending key. It does copy the service's behaviour where
- * that behaviour shapes the screens: one pending key at a time, [TotpApi.MAX_KEYS] confirmed keys, ids drawn from the
- * lowest free slot, and name length enforced.
- */
-class InMemoryTotpApi : TotpApi {
-
-  companion object {
-    /** What the service generates: a 256-bit key. */
-    private const val KEY_LENGTH_BYTES = 32
-  }
-
-  private val lock = Any()
-  private val secureRandom = SecureRandom()
-
-  private var hasPendingKey = false
-  private val confirmedKeys = mutableMapOf<Int, TotpApi.Metadata>()
-
-  override suspend fun generateKey(): RequestResult<TotpApi.GeneratedKey, TotpApi.GenerateKeyError> = synchronized(lock) {
-    if (confirmedKeys.size >= TotpApi.MAX_KEYS) {
-      return RequestResult.NonSuccess(TotpApi.GenerateKeyError.TooManyKeys)
-    }
-
-    hasPendingKey = true
-
-    RequestResult.Success(TotpApi.GeneratedKey(key = ByteArray(KEY_LENGTH_BYTES).also { secureRandom.nextBytes(it) }))
-  }
-
-  override suspend fun confirmKey(oneTimePassword: Int, metadata: TotpApi.Metadata): RequestResult<Int, TotpApi.ConfirmKeyError> = synchronized(lock) {
-    requireNameFits(metadata)
-
-    if (!hasPendingKey) {
-      return RequestResult.NonSuccess(TotpApi.ConfirmKeyError.NotVerified)
-    }
-
-    if (confirmedKeys.size >= TotpApi.MAX_KEYS) {
-      return RequestResult.NonSuccess(TotpApi.ConfirmKeyError.TooManyKeys)
-    }
-
-    val keyId = nextKeyId()
-    confirmedKeys[keyId] = metadata
-    hasPendingKey = false
-
-    RequestResult.Success(keyId)
-  }
-
-  override suspend fun listKeys(): RequestResult<List<TotpApi.RemoteKey>, Nothing> = synchronized(lock) {
-    RequestResult.Success(
-      confirmedKeys.entries
-        .sortedBy { it.key }
-        .map { (keyId, metadata) -> TotpApi.RemoteKey(keyId = keyId, metadata = metadata) }
-    )
-  }
-
-  override suspend fun setKeyMetadata(keyId: Int, metadata: TotpApi.Metadata): RequestResult<Unit, TotpApi.SetKeyMetadataError> = synchronized(lock) {
-    requireNameFits(metadata)
-
-    if (keyId !in confirmedKeys) {
-      return RequestResult.NonSuccess(TotpApi.SetKeyMetadataError.KeyNotFound)
-    }
-    confirmedKeys[keyId] = metadata
-
-    RequestResult.Success(Unit)
-  }
-
-  override suspend fun removeKey(keyId: Int): RequestResult<Unit, Nothing> = synchronized(lock) {
-    confirmedKeys.remove(keyId)
-    RequestResult.Success(Unit)
-  }
-
-  /** The service hands out the lowest free id rather than counting upwards, so removing a key frees its id for reuse. */
-  private fun nextKeyId(): Int = TotpApi.KEY_ID_RANGE.first { it !in confirmedKeys }
-
-  private fun requireNameFits(metadata: TotpApi.Metadata) {
-    require(metadata.name.toByteArray(Charsets.UTF_8).size <= TotpApi.Metadata.NAME_MAX_LENGTH) {
-      "Name must be at most ${TotpApi.Metadata.NAME_MAX_LENGTH} bytes of UTF-8"
-    }
-  }
-}

app/src/main/java/org/thoughtcrime/securesms/components/settings/app/account/authenticator/TotpApi.kt

@@ -1,86 +0,0 @@
-/*
- * Copyright 2026 Signal Messenger, LLC
- * SPDX-License-Identifier: AGPL-3.0-only
- */
-
-package org.thoughtcrime.securesms.components.settings.app.account.authenticator
-
-import org.signal.core.util.censor
-import org.signal.libsignal.net.BadRequestError
-import org.signal.libsignal.net.RequestResult
-import java.time.Instant
-
-/**
- * All TOTP operations.
- */
-interface TotpApi {
-
-  companion object {
-    /** How many confirmed keys an account may have, which the service enforces. */
-    const val MAX_KEYS = 2
-
-    /** The ids the service will assign, which fit in a byte with the sign bit clear. */
-    val KEY_ID_RANGE = 0..127
-  }
-
-  /**
-   * Generates a new pending key, replacing any pending key already on the account. The key doesn't take effect, or show
-   * up in [listKeys], until [confirmKey] proves the caller kept a copy of it.
-   */
-  suspend fun generateKey(): RequestResult<GeneratedKey, GenerateKeyError>
-
-  /**
-   * Confirms the pending key by proving a one-time password can be derived from it, and attaches [metadata] to it,
-   * returning the id the service assigned.
-   */
-  suspend fun confirmKey(oneTimePassword: Int, metadata: Metadata): RequestResult<Int, ConfirmKeyError>
-
-  /** The confirmed keys on the account. Key material is never returned, only metadata and parameters. */
-  suspend fun listKeys(): RequestResult<List<RemoteKey>, Nothing>
-
-  /** Replaces the metadata attached to a confirmed key. */
-  suspend fun setKeyMetadata(keyId: Int, metadata: Metadata): RequestResult<Unit, SetKeyMetadataError>
-
-  /** Removes a key, which also succeeds when there's no key with that id, so retries look the same as the first try. */
-  suspend fun removeKey(keyId: Int): RequestResult<Unit, Nothing>
-
-  data class Metadata(val name: String, val createdAt: Instant) {
-    override fun toString(): String = "Metadata(name=${name.censor()}, createdAt=$createdAt)"
-
-    companion object {
-      /** How long [name] may be, in bytes of UTF-8, which is the room the service's metadata blob leaves for it. */
-      const val NAME_MAX_LENGTH = 98
-    }
-  }
-
-  data class GeneratedKey(val key: ByteArray) {
-    override fun equals(other: Any?): Boolean = other is GeneratedKey && key.contentEquals(other.key)
-    override fun hashCode(): Int = key.contentHashCode()
-    override fun toString(): String = "GeneratedKey()"
-  }
-
-  data class RemoteKey(
-    /** The account-specific id, in [KEY_ID_RANGE]. */
-    val keyId: Int,
-    val metadata: Metadata
-  ) {
-    override fun toString(): String = "RemoteKey(keyId=$keyId)"
-  }
-
-  sealed interface GenerateKeyError : BadRequestError {
-    /** The account already has [MAX_KEYS] keys. */
-    data object TooManyKeys : GenerateKeyError
-  }
-
-  sealed interface ConfirmKeyError : BadRequestError {
-    /** The password was wrong, the clocks are too far apart, or there was no pending key. The service can't tell us which. */
-    data object NotVerified : ConfirmKeyError
-
-    /** The account filled up with keys between generating this one and confirming it. */
-    data object TooManyKeys : ConfirmKeyError
-  }
-
-  sealed interface SetKeyMetadataError : BadRequestError {
-    data object KeyNotFound : SetKeyMetadataError
-  }
-}

app/src/main/java/org/thoughtcrime/securesms/components/settings/app/account/authenticator/TotpRepository.kt

@@ -6,42 +6,62 @@
 package org.thoughtcrime.securesms.components.settings.app.account.authenticator
 
 import org.signal.appsettings.totpapplist.TotpApp
+import org.signal.core.models.MasterKey
 import org.signal.core.util.Base32
 import org.signal.core.util.logging.Log
+import org.signal.libsignal.net.MfaKeyNotFoundException
+import org.signal.libsignal.net.MfaMetadata
+import org.signal.libsignal.net.OneTimePasswordNotVerifiedException
 import org.signal.libsignal.net.RequestResult
+import org.signal.libsignal.net.TooManyMfaKeysException
+import org.signal.libsignal.net.TooManyTotpKeysException
+import org.signal.libsignal.net.TotpParameters
+import org.signal.network.api.AccountApiV2
+import org.thoughtcrime.securesms.keyvalue.SignalStore
+import org.thoughtcrime.securesms.net.SignalNetwork
 import java.net.URLEncoder
 import java.time.Instant
 
 /**
- * Everything the authenticator app screens need, sitting between them and the TOTP operations on [TotpApi].
+ * Everything the authenticator app screens need, sitting between them and the TOTP endpoints on [AccountApiV2].
+ *
+ * The name the user gives an app and the time they set it up live in the metadata the service stores against each
+ * key. libsignal encrypts that metadata under a key derived from the master key, so the service never reads it --
+ * all this layer does is hand the master key over and map the results into what the screens show.
  */
 class TotpRepository(
-  private val api: TotpApi = SHARED_API,
+  private val api: AccountApiV2 = SignalNetwork.accountV2,
+  private val masterKeyProvider: () -> MasterKey = { SignalStore.svr.masterKey },
   private val clock: () -> Long = System::currentTimeMillis
 ) {
 
   companion object {
     private val TAG = Log.tag(TotpRepository::class)
 
-    /** Shared so that every screen in the flow sees the same state until there's a service behind this. */
-    private val SHARED_API: TotpApi = InMemoryTotpApi()
+    /**
+     * How many authenticator apps an account may have, which the service enforces. libsignal reports hitting the
+     * limit but doesn't expose the number, so the screens that want to show it get it from here.
+     */
+    const val MAX_APPS = 2
 
     private const val ISSUER = "Signal"
 
-    /** What the service uses, and what every authenticator app supports without reading a single URI parameter. */
-    private const val ALGORITHM = "SHA1"
-    private const val DIGITS = 6
-    private const val PERIOD_SECONDS = 30
+    /** The algorithm names the Key Uri Format defines, keyed by what [TotpParameters.algorithm] calls them. */
+    private val URI_ALGORITHMS = mapOf(
+      "HmacSHA1" to "SHA1",
+      "HmacSHA256" to "SHA256",
+      "HmacSHA512" to "SHA512"
+    )
 
     /** How many characters of the display form go between spaces. */
     private const val DISPLAY_GROUP_SIZE = 4
 
-    const val MAX_NAME_LENGTH_BYTES = TotpApi.Metadata.NAME_MAX_LENGTH
+    const val MAX_NAME_LENGTH_BYTES = MfaMetadata.NAME_MAX_LENGTH
     const val MAX_NAME_LENGTH_GRAPHEMES = 30
   }
 
   fun getMaxApps(): Int {
-    return TotpApi.MAX_KEYS
+    return MAX_APPS
   }
 
   /**
@@ -49,19 +69,25 @@ class TotpRepository(
    * service holds the pending key from here until [confirmPendingApp], so nothing is kept on this side.
    */
   suspend fun beginSetup(accountName: String): BeginSetupResult {
-    return when (val result = api.generateKey()) {
+    return when (val result = api.generateTotpKey()) {
       is RequestResult.Success -> {
-        val key = result.result.key
+        val generated = result.result
+
+        val setupUri = buildSetupUri(key = generated.key, parameters = generated.parameters, accountName = accountName)
+        if (setupUri == null) {
+          Log.w(TAG, "The service generated a key with parameters a setup link can't describe: ${generated.parameters}")
+          return BeginSetupResult.NetworkFailure
+        }
 
         BeginSetupResult.Success(
-          setupUri = buildSetupUri(key = key, accountName = accountName),
-          displayKey = Base32.encode(key).chunked(DISPLAY_GROUP_SIZE).joinToString(" "),
-          clipboardKey = Base32.encode(key)
+          setupUri = setupUri,
+          displayKey = Base32.encode(generated.key).chunked(DISPLAY_GROUP_SIZE).joinToString(" "),
+          clipboardKey = Base32.encode(generated.key)
         )
       }
       is RequestResult.NonSuccess -> {
         when (result.error) {
-          TotpApi.GenerateKeyError.TooManyKeys -> BeginSetupResult.TooManyApps
+          is TooManyTotpKeysException, is TooManyMfaKeysException -> BeginSetupResult.TooManyApps
         }
       }
       is RequestResult.RetryableNetworkError -> {
@@ -85,15 +111,15 @@ class TotpRepository(
   suspend fun confirmPendingApp(code: String): ConfirmResult {
     val oneTimePassword = code.toIntOrNull() ?: return ConfirmResult.IncorrectCode
 
-    val metadata = TotpApi.Metadata(name = "", createdAt = Instant.ofEpochMilli(clock()))
+    val metadata = MfaMetadata(name = "", createdAt = Instant.ofEpochMilli(clock()))
 
-    return when (val result = api.confirmKey(oneTimePassword = oneTimePassword, metadata = metadata)) {
+    return when (val result = api.confirmTotpKey(oneTimePassword = oneTimePassword, metadata = metadata, masterKey = masterKeyProvider())) {
       is RequestResult.Success -> {
         ConfirmResult.Success(appId = result.result.toLong())
       }
       is RequestResult.NonSuccess -> when (result.error) {
-        TotpApi.ConfirmKeyError.NotVerified -> ConfirmResult.IncorrectCode
-        TotpApi.ConfirmKeyError.TooManyKeys -> {
+        is OneTimePasswordNotVerifiedException -> ConfirmResult.IncorrectCode
+        is TooManyMfaKeysException -> {
           Log.w(TAG, "The account filled up with keys between generating this one and confirming it.")
           ConfirmResult.TooManyApps
         }
@@ -109,44 +135,50 @@ class TotpRepository(
     }
   }
 
-  /** The authenticator apps on the account, newest id last. */
+  /** The authenticator apps on the account, newest id last, with anything we can't read left out. */
   suspend fun getTotpApps(): AppsResult {
-    return when (val result = api.listKeys()) {
-      is RequestResult.Success -> {
-        AppsResult.Success(
-          result.result.map { key ->
-            TotpApp(
-              id = key.keyId.toLong(),
-              name = key.metadata.name,
-              createdAt = key.metadata.createdAt.toEpochMilli()
-            )
-          }
-        )
-      }
+    val keys = when (val result = api.listMfaKeys(masterKeyProvider())) {
+      is RequestResult.Success -> result.result
       is RequestResult.RetryableNetworkError -> {
         Log.w(TAG, "Couldn't list keys.", result.networkError)
-        AppsResult.NetworkFailure
+        return AppsResult.NetworkFailure
       }
       is RequestResult.ApplicationError -> {
         Log.w(TAG, "Couldn't list keys.", result.cause)
-        AppsResult.NetworkFailure
+        return AppsResult.NetworkFailure
       }
       is RequestResult.NonSuccess -> error("Code branch is unreachable")
     }
+
+    val apps = keys.mapNotNull { key ->
+      val metadata = key.metadata
+      if (metadata == null) {
+        Log.w(TAG, "Couldn't read the metadata for key ${key.id}. Leaving it out of the list.")
+        null
+      } else {
+        TotpApp(
+          id = key.id.toLong(),
+          name = metadata.name,
+          createdAt = metadata.createdAt.toEpochMilli()
+        )
+      }
+    }
+
+    return AppsResult.Success(apps)
   }
 
-  /** Renames [app], which means handing the whole metadata blob back to the service. */
+  /** Renames [app], which means re-encrypting its metadata and handing the whole blob back to the service. */
   suspend fun renameTotpApp(app: TotpApp, name: String): UpdateResult {
-    return setMetadata(app.id, TotpApi.Metadata(name = name, createdAt = Instant.ofEpochMilli(app.createdAt)))
+    return setMetadata(app.id, MfaMetadata(name = name, createdAt = Instant.ofEpochMilli(app.createdAt)))
   }
 
   /** Names a newly confirmed app, which was confirmed without one moments ago. */
   suspend fun nameNewTotpApp(appId: Long, name: String): UpdateResult {
-    return setMetadata(appId, TotpApi.Metadata(name = name, createdAt = Instant.ofEpochMilli(clock())))
+    return setMetadata(appId, MfaMetadata(name = name, createdAt = Instant.ofEpochMilli(clock())))
   }
 
   suspend fun removeTotpApp(appId: Long): UpdateResult {
-    return when (val result = api.removeKey(appId.toInt())) {
+    return when (val result = api.removeMfaKey(appId.toInt())) {
       is RequestResult.Success -> UpdateResult.Success
       is RequestResult.RetryableNetworkError -> {
         Log.w(TAG, "Couldn't remove the key.", result.networkError)
@@ -160,11 +192,11 @@ class TotpRepository(
     }
   }
 
-  private suspend fun setMetadata(appId: Long, metadata: TotpApi.Metadata): UpdateResult {
-    return when (val result = api.setKeyMetadata(keyId = appId.toInt(), metadata = metadata)) {
+  private suspend fun setMetadata(appId: Long, metadata: MfaMetadata): UpdateResult {
+    return when (val result = api.setMfaKeyMetadata(keyId = appId.toInt(), metadata = metadata, masterKey = masterKeyProvider())) {
       is RequestResult.Success -> UpdateResult.Success
       is RequestResult.NonSuccess -> when (result.error) {
-        TotpApi.SetKeyMetadataError.KeyNotFound -> UpdateResult.AppNotFound
+        is MfaKeyNotFoundException -> UpdateResult.AppNotFound
       }
       is RequestResult.RetryableNetworkError -> {
         Log.w(TAG, "Couldn't set key metadata.", result.networkError)
@@ -179,17 +211,20 @@ class TotpRepository(
 
   /**
    * The `otpauth://` URI that hands the key to an authenticator app, following the de facto Key Uri Format every app
-   * implements. Note that a lot of apps ignore params like "algorithm", but we set them just in case.
+   * implements, or null for parameters the format can't describe. Note that a lot of apps ignore params like
+   * "algorithm", but we set them just in case.
    */
-  private fun buildSetupUri(key: ByteArray, accountName: String): String {
+  private fun buildSetupUri(key: ByteArray, parameters: TotpParameters, accountName: String): String? {
+    val algorithm = URI_ALGORITHMS[parameters.algorithm] ?: return null
+
     val label = if (accountName.isBlank()) encode(ISSUER) else "${encode(ISSUER)}:${encode(accountName)}"
 
     val query = listOf(
       "secret" to Base32.encode(key),
       "issuer" to ISSUER,
-      "algorithm" to ALGORITHM,
-      "digits" to DIGITS.toString(),
-      "period" to PERIOD_SECONDS.toString()
+      "algorithm" to algorithm,
+      "digits" to parameters.passwordLength.toString(),
+      "period" to parameters.timeStep.seconds.toString()
     ).joinToString("&") { (name, value) -> "$name=${encode(value)}" }
 
     return "otpauth://totp/$label?$query"

app/src/test/java/org/thoughtcrime/securesms/components/settings/app/account/authenticator/InMemoryTotpApiTest.kt

@@ -1,129 +0,0 @@
-/*
- * Copyright 2026 Signal Messenger, LLC
- * SPDX-License-Identifier: AGPL-3.0-only
- */
-
-package org.thoughtcrime.securesms.components.settings.app.account.authenticator
-
-import assertk.assertFailure
-import assertk.assertThat
-import assertk.assertions.hasSize
-import assertk.assertions.isEqualTo
-import assertk.assertions.isInstanceOf
-import kotlinx.coroutines.test.runTest
-import org.junit.Test
-import org.signal.libsignal.net.RequestResult
-import java.time.Instant
-
-/**
- * Covers the parts of the stand-in that copy behaviour the service is strict about, since those are the parts most
- * likely to be wrong once there's a real service behind [TotpApi].
- */
-class InMemoryTotpApiTest {
-
-  companion object {
-    private const val NOW = 1_700_000_000_000L
-    private const val CODE = 123456
-    private val METADATA = TotpApi.Metadata(name = "Aegis", createdAt = Instant.ofEpochMilli(NOW))
-    private val OTHER_METADATA = TotpApi.Metadata(name = "Aegis on my tablet", createdAt = Instant.ofEpochMilli(NOW))
-  }
-
-  private val api = InMemoryTotpApi()
-
-  @Test
-  fun `a generated key is 32 bytes`() = runTest {
-    val result = api.generateKey()
-
-    assertThat(result).isInstanceOf(RequestResult.Success::class)
-    assertThat((result as RequestResult.Success).result.key.size).isEqualTo(32)
-  }
-
-  @Test
-  fun `a pending key doesn't show up until it's confirmed`() = runTest {
-    api.generateKey()
-
-    assertThat(listedKeys()).hasSize(0)
-  }
-
-  @Test
-  fun `a confirmed key is assigned the lowest free id`() = runTest {
-    val first = confirmNewKey()
-    val second = confirmNewKey()
-
-    assertThat(first).isEqualTo(0)
-    assertThat(second).isEqualTo(1)
-  }
-
-  @Test
-  fun `an id freed by a removal is handed out again rather than counting upwards`() = runTest {
-    confirmNewKey()
-    val second = confirmNewKey()
-    api.removeKey(second)
-
-    assertThat(confirmNewKey()).isEqualTo(second)
-  }
-
-  @Test
-  fun `confirming with no pending key doesn't confirm anything`() = runTest {
-    assertThat(api.confirmKey(oneTimePassword = CODE, metadata = METADATA)).isEqualTo(RequestResult.NonSuccess(TotpApi.ConfirmKeyError.NotVerified))
-    assertThat(listedKeys()).hasSize(0)
-  }
-
-  @Test
-  fun `an account at its limit can't generate another key`() = runTest {
-    repeat(TotpApi.MAX_KEYS) { confirmNewKey() }
-
-    assertThat(api.generateKey()).isEqualTo(RequestResult.NonSuccess(TotpApi.GenerateKeyError.TooManyKeys))
-  }
-
-  @Test
-  fun `removing a key makes room for another`() = runTest {
-    repeat(TotpApi.MAX_KEYS) { confirmNewKey() }
-    api.removeKey(0)
-
-    assertThat(api.generateKey()).isInstanceOf(RequestResult.Success::class)
-  }
-
-  @Test
-  fun `metadata can be replaced on a confirmed key`() = runTest {
-    val keyId = confirmNewKey()
-
-    assertThat(api.setKeyMetadata(keyId, OTHER_METADATA)).isEqualTo(RequestResult.Success(Unit))
-    assertThat(listedKeys().first().metadata).isEqualTo(OTHER_METADATA)
-  }
-
-  @Test
-  fun `metadata for a key that isn't there is reported rather than created`() = runTest {
-    assertThat(api.setKeyMetadata(7, METADATA)).isEqualTo(RequestResult.NonSuccess(TotpApi.SetKeyMetadataError.KeyNotFound))
-  }
-
-  /** The service leaves a fixed amount of room for the name, so a name that doesn't fit is the caller's bug. */
-  @Test
-  fun `a name longer than the room the service leaves is refused`() = runTest {
-    val keyId = confirmNewKey()
-    val tooLong = METADATA.copy(name = "a".repeat(TotpApi.Metadata.NAME_MAX_LENGTH + 1))
-
-    assertFailure { api.setKeyMetadata(keyId, tooLong) }.isInstanceOf(IllegalArgumentException::class)
-  }
-
-  /** The service reports success either way, so a retried removal looks like the original. */
-  @Test
-  fun `removing a key that isn't there still succeeds`() = runTest {
-    assertThat(api.removeKey(7)).isEqualTo(RequestResult.Success(Unit))
-  }
-
-  @Test
-  fun `keys are listed in ascending id order`() = runTest {
-    confirmNewKey()
-    confirmNewKey()
-
-    assertThat(listedKeys().map { it.keyId }).isEqualTo(listOf(0, 1))
-  }
-
-  private suspend fun confirmNewKey(): Int {
-    api.generateKey()
-    return (api.confirmKey(oneTimePassword = CODE, metadata = METADATA) as RequestResult.Success).result
-  }
-
-  private suspend fun listedKeys(): List<TotpApi.RemoteKey> = (api.listKeys() as RequestResult.Success).result
-}

app/src/test/java/org/thoughtcrime/securesms/components/settings/app/account/authenticator/TotpRepositoryTest.kt

@@ -12,13 +12,33 @@ import assertk.assertions.isEmpty
 import assertk.assertions.isEqualTo
 import assertk.assertions.isInstanceOf
 import assertk.assertions.startsWith
+import io.mockk.coEvery
+import io.mockk.coVerify
+import io.mockk.mockk
+import io.mockk.slot
 import kotlinx.coroutines.test.runTest
+import org.junit.Before
 import org.junit.Test
 import org.signal.appsettings.totpapplist.TotpApp
+import org.signal.core.models.MasterKey
+import org.signal.libsignal.net.ConfirmedMfaKey
+import org.signal.libsignal.net.MfaKeyKind
+import org.signal.libsignal.net.MfaKeyNotFoundException
+import org.signal.libsignal.net.MfaMetadata
+import org.signal.libsignal.net.OneTimePasswordNotVerifiedException
+import org.signal.libsignal.net.PendingTotpKey
+import org.signal.libsignal.net.RequestResult
+import org.signal.libsignal.net.TooManyMfaKeysException
+import org.signal.libsignal.net.TooManyTotpKeysException
+import org.signal.libsignal.net.TotpParameters
+import org.signal.network.api.AccountApiV2
 import org.thoughtcrime.securesms.components.settings.app.account.authenticator.TotpRepository.AppsResult
 import org.thoughtcrime.securesms.components.settings.app.account.authenticator.TotpRepository.BeginSetupResult
 import org.thoughtcrime.securesms.components.settings.app.account.authenticator.TotpRepository.ConfirmResult
 import org.thoughtcrime.securesms.components.settings.app.account.authenticator.TotpRepository.UpdateResult
+import java.io.IOException
+import java.time.Duration
+import java.time.Instant
 
 class TotpRepositoryTest {
 
@@ -26,11 +46,28 @@ class TotpRepositoryTest {
     private const val NOW = 1_700_000_000_000L
     private const val ACCOUNT_NAME = "8B4A1F0C"
     private const val CODE = "123456"
+    private const val KEY_ID = 1
+
+    private val KEY = ByteArray(32) { it.toByte() }
+    private val MASTER_KEY = MasterKey(ByteArray(32) { (it + 100).toByte() })
+
+    /** What the service generates: a 256-bit key with HMAC-SHA1, six digits, thirty second steps. */
+    private val PARAMETERS = TotpParameters(algorithm = "HmacSHA1", passwordLength = 6, timeStep = Duration.ofSeconds(30))
+    private val PENDING_KEY = PendingTotpKey(key = KEY, parameters = PARAMETERS)
   }
 
   private var now = NOW
-  private val api = InMemoryTotpApi()
-  private val repository = TotpRepository(api = api, clock = { now })
+  private val api = mockk<AccountApiV2>()
+  private val repository = TotpRepository(api = api, masterKeyProvider = { MASTER_KEY }, clock = { now })
+
+  @Before
+  fun setUp() {
+    coEvery { api.generateTotpKey() } returns RequestResult.Success(PENDING_KEY)
+    coEvery { api.confirmTotpKey(any(), any(), any()) } returns RequestResult.Success(KEY_ID)
+    coEvery { api.listMfaKeys(any()) } returns RequestResult.Success(emptyList())
+    coEvery { api.setMfaKeyMetadata(any(), any(), any()) } returns RequestResult.Success(Unit)
+    coEvery { api.removeMfaKey(any()) } returns RequestResult.Success(Unit)
+  }
 
   @Test
   fun `beginSetup returns a link and a key in both the forms the screen needs`() = runTest {
@@ -74,7 +111,7 @@ class TotpRepositoryTest {
   }
 
   @Test
-  fun `setup asks for the parameters every authenticator app supports`() = runTest {
+  fun `beginSetup writes out the parameters the service chose`() = runTest {
     val result = repository.beginSetup(ACCOUNT_NAME) as BeginSetupResult.Success
 
     assertThat(result.setupUri).contains("algorithm=SHA1")
@@ -82,97 +119,164 @@ class TotpRepositoryTest {
     assertThat(result.setupUri).contains("period=30")
   }
 
+  /** A link naming an algorithm the format doesn't define would pair an app whose codes never confirm, which is worse than not starting. */
   @Test
-  fun `an account at its limit is told rather than handed a key`() = runTest {
-    repeat(TotpApi.MAX_KEYS) { confirmNewApp() }
+  fun `a key with parameters a setup link can't describe fails setup rather than being handed out`() = runTest {
+    coEvery { api.generateTotpKey() } returns RequestResult.Success(
+      PendingTotpKey(key = KEY, parameters = TotpParameters(algorithm = "HmacMD5", passwordLength = 6, timeStep = Duration.ofSeconds(30)))
+    )
+
+    assertThat(repository.beginSetup(ACCOUNT_NAME)).isEqualTo(BeginSetupResult.NetworkFailure)
+  }
+
+  @Test
+  fun `an account at its TOTP key limit is told rather than handed a key`() = runTest {
+    coEvery { api.generateTotpKey() } returns RequestResult.NonSuccess(TooManyTotpKeysException("full"))
 
     assertThat(repository.beginSetup(ACCOUNT_NAME)).isEqualTo(BeginSetupResult.TooManyApps)
   }
 
   @Test
-  fun `a confirmed app shows up in the list with the time it was confirmed`() = runTest {
-    val appId = confirmNewApp()
+  fun `an account at its overall MFA key limit is told rather than handed a key`() = runTest {
+    coEvery { api.generateTotpKey() } returns RequestResult.NonSuccess(TooManyMfaKeysException("full"))
 
-    val apps = (repository.getTotpApps() as AppsResult.Success).apps
+    assertThat(repository.beginSetup(ACCOUNT_NAME)).isEqualTo(BeginSetupResult.TooManyApps)
+  }
 
-    assertThat(apps).hasSize(1)
-    assertThat(apps.first().id).isEqualTo(appId)
-    assertThat(apps.first().createdAt).isEqualTo(NOW)
+  @Test
+  fun `a service we couldn't reach fails setup`() = runTest {
+    coEvery { api.generateTotpKey() } returns RequestResult.RetryableNetworkError(IOException("offline"))
+
+    assertThat(repository.beginSetup(ACCOUNT_NAME)).isEqualTo(BeginSetupResult.NetworkFailure)
+  }
+
+  @Test
+  fun `a code that isn't a number is reported as a wrong code without asking the service`() = runTest {
+    assertThat(repository.confirmPendingApp("abcdef")).isEqualTo(ConfirmResult.IncorrectCode)
+
+    coVerify(exactly = 0) { api.confirmTotpKey(any(), any(), any()) }
+  }
+
+  @Test
+  fun `a code the service rejects is reported as a wrong code`() = runTest {
+    coEvery { api.confirmTotpKey(any(), any(), any()) } returns RequestResult.NonSuccess(OneTimePasswordNotVerifiedException("nope"))
+
+    assertThat(repository.confirmPendingApp(CODE)).isEqualTo(ConfirmResult.IncorrectCode)
+  }
+
+  @Test
+  fun `an account that filled up while the key was pending is told so`() = runTest {
+    coEvery { api.confirmTotpKey(any(), any(), any()) } returns RequestResult.NonSuccess(TooManyMfaKeysException("full"))
+
+    assertThat(repository.confirmPendingApp(CODE)).isEqualTo(ConfirmResult.TooManyApps)
+  }
+
+  @Test
+  fun `a code the service accepts confirms the app`() = runTest {
+    val result = repository.confirmPendingApp(CODE)
+
+    assertThat(result).isInstanceOf(ConfirmResult.Success::class)
+    assertThat((result as ConfirmResult.Success).appId).isEqualTo(KEY_ID.toLong())
   }
 
   /** The service wants metadata at confirmation time, and the user hasn't been asked for a name yet. */
   @Test
-  fun `a newly confirmed app starts out with no name`() = runTest {
-    confirmNewApp()
+  fun `a key is confirmed without a name, stamped with the time it was confirmed`() = runTest {
+    val metadata = slot<MfaMetadata>()
+    coEvery { api.confirmTotpKey(any(), capture(metadata), any()) } returns RequestResult.Success(KEY_ID)
 
-    val apps = (repository.getTotpApps() as AppsResult.Success).apps
+    repository.confirmPendingApp(CODE)
 
-    assertThat(apps.first().name).isEqualTo("")
+    assertThat(metadata.captured.name).isEqualTo("")
+    assertThat(metadata.captured.createdAt).isEqualTo(Instant.ofEpochMilli(NOW))
   }
 
   @Test
-  fun `naming a new app names it`() = runTest {
-    val appId = confirmNewApp()
+  fun `the confirmed keys on the account come back as apps`() = runTest {
+    coEvery { api.listMfaKeys(any()) } returns RequestResult.Success(
+      listOf(
+        ConfirmedMfaKey(id = KEY_ID, metadata = MfaMetadata(name = "Aegis", createdAt = Instant.ofEpochMilli(NOW)), kind = MfaKeyKind.TOTP)
+      )
+    )
 
-    assertThat(repository.nameNewTotpApp(appId, "Aegis")).isEqualTo(UpdateResult.Success)
+    val apps = (repository.getTotpApps() as AppsResult.Success).apps
 
-    assertThat(listedApp(appId)?.name).isEqualTo("Aegis")
+    assertThat(apps).hasSize(1)
+    assertThat(apps.first().id).isEqualTo(KEY_ID.toLong())
+    assertThat(apps.first().name).isEqualTo("Aegis")
+    assertThat(apps.first().createdAt).isEqualTo(NOW)
+  }
+
+  /** Metadata we can't read was written under some other key, so listing it as a nameless app would be worse than omitting it. */
+  @Test
+  fun `a key whose metadata can't be read is left out of the list`() = runTest {
+    coEvery { api.listMfaKeys(any()) } returns RequestResult.Success(
+      listOf(ConfirmedMfaKey(id = KEY_ID, metadata = null, kind = MfaKeyKind.TOTP))
+    )
+
+    assertThat((repository.getTotpApps() as AppsResult.Success).apps).isEmpty()
+  }
+
+  /** The list is about what's on the account, not what this client understands, so a newer device's key still shows. */
+  @Test
+  fun `a key of a kind this client doesn't know still shows up as an app`() = runTest {
+    coEvery { api.listMfaKeys(any()) } returns RequestResult.Success(
+      listOf(ConfirmedMfaKey(id = KEY_ID, metadata = MfaMetadata(name = "Future", createdAt = Instant.ofEpochMilli(NOW)), kind = MfaKeyKind.UNKNOWN))
+    )
+
+    assertThat((repository.getTotpApps() as AppsResult.Success).apps).hasSize(1)
+  }
+
+  @Test
+  fun `a list we couldn't fetch is a network failure`() = runTest {
+    coEvery { api.listMfaKeys(any()) } returns RequestResult.RetryableNetworkError(IOException("offline"))
+
+    assertThat(repository.getTotpApps()).isEqualTo(AppsResult.NetworkFailure)
+  }
+
+  @Test
+  fun `naming a new app stamps it with the current time`() = runTest {
+    val metadata = slot<MfaMetadata>()
+    coEvery { api.setMfaKeyMetadata(eq(KEY_ID), capture(metadata), any()) } returns RequestResult.Success(Unit)
+
+    assertThat(repository.nameNewTotpApp(KEY_ID.toLong(), "Aegis")).isEqualTo(UpdateResult.Success)
+
+    assertThat(metadata.captured.name).isEqualTo("Aegis")
+    assertThat(metadata.captured.createdAt).isEqualTo(Instant.ofEpochMilli(NOW))
   }
 
   @Test
   fun `renaming keeps the time the app was confirmed`() = runTest {
-    val appId = confirmNewApp()
-    repository.nameNewTotpApp(appId, "Aegis")
+    val metadata = slot<MfaMetadata>()
+    coEvery { api.setMfaKeyMetadata(eq(KEY_ID), capture(metadata), any()) } returns RequestResult.Success(Unit)
+    val app = TotpApp(id = KEY_ID.toLong(), name = "Aegis", createdAt = NOW)
     now += 60_000
 
-    assertThat(repository.renameTotpApp(listedApp(appId)!!, "Aegis on my tablet")).isEqualTo(UpdateResult.Success)
+    assertThat(repository.renameTotpApp(app, "Aegis on my tablet")).isEqualTo(UpdateResult.Success)
 
-    val app = listedApp(appId)
-    assertThat(app?.name).isEqualTo("Aegis on my tablet")
-    assertThat(app?.createdAt).isEqualTo(NOW)
+    assertThat(metadata.captured.name).isEqualTo("Aegis on my tablet")
+    assertThat(metadata.captured.createdAt).isEqualTo(Instant.ofEpochMilli(NOW))
   }
 
   @Test
-  fun `renaming an app that isn't there is reported rather than creating one`() = runTest {
+  fun `renaming an app the service no longer has is reported as not found`() = runTest {
+    coEvery { api.setMfaKeyMetadata(any(), any(), any()) } returns RequestResult.NonSuccess(MfaKeyNotFoundException("gone"))
     val gone = TotpApp(id = 7, name = "Aegis", createdAt = NOW)
 
     assertThat(repository.renameTotpApp(gone, "Aegis on my tablet")).isEqualTo(UpdateResult.AppNotFound)
   }
 
   @Test
-  fun `a removed app leaves the list`() = runTest {
-    val appId = confirmNewApp()
+  fun `removing an app removes its key`() = runTest {
+    assertThat(repository.removeTotpApp(KEY_ID.toLong())).isEqualTo(UpdateResult.Success)
 
-    assertThat(repository.removeTotpApp(appId)).isEqualTo(UpdateResult.Success)
-    assertThat((repository.getTotpApps() as AppsResult.Success).apps).isEmpty()
-  }
-
-  /** The service can't tell a wrong code from a missing pending key, so neither can we. */
-  @Test
-  fun `confirming with nothing pending is just a wrong code`() = runTest {
-    assertThat(repository.confirmPendingApp(CODE)).isEqualTo(ConfirmResult.IncorrectCode)
+    coVerify { api.removeMfaKey(KEY_ID) }
   }
 
   @Test
-  fun `a code that isn't a number is reported as a wrong code`() = runTest {
-    repository.beginSetup(ACCOUNT_NAME)
+  fun `a removal we couldn't send is a network failure`() = runTest {
+    coEvery { api.removeMfaKey(any()) } returns RequestResult.RetryableNetworkError(IOException("offline"))
 
-    assertThat(repository.confirmPendingApp("abcdef")).isEqualTo(ConfirmResult.IncorrectCode)
-  }
-
-  @Test
-  fun `a code confirms the app`() = runTest {
-    repository.beginSetup(ACCOUNT_NAME)
-
-    assertThat(repository.confirmPendingApp(CODE)).isInstanceOf(ConfirmResult.Success::class)
-  }
-
-  private suspend fun listedApp(appId: Long): TotpApp? {
-    return (repository.getTotpApps() as AppsResult.Success).apps.firstOrNull { it.id == appId }
-  }
-
-  private suspend fun confirmNewApp(): Long {
-    repository.beginSetup(ACCOUNT_NAME)
-    return (repository.confirmPendingApp(CODE) as ConfirmResult.Success).appId
+    assertThat(repository.removeTotpApp(KEY_ID.toLong())).isEqualTo(UpdateResult.NetworkFailure)
   }
 }

lib/network/src/main/java/org/signal/network/api/AccountApiV2.kt

@@ -17,6 +17,12 @@ import org.signal.libsignal.net.AuthAccountsService
 import org.signal.libsignal.net.AuthDevicesService
 import org.signal.libsignal.net.AuthUsernamesService
 import org.signal.libsignal.net.BadRequestError
+import org.signal.libsignal.net.ConfirmTotpKeyError
+import org.signal.libsignal.net.ConfirmedMfaKey
+import org.signal.libsignal.net.GenerateTotpKeyError
+import org.signal.libsignal.net.MfaKeyNotFoundException
+import org.signal.libsignal.net.MfaMetadata
+import org.signal.libsignal.net.PendingTotpKey
 import org.signal.libsignal.net.RequestResult
 import org.signal.libsignal.net.SvrKey
 import org.signal.libsignal.net.UsernameNotAvailableException
@@ -184,6 +190,64 @@ class AccountApiV2(private val authWebSocket: SignalWebSocket.AuthenticatedWebSo
     }
   }
 
+  /**
+   * Generates and stores a new pending TOTP key for the account, replacing any pending key already there. The key
+   * doesn't take effect, or show up in [listMfaKeys], until [confirmTotpKey] proves the caller kept a copy of it,
+   * which must happen within 24 hours. This is the only time the key material is ever handed out.
+   *
+   * A [TooManyTotpKeysException][org.signal.libsignal.net.TooManyTotpKeysException] or
+   * [TooManyMfaKeysException][org.signal.libsignal.net.TooManyMfaKeysException] means the account is at its limit,
+   * and a key has to be removed before another can be added.
+   */
+  suspend fun generateTotpKey(): RequestResult<PendingTotpKey, GenerateTotpKeyError> {
+    return authWebSocket.runCatchingWithChatConnection { connection ->
+      AuthAccountsService(connection).generateTotpKey()
+    }
+  }
+
+  /**
+   * Confirms the pending TOTP key by proving a one-time password can be derived from it, and attaches [metadata] to
+   * it, returning the id the service assigned. The metadata is encrypted under a key derived from [masterKey], so the
+   * service never sees it.
+   *
+   * A [OneTimePasswordNotVerifiedException][org.signal.libsignal.net.OneTimePasswordNotVerifiedException] means the
+   * password was wrong, the clocks are too far apart, or there was no pending key -- the service can't tell us which.
+   */
+  suspend fun confirmTotpKey(oneTimePassword: Int, metadata: MfaMetadata, masterKey: MasterKey): RequestResult<Int, ConfirmTotpKeyError> {
+    return authWebSocket.runCatchingWithChatConnection { connection ->
+      AuthAccountsService(connection).confirmTotpKey(oneTimePassword = oneTimePassword, metadata = metadata, svrKey = SvrKey(masterKey.serialize()))
+    }
+  }
+
+  /**
+   * The confirmed MFA keys on the account. Key material is never returned, only metadata and parameters. A key whose
+   * metadata can't be decrypted under [masterKey] comes back with null metadata rather than being dropped.
+   */
+  suspend fun listMfaKeys(masterKey: MasterKey): RequestResult<List<ConfirmedMfaKey>, Nothing> {
+    return authWebSocket.runCatchingWithChatConnection { connection ->
+      AuthAccountsService(connection).listMfaKeys(SvrKey(masterKey.serialize()))
+    }
+  }
+
+  /**
+   * Replaces the metadata attached to a confirmed MFA key, encrypted under a key derived from [masterKey].
+   */
+  suspend fun setMfaKeyMetadata(keyId: Int, metadata: MfaMetadata, masterKey: MasterKey): RequestResult<Unit, MfaKeyNotFoundException> {
+    return authWebSocket.runCatchingWithChatConnection { connection ->
+      AuthAccountsService(connection).setMfaKeyMetadata(keyId = keyId, metadata = metadata, svrKey = SvrKey(masterKey.serialize()))
+    }
+  }
+
+  /**
+   * Removes an MFA key, which also succeeds when there's no key with that id, so retries look the same as the
+   * first try.
+   */
+  suspend fun removeMfaKey(keyId: Int): RequestResult<Unit, Nothing> {
+    return authWebSocket.runCatchingWithChatConnection { connection ->
+      AuthAccountsService(connection).removeMfaKey(keyId)
+    }
+  }
+
   /**
    * Deletes the account off the service.
    *