Reset backup id on zk verification failure during restore attempts.

2026-02-24 03:35:58 +00:00 · 2025-10-08 13:49:53 -04:00
parent 6e8f982e7b
commit a5cca5b0fd
13 changed files with 249 additions and 86 deletions
--- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/BackupRepository.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/BackupRepository.kt
@@ -177,6 +177,7 @@ import java.util.concurrent.CountDownLatch
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicLong
 import kotlin.jvm.optionals.getOrNull
+import kotlin.time.Duration
 import kotlin.time.Duration.Companion.days
 import kotlin.time.Duration.Companion.milliseconds
 import kotlin.time.Duration.Companion.seconds
@@ -285,6 +286,19 @@ object BackupRepository {
    val messageBackupKey = SignalStore.backup.messageBackupKey
    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
    return SignalNetwork.archive.triggerBackupIdReservation(messageBackupKey, mediaRootBackupKey, SignalStore.account.requireAci())
+      .runIfSuccessful {
+        SignalStore.backup.messageCredentials.clearAll()
+        SignalStore.backup.mediaCredentials.clearAll()
+      }
+  }
+
+  @WorkerThread
+  fun triggerBackupIdReservationForRestore(): NetworkResult<Unit> {
+    val messageBackupKey = SignalStore.backup.messageBackupKey
+    return SignalNetwork.archive.triggerBackupIdReservation(messageBackupKey, null, SignalStore.account.requireAci())
+      .runIfSuccessful {
+        SignalStore.backup.messageCredentials.clearAll()
+      }
  }

  /**
@@ -1818,13 +1832,25 @@ object BackupRepository {
        return RestoreTimestampResult.Success(SignalStore.backup.lastBackupTime)
      }

-      timestampResult is NetworkResult.StatusCodeError && (timestampResult.code == 401 || timestampResult.code == 404) -> {
+      timestampResult is NetworkResult.StatusCodeError && timestampResult.code == 404 -> {
        Log.i(TAG, "No backup file exists")
        SignalStore.backup.lastBackupTime = 0L
        SignalStore.backup.isBackupTimestampRestored = true
        return RestoreTimestampResult.NotFound
      }

+      timestampResult is NetworkResult.StatusCodeError && timestampResult.code == 401 -> {
+        Log.i(TAG, "Backups not enabled")
+        SignalStore.backup.lastBackupTime = 0L
+        SignalStore.backup.isBackupTimestampRestored = true
+        return RestoreTimestampResult.BackupsNotEnabled
+      }
+
+      timestampResult is NetworkResult.ApplicationError && timestampResult.getCause() is VerificationFailedException -> {
+        Log.w(TAG, "Entered AEP fails zk verification", timestampResult.getCause())
+        return RestoreTimestampResult.VerificationFailure
+      }
+
      else -> {
        Log.w(TAG, "Could not check for backup file.", timestampResult.getCause())
        return RestoreTimestampResult.Failure
@@ -1832,11 +1858,40 @@ object BackupRepository {
    }
  }

-  fun verifyBackupKeyAssociatedWithAccount(aci: ACI, aep: AccountEntropyPool): MessageBackupTier? {
+  fun verifyBackupKeyAssociatedWithAccount(aci: ACI, aep: AccountEntropyPool): RestoreTimestampResult {
+    Log.i(TAG, "Verifying enter aep is associated with account")
+    var result: RestoreTimestampResult = getBackupTimestampToVerifyAepAssociatedWithAccountAndHasBackup(aci, aep)
+
+    if (result is RestoreTimestampResult.VerificationFailure) {
+      Log.w(TAG, "Resetting backup id reservation due to zk verification failure")
+      val triggerResult = SignalNetwork.archive.triggerBackupIdReservation(aep.deriveMessageBackupKey(), null, aci)
+      result = when {
+        triggerResult is NetworkResult.Success -> {
+          Log.i(TAG, "Reset successful, retrying aep verification")
+          SignalStore.backup.messageCredentials.clearAll()
+          getBackupTimestampToVerifyAepAssociatedWithAccountAndHasBackup(aci, aep)
+        }
+
+        triggerResult is NetworkResult.StatusCodeError && triggerResult.code == 429 -> {
+          Log.w(TAG, "Rate limited when resetting backup id, failing operation $triggerResult")
+          RestoreTimestampResult.RateLimited(triggerResult.retryAfter())
+        }
+
+        else -> {
+          Log.w(TAG, "Reset backup id failed, failing operation", triggerResult.getCause())
+          result
+        }
+      }
+    }
+
+    return result
+  }
+
+  private fun getBackupTimestampToVerifyAepAssociatedWithAccountAndHasBackup(aci: ACI, aep: AccountEntropyPool): RestoreTimestampResult {
    val currentTime = System.currentTimeMillis()
    val messageBackupKey = aep.deriveMessageBackupKey()

-    val result: NetworkResult<MessageBackupTier> = SignalNetwork.archive.getServiceCredentials(currentTime)
+    val result: NetworkResult<ZonedDateTime> = SignalNetwork.archive.getServiceCredentials(currentTime)
      .then { result ->
        val credential: ArchiveServiceCredential? = ArchiveServiceCredentials(result.messageCredentials.associateBy { it.redemptionTime }).getForCurrentTime(currentTime.milliseconds)

@@ -1851,20 +1906,41 @@ object BackupRepository {
          )
        }
      }
-      .map { messageAccess ->
-        val zkCredential = SignalNetwork.archive.getZkCredential(aci, messageAccess)
-        if (zkCredential.backupLevel == BackupLevel.PAID) {
-          MessageBackupTier.PAID
-        } else {
-          MessageBackupTier.FREE
-        }
+      .then { messageAccess ->
+        SignalNetwork.archive.getBackupInfo(SignalStore.account.requireAci(), messageAccess)
+          .then { info -> SignalNetwork.archive.getCdnReadCredentials(info.cdn ?: RemoteConfig.backupFallbackArchiveCdn, aci, messageAccess).map { it.headers to info } }
+          .then { pair ->
+            val (cdnCredentials, info) = pair
+            NetworkResult.fromFetch {
+              AppDependencies.signalServiceMessageReceiver.getCdnLastModifiedTime(info.cdn!!, cdnCredentials, "backups/${info.backupDir}/${info.backupName}")
+            }
+          }
      }

-    return if (result is NetworkResult.Success) {
-      result.result
-    } else {
-      Log.i(TAG, "Unable to verify backup key", result.getCause())
-      null
+    return when {
+      result is NetworkResult.Success -> {
+        RestoreTimestampResult.Success(result.result.toMillis())
+      }
+
+      result is NetworkResult.StatusCodeError && result.code == 404 -> {
+        Log.i(TAG, "No backup file exists")
+        RestoreTimestampResult.NotFound
+      }
+
+      result is NetworkResult.StatusCodeError && result.code == 401 -> {
+        Log.i(TAG, "Backups not enabled")
+        RestoreTimestampResult.BackupsNotEnabled
+      }
+
+      result is NetworkResult.ApplicationError && result.getCause() is VerificationFailedException -> {
+        Log.w(TAG, "Entered AEP fails zk verification", result.getCause())
+        RestoreTimestampResult.VerificationFailure
+      }
+
+      else -> {
+        Log.w(TAG, "Could not check for backup file.", result.getCause())
+        RestoreTimestampResult.Failure
+      }
    }
  }

@@ -1993,7 +2069,11 @@ object BackupRepository {

      return SignalNetwork.archive
        .triggerBackupIdReservation(messageBackupKey, mediaRootBackupKey, SignalStore.account.requireAci())
-        .then { getArchiveServiceAccessPair() }
+        .then {
+          SignalStore.backup.messageCredentials.clearAll()
+          SignalStore.backup.mediaCredentials.clearAll()
+          getArchiveServiceAccessPair()
+        }
        .then { credential -> SignalNetwork.archive.setPublicKey(SignalStore.account.requireAci(), credential.messageBackupAccess).map { credential } }
        .then { credential -> SignalNetwork.archive.setPublicKey(SignalStore.account.requireAci(), credential.mediaBackupAccess).map { credential } }
        .runIfSuccessful { SignalStore.backup.backupsInitialized = true }
@@ -2145,20 +2225,24 @@ object BackupRepository {
        SignalStore.backup.nextBackupSecretData = result.data.nextBackupSecretData
        result.data.forwardSecrecyToken
      }
+
      is SvrBApi.RestoreResult.NetworkError -> {
        Log.w(TAG, "[remoteRestore] Network error during SVRB.", result.exception)
        return RemoteRestoreResult.NetworkError
      }
+
      is SvrBApi.RestoreResult.RestoreFailedError,
      SvrBApi.RestoreResult.InvalidDataError -> {
        Log.w(TAG, "[remoteRestore] Permanent SVRB error! $result")
        return RemoteRestoreResult.PermanentSvrBFailure
      }
+
      SvrBApi.RestoreResult.DataMissingError,
      is SvrBApi.RestoreResult.SvrError -> {
        Log.w(TAG, "[remoteRestore] Failed to fetch SVRB data: $result")
        return RemoteRestoreResult.Failure
      }
+
      is SvrBApi.RestoreResult.UnknownError -> {
        Log.e(TAG, "[remoteRestore] Unknown SVRB result! Crashing.", result.throwable)
        throw result.throwable
@@ -2383,6 +2467,9 @@ sealed interface RemoteRestoreResult {
 sealed interface RestoreTimestampResult {
  data class Success(val timestamp: Long) : RestoreTimestampResult
  data object NotFound : RestoreTimestampResult
+  data object BackupsNotEnabled : RestoreTimestampResult
+  data object VerificationFailure : RestoreTimestampResult
+  data class RateLimited(val retryAfter: Duration?) : RestoreTimestampResult
  data object Failure : RestoreTimestampResult
 }

--- a/app/src/main/java/org/thoughtcrime/securesms/keyvalue/AccountValues.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/keyvalue/AccountValues.kt
@@ -165,6 +165,7 @@ class AccountValues internal constructor(store: KeyValueStore, context: Context)

  fun restoreAccountEntropyPool(aep: AccountEntropyPool) {
    AEP_LOCK.withLock {
+      Log.i(TAG, "Restoring AEP from registration source", Throwable())
      store
        .beginWrite()
        .putString(KEY_ACCOUNT_ENTROPY_POOL, aep.value)
--- a/app/src/main/java/org/thoughtcrime/securesms/keyvalue/BackupValues.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/keyvalue/BackupValues.kt
@@ -92,6 +92,8 @@ class BackupValues(store: KeyValueStore) : SignalStoreValues(store) {

    private const val KEY_NEXT_BACKUP_SECRET_DATA = "backup.next_backup_secret_data"

+    private const val KEY_RESTORING_VIA_QR = "backup.restore_via_qr"
+
    private val cachedCdnCredentialsExpiresIn: Duration = 12.hours

    private val lock = ReentrantLock()
@@ -387,6 +389,9 @@ class BackupValues(store: KeyValueStore) : SignalStoreValues(store) {
  /** The value from the last successful SVRB operation that must be passed to the next SVRB operation. */
  var nextBackupSecretData by nullableBlobValue(KEY_NEXT_BACKUP_SECRET_DATA, null)

