Reset backup id on zk verification failure during restore attempts.

2025-12-24 13:08:46 +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())
--- a/app/src/main/res/values/strings.xml
+++ b/app/src/main/res/values/strings.xml
@@ -8654,6 +8654,10 @@
    <string name="EnterBackupKey_backup_not_found">Backup not found</string>
    <!-- Dialog message shown during registration and trying to restore a remote backup -->
    <string name="EnterBackupKey_backup_key_you_entered_is_correct_but_no_backup">The recovery key you entered is correct, but there is no backup associated with it. If you still have your old phone, make sure backups are enabled and that a backup has been completed and try again.</string>
+    <!-- Dialog message shown during registration and trying to restore a remote backup -->
+    <string name="EnterBackupKey_backup_key_incorrect_or_backups_not_enabled">You may have entered your recovery key incorrectly, or you have not completed at least one backup on your old device.\n\nMake sure you’re registering with the same phone number and recovery key you saved when enabling Signal Secure Backups, and confirm that your old device has completed a backup.</string>
+    <!-- Dialog message shown during registration and trying to restore a remote backup -->
+    <string name="EnterBackupKey_backup_key_check_rate_limited">You have run out of recovery key attempts. For your security, the number of attempts you can make is limited. Wait one hour and try again.</string>
    <!-- Dialog button shown during registration and trying to restore a remote backup to instead skip said backup process -->
    <string name="EnterBackupKey_skip_restore">Skip restore</string>
    <!-- Dialog title text when encountering backup restore network error -->
--- a/libsignal-service/src/main/java/org/whispersystems/signalservice/api/archive/ArchiveApi.kt
+++ b/libsignal-service/src/main/java/org/whispersystems/signalservice/api/archive/ArchiveApi.kt
@@ -95,19 +95,21 @@ class ArchiveApi(
   * Ensures that you reserve backupIds for both messages and media on the service. This must be done before any other
   * backup-related calls. You only need to do it once, but repeated calls are safe.
   *
+   * Passing null for either key will skip reserving for that backup and not cost a rate limit permit.
+   *
   * PUT /v1/archives/backupid
   *
   * - 204: Success
   * - 400: Invalid credential
   * - 429: Rate-limited
   */
-  fun triggerBackupIdReservation(messageBackupKey: MessageBackupKey, mediaRootBackupKey: MediaRootBackupKey, aci: ACI): NetworkResult<Unit> {
-    val messageBackupRequestContext = BackupAuthCredentialRequestContext.create(messageBackupKey.value, aci.rawUuid)
-    val mediaBackupRequestContext = BackupAuthCredentialRequestContext.create(mediaRootBackupKey.value, aci.rawUuid)
+  fun triggerBackupIdReservation(messageBackupKey: MessageBackupKey?, mediaRootBackupKey: MediaRootBackupKey?, aci: ACI): NetworkResult<Unit> {
+    val messageBackupRequestContext = messageBackupKey?.let { BackupAuthCredentialRequestContext.create(messageBackupKey.value, aci.rawUuid) }
+    val mediaBackupRequestContext = mediaRootBackupKey?.let { BackupAuthCredentialRequestContext.create(mediaRootBackupKey.value, aci.rawUuid) }

    val request = WebSocketRequestMessage.put(
      "/v1/archives/backupid",
-      ArchiveSetBackupIdRequest(messageBackupRequestContext.request, mediaBackupRequestContext.request)
+      ArchiveSetBackupIdRequest(messageBackupRequestContext?.request, mediaBackupRequestContext?.request)
    )

    return NetworkResult.fromWebSocketRequest(authWebSocket, request)
--- a/libsignal-service/src/main/java/org/whispersystems/signalservice/api/archive/ArchiveSetBackupIdRequest.kt
+++ b/libsignal-service/src/main/java/org/whispersystems/signalservice/api/archive/ArchiveSetBackupIdRequest.kt
@@ -19,10 +19,10 @@ import org.signal.libsignal.zkgroup.backups.BackupAuthCredentialRequest
 class ArchiveSetBackupIdRequest(
  @JsonProperty
  @JsonSerialize(using = BackupAuthCredentialRequestSerializer::class)
-  val messagesBackupAuthCredentialRequest: BackupAuthCredentialRequest,
+  val messagesBackupAuthCredentialRequest: BackupAuthCredentialRequest?,
  @JsonProperty
  @JsonSerialize(using = BackupAuthCredentialRequestSerializer::class)
-  val mediaBackupAuthCredentialRequest: BackupAuthCredentialRequest
+  val mediaBackupAuthCredentialRequest: BackupAuthCredentialRequest?
 ) {
  class BackupAuthCredentialRequestSerializer : JsonSerializer<BackupAuthCredentialRequest>() {
    override fun serialize(value: BackupAuthCredentialRequest, gen: JsonGenerator, serializers: SerializerProvider) {