Fix local backup v2 restore pre-registration in regV5.

app/src/main/java/org/thoughtcrime/securesms/backup/v2/local/LocalArchiver.kt

@@ -26,6 +26,7 @@ import org.signal.core.util.toJson
 import org.signal.libsignal.crypto.Aes256Ctr32
 import org.thoughtcrime.securesms.backup.LocalExportProgress
 import org.thoughtcrime.securesms.backup.v2.BackupRepository
+import org.thoughtcrime.securesms.backup.v2.ImportResult
 import org.thoughtcrime.securesms.database.AttachmentTable
 import org.thoughtcrime.securesms.database.SignalDatabase
 import org.thoughtcrime.securesms.keyvalue.SignalStore
@@ -269,13 +270,18 @@ object LocalArchiver {
 
       val mainStreamLength = snapshotFileSystem.mainLength() ?: return ArchiveResult.failure(RestoreFailure.MainStream)
 
-      BackupRepository.importLocal(
+      val importResult = BackupRepository.importLocal(
         mainStreamFactory = { snapshotFileSystem.mainInputStream()!! },
         mainStreamLength = mainStreamLength,
         selfData = selfData,
         backupId = backupId,
         messageBackupKey = messageBackupKey
       )
+
+      if (importResult is ImportResult.Failure) {
+        Log.w(TAG, "Local backup import failed")
+        return RestoreResult.failure(RestoreFailure.ImportFailed)
+      }
     } finally {
       metadataStream?.close()
     }
@@ -350,6 +356,7 @@ object LocalArchiver {
     data object MainStream : RestoreFailure
     data object Cancelled : RestoreFailure
     data object BackupIdMissing : RestoreFailure
+    data object ImportFailed : RestoreFailure
     data class VersionMismatch(val backupVersion: Int, val supportedVersion: Int) : RestoreFailure
   }
 

app/src/main/java/org/thoughtcrime/securesms/registration/v2/AppRegistrationStorageController.kt

@@ -353,6 +353,16 @@ class AppRegistrationStorageController(private val context: Context) : StorageCo
     }
   }
 