+  /** True if attempting to restore backup from quick restore/QR code */
+  var restoringViaQr by booleanValue(KEY_RESTORING_VIA_QR, false)
+
  /**
   * If true, it means we have been told that remote storage is full, but we have not yet run any of our "garbage collection" tasks, like committing deletes
   * or pruning orphaned media.
--- a/app/src/main/java/org/thoughtcrime/securesms/keyvalue/RestoreDecisionStateExt.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/keyvalue/RestoreDecisionStateExt.kt
@@ -37,7 +37,7 @@ val RestoreDecisionState.includeDeviceToDeviceTransfer: Boolean
    RestoreDecisionState.State.INTEND_TO_RESTORE -> {
      this.intendToRestoreData?.hasOldDevice == true
    }
-    else -> true
+    else -> false
  }

 /** Has a final decision been made regarding restoring. */
--- a/app/src/main/java/org/thoughtcrime/securesms/registration/ui/restore/RemoteRestoreViewModel.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/registration/ui/restore/RemoteRestoreViewModel.kt
@@ -56,7 +56,23 @@ class RemoteRestoreViewModel(isOnlyRestoreOption: Boolean) : ViewModel() {
  fun reload() {
    viewModelScope.launch(Dispatchers.IO) {
      store.update { it.copy(loadState = ScreenState.LoadState.LOADING, loadAttempts = it.loadAttempts + 1) }
-      val result = BackupRepository.restoreBackupFileTimestamp()
+      Log.i(TAG, "Fetching remote backup information")
+      var result: RestoreTimestampResult = BackupRepository.restoreBackupFileTimestamp()
+
+      if (result is RestoreTimestampResult.VerificationFailure && SignalStore.account.restoredAccountEntropyPool) {
+        Log.w(TAG, "Resetting backup id reservation due to zk verification failure with restored AEP")
+        result = when (val triggerResult = BackupRepository.triggerBackupIdReservationForRestore()) {
+          is NetworkResult.Success -> {
+            Log.i(TAG, "Reset successful, trying to restore timestamp")
+            BackupRepository.restoreBackupFileTimestamp()
+          }
+          else -> {
+            Log.w(TAG, "Reset unsuccessful, failing", triggerResult.getCause())
+            result
+          }
+        }
+      }
+
      store.update {
        when (result) {
          is RestoreTimestampResult.Success -> {
@@ -68,6 +84,7 @@ class RemoteRestoreViewModel(isOnlyRestoreOption: Boolean) : ViewModel() {
            )
          }

+          is RestoreTimestampResult.BackupsNotEnabled,
          is RestoreTimestampResult.NotFound -> {
            it.copy(loadState = ScreenState.LoadState.NOT_FOUND)
          }
--- a/app/src/main/java/org/thoughtcrime/securesms/registration/ui/restore/RestoreViaQrViewModel.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/registration/ui/restore/RestoreViaQrViewModel.kt
@@ -165,6 +165,7 @@ class RestoreViaQrViewModel : ViewModel() {

        SignalStore.backup.lastBackupTime = result.message.backupTimestampMs ?: 0
        SignalStore.backup.isBackupTimestampRestored = true
+        SignalStore.backup.restoringViaQr = true
        SignalStore.backup.backupTier = when (result.message.tier) {
          RegistrationProvisionMessage.Tier.FREE -> MessageBackupTier.FREE
          RegistrationProvisionMessage.Tier.PAID -> MessageBackupTier.PAID
--- a/app/src/main/java/org/thoughtcrime/securesms/restore/RestoreViewModel.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/restore/RestoreViewModel.kt
@@ -25,7 +25,6 @@ import org.thoughtcrime.securesms.keyvalue.skippedRestoreChoice
 import org.thoughtcrime.securesms.registration.data.QuickRegistrationRepository
 import org.thoughtcrime.securesms.registration.ui.restore.RestoreMethod
 import org.thoughtcrime.securesms.registration.ui.restore.StorageServiceRestore
-import org.thoughtcrime.securesms.restore.transferorrestore.BackupRestorationType
 import org.whispersystems.signalservice.api.provisioning.RestoreMethod as ApiRestoreMethod

 /**
@@ -44,22 +43,6 @@ class RestoreViewModel : ViewModel() {
    }
  }

-  fun onTransferFromAndroidDeviceSelected() {
-    store.update {
-      it.copy(restorationType = BackupRestorationType.DEVICE_TRANSFER)
-    }
-  }
-
-  fun onRestoreFromLocalBackupSelected() {
-    store.update {
-      it.copy(restorationType = BackupRestorationType.LOCAL_BACKUP)
-    }
-  }
-
-  fun getBackupRestorationType(): BackupRestorationType {
-    return store.value.restorationType
-  }
-
  fun setBackupFileUri(backupFileUri: Uri) {
    store.update {
      it.copy(backupFile = backupFileUri)
@@ -75,17 +58,17 @@ class RestoreViewModel : ViewModel() {
  }

  fun getAvailableRestoreMethods(): List<RestoreMethod> {
-    if (SignalStore.registration.isOtherDeviceAndroid || SignalStore.registration.restoreDecisionState.skippedRestoreChoice || !SignalStore.backup.isBackupTimestampRestored) {
+    if (SignalStore.registration.isOtherDeviceAndroid || SignalStore.registration.restoreDecisionState.skippedRestoreChoice) {
      val methods = mutableListOf(RestoreMethod.FROM_LOCAL_BACKUP_V1)

-      if (SignalStore.registration.restoreDecisionState.includeDeviceToDeviceTransfer) {
+      if (SignalStore.registration.isOtherDeviceAndroid && SignalStore.registration.restoreDecisionState.includeDeviceToDeviceTransfer) {
        methods.add(0, RestoreMethod.FROM_OLD_DEVICE)
      }

      when (SignalStore.backup.backupTier) {
        MessageBackupTier.FREE -> methods.add(1, RestoreMethod.FROM_SIGNAL_BACKUPS)
        MessageBackupTier.PAID -> methods.add(0, RestoreMethod.FROM_SIGNAL_BACKUPS)
-        null -> if (!SignalStore.backup.isBackupTimestampRestored) {
+        null -> if (!SignalStore.backup.restoringViaQr) {
          methods.add(1, RestoreMethod.FROM_SIGNAL_BACKUPS)
        }
      }
@@ -93,7 +76,7 @@ class RestoreViewModel : ViewModel() {
      return methods
    }

-    if (SignalStore.backup.backupTier != null || !SignalStore.backup.isBackupTimestampRestored) {
+    if (SignalStore.backup.restoringViaQr && SignalStore.backup.backupTier != null) {
      return listOf(RestoreMethod.FROM_SIGNAL_BACKUPS)
    }

@@ -104,6 +87,10 @@ class RestoreViewModel : ViewModel() {
    return SignalStore.account.restoredAccountEntropyPool
  }

+  fun hasRestoredBackupDataFromQr(): Boolean {
+    return SignalStore.backup.restoringViaQr && SignalStore.backup.backupTier != null
+  }
+
  fun skipRestore() {
    SignalStore.registration.restoreDecisionState = RestoreDecisionState.Skipped

--- a/app/src/main/java/org/thoughtcrime/securesms/restore/enterbackupkey/PostRegistrationEnterBackupKeyFragment.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/restore/enterbackupkey/PostRegistrationEnterBackupKeyFragment.kt
@@ -24,7 +24,6 @@ import org.signal.core.ui.compose.Dialogs
 import org.signal.core.util.logging.Log
 import org.thoughtcrime.securesms.R
 import org.thoughtcrime.securesms.compose.ComposeFragment
-import org.thoughtcrime.securesms.registration.ui.restore.AccountEntropyPoolVerification
 import org.thoughtcrime.securesms.registration.ui.restore.EnterBackupKeyScreen
 import org.thoughtcrime.securesms.util.CommunicationActions
 import org.thoughtcrime.securesms.util.navigation.safeNavigate
@@ -72,17 +71,14 @@ class PostRegistrationEnterBackupKeyFragment : ComposeFragment() {
      chunkLength = 4,
      aepValidationError = state.aepValidationError,
      onBackupKeyChanged = viewModel::updateBackupKey,
-      onNextClicked = {
-        viewModel.restoreBackupTier()
-      },
+      onNextClicked = { viewModel.restoreBackupTimestamp() },
      onLearnMore = { CommunicationActions.openBrowserLink(requireContext(), LEARN_MORE_URL) },
      onSkip = { findNavController().popBackStack() }
    ) {
      ErrorContent(
-        showBackupTierNotRestoreError = state.showBackupTierNotRestoreError,
-        aepError = state.aepValidationError,
+        errorDialog = state.errorDialog,
        onBackupKeyHelp = { CommunicationActions.openBrowserLink(requireContext(), LEARN_MORE_URL) },
-        onBackupTierNotRestoredDismiss = viewModel::hideRestoreBackupTierFailed
+        onDismiss = viewModel::hideErrorDialog
      )
    }
  }
@@ -90,20 +86,71 @@ class PostRegistrationEnterBackupKeyFragment : ComposeFragment() {

@Composable
 private fun ErrorContent(
-  showBackupTierNotRestoreError: Boolean,
-  aepError: AccountEntropyPoolVerification.AEPValidationError?,
+  errorDialog: PostRegistrationEnterBackupKeyViewModel.ErrorDialog?,
  onBackupKeyHelp: () -> Unit = {},
-  onBackupTierNotRestoredDismiss: () -> Unit = {}
+  onDismiss: () -> Unit = {}
 ) {
-  if (aepError == AccountEntropyPoolVerification.AEPValidationError.Incorrect && showBackupTierNotRestoreError) {
-    Dialogs.SimpleAlertDialog(
-      title = stringResource(R.string.EnterBackupKey_incorrect_backup_key_title),
-      body = stringResource(R.string.EnterBackupKey_incorrect_backup_key_message),
-      confirm = stringResource(R.string.EnterBackupKey_try_again),
-      dismiss = stringResource(R.string.EnterBackupKey_backup_key_help),
-      onConfirm = {},
-      onDeny = onBackupKeyHelp,
-      onDismiss = onBackupTierNotRestoredDismiss
-    )
+  if (errorDialog == null) {
+    return
+  }
+
+  when (errorDialog) {
+    PostRegistrationEnterBackupKeyViewModel.ErrorDialog.AEP_INVALID -> {
+      Dialogs.SimpleAlertDialog(
+        title = stringResource(R.string.EnterBackupKey_incorrect_backup_key_title),
+        body = stringResource(R.string.EnterBackupKey_incorrect_backup_key_message),
+        confirm = stringResource(R.string.EnterBackupKey_try_again),
+        dismiss = stringResource(R.string.EnterBackupKey_backup_key_help),
+        onConfirm = {},
+        onDeny = onBackupKeyHelp,
+        onDismiss = onDismiss
+      )
+    }
+
+    PostRegistrationEnterBackupKeyViewModel.ErrorDialog.BACKUP_NOT_FOUND -> {
+      Dialogs.SimpleAlertDialog(
+        title = stringResource(R.string.EnterBackupKey_backup_not_found),
+        body = stringResource(R.string.EnterBackupKey_backup_key_you_entered_is_correct_but_no_backup),
+        confirm = stringResource(R.string.EnterBackupKey_try_again),
+        dismiss = stringResource(R.string.EnterBackupKey_backup_key_help),
+        onConfirm = {},
+        onDeny = onBackupKeyHelp,
+        onDismiss = onDismiss
+      )
+    }
+
+    PostRegistrationEnterBackupKeyViewModel.ErrorDialog.UNKNOWN_ERROR -> {
+      Dialogs.SimpleAlertDialog(
+        title = stringResource(R.string.EnterBackupKey_cant_restore_backup),
+        body = stringResource(R.string.EnterBackupKey_your_backup_cant_be_restored_right_now),
+        confirm = stringResource(R.string.EnterBackupKey_try_again),
+        onConfirm = {},
+        onDismiss = onDismiss
+      )
+    }
+
+    PostRegistrationEnterBackupKeyViewModel.ErrorDialog.BACKUPS_NOT_ENABLED -> {
+      Dialogs.SimpleAlertDialog(
+        title = stringResource(R.string.EnterBackupKey_backup_not_found),
+        body = stringResource(R.string.EnterBackupKey_backup_key_incorrect_or_backups_not_enabled),
+        confirm = stringResource(R.string.EnterBackupKey_try_again),
+        dismiss = stringResource(R.string.EnterBackupKey_backup_key_help),
+        onConfirm = {},
+        onDeny = onBackupKeyHelp,
+        onDismiss = onDismiss
+      )
+    }
+
+    PostRegistrationEnterBackupKeyViewModel.ErrorDialog.RATE_LIMITED -> {
+      Dialogs.SimpleAlertDialog(
+        title = stringResource(R.string.EnterBackupKey_backup_not_found),
+        body = stringResource(R.string.EnterBackupKey_backup_key_check_rate_limited),
+        confirm = stringResource(R.string.EnterBackupKey_try_again),
+        dismiss = stringResource(R.string.EnterBackupKey_backup_key_help),
+        onConfirm = {},
+        onDeny = onBackupKeyHelp,
+        onDismiss = onDismiss
+      )
+    }
  }
 }
--- a/app/src/main/java/org/thoughtcrime/securesms/restore/enterbackupkey/PostRegistrationEnterBackupKeyViewModel.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/restore/enterbackupkey/PostRegistrationEnterBackupKeyViewModel.kt
@@ -15,9 +15,9 @@ import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.withContext
 import org.signal.core.util.logging.Log
 import org.thoughtcrime.securesms.backup.v2.BackupRepository
+import org.thoughtcrime.securesms.backup.v2.RestoreTimestampResult
 import org.thoughtcrime.securesms.keyvalue.SignalStore
 import org.thoughtcrime.securesms.registration.ui.restore.AccountEntropyPoolVerification
 import org.thoughtcrime.securesms.registration.ui.restore.AccountEntropyPoolVerification.AEPValidationError
@@ -51,33 +51,37 @@ class PostRegistrationEnterBackupKeyViewModel : ViewModel() {
    }
  }

-  fun restoreBackupTier() {
+  fun restoreBackupTimestamp() {
    store.update { it.copy(inProgress = true) }
    viewModelScope.launch(Dispatchers.IO) {
      val aep = AccountEntropyPool.parseOrNull(backupKey)
-      val backupTier = withContext(Dispatchers.IO) {
-        if (aep != null) {
-          BackupRepository.verifyBackupKeyAssociatedWithAccount(SignalStore.account.requireAci(), aep)
-        } else {
-          Log.w(TAG, "Parsed AEP is null, failing")
-          null
-        }
+
+      val restoreTimestampResult = if (aep != null) {
+        BackupRepository.verifyBackupKeyAssociatedWithAccount(SignalStore.account.requireAci(), aep)
+      } else {
+        Log.w(TAG, "Parsed AEP is null, failing")
+        store.update { it.copy(aepValidationError = AEPValidationError.Invalid, errorDialog = ErrorDialog.AEP_INVALID) }
+        return@launch
      }

-      if (backupTier != null) {
-        Log.i(TAG, "Backup tier found with entered AEP, migrating to new AEP and moving on to restore")
-        SignalStore.account.restoreAccountEntropyPool(aep!!)
-        store.update { it.copy(restoreBackupTierSuccessful = true) }
-      } else {
-        Log.w(TAG, "Unable to validate AEP against currently registered account")
-        store.update { it.copy(aepValidationError = AEPValidationError.Incorrect, showBackupTierNotRestoreError = true) }
+      when (restoreTimestampResult) {
+        is RestoreTimestampResult.Success -> {
+          Log.i(TAG, "Backup timestamp found with entered AEP, migrating to new AEP and moving on to restore")
+          SignalStore.account.restoreAccountEntropyPool(aep)
+          store.update { it.copy(restoreBackupTierSuccessful = true) }
+        }
+
+        RestoreTimestampResult.NotFound -> store.update { it.copy(errorDialog = ErrorDialog.BACKUP_NOT_FOUND) }
+        RestoreTimestampResult.BackupsNotEnabled -> store.update { it.copy(errorDialog = ErrorDialog.BACKUPS_NOT_ENABLED) }
+        is RestoreTimestampResult.RateLimited -> store.update { it.copy(errorDialog = ErrorDialog.RATE_LIMITED) }
+        else -> store.update { it.copy(errorDialog = ErrorDialog.UNKNOWN_ERROR) }
      }
    }
  }

-  fun hideRestoreBackupTierFailed() {
+  fun hideErrorDialog() {
    store.update {
-      it.copy(showBackupTierNotRestoreError = false, inProgress = false)
+      it.copy(errorDialog = null, inProgress = false)
    }
  }

@@ -85,7 +89,15 @@ class PostRegistrationEnterBackupKeyViewModel : ViewModel() {
    val backupKeyValid: Boolean = false,
    val inProgress: Boolean = false,
    val restoreBackupTierSuccessful: Boolean = false,
-    val showBackupTierNotRestoreError: Boolean = false,
-    val aepValidationError: AEPValidationError? = null
+    val aepValidationError: AEPValidationError? = null,
+    val errorDialog: ErrorDialog? = null
  )
+
+  enum class ErrorDialog {
+    AEP_INVALID,
+    BACKUPS_NOT_ENABLED,
+    BACKUP_NOT_FOUND,
+    RATE_LIMITED,
+    UNKNOWN_ERROR
+  }
 }
--- a/app/src/main/java/org/thoughtcrime/securesms/restore/selection/SelectRestoreMethodFragment.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/restore/selection/SelectRestoreMethodFragment.kt
@@ -88,7 +88,7 @@ class SelectRestoreMethodFragment : ComposeFragment() {

    when (method) {
      RestoreMethod.FROM_SIGNAL_BACKUPS -> {
-        if (viewModel.hasRestoredAccountEntropyPool()) {
+        if (viewModel.hasRestoredBackupDataFromQr()) {
          startActivity(RemoteRestoreActivity.getIntent(requireContext()))
        } else {
          findNavController().safeNavigate(SelectRestoreMethodFragmentDirections.goToPostRestoreEnterBackupKey())