+  override suspend fun verifyLocalBackupKey(backupUri: Uri, aep: AccountEntropyPool): Boolean = withContext(Dispatchers.IO) {
+    val backupDir = DocumentFile.fromTreeUri(context, backupUri)
+    if (backupDir == null || !backupDir.canRead()) {
+      Log.w(TAG, "[verifyLocalBackupKey] Could not open backup directory.")
+      return@withContext false
+    }
+
+    LocalArchiver.canDecryptMainArchive(SnapshotFileSystem(context, backupDir), aep.deriveMessageBackupKey())
+  }
+
   override fun restoreLocalBackupV2(rootUri: Uri, backupUri: Uri, aep: AccountEntropyPool): Flow<LocalBackupRestoreProgress> = callbackFlow {
     Log.d(TAG, "Starting V2 local backup restore from backup=$backupUri, root=$rootUri")
 

demo/registration/src/main/java/org/signal/registration/sample/dependencies/DemoStorageController.kt

@@ -318,6 +318,8 @@ class DemoStorageController(private val context: Context) : StorageController {
     Log.d(TAG, "Simulated V1 restore complete.")
   }.flowOn(Dispatchers.IO)
 
+  override suspend fun verifyLocalBackupKey(backupUri: Uri, aep: AccountEntropyPool): Boolean = true
+
   override fun restoreLocalBackupV2(rootUri: Uri, backupUri: Uri, aep: AccountEntropyPool): Flow<LocalBackupRestoreProgress> = flow {
     Log.d(TAG, "Starting simulated V2 local backup restore from backup=$backupUri, root=$rootUri")
 

feature/registration/src/main/java/org/signal/registration/RegistrationNavigation.kt

@@ -46,6 +46,7 @@ import org.signal.core.util.serialization.AccountEntropyPoolSerializer
 import org.signal.registration.screens.accountlocked.AccountLockedScreen
 import org.signal.registration.screens.accountlocked.AccountLockedScreenEvents
 import org.signal.registration.screens.accountlocked.AccountLockedState
+import org.signal.registration.screens.aepentry.EnterAepForLocalBackupResult
 import org.signal.registration.screens.aepentry.EnterAepForLocalBackupViewModel
 import org.signal.registration.screens.aepentry.EnterAepForRemoteBackupPostRegistrationViewModel
 import org.signal.registration.screens.aepentry.EnterAepForRemoteBackupPreRegistrationViewModel
@@ -231,8 +232,19 @@ sealed interface RegistrationRoute : NavKey, Parcelable {
   @Serializable
   data object EnterLocalBackupV1Passphrase : RegistrationRoute
 
+  /**
+   * Recovery key entry for a local V2 backup.
+   *
+   * When [isPreRegistration] is true (pre-registration manual restore), submitting the key first verifies it can
+   * decrypt the backup at [backupUri], then registers the account via the recovery password derived from it. A backup
+   * belonging to a different account is surfaced to the user, who can choose to restore it after verifying over SMS.
+   * When false (already registered), the key is simply handed back to the restore screen to decrypt the backup.
+   */
   @Serializable
-  data object EnterAepForLocalBackup : RegistrationRoute
+  data class EnterAepForLocalBackup(
+    val isPreRegistration: Boolean = false,
+    val backupUri: String? = null
+  ) : RegistrationRoute
 
   @Serializable
   data class EnterAepForRemoteBackupPreRegistration(val e164: String) : RegistrationRoute
@@ -275,6 +287,7 @@ sealed interface RegistrationRoute : NavKey, Parcelable {
 private const val CAPTCHA_RESULT = "captcha_token"
 private const val COUNTRY_CODE_RESULT = "country_code_result"
 private const val BACKUP_CREDENTIAL_RESULT = "backup_credential_result"
+private const val AEP_FOR_LOCAL_BACKUP_RESULT = "aep_for_local_backup_result"
 private const val LOCAL_BACKUP_RESTORE_RESULT = "local_backup_restore_result"
 private const val PHONE_NUMBER_DISCOVERABILITY_RESULT = "phone_number_discoverability_result"
 private const val PIN_LEARN_MORE_URL = "https://support.signal.org/hc/articles/360007059792"
@@ -349,7 +362,7 @@ fun RegistrationNavHost(
     transitionSpec = { TransitionSpecs.HorizontalSlide.transitionSpec },
     popTransitionSpec = {
       when {
-        initialState.key == RegistrationRoute.EnterAepForLocalBackup.toString() || initialState.key == RegistrationRoute.EnterAepForRemoteBackupPreRegistration.toString() -> {
+        initialState.key.toString().startsWith("EnterAepForLocalBackup") || initialState.key == RegistrationRoute.EnterAepForRemoteBackupPreRegistration.toString() -> {
           TransitionSpecs.HorizontalSlide.transitionSpec
         }
 
@@ -755,6 +768,13 @@ private fun EntryProviderScope<NavKey>.navigationEntries(
       }
     }
 
+    ResultEffect<EnterAepForLocalBackupResult>(registrationViewModel.resultBus, AEP_FOR_LOCAL_BACKUP_RESULT) { result ->
+      when (result) {
+        is EnterAepForLocalBackupResult.RestoreReady -> viewModel.onEvent(LocalBackupRestoreEvents.PassphraseSubmitted(result.key))
+        is EnterAepForLocalBackupResult.RegistrationDeferredToSms -> viewModel.onEvent(LocalBackupRestoreEvents.RegistrationDeferredToSms)
+      }
+    }
+
     LocalBackupRestoreScreen(
       state = state,
       onEvent = { viewModel.onEvent(it) }
@@ -777,13 +797,17 @@ private fun EntryProviderScope<NavKey>.navigationEntries(
   // TODO I think we can re-use the screen but attach different viewmodels to progress forward rather than do for-result flows?
 
   // -- Enter AEP
-  entry<RegistrationRoute.EnterAepForLocalBackup> {
+  entry<RegistrationRoute.EnterAepForLocalBackup> { key ->
     val context = LocalContext.current
     val viewModel: EnterAepForLocalBackupViewModel = viewModel(
       factory = EnterAepForLocalBackupViewModel.Factory(
+        isPreRegistration = key.isPreRegistration,
+        backupUri = key.backupUri,
+        repository = registrationRepository,
+        parentState = registrationViewModel.state,
         parentEventEmitter = registrationViewModel::onEvent,
         resultBus = registrationViewModel.resultBus,
-        resultKey = BACKUP_CREDENTIAL_RESULT,
+        resultKey = AEP_FOR_LOCAL_BACKUP_RESULT,
         isPasswordManagerAvailable = RegistrationCredentialManager.isSupported(context)
       )
     )

feature/registration/src/main/java/org/signal/registration/RegistrationRepository.kt

@@ -906,6 +906,11 @@ class RegistrationRepository(val context: Context, val networkController: Networ
     return storageController.restoreLocalBackupV2(rootUri, backupUri, aep)
   }
 
+  /** Verifies that [aep] can decrypt the V2 local backup at [backupUri] without restoring anything. */
+  suspend fun verifyLocalBackupKey(backupUri: Uri, aep: AccountEntropyPool): Boolean = withContext(Dispatchers.IO) {
+    storageController.verifyLocalBackupKey(backupUri, aep)
+  }
+
   suspend fun scanLocalBackupFolder(folderUri: Uri): List<LocalBackupInfo> = withContext(Dispatchers.IO) {
     storageController.scanLocalBackupFolder(folderUri)
   }

feature/registration/src/main/java/org/signal/registration/StorageController.kt

@@ -125,6 +125,13 @@ interface StorageController {
    */
   fun restoreLocalBackupV2(rootUri: Uri, backupUri: Uri, aep: AccountEntropyPool): Flow<LocalBackupRestoreProgress>
 
+  /**
+   * Verifies that [aep] can decrypt the V2 (folder-based) backup at [backupUri], without restoring anything.
+   * Used to distinguish a mistyped recovery key from a key that belongs to a different account before
+   * attempting recovery-password registration.
+   */
+  suspend fun verifyLocalBackupKey(backupUri: Uri, aep: AccountEntropyPool): Boolean
+
   /**
    * Begins restoring from a remote (server-hosted) backup.
    *

feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepEvents.kt

@@ -21,4 +21,10 @@ sealed class EnterAepEvents {
 
   /** Dismiss a registration error dialog. */
   data object DismissError : EnterAepEvents()
+
+  /** User confirmed restoring a backup that belongs to a different account, deferring the restore until after SMS verification. */
+  data object ConfirmDifferentAccountRestore : EnterAepEvents()
+
+  /** User dismissed the different-account warning dialog without restoring. */
+  data object DismissDifferentAccountDialog : EnterAepEvents()
 }

feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForLocalBackupViewModel.kt

@@ -5,6 +5,8 @@
 
 package org.signal.registration.screens.aepentry
 
+import androidx.annotation.VisibleForTesting
+import androidx.core.net.toUri
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.viewModelScope
@@ -13,14 +15,35 @@ import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.flow.update
+import org.signal.core.models.AccountEntropyPool
 import org.signal.core.ui.compose.EventDrivenViewModel
 import org.signal.core.ui.navigation.ResultEventBus
 import org.signal.core.util.logging.Log
+import org.signal.libsignal.net.RequestResult
+import org.signal.registration.NetworkController
 import org.signal.registration.RegistrationFlowEvent
+import org.signal.registration.RegistrationFlowState
+import org.signal.registration.RegistrationRepository
+import org.signal.registration.RegistrationRoute
 import org.signal.registration.screens.util.navigateBack
+import org.signal.registration.screens.util.navigateTo
 
+/**
+ * Recovery key entry for a local V2 backup restore.
+ *
+ * When [isPreRegistration] is true (pre-registration manual restore), submitting the key first verifies it can
+ * decrypt the selected backup, then attempts recovery-password registration with it -- the restore itself only ever
+ * runs against a registered account. If the server rejects the recovery password, the key is valid but the backup
+ * belongs to a different account: the user is warned and can choose to restore it anyway, which defers the import
+ * until after they verify their number over SMS.
+ *
+ * When [isPreRegistration] is false (already registered), the key is simply handed back to the restore screen.
+ */
 class EnterAepForLocalBackupViewModel(
+  private val isPreRegistration: Boolean,
+  private val backupUri: String?,
+  private val repository: RegistrationRepository,
+  private val parentState: StateFlow<RegistrationFlowState>,
   private val parentEventEmitter: (RegistrationFlowEvent) -> Unit,
   private val resultBus: ResultEventBus,
   private val resultKey: String,
@@ -41,33 +64,167 @@ class EnterAepForLocalBackupViewModel(
   }
 
   override suspend fun processEvent(event: EnterAepEvents) {
+    applyEvent(_state.value, event) { _state.value = it }
+  }
+
+  @VisibleForTesting
+  suspend fun applyEvent(inputState: EnterAepState, event: EnterAepEvents, stateEmitter: (EnterAepState) -> Unit) {
     when (event) {
       is EnterAepEvents.BackupKeyChanged -> {
-        _state.update { EnterAepScreenEventHandler.applyEvent(it, event) }
+        stateEmitter(EnterAepScreenEventHandler.applyEvent(inputState, event))
       }
       is EnterAepEvents.Submit -> {
-        if (_state.value.isBackupKeyValid) {
-          resultBus.sendResult(resultKey, _state.value.backupKey)
-          parentEventEmitter.navigateBack()
-        }
+        applySubmit(inputState, stateEmitter)
       }
       is EnterAepEvents.Cancel -> {
         parentEventEmitter.navigateBack()
       }
       is EnterAepEvents.DismissError -> {
-        _state.update { EnterAepScreenEventHandler.applyEvent(it, event) }
+        stateEmitter(EnterAepScreenEventHandler.applyEvent(inputState, event))
+      }
+      is EnterAepEvents.ConfirmDifferentAccountRestore -> {
+        applyConfirmDifferentAccountRestore(inputState, stateEmitter)
+      }
+      is EnterAepEvents.DismissDifferentAccountDialog -> {
+        stateEmitter(inputState.copy(showDifferentAccountDialog = false))
       }
     }
   }
 
+  private suspend fun applySubmit(inputState: EnterAepState, stateEmitter: (EnterAepState) -> Unit) {
+    check(inputState.isBackupKeyValid) { "AEP is not valid, should not have gotten here." }
+
+    if (!isPreRegistration) {
+      resultBus.sendResult(resultKey, EnterAepForLocalBackupResult.RestoreReady(inputState.backupKey))
+      parentEventEmitter.navigateBack()
+      return
+    }
+
+    val aep = AccountEntropyPool(inputState.backupKey)
+
+    stateEmitter(inputState.copy(isRegistering = true))
+
+    // Confirm the key actually decrypts the backup before going to the server, so a recovery-password rejection can
+    // only mean the backup belongs to a different account rather than a mistyped key.
+    if (!repository.verifyLocalBackupKey(checkNotNull(backupUri).toUri(), aep)) {
+      Log.w(TAG, "[Submit] Entered key cannot decrypt the selected backup.")
+      stateEmitter(inputState.copy(isRegistering = false, aepValidationError = AepValidationError.Incorrect))
+      return
+    }
+
+    parentEventEmitter(RegistrationFlowEvent.UserSuppliedAepSubmitted(aep))
+
+    Log.i(TAG, "[Submit] Attempting registration with RRP derived from user-supplied AEP.")
+
+    attemptToRegister(inputState, aep, provideRegistrationLock = false, stateEmitter)
+  }
+
+  private suspend fun attemptToRegister(inputState: EnterAepState, aep: AccountEntropyPool, provideRegistrationLock: Boolean, stateEmitter: (EnterAepState) -> Unit) {
+    val e164 = checkNotNull(parentState.value.sessionE164) { "No e164 present in the flow state, should not have gotten here." }
+    val masterKey = aep.deriveMasterKey()
+    val recoveryPassword = masterKey.deriveRegistrationRecoveryPassword()
+    val registrationLock = masterKey.deriveRegistrationLock().takeIf { provideRegistrationLock }
+
+    when (val result = repository.registerAccountWithRecoveryPassword(e164, recoveryPassword, registrationLock, existingAccountEntropyPool = aep)) {
+      is RequestResult.Success -> {
+        Log.i(TAG, "[Submit] Successfully registered using RRP from user-supplied AEP. Proceeding with the restore.")
+        val (response, keyMaterial) = result.result
+
+        stateEmitter(inputState.copy(isRegistering = false))
+        parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool, response.storageCapable))
+        resultBus.sendResult(resultKey, EnterAepForLocalBackupResult.RestoreReady(inputState.backupKey))
+        parentEventEmitter.navigateBack()
+      }
+      is RequestResult.NonSuccess -> {
+        when (val error = result.error) {
+          is NetworkController.RegisterAccountError.RegistrationRecoveryPasswordIncorrect -> {
+            Log.w(TAG, "[Submit] RRP incorrect, but the key decrypts the backup. The backup belongs to a different account. Message: ${error.message}")
+            stateEmitter(inputState.copy(isRegistering = false, showDifferentAccountDialog = true))
+          }
+          is NetworkController.RegisterAccountError.InvalidRequest -> {
+            Log.w(TAG, "[Submit] Invalid request. Message: ${error.message}")
+            stateEmitter(
+              inputState.copy(
+                isRegistering = false,
+                registrationError = RegistrationError.UnknownError
+              )
+            )
+          }
+          is NetworkController.RegisterAccountError.RegistrationLock -> {
+            if (provideRegistrationLock) {
+              Log.w(TAG, "[Submit] Still registration locked after providing the reglock token derived from the AEP. Falling back to PIN entry.")
+              stateEmitter(inputState.copy(isRegistering = false))
+              parentEventEmitter.navigateTo(
+                RegistrationRoute.PinEntryForRegistrationLock(
+                  timeRemaining = error.data.timeRemaining,
+                  svrCredentials = error.data.svr2Credentials
+                )
+              )
+            } else {
+              Log.w(TAG, "[Submit] Registration locked. Retrying with the reglock token derived from the AEP.")
+              attemptToRegister(inputState, aep, provideRegistrationLock = true, stateEmitter)
+            }
+          }
+          is NetworkController.RegisterAccountError.RateLimited -> {
+            Log.w(TAG, "[Submit] Rate limited (retryAfter: ${error.retryAfter}).")
+            stateEmitter(inputState.copy(isRegistering = false, registrationError = RegistrationError.RateLimited))
+          }
+          is NetworkController.RegisterAccountError.SessionNotFoundOrNotVerified -> {
+            error("[Submit] Session not found or not verified. This should not happen with RRP-based registration.")
+          }
+          is NetworkController.RegisterAccountError.DeviceTransferPossible -> {
+            error("[Submit] Device transfer possible. This should not happen with RRP-based registration.")
+          }
+        }
+      }
+      is RequestResult.RetryableNetworkError -> {
+        Log.w(TAG, "[Submit] Network error.", result.networkError)
+        stateEmitter(inputState.copy(isRegistering = false, registrationError = RegistrationError.NetworkError))
+      }
+      is RequestResult.ApplicationError -> {
+        Log.w(TAG, "[Submit] Application error.", result.cause)
+        stateEmitter(inputState.copy(isRegistering = false, registrationError = RegistrationError.UnknownError))
+      }
+    }
+  }
+
+  /**
+   * The user chose to restore the different-account backup anyway. Force the session/SMS path (the recovery password
+   * won't work), then hand control back so the flow verifies the number over SMS. The restore is resumed once
+   * registration completes, keyed off the still-set [RegistrationFlowState.pendingRestoreOption] and the entered
+   * [RegistrationFlowState.unverifiedRestoredAep].
+   */
+  private fun applyConfirmDifferentAccountRestore(inputState: EnterAepState, stateEmitter: (EnterAepState) -> Unit) {
+    Log.i(TAG, "[ConfirmDifferentAccountRestore] Deferring the restore until after SMS verification.")
+
+    stateEmitter(inputState.copy(showDifferentAccountDialog = false))
+
+    parentEventEmitter(RegistrationFlowEvent.RecoveryPasswordInvalid)
+    resultBus.sendResult(resultKey, EnterAepForLocalBackupResult.RegistrationDeferredToSms)
+    parentEventEmitter.navigateBack()
+  }
+
   class Factory(
+    private val isPreRegistration: Boolean,
+    private val backupUri: String?,
+    private val repository: RegistrationRepository,
+    private val parentState: StateFlow<RegistrationFlowState>,
     private val parentEventEmitter: (RegistrationFlowEvent) -> Unit,
     private val resultBus: ResultEventBus,
     private val resultKey: String,
     private val isPasswordManagerAvailable: Boolean = false
   ) : ViewModelProvider.Factory {
     override fun <T : ViewModel> create(modelClass: Class<T>): T {
-      return EnterAepForLocalBackupViewModel(parentEventEmitter, resultBus, resultKey, isPasswordManagerAvailable) as T
+      return EnterAepForLocalBackupViewModel(isPreRegistration, backupUri, repository, parentState, parentEventEmitter, resultBus, resultKey, isPasswordManagerAvailable) as T
     }
   }
 }
+
+/** Result sent back to the local backup restore screen from [EnterAepForLocalBackupViewModel]. */
+sealed interface EnterAepForLocalBackupResult {
+  /** The account is registered (either it already was, or RRP registration just succeeded) and the backup can be restored with [key]. */
+  data class RestoreReady(val key: String) : EnterAepForLocalBackupResult
+
+  /** The backup belongs to a different account and the user chose to restore it anyway. Registration must happen over SMS first. */
+  data object RegistrationDeferredToSms : EnterAepForLocalBackupResult
+}

feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForRemoteBackupPostRegistrationViewModel.kt

@@ -53,6 +53,10 @@ class EnterAepForRemoteBackupPostRegistrationViewModel(
       is EnterAepEvents.DismissError -> {
         stateEmitter(EnterAepScreenEventHandler.applyEvent(inputState, event))
       }
+      is EnterAepEvents.ConfirmDifferentAccountRestore,
+      is EnterAepEvents.DismissDifferentAccountDialog -> {
+        error("Different-account handling only exists for local backup restores.")
+      }
     }
   }
 

feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForRemoteBackupPreRegistrationViewModel.kt

@@ -63,6 +63,10 @@ class EnterAepForRemoteBackupPreRegistrationViewModel(
       is EnterAepEvents.DismissError -> {
         stateEmitter(EnterAepScreenEventHandler.applyEvent(inputState, event))
       }
+      is EnterAepEvents.ConfirmDifferentAccountRestore,
+      is EnterAepEvents.DismissDifferentAccountDialog -> {
+        error("Different-account handling only exists for local backup restores.")
+      }
     }
   }
 

feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepScreen.kt

@@ -74,12 +74,32 @@ fun EnterAepScreen(
 ) {
   RegistrationErrorDialog(state.registrationError, onEvent)
 
+  if (state.showDifferentAccountDialog) {
+    DifferentAccountDialog(onEvent)
+  }
+
   when (val layoutParams = RegistrationScaffold.rememberLayoutParams()) {
     is RegistrationScaffold.Params.OnePane -> OnePaneLayout(layoutParams, state, onEvent, modifier)
     is RegistrationScaffold.Params.TwoPane -> TwoPaneLayout(layoutParams, state, onEvent, modifier)
   }
 }
 
+/**
+ * Warns that the entered key decrypts the backup but the backup was created by a different account, offering to
+ * restore it anyway (which requires verifying the phone number over SMS first).
+ */
+@Composable
+private fun DifferentAccountDialog(onEvent: (EnterAepEvents) -> Unit) {
+  Dialogs.SimpleAlertDialog(
+    title = stringResource(R.string.EnterAepScreen__restore_to_new_account),
+    body = stringResource(R.string.EnterAepScreen__restore_to_new_account_body),
+    confirm = stringResource(R.string.EnterAepScreen__restore),
+    dismiss = stringResource(android.R.string.cancel),
+    onConfirm = { onEvent(EnterAepEvents.ConfirmDifferentAccountRestore) },
+    onDismiss = { onEvent(EnterAepEvents.DismissDifferentAccountDialog) }
+  )
+}
+
 /**
  * Shows a dismissable dialog for generic registration errors (network/rate-limit/unknown). Incorrect-key errors are
  * surfaced inline on the text field instead, so they are intentionally not shown here.

feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepState.kt

@@ -17,10 +17,12 @@ data class EnterAepState(
   val chunkLength: Int = 4,
   val isRegistering: Boolean = false,
   val registrationError: RegistrationError? = null,
+  /** The entered key decrypts the backup, but the backup belongs to a different account. Asks whether to restore it anyway after SMS verification. */
+  val showDifferentAccountDialog: Boolean = false,
   /** Whether a password manager / credential provider is available to fill the recovery key. */
   val isPasswordManagerAvailable: Boolean = false
 ) {
-  override fun toString(): String = "EnterAepState(enteredText=${enteredText.censor()}, backupKey=${backupKey.censor()}, isBackupKeyValid=$isBackupKeyValid, aepValidationError=$aepValidationError, chunkLength=$chunkLength, isRegistering=$isRegistering, registrationError=$registrationError, isPasswordManagerAvailable=$isPasswordManagerAvailable)"
+  override fun toString(): String = "EnterAepState(enteredText=${enteredText.censor()}, backupKey=${backupKey.censor()}, isBackupKeyValid=$isBackupKeyValid, aepValidationError=$aepValidationError, chunkLength=$chunkLength, isRegistering=$isRegistering, registrationError=$registrationError, showDifferentAccountDialog=$showDifferentAccountDialog, isPasswordManagerAvailable=$isPasswordManagerAvailable)"
 }
 
 sealed interface AepValidationError {

feature/registration/src/main/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreEvents.kt

@@ -34,6 +34,9 @@ sealed interface LocalBackupRestoreEvents {
     override fun toString(): String = "PassphraseSubmitted(credential=${credential.censor()})"
   }
 
+  /** The backup belongs to a different account and the user chose to restore it after verifying over SMS. Hands control back to the phone number screen. */
+  data object RegistrationDeferredToSms : LocalBackupRestoreEvents
+
   /** The folder picker was dismissed without selecting a folder. */
   data object FolderPickerDismissed : LocalBackupRestoreEvents
 

feature/registration/src/main/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreResult.kt

@@ -6,6 +6,7 @@
 package org.signal.registration.screens.localbackuprestore
 
 import org.signal.core.models.AccountEntropyPool
+import org.signal.registration.RegistrationFlowState
 
 /**
  * Result communicated back from the pre-registration local backup restore flow
@@ -15,6 +16,12 @@ sealed interface LocalBackupRestoreResult {
   /** The restore completed successfully. Contains the AEP if V2 backup, or the restored AEP for V1. */
   data class Success(val aep: AccountEntropyPool?) : LocalBackupRestoreResult
 
+  /**
+   * The backup belongs to a different account and the user chose to restore it anyway. The phone number must be
+   * verified over SMS, after which the restore recorded in [RegistrationFlowState.pendingLocalBackupRestore] runs.
+   */
+  data object DeferredToSms : LocalBackupRestoreResult
+
   /** The user canceled the restore flow. The pending restore option should be cleared. */
   data object Canceled : LocalBackupRestoreResult
 }

feature/registration/src/main/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreViewModel.kt

@@ -81,6 +81,10 @@ class LocalBackupRestoreViewModel(
       is LocalBackupRestoreEvents.PassphraseSubmitted -> {
         applyPassphraseSubmitted(state, event.credential, stateEmitter)
       }
+      is LocalBackupRestoreEvents.RegistrationDeferredToSms -> {
+        resultBus.sendResult(resultKey, LocalBackupRestoreResult.DeferredToSms)
+        parentEventEmitter.navigateBack()
+      }
       is LocalBackupRestoreEvents.ChooseDifferentFolder -> {
         stateEmitter(LocalBackupRestoreState(launchFolderPicker = true, storageCapable = state.storageCapable))
       }
@@ -116,7 +120,10 @@ class LocalBackupRestoreViewModel(
 
     val credentialRoute = when (backup.type) {
       LocalBackupInfo.BackupType.V1 -> RegistrationRoute.EnterLocalBackupV1Passphrase
-      LocalBackupInfo.BackupType.V2 -> RegistrationRoute.EnterAepForLocalBackup
+      LocalBackupInfo.BackupType.V2 -> RegistrationRoute.EnterAepForLocalBackup(
+        isPreRegistration = isPreRegistration,
+        backupUri = backup.uri.toString()
+      )
     }
     parentEventEmitter.navigateTo(credentialRoute)
   }
@@ -132,10 +139,16 @@ class LocalBackupRestoreViewModel(
     startRestore(backup, state.selectedFolderUri, credential, aep)
   }
 
-  private suspend fun onRestoreComplete(state: LocalBackupRestoreState, progress: LocalBackupRestoreProgress.Complete) {
+  private suspend fun onRestoreComplete(state: LocalBackupRestoreState, progress: LocalBackupRestoreProgress.Complete, backupType: LocalBackupInfo.BackupType) {
     repository.persistRestoredBackupState(progress.restoredSvrPin, progress.restoredProfileKey)
 
-    if (isPreRegistration) {
+    // The restore ran, so clear the pending-restore signal. Otherwise a post-registration screen (verification code or
+    // reglock PIN) would see it still set and route back here to restore again.
+    parentEventEmitter(RegistrationFlowEvent.PendingRestoreOptionSelected(null))
+
+    // V1 backups restore before registration, then the phone number screen registers with the restored data.
+    // V2 backups only restore against a registered account, so completion always continues the post-registration flow.
+    if (isPreRegistration && backupType == LocalBackupInfo.BackupType.V1) {
       repository.persistRestoredIdentityKeys(progress.restoredAciIdentityKey, progress.restoredPniIdentityKey)
       repository.setRestoreDecision(RestoreDecision.COMPLETED)
       resultBus.sendResult(resultKey, LocalBackupRestoreResult.Success(state.aep ?: progress.restoredAccountEntropyPool))
@@ -218,7 +231,7 @@ class LocalBackupRestoreViewModel(
             storageCapable = currentState.storageCapable
           )
           is LocalBackupRestoreProgress.Complete -> {
-            onRestoreComplete(_state.value.copy(aep = currentState.aep, v1Passphrase = currentState.v1Passphrase, storageCapable = currentState.storageCapable), progress)
+            onRestoreComplete(_state.value.copy(aep = currentState.aep, v1Passphrase = currentState.v1Passphrase, storageCapable = currentState.storageCapable), progress, backup.type)
             _state.value
           }
           is LocalBackupRestoreProgress.IncorrectCredential -> {

feature/registration/src/main/java/org/signal/registration/screens/phonenumber/PhoneNumberEntryViewModel.kt

@@ -142,6 +142,13 @@ class PhoneNumberEntryViewModel(
             localState = applyLocalBackupRestoreCompleted(localState, event.result.aep, parentEventEmitter)
             stateEmitter(localState.copy(showSpinner = false))
           }
+          is LocalBackupRestoreResult.DeferredToSms -> {
+            Log.i(TAG, "[LocalRestore] Backup belongs to a different account. Verifying the number over SMS before restoring.")
+            var localState = state.copy(showSpinner = true)
+            stateEmitter(localState)
+            localState = applySessionBasedRegistration(localState, localState.sessionE164 ?: "+${localState.countryCode}${localState.nationalNumber}", parentEventEmitter)
+            stateEmitter(localState.copy(showSpinner = false))
+          }
           is LocalBackupRestoreResult.Canceled -> {
             parentEventEmitter(RegistrationFlowEvent.PendingRestoreOptionSelected(null))
           }
@@ -344,8 +351,8 @@ class PhoneNumberEntryViewModel(
   }
 
   /**
-   * Handles the result of a pre-registration local backup restore.
-   * If an AEP was obtained (V2 backup), attempts RRP-based registration.
+   * Handles the result of a pre-registration V1 local backup restore (V2 backups register before restoring instead).
+   * If the restored database contained an AEP, attempts RRP-based registration with it.
    * Falls back to SVR check and SMS verification if RRP fails or no AEP is available.
    */
   private suspend fun applyLocalBackupRestoreCompleted(

feature/registration/src/main/java/org/signal/registration/screens/pinentry/PinEntryForRegistrationLockViewModel.kt

@@ -19,6 +19,7 @@ import org.signal.core.ui.compose.EventDrivenViewModel
 import org.signal.core.util.logging.Log
 import org.signal.libsignal.net.RequestResult
 import org.signal.registration.NetworkController
+import org.signal.registration.PendingRestoreOption
 import org.signal.registration.RegistrationFlowEvent
 import org.signal.registration.RegistrationFlowState
 import org.signal.registration.RegistrationRepository
@@ -163,7 +164,12 @@ class PinEntryForRegistrationLockViewModel(
         parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool, response.storageCapable))
         repository.enqueueSvrResetGuessCountJob()
         repository.restoreAccountRecord()
+        val pendingRestore = pendingRestoreNavigation()
         when {
+          pendingRestore != null -> {
+            Log.i(TAG, "[PinEntered] A restore was pending behind the registration lock. Resuming it now.")
+            parentEventEmitter.navigateTo(pendingRestore)
+          }
           response.reregistration && parentState.value.pendingRestoreOption == null -> parentEventEmitter.navigateTo(RegistrationRoute.ArchiveRestoreSelection.forPostRegisterWithPinKnown())
           else -> parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
         }
@@ -212,6 +218,20 @@ class PinEntryForRegistrationLockViewModel(
     }
   }
 
+  /**
+   * If the user pre-selected a restore (see [RegistrationFlowState.pendingRestoreOption]) and it hasn't run yet,
+   * returns the restore screen to resume it now that the account is registered; otherwise null. Used to pick a
+   * restore back up when it was blocked behind a registration lock that had to be cleared with the PIN first.
+   */
+  private fun pendingRestoreNavigation(): RegistrationRoute? {
+    val aep = parentState.value.unverifiedRestoredAep ?: return null
+    return when (parentState.value.pendingRestoreOption) {
+      PendingRestoreOption.LocalBackup -> RegistrationRoute.LocalBackupRestore(isPreRegistration = false, aep = aep)
+      PendingRestoreOption.RemoteBackup -> RegistrationRoute.RemoteRestore(aep)
+      null -> null
+    }
+  }
+
   class Factory(
     private val repository: RegistrationRepository,
     private val parentState: StateFlow<RegistrationFlowState>,

feature/registration/src/main/java/org/signal/registration/screens/verificationcode/VerificationCodeViewModel.kt

@@ -31,6 +31,7 @@ import org.signal.core.ui.compose.EventDrivenViewModel
 import org.signal.core.util.logging.Log
 import org.signal.libsignal.net.RequestResult
 import org.signal.registration.NetworkController
+import org.signal.registration.PendingRestoreOption
 import org.signal.registration.RegistrationFlowEvent
 import org.signal.registration.RegistrationFlowState
 import org.signal.registration.RegistrationRepository
@@ -364,7 +365,12 @@ class VerificationCodeViewModel(
 
         parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool, response.storageCapable))
 
+        val pendingRestore = pendingRestoreNavigation()
         when {
+          pendingRestore != null -> {
+            Log.i(TAG, "[Register] A restore was deferred until after SMS verification. Resuming it now.")
+            parentEventEmitter.navigateTo(pendingRestore)
+          }
           response.reregistration && parentState.value.pendingRestoreOption == null -> parentEventEmitter.navigateTo(RegistrationRoute.ArchiveRestoreSelection.forPostRegisterWithPinUnknown())
           response.storageCapable -> parentEventEmitter.navigateTo(RegistrationRoute.PinEntryForSvrRestore)
           else -> parentEventEmitter.navigateTo(RegistrationRoute.PinCreate)
@@ -415,6 +421,20 @@ class VerificationCodeViewModel(
     }
   }
 
+  /**
+   * If the user pre-selected a restore (see [RegistrationFlowState.pendingRestoreOption]) and it hasn't run yet,
+   * returns the restore screen to resume it now that the account is registered; otherwise null. Used to pick a
+   * restore back up after it was deferred to SMS verification (e.g. a local backup that belongs to a different account).
+   */
+  private fun pendingRestoreNavigation(): RegistrationRoute? {
+    val aep = parentState.value.unverifiedRestoredAep ?: return null
+    return when (parentState.value.pendingRestoreOption) {
+      PendingRestoreOption.LocalBackup -> RegistrationRoute.LocalBackupRestore(isPreRegistration = false, aep = aep)
+      PendingRestoreOption.RemoteBackup -> RegistrationRoute.RemoteRestore(aep)
+      null -> null
+    }
+  }
+
   private suspend fun applyResendCode(
     state: VerificationCodeState,
     transport: NetworkController.VerificationCodeTransport

feature/registration/src/main/res/values/strings.xml

@@ -228,6 +228,12 @@
     <string name="EnterAepScreen__incorrect_recovery_key">Incorrect recovery key</string>
     <!-- Button that prompts the device password manager to fill in the recovery key -->
     <string name="EnterAepScreen__fill_from_password_manager">Fill from password manager</string>
+    <!-- Title of the dialog warning that the backup being restored was created with a different account -->
+    <string name="EnterAepScreen__restore_to_new_account">Restore to new account?</string>
+    <!-- Body of the dialog warning the user that the backup was created with a different account and they will lose group memberships -->
+    <string name="EnterAepScreen__restore_to_new_account_body">The backup you are restoring was created using a different Signal account. You can restore this backup to a new account, but you will no longer be a member of any groups that were restored.</string>
+    <!-- Confirm button label for the restore to new account dialog -->
+    <string name="EnterAepScreen__restore">Restore</string>
     <!-- Text shown while preparing the restore -->
     <string name="LocalBackupRestoreScreen__preparing_restore">Preparing restore…</string>
     <!-- Title shown while restore is in progress -->

feature/registration/src/test/java/org/signal/registration/RegistrationEndToEndTest.kt

@@ -233,7 +233,7 @@ class RegistrationEndToEndTest {
   }
 
   @Test
-  fun `restoring a remote backup for a reglocked account whose reglock is not derived from the aep falls back to pin entry and registers without a session`() {
+  fun `restoring a remote backup for a reglocked account whose reglock is not derived from the aep falls back to pin entry then resumes the restore`() {
     val aep = AccountEntropyPool.generate()
     val svrMasterKey = MasterKey(ByteArray(32) { it.toByte() })
 
@@ -261,6 +261,11 @@ class RegistrationEndToEndTest {
       }
     }
 
+    // The backup contains the user's PIN, so no PIN screens are needed after the restore
+    storageController.onRestoreRemoteBackup = {
+      flowOf(RemoteBackupRestoreProgress.Complete(restoredSvrPin = PIN, restoredProfileKey = null))
+    }
+
     var registrationComplete = false
     launchRegistrationFlow(onRegistrationComplete = { registrationComplete = true })
 
@@ -274,6 +279,9 @@ class RegistrationEndToEndTest {
     composeTestRule.onNodeWithTag(TestTags.PIN_ENTRY_INPUT).performTextInput(PIN)
     composeTestRule.onNodeWithTag(TestTags.PIN_ENTRY_CONTINUE_BUTTON).performClick()
 
+    // With the reglock cleared, the remote restore the user chose resumes rather than being dropped
+    startRemoteRestore()
+
     waitFor("registration to complete") { registrationComplete }
 
     assert(networkController.lastRestoreMasterKeyRequest?.pin == PIN) { "Expected master key restore with pin $PIN but was ${networkController.lastRestoreMasterKeyRequest}" }
@@ -284,6 +292,7 @@ class RegistrationEndToEndTest {
     val committed = storageController.committedData
     assert(committed != null) { "Expected registration data to be committed" }
     assert(committed!!.accountData?.e164 == E164) { "Expected committed e164 $E164 but was ${committed.accountData?.e164}" }
+    assert(storageController.restoreDecision == RestoreDecision.COMPLETED) { "Expected COMPLETED restore decision (the remote backup was restored) but was ${storageController.restoreDecision}" }
   }
 
   @Test
@@ -670,11 +679,13 @@ class RegistrationEndToEndTest {
   }
 
   @Test
-  fun `restoring a local backup whose aep the server rejects falls back to sms verification`() {
+  fun `restoring a local backup for a different account warns the user, verifies over sms, then imports the backup`() {
     val aep = AccountEntropyPool.generate()
 
     // The AEP decrypts the backup fine, but it doesn't belong to the account for this phone number
+    val registerRequests = mutableListOf<FakeNetworkController.RegisterAccountRequest>()
     networkController.onRegisterAccount = { request ->
+      registerRequests += request
       if (request.recoveryPassword != null) {
         RequestResult.NonSuccess(RegisterAccountError.RegistrationRecoveryPasswordIncorrect("wrong recovery password"))
       } else {
@@ -682,6 +693,14 @@ class RegistrationEndToEndTest {
       }
     }
 
+    // The import only ever runs against a registered account, so it must come after SMS verification
+    var registerAttemptsWhenRestoreRan = -1
+    val defaultLocalRestore = storageController.onRestoreLocalBackupV2
+    storageController.onRestoreLocalBackupV2 = { backupUri, restoreAep ->
+      registerAttemptsWhenRestoreRan = registerRequests.size
+      defaultLocalRestore(backupUri, restoreAep)
+    }
+
     var registrationComplete = false
     launchRegistrationFlow(folderPickerResult = backupFolderUri, onRegistrationComplete = { registrationComplete = true })
 
@@ -690,19 +709,146 @@ class RegistrationEndToEndTest {
     enterPhoneNumber()
     restoreLocalBackup(aep)
 
-    // The backup was restored locally, but recovery-password registration was rejected, so the flow
-    // falls back to verifying the phone number over SMS
+    // The key decrypts the backup but the recovery password is rejected, meaning the backup belongs to a different
+    // account. The user is warned and confirms restoring it to this account anyway.
+    waitForTag(Dialogs.TEST_TAG_ALERT_DIALOG_CONFIRM_BUTTON)
+    composeTestRule.onNodeWithTag(Dialogs.TEST_TAG_ALERT_DIALOG_CONFIRM_BUTTON).performClick()
+
+    // Confirming requires verifying the phone number over SMS. Afterwards the restore resumes on the standard
+    // post-registration local restore screen, where the folder is re-selected (the entered key is reused).
     submitVerificationCode(VERIFICATION_CODE)
+    restoreFoundLocalBackup()
     createPin(PIN)
 
     waitFor("registration to complete") { registrationComplete }
 
+    assert(registerRequests.first().recoveryPassword == aep.deriveMasterKey().deriveRegistrationRecoveryPassword()) {
+      "Expected the first registration attempt to use the recovery password derived from the backup's AEP"
+    }
     assert(networkController.lastRegisterAccountRequest?.sessionId != null) { "Expected the final registration to use a verified session" }
+    assert(registerAttemptsWhenRestoreRan == registerRequests.size) {
+      "Expected the backup to be imported only after registration completed, but the import ran after $registerAttemptsWhenRestoreRan of ${registerRequests.size} registration attempts"
+    }
 
     val committed = storageController.committedData
     assert(committed != null) { "Expected registration data to be committed" }
     assert(committed!!.accountData?.e164 == E164) { "Expected committed e164 $E164 but was ${committed.accountData?.e164}" }
     assert(committed.pin == PIN) { "Expected committed pin $PIN but was ${committed.pin}" }
+    assert(committed.accountEntropyPool.isNotEmpty() && committed.accountEntropyPool != aep.value) {
+      "Expected a fresh AEP to be committed rather than the foreign backup's AEP"
+    }
+    assert(storageController.restoreDecision == RestoreDecision.COMPLETED) { "Expected COMPLETED restore decision but was ${storageController.restoreDecision}" }
+  }
+
+  @Test
+  fun `declining to restore a local backup for a different account returns to the recovery key entry screen`() {
+    val aep = AccountEntropyPool.generate()
+
+    networkController.onRegisterAccount = { request ->
+      if (request.recoveryPassword != null) {
+        RequestResult.NonSuccess(RegisterAccountError.RegistrationRecoveryPasswordIncorrect("wrong recovery password"))
+      } else {
+        RequestResult.Success(networkController.registerAccountResponse(request.e164))
+      }
+    }
+
+    launchRegistrationFlow(folderPickerResult = backupFolderUri)
+
+    startManualRestore()
+    chooseRestoreOption(TestTags.ARCHIVE_RESTORE_SELECTION_FROM_BACKUP_FOLDER)
+    enterPhoneNumber()
+    restoreLocalBackup(aep)
+
+    // The user is warned that the backup belongs to a different account and declines
+    waitForTag(Dialogs.TEST_TAG_ALERT_DIALOG_DISMISS_BUTTON)
+    composeTestRule.onNodeWithTag(Dialogs.TEST_TAG_ALERT_DIALOG_DISMISS_BUTTON).performClick()
+
+    // The user stays on the recovery key entry screen and nothing was restored or registered
+    waitForTag(TestTags.ENTER_AEP_SCREEN)
+    assert(storageController.committedData == null) { "Expected no registration data to be committed" }
+    assert(storageController.restoreDecision == null) { "Expected no restore decision to have been made" }
+  }
+
+  @Test
+  fun `restoring a local backup for a reglocked account whose reglock is not derived from the aep verifies the pin then restores`() {
+    val aep = AccountEntropyPool.generate()
+    val svrMasterKey = MasterKey(ByteArray(32) { it.toByte() })
+
+    // The AEP is valid for the account (its recovery password is accepted), but the account's reglock is governed by a
+    // separate master key held in SVR, so the reglock proof derived from the AEP is rejected.
+    networkController.onRegisterAccount = { request ->
+      when {
+        request.registrationLock == svrMasterKey.deriveRegistrationLock() -> RequestResult.Success(networkController.registerAccountResponse(request.e164))
+        else -> RequestResult.NonSuccess(
+          RegisterAccountError.RegistrationLock(
+            RegistrationLockResponse(
+              timeRemaining = 14.days.inWholeMilliseconds,
+              svr2Credentials = SvrCredentials(username = "svr-user", password = "svr-pass")
+            )
+          )
+        )
+      }
+    }
+
+    networkController.onRestoreMasterKeyFromSvr = { request ->
+      if (request.pin == PIN) {
+        RequestResult.Success(MasterKeyResponse(svrMasterKey))
+      } else {
+        RequestResult.NonSuccess(RestoreMasterKeyError.WrongPin(triesRemaining = 3))
+      }
+    }
+
+    // The backup contains the user's PIN, so the flow finishes without any PIN screens after the restore
+    storageController.onRestoreLocalBackupV2 = { _, _ ->
+      flowOf(LocalBackupRestoreProgress.Complete(restoredSvrPin = PIN, restoredProfileKey = null))
+    }
+
+    var registrationComplete = false
+    launchRegistrationFlow(folderPickerResult = backupFolderUri, onRegistrationComplete = { registrationComplete = true })
+
+    startManualRestore()
+    chooseRestoreOption(TestTags.ARCHIVE_RESTORE_SELECTION_FROM_BACKUP_FOLDER)
+    enterPhoneNumber()
+    restoreLocalBackup(aep)
+
+    // The AEP-derived reglock proof was rejected, so the user must prove their PIN to clear the registration lock
+    waitForTag(TestTags.PIN_ENTRY_SCREEN)
+    composeTestRule.onNodeWithTag(TestTags.PIN_ENTRY_INPUT).performTextInput(PIN)
+    composeTestRule.onNodeWithTag(TestTags.PIN_ENTRY_CONTINUE_BUTTON).performClick()
+
+    // With the reglock cleared, the pending backup restore resumes on the post-registration local restore screen
+    restoreFoundLocalBackup()
+
+    waitFor("registration to complete") { registrationComplete }
+
+    assert(networkController.lastRestoreMasterKeyRequest?.pin == PIN) { "Expected the PIN to be used to restore the master key from SVR" }
+    assert(networkController.lastRegisterAccountRequest?.registrationLock == svrMasterKey.deriveRegistrationLock()) { "Expected registration with the reglock proof derived from the SVR master key" }
+
+    val committed = storageController.committedData
+    assert(committed != null) { "Expected registration data to be committed" }
+    assert(committed!!.accountData?.e164 == E164) { "Expected committed e164 $E164 but was ${committed.accountData?.e164}" }
+    assert(storageController.restoreDecision == RestoreDecision.COMPLETED) { "Expected COMPLETED restore decision (the backup was restored) but was ${storageController.restoreDecision}" }
+  }
+
+  @Test
+  fun `entering a recovery key that cannot decrypt the local backup shows an inline error without registering`() {
+    val aep = AccountEntropyPool.generate()
+
+    storageController.onVerifyLocalBackupKey = { _, _ -> false }
+
+    launchRegistrationFlow(folderPickerResult = backupFolderUri)
+
+    startManualRestore()
+    chooseRestoreOption(TestTags.ARCHIVE_RESTORE_SELECTION_FROM_BACKUP_FOLDER)
+    enterPhoneNumber()
+    restoreLocalBackup(aep)
+
+    // The key can't decrypt the backup, so submission is rejected inline before any registration attempt
+    waitFor("the incorrect key to be rejected") {
+      composeTestRule.onAllNodesWithTag(TestTags.ENTER_AEP_NEXT_BUTTON).fetchSemanticsNodes().firstOrNull()
+        ?.config?.getOrNull(SemanticsProperties.Disabled) != null
+    }
+    assert(networkController.lastRegisterAccountRequest == null) { "Expected no registration attempt for a key that cannot decrypt the backup" }
   }
 
   @Test

feature/registration/src/test/java/org/signal/registration/fakes/FakeStorageController.kt

@@ -64,6 +64,8 @@ class FakeStorageController : StorageController {
     flowOf(LocalBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null))
   }
 
+  var onVerifyLocalBackupKey: suspend (backupUri: Uri, aep: AccountEntropyPool) -> Boolean = { _, _ -> true }
+
   var onRestoreRemoteBackup: (aep: AccountEntropyPool) -> Flow<RemoteBackupRestoreProgress> = {
     flowOf(RemoteBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null))
   }
@@ -92,7 +94,10 @@ class FakeStorageController : StorageController {
   }
 
   override suspend fun setRestoreDecision(decision: RestoreDecision) {
-    restoreDecision = decision
+    // Mirrors the real controller: only the first decision sticks, later ones are ignored
+    if (restoreDecision == null) {
+      restoreDecision = decision
+    }
   }
 
   override fun restoreLocalBackupV1(rootUri: Uri, backupUri: Uri, passphrase: String): Flow<LocalBackupRestoreProgress> {
@@ -103,6 +108,10 @@ class FakeStorageController : StorageController {
     return onRestoreLocalBackupV2(backupUri, aep)
   }
 
+  override suspend fun verifyLocalBackupKey(backupUri: Uri, aep: AccountEntropyPool): Boolean {
+    return onVerifyLocalBackupKey(backupUri, aep)
+  }
+
   override fun restoreRemoteBackup(aep: AccountEntropyPool): Flow<RemoteBackupRestoreProgress> {
     return onRestoreRemoteBackup(aep)
   }

feature/registration/src/test/java/org/signal/registration/screens/aepentry/EnterAepForLocalBackupViewModelTest.kt

@@ -5,114 +5,280 @@
 
 package org.signal.registration.screens.aepentry
 
+import android.net.Uri
 import assertk.assertThat
 import assertk.assertions.hasSize
-import assertk.assertions.isEmpty
 import assertk.assertions.isEqualTo
+import assertk.assertions.isInstanceOf
 import assertk.assertions.isNull
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.advanceUntilIdle
-import kotlinx.coroutines.test.resetMain
+import assertk.assertions.prop
+import io.mockk.coEvery
+import io.mockk.coVerify
+import io.mockk.every
+import io.mockk.mockk
+import io.mockk.mockkStatic
+import io.mockk.unmockkStatic
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.runTest
-import kotlinx.coroutines.test.setMain
 import org.junit.After
 import org.junit.Before
 import org.junit.Test
+import org.signal.core.models.AccountEntropyPool
 import org.signal.core.ui.navigation.ResultEventBus
+import org.signal.libsignal.net.RequestResult
+import org.signal.registration.KeyMaterial
+import org.signal.registration.NetworkController
 import org.signal.registration.RegistrationFlowEvent
+import org.signal.registration.RegistrationFlowState
+import org.signal.registration.RegistrationRepository
+import org.signal.registration.RegistrationRoute
 
-@OptIn(ExperimentalCoroutinesApi::class)
 class EnterAepForLocalBackupViewModelTest {
 
-  private lateinit var viewModel: EnterAepForLocalBackupViewModel
+  private lateinit var mockRepository: RegistrationRepository
   private lateinit var resultBus: ResultEventBus
+  private lateinit var parentState: MutableStateFlow<RegistrationFlowState>
   private lateinit var emittedParentEvents: MutableList<RegistrationFlowEvent>
   private lateinit var parentEventEmitter: (RegistrationFlowEvent) -> Unit
-
-  private val resultKey = "test-result-key"
-
-  private val testDispatcher = StandardTestDispatcher()
+  private lateinit var emittedStates: MutableList<EnterAepState>
+  private lateinit var stateEmitter: (EnterAepState) -> Unit
 
   @Before
   fun setup() {
-    Dispatchers.setMain(testDispatcher)
+    mockkStatic(Uri::class)
+    every { Uri.parse(any()) } answers { mockk(relaxed = true) }
+
+    mockRepository = mockk(relaxed = true)
     resultBus = ResultEventBus()
+    parentState = MutableStateFlow(RegistrationFlowState(sessionE164 = E164))
     emittedParentEvents = mutableListOf()
     parentEventEmitter = { event -> emittedParentEvents.add(event) }
-    viewModel = EnterAepForLocalBackupViewModel(
-      parentEventEmitter = parentEventEmitter,
-      resultBus = resultBus,
-      resultKey = resultKey
-    )
+    emittedStates = mutableListOf()
+    stateEmitter = { state -> emittedStates.add(state) }
   }
 
   @After
   fun tearDown() {
-    Dispatchers.resetMain()
+    unmockkStatic(Uri::class)
   }
 
-  // ==================== BackupKeyChanged Tests ====================
-
-  @Test
-  fun `BackupKeyChanged updates backup key in state`() = runTest {
-    val testKey = VALID_AEP
-
-    viewModel.onEvent(EnterAepEvents.BackupKeyChanged(testKey))
-    advanceUntilIdle()
-
-    assertThat(viewModel.state.value.backupKey).isEqualTo(testKey)
+  private fun createViewModel(isPreRegistration: Boolean = true): EnterAepForLocalBackupViewModel {
+    return EnterAepForLocalBackupViewModel(
+      isPreRegistration = isPreRegistration,
+      backupUri = BACKUP_URI,
+      repository = mockRepository,
+      parentState = parentState,
+      parentEventEmitter = parentEventEmitter,
+      resultBus = resultBus,
+      resultKey = RESULT_KEY
+    )
   }
 
-  // ==================== Submit Tests ====================
+  private fun latestResult(): EnterAepForLocalBackupResult? {
+    return resultBus.channelMap[RESULT_KEY]?.tryReceive()?.getOrNull() as EnterAepForLocalBackupResult?
+  }
+
+  // ==================== Already-registered mode ====================
 
   @Test
-  fun `Submit with valid key sends result via resultBus and emits NavigateBack`() = runTest {
-    viewModel.onEvent(EnterAepEvents.BackupKeyChanged(VALID_AEP))
-    viewModel.onEvent(EnterAepEvents.Submit)
-    advanceUntilIdle()
+  fun `Submit without requiring registration hands the key back and navigates back`() = runTest {
+    val viewModel = createViewModel(isPreRegistration = false)
+    val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true)
 
-    val result = resultBus.channelMap[resultKey]?.tryReceive()?.getOrNull()
-    assertThat(result).isEqualTo(VALID_AEP)
+    viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter)
+
+    assertThat(latestResult()).isEqualTo(EnterAepForLocalBackupResult.RestoreReady(VALID_AEP))
     assertThat(emittedParentEvents).hasSize(1)
     assertThat(emittedParentEvents.first()).isEqualTo(RegistrationFlowEvent.NavigateBack)
+    coVerify(exactly = 0) { mockRepository.registerAccountWithRecoveryPassword(any(), any(), any(), any(), any(), any()) }
+  }
+
+  // ==================== Register-first mode ====================
+
+  @Test
+  fun `Submit with a key that cannot decrypt the backup shows an inline error without registering`() = runTest {
+    val viewModel = createViewModel()
+    val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true)
+
+    coEvery { mockRepository.verifyLocalBackupKey(any(), any()) } returns false
+
+    viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter)
+
+    assertThat(emittedStates.last().aepValidationError).isEqualTo(AepValidationError.Incorrect)
+    assertThat(emittedStates.last().isRegistering).isEqualTo(false)
+    coVerify(exactly = 0) { mockRepository.registerAccountWithRecoveryPassword(any(), any(), any(), any(), any(), any()) }
   }
 
   @Test
-  fun `Submit with invalid key does not send result or navigate`() = runTest {
-    viewModel.onEvent(EnterAepEvents.BackupKeyChanged("too-short"))
-    viewModel.onEvent(EnterAepEvents.Submit)
-    advanceUntilIdle()
+  fun `Submit with successful registration emits UserSuppliedAepSubmitted and Registered, then hands the key back`() = runTest {
+    val viewModel = createViewModel()
+    val aep = AccountEntropyPool(VALID_AEP)
+    val mockKeyMaterial = mockk<KeyMaterial>(relaxed = true) {
+      io.mockk.every { accountEntropyPool } returns aep
+    }
+    val mockResponse = mockk<NetworkController.RegisterAccountResponse>(relaxed = true)
+    val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true)
 
-    assertThat(resultBus.channelMap[resultKey]).isNull()
-    assertThat(emittedParentEvents).isEmpty()
+    coEvery { mockRepository.verifyLocalBackupKey(any(), any()) } returns true
+    coEvery { mockRepository.registerAccountWithRecoveryPassword(any(), any(), any(), any(), any(), any()) } returns
+      RequestResult.Success(mockResponse to mockKeyMaterial)
+
+    viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter)
+
+    assertThat(emittedParentEvents).hasSize(3)
+    assertThat(emittedParentEvents[0]).isInstanceOf<RegistrationFlowEvent.UserSuppliedAepSubmitted>()
+    assertThat(emittedParentEvents[1]).isInstanceOf<RegistrationFlowEvent.Registered>()
+    assertThat(emittedParentEvents[2]).isEqualTo(RegistrationFlowEvent.NavigateBack)
+    assertThat(latestResult()).isEqualTo(EnterAepForLocalBackupResult.RestoreReady(VALID_AEP))
+  }
+
+  @Test
+  fun `Submit with RegistrationRecoveryPasswordIncorrect shows the different account dialog`() = runTest {
+    val viewModel = createViewModel()
+    val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true)
+
+    coEvery { mockRepository.verifyLocalBackupKey(any(), any()) } returns true
+    coEvery { mockRepository.registerAccountWithRecoveryPassword(any(), any(), any(), any(), any(), any()) } returns
+      RequestResult.NonSuccess(
+        NetworkController.RegisterAccountError.RegistrationRecoveryPasswordIncorrect("Incorrect")
+      )
+
+    viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter)
+
+    assertThat(emittedStates.last().showDifferentAccountDialog).isEqualTo(true)
+    assertThat(emittedStates.last().isRegistering).isEqualTo(false)
+    assertThat(latestResult()).isNull()
+  }
+
+  @Test
+  fun `ConfirmDifferentAccountRestore forces the session path and hands control back for SMS verification`() = runTest {
+    val viewModel = createViewModel()
+    val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true, showDifferentAccountDialog = true)
+
+    viewModel.applyEvent(initialState, EnterAepEvents.ConfirmDifferentAccountRestore, stateEmitter)
+
+    assertThat(emittedStates.last().showDifferentAccountDialog).isEqualTo(false)
+    assertThat(emittedParentEvents).hasSize(2)
+    assertThat(emittedParentEvents[0]).isEqualTo(RegistrationFlowEvent.RecoveryPasswordInvalid)
+    assertThat(emittedParentEvents[1]).isEqualTo(RegistrationFlowEvent.NavigateBack)
+    assertThat(latestResult()).isEqualTo(EnterAepForLocalBackupResult.RegistrationDeferredToSms)
+  }
+
+  @Test
+  fun `DismissDifferentAccountDialog clears the dialog`() = runTest {
+    val viewModel = createViewModel()
+    val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true, showDifferentAccountDialog = true)
+
+    viewModel.applyEvent(initialState, EnterAepEvents.DismissDifferentAccountDialog, stateEmitter)
+
+    assertThat(emittedStates.last().showDifferentAccountDialog).isEqualTo(false)
+    assertThat(emittedParentEvents).hasSize(0)
+  }
+
+  @Test
+  fun `Submit with RegistrationLock retries with the reglock token derived from the AEP`() = runTest {
+    val viewModel = createViewModel()
+    val aep = AccountEntropyPool(VALID_AEP)
+    val mockKeyMaterial = mockk<KeyMaterial>(relaxed = true) {
+      io.mockk.every { accountEntropyPool } returns aep
+    }
+    val mockResponse = mockk<NetworkController.RegisterAccountResponse>(relaxed = true)
+    val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true)
+    val registrationLockData = NetworkController.RegistrationLockResponse(
+      timeRemaining = 86400000L,
+      svr2Credentials = NetworkController.SvrCredentials(username = "test-username", password = "test-password")
+    )
+
+    coEvery { mockRepository.verifyLocalBackupKey(any(), any()) } returns true
+    coEvery { mockRepository.registerAccountWithRecoveryPassword(any(), any(), registrationLock = any<String>(), any(), any(), any()) } returns
+      RequestResult.Success(mockResponse to mockKeyMaterial)
+    coEvery { mockRepository.registerAccountWithRecoveryPassword(any(), any(), registrationLock = null, any(), any(), any()) } returns
+      RequestResult.NonSuccess(
+        NetworkController.RegisterAccountError.RegistrationLock(registrationLockData)
+      )
+
+    viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter)
+
+    coVerify {
+      mockRepository.registerAccountWithRecoveryPassword(any(), any(), registrationLock = aep.deriveMasterKey().deriveRegistrationLock(), any(), any(), any())
+    }
+    assertThat(latestResult()).isEqualTo(EnterAepForLocalBackupResult.RestoreReady(VALID_AEP))
+  }
+
+  @Test
+  fun `Submit with RegistrationLock when already providing the reglock token navigates to PinEntryForRegistrationLock`() = runTest {
+    val viewModel = createViewModel()
+    val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true)
+    val registrationLockData = NetworkController.RegistrationLockResponse(
+      timeRemaining = 86400000L,
+      svr2Credentials = NetworkController.SvrCredentials(username = "test-username", password = "test-password")
+    )
+
+    coEvery { mockRepository.verifyLocalBackupKey(any(), any()) } returns true
+    coEvery { mockRepository.registerAccountWithRecoveryPassword(any(), any(), any(), any(), any(), any()) } returns
+      RequestResult.NonSuccess(
+        NetworkController.RegisterAccountError.RegistrationLock(registrationLockData)
+      )
+
+    viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter)
+
+    assertThat(emittedStates.last().isRegistering).isEqualTo(false)
+    assertThat(emittedParentEvents.last())
+      .isInstanceOf<RegistrationFlowEvent.NavigateToScreen>()
+      .prop(RegistrationFlowEvent.NavigateToScreen::route)
+      .isInstanceOf<RegistrationRoute.PinEntryForRegistrationLock>()
+  }
+
+  @Test
+  fun `Submit with RateLimited sets registrationError to RateLimited`() = runTest {
+    val viewModel = createViewModel()
+    val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true)
+
+    coEvery { mockRepository.verifyLocalBackupKey(any(), any()) } returns true
+    coEvery { mockRepository.registerAccountWithRecoveryPassword(any(), any(), any(), any(), any(), any()) } returns
+      RequestResult.NonSuccess(
+        NetworkController.RegisterAccountError.RateLimited(kotlin.time.Duration.parse("1m"))
+      )
+
+    viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter)
+
+    assertThat(emittedStates.last().registrationError).isEqualTo(RegistrationError.RateLimited)
+    assertThat(emittedStates.last().isRegistering).isEqualTo(false)
+  }
+
+  @Test
+  fun `Submit with RetryableNetworkError sets registrationError to NetworkError`() = runTest {
+    val viewModel = createViewModel()
+    val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true)
+
+    coEvery { mockRepository.verifyLocalBackupKey(any(), any()) } returns true
+    coEvery { mockRepository.registerAccountWithRecoveryPassword(any(), any(), any(), any(), any(), any()) } returns
+      RequestResult.RetryableNetworkError(java.io.IOException("Network error"))
+
+    viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter)
+
+    assertThat(emittedStates.last().registrationError).isEqualTo(RegistrationError.NetworkError)
+    assertThat(emittedStates.last().isRegistering).isEqualTo(false)
   }
 
   // ==================== Cancel Tests ====================
 
   @Test
   fun `Cancel emits NavigateBack`() = runTest {
-    viewModel.onEvent(EnterAepEvents.Cancel)
-    advanceUntilIdle()
+    val viewModel = createViewModel()
+
+    viewModel.applyEvent(EnterAepState(), EnterAepEvents.Cancel, stateEmitter)
 
     assertThat(emittedParentEvents).hasSize(1)
     assertThat(emittedParentEvents.first()).isEqualTo(RegistrationFlowEvent.NavigateBack)
   }
 
-  // ==================== DismissError Tests ====================
-
-  @Test
-  fun `DismissError clears registrationError from state`() = runTest {
-    viewModel.onEvent(EnterAepEvents.DismissError)
-    advanceUntilIdle()
-
-    assertThat(viewModel.state.value.registrationError).isNull()
-  }
-
   // ==================== Constants ====================
 
   companion object {
     private const val VALID_AEP = "uy38jh2778hjjhj8lk19ga61s672jsj089r023s6a57809bap92j2yh5t326vv7t"
+    private const val E164 = "+15551234567"
+    private const val BACKUP_URI = "content://test/backups/signal-backup-2026-01-01-12-00-00"
+    private const val RESULT_KEY = "test-result-key"
   }
 }

feature/registration/src/test/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreViewModelTest.kt

@@ -72,14 +72,20 @@ class LocalBackupRestoreViewModelTest {
     Dispatchers.resetMain()
   }
 
-  private fun createViewModel(isPreRegistration: Boolean, storageCapable: Boolean = true): LocalBackupRestoreViewModel {
+  private fun createViewModel(
+    isPreRegistration: Boolean,
+    storageCapable: Boolean = true,
+    knownAep: AccountEntropyPool? = null,
+    parentState: RegistrationFlowState = RegistrationFlowState(storageCapable = storageCapable)
+  ): LocalBackupRestoreViewModel {
     return LocalBackupRestoreViewModel(
       repository = mockRepository,
-      parentState = flowOf(RegistrationFlowState(storageCapable = storageCapable)),
+      parentState = flowOf(parentState),
       parentEventEmitter = parentEventEmitter,
       isPreRegistration = isPreRegistration,
       resultBus = resultBus,
-      resultKey = resultKey
+      resultKey = resultKey,
+      knownAep = knownAep
     )
   }
 
@@ -149,7 +155,7 @@ class LocalBackupRestoreViewModelTest {
   // ==================== RestoreBackup with V2 Tests ====================
 
   @Test
-  fun `RestoreBackup with V2 backup navigates to EnterAepForLocalBackup`() = runTest {
+  fun `RestoreBackup with V2 backup post-registration navigates to EnterAepForLocalBackup without requiring registration`() = runTest {
     val viewModel = createViewModel(isPreRegistration = false)
     val backupInfo = LocalBackupInfo(
       type = LocalBackupInfo.BackupType.V2,
@@ -162,10 +168,70 @@ class LocalBackupRestoreViewModelTest {
     viewModel.applyEvent(initialState, LocalBackupRestoreEvents.RestoreBackup, stateEmitter)
 
     assertThat(emittedParentEvents).hasSize(1)
-    assertThat(emittedParentEvents.first())
+    val route = assertThat(emittedParentEvents.first())
       .isInstanceOf<RegistrationFlowEvent.NavigateToScreen>()
       .prop(RegistrationFlowEvent.NavigateToScreen::route)
-      .isEqualTo(RegistrationRoute.EnterAepForLocalBackup)
+    route.isInstanceOf<RegistrationRoute.EnterAepForLocalBackup>().prop(RegistrationRoute.EnterAepForLocalBackup::isPreRegistration).isEqualTo(false)
+  }
+
+  @Test
+  fun `RestoreBackup with V2 backup pre-registration navigates to EnterAepForLocalBackup requiring registration`() = runTest {
+    val viewModel = createViewModel(isPreRegistration = true)
+    val backupInfo = LocalBackupInfo(
+      type = LocalBackupInfo.BackupType.V2,
+      date = LocalDateTime.now(),
+      name = "backup.bin",
+      uri = mockk(relaxed = true)
+    )
+    val initialState = LocalBackupRestoreState(backupInfo = backupInfo, selectedFolderUri = mockk(relaxed = true))
+
+    viewModel.applyEvent(initialState, LocalBackupRestoreEvents.RestoreBackup, stateEmitter)
+
+    assertThat(emittedParentEvents).hasSize(1)
+    val route = assertThat(emittedParentEvents.first())
+      .isInstanceOf<RegistrationFlowEvent.NavigateToScreen>()
+      .prop(RegistrationFlowEvent.NavigateToScreen::route)
+    route.isInstanceOf<RegistrationRoute.EnterAepForLocalBackup>().prop(RegistrationRoute.EnterAepForLocalBackup::isPreRegistration).isEqualTo(true)
+  }
+
+  // ==================== Deferred restore Tests ====================
+
+  @Test
+  fun `RegistrationDeferredToSms forwards the deferral to the phone number screen and navigates back`() = runTest {
+    val viewModel = createViewModel(isPreRegistration = true)
+    val initialState = LocalBackupRestoreState()
+
+    viewModel.applyEvent(initialState, LocalBackupRestoreEvents.RegistrationDeferredToSms, stateEmitter)
+
+    val result = resultBus.channelMap[resultKey]?.tryReceive()?.getOrNull()
+    assertThat(result).isNotNull().isEqualTo(LocalBackupRestoreResult.DeferredToSms)
+    assertThat(emittedParentEvents).hasSize(1)
+    assertThat(emittedParentEvents.first()).isEqualTo(RegistrationFlowEvent.NavigateBack)
+  }
+
+  @Test
+  fun `a completed restore clears the pending restore option so it is not resumed again post-registration`() = runTest(testDispatcher) {
+    val viewModel = createViewModel(isPreRegistration = false, storageCapable = false, knownAep = AccountEntropyPool(VALID_AEP))
+    val backupInfo = LocalBackupInfo(
+      type = LocalBackupInfo.BackupType.V2,
+      date = LocalDateTime.now(),
+      name = "signal-backup",
+      uri = mockk()
+    )
+    val initialState = LocalBackupRestoreState(backupInfo = backupInfo, selectedFolderUri = mockk())
+
+    every { mockRepository.restoreV2Backup(any(), any(), any()) } returns flowOf(
+      LocalBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null)
+    )
+
+    viewModel.applyEvent(initialState, LocalBackupRestoreEvents.PassphraseSubmitted(VALID_AEP), stateEmitter)
+
+    coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) }
+    assertThat(emittedParentEvents).contains(RegistrationFlowEvent.PendingRestoreOptionSelected(null))
+    assertThat(emittedParentEvents.last())
+      .isInstanceOf<RegistrationFlowEvent.NavigateToScreen>()
+      .prop(RegistrationFlowEvent.NavigateToScreen::route)
+      .isEqualTo(RegistrationRoute.PinCreate)
   }
 
   // ==================== RestoreBackup with no backup Tests ====================
@@ -312,8 +378,7 @@ class LocalBackupRestoreViewModelTest {
 
     coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) }
     coVerify(exactly = 0) { mockRepository.restoreAccountRecord(any()) }
-    assertThat(emittedParentEvents).hasSize(1)
-    assertThat(emittedParentEvents.first())
+    assertThat(emittedParentEvents.last())
       .isInstanceOf<RegistrationFlowEvent.NavigateToScreen>()
       .prop(RegistrationFlowEvent.NavigateToScreen::route)
       .isEqualTo(RegistrationRoute.PinEntryForSvrRestore)
@@ -340,8 +405,7 @@ class LocalBackupRestoreViewModelTest {
 
     coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) }
     coVerify(exactly = 0) { mockRepository.restoreAccountRecord(any()) }
-    assertThat(emittedParentEvents).hasSize(1)
-    assertThat(emittedParentEvents.first())
+    assertThat(emittedParentEvents.last())
       .isInstanceOf<RegistrationFlowEvent.NavigateToScreen>()
       .prop(RegistrationFlowEvent.NavigateToScreen::route)
       .isEqualTo(RegistrationRoute.PinCreate)

feature/registration/src/test/java/org/signal/registration/screens/pinentry/PinEntryForRegistrationLockViewModelTest.kt

@@ -19,10 +19,12 @@ import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
+import org.signal.core.models.AccountEntropyPool
 import org.signal.core.models.MasterKey
 import org.signal.libsignal.net.RequestResult
 import org.signal.registration.KeyMaterial
 import org.signal.registration.NetworkController
+import org.signal.registration.PendingRestoreOption
 import org.signal.registration.RegistrationFlowEvent
 import org.signal.registration.RegistrationFlowState
 import org.signal.registration.RegistrationRepository
@@ -116,6 +118,60 @@ class PinEntryForRegistrationLockViewModelTest {
     assertThat(emittedStates.last().loading).isEqualTo(true)
   }
 
+  @Test
+  fun `PinEntered resumes a pending local backup restore after clearing the registration lock`() = runTest {
+    val masterKey = mockk<MasterKey>(relaxed = true)
+    val keyMaterial = mockk<KeyMaterial>(relaxed = true)
+    val restoreAep = AccountEntropyPool.generate()
+    val initialState = PinEntryState(mode = PinEntryState.Mode.RegistrationLock)
+
+    parentState.value = parentState.value.copy(
+      pendingRestoreOption = PendingRestoreOption.LocalBackup,
+      unverifiedRestoredAep = restoreAep
+    )
+
+    coEvery { mockRepository.restoreMasterKeyFromSvr(any(), any(), forRegistrationLock = true) } returns
+      RequestResult.Success(NetworkController.MasterKeyResponse(masterKey))
+    coEvery { mockRepository.registerAccountWithSession(any(), any(), any(), any()) } returns
+      RequestResult.Success(createRegisterAccountResponse() to keyMaterial)
+
+    viewModel.applyEvent(initialState, PinEntryScreenEvents.PinEntered("123456"), parentEventEmitter, stateEmitter)
+
+    assertThat(emittedParentEvents.last())
+      .isInstanceOf<RegistrationFlowEvent.NavigateToScreen>()
+      .prop(RegistrationFlowEvent.NavigateToScreen::route)
+      .isInstanceOf<RegistrationRoute.LocalBackupRestore>()
+      .prop(RegistrationRoute.LocalBackupRestore::aep)
+      .isEqualTo(restoreAep)
+  }
+
+  @Test
+  fun `PinEntered resumes a pending remote backup restore after clearing the registration lock`() = runTest {
+    val masterKey = mockk<MasterKey>(relaxed = true)
+    val keyMaterial = mockk<KeyMaterial>(relaxed = true)
+    val restoreAep = AccountEntropyPool.generate()
+    val initialState = PinEntryState(mode = PinEntryState.Mode.RegistrationLock)
+
+    parentState.value = parentState.value.copy(
+      pendingRestoreOption = PendingRestoreOption.RemoteBackup,
+      unverifiedRestoredAep = restoreAep
+    )
+
+    coEvery { mockRepository.restoreMasterKeyFromSvr(any(), any(), forRegistrationLock = true) } returns
+      RequestResult.Success(NetworkController.MasterKeyResponse(masterKey))
+    coEvery { mockRepository.registerAccountWithSession(any(), any(), any(), any()) } returns
+      RequestResult.Success(createRegisterAccountResponse() to keyMaterial)
+
+    viewModel.applyEvent(initialState, PinEntryScreenEvents.PinEntered("123456"), parentEventEmitter, stateEmitter)
+
+    assertThat(emittedParentEvents.last())
+      .isInstanceOf<RegistrationFlowEvent.NavigateToScreen>()
+      .prop(RegistrationFlowEvent.NavigateToScreen::route)
+      .isInstanceOf<RegistrationRoute.RemoteRestore>()
+      .prop(RegistrationRoute.RemoteRestore::aep)
+      .isEqualTo(restoreAep)
+  }
+
   @Test
   fun `PinEntered with wrong PIN returns state with tries remaining`() = runTest {
     val triesRemaining = 3