Fix some reglock bugs in regV5.

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

@@ -100,6 +100,7 @@ import org.signal.registration.screens.remotebackuprestore.RemoteRestoreScreen
 import org.signal.registration.screens.restoreselection.ArchiveRestoreOption
 import org.signal.registration.screens.restoreselection.ArchiveRestoreSelectionScreen
 import org.signal.registration.screens.restoreselection.ArchiveRestoreSelectionViewModel
+import org.signal.registration.screens.restoreselection.RegisteredState
 import org.signal.registration.screens.util.navigateBack
 import org.signal.registration.screens.util.navigateTo
 import org.signal.registration.screens.verificationcode.VerificationCodeScreen
@@ -162,7 +163,7 @@ sealed interface RegistrationRoute : NavKey, Parcelable {
   data object PinCreate : RegistrationRoute
 
   @Serializable
-  data class ArchiveRestoreSelection(val restoreOptions: List<ArchiveRestoreOption>, val isPreRegistration: Boolean) : RegistrationRoute {
+  data class ArchiveRestoreSelection(val restoreOptions: List<ArchiveRestoreOption>, val registeredState: RegisteredState) : RegistrationRoute {
     companion object {
       fun forQuickRestore(hasRemoteBackup: Boolean): ArchiveRestoreSelection {
         return ArchiveRestoreSelection(
@@ -174,7 +175,7 @@ sealed interface RegistrationRoute : NavKey, Parcelable {
             add(ArchiveRestoreOption.DeviceTransfer)
             add(ArchiveRestoreOption.None)
           },
-          isPreRegistration = true
+          registeredState = RegisteredState.NotRegistered
         )
       }
 
@@ -185,18 +186,29 @@ sealed interface RegistrationRoute : NavKey, Parcelable {
             add(ArchiveRestoreOption.LocalBackup)
             add(ArchiveRestoreOption.None)
           },
-          isPreRegistration = true
+          registeredState = RegisteredState.NotRegistered
         )
       }
 
-      fun forPostRegister(): ArchiveRestoreSelection {
+      fun forPostRegisterWithPinUnknown(): ArchiveRestoreSelection {
         return ArchiveRestoreSelection(
           restoreOptions = buildList {
             add(ArchiveRestoreOption.SignalSecureBackup)
             add(ArchiveRestoreOption.LocalBackup)
             add(ArchiveRestoreOption.None)
           },
-          isPreRegistration = false
+          registeredState = RegisteredState.RegisteredAndPinUnknown
+        )
+      }
+
+      fun forPostRegisterWithPinKnown(): ArchiveRestoreSelection {
+        return ArchiveRestoreSelection(
+          restoreOptions = buildList {
+            add(ArchiveRestoreOption.SignalSecureBackup)
+            add(ArchiveRestoreOption.LocalBackup)
+            add(ArchiveRestoreOption.None)
+          },
+          registeredState = RegisteredState.RegisteredAndPinKnown
         )
       }
     }
@@ -694,7 +706,7 @@ private fun EntryProviderScope<NavKey>.navigationEntries(
     val viewModel: ArchiveRestoreSelectionViewModel = viewModel(
       factory = ArchiveRestoreSelectionViewModel.Factory(
         restoreOptions = key.restoreOptions,
-        isPreRegistration = key.isPreRegistration,
+        registeredState = key.registeredState,
         repository = registrationRepository,
         parentState = registrationViewModel.state,
         parentEventEmitter = registrationViewModel::onEvent

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

@@ -648,43 +648,11 @@ class RegistrationRepository(val context: Context, val networkController: Networ
    * The work continues in the background even if [timeout] elapses. See [NetworkController.restoreAccountRecord].
    */
   suspend fun restoreAccountRecord(
-    timeout: Duration
+    timeout: Duration = 10.seconds
   ): RequestResult<Unit, NetworkController.RestoreAccountRecordError> = withContext(Dispatchers.IO) {
     networkController.restoreAccountRecord(timeout)
   }
 
-  /**
-   * Best-effort restore the AccountRecord (when local profile data is incomplete) and then signal
-   * registration completion on [parentEventEmitter]. The Profile screen is intentionally not
-   * routed to from here for now — even when the restore doesn't fully populate profile data, we
-   * emit [RegistrationFlowEvent.RegistrationComplete].
-   *
-   * Intended for any screen that, in the legacy flow, would have signalled "we're done". Pre-
-   * existing-data callers (re-registration, device transfer, backup restore) won't pay the
-   * restore-record cost.
-   */
-  suspend fun finishRegistrationOrCreateProfile(
-    parentEventEmitter: (RegistrationFlowEvent) -> Unit,
-    restoreTimeout: Duration = 10.seconds
-  ) {
-    if (hasProfileNameAndAvatar()) {
-      Log.i(TAG, "[finishRegistrationOrCreateProfile] Profile name + avatar already on disk; finishing.")
-      parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
-      return
-    }
-
-    Log.i(TAG, "[finishRegistrationOrCreateProfile] Profile data incomplete; attempting best-effort account-record restore (timeout=${restoreTimeout.inWholeSeconds}s).")
-    restoreAccountRecord(restoreTimeout)
-
-    Log.i(TAG, "[finishRegistrationOrCreateProfile] Account-record restore finished; finishing without routing to Profile screen.")
-    parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
-  }
-
-  private suspend fun hasProfileNameAndAvatar(): Boolean {
-    val stored = getStoredProfileData()
-    return stored.givenName.isNotEmpty() && stored.avatar != null
-  }
-
   /**
    * Persists the freshly-created profile to local storage and arranges for it to be uploaded.
    * See [NetworkController.setProfile].

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

@@ -24,6 +24,7 @@ import kotlinx.coroutines.launch
 import org.signal.core.ui.navigation.ResultEventBus
 import org.signal.core.util.logging.Log
 import org.signal.registration.screens.EventDrivenViewModel
+import org.signal.registration.screens.restoreselection.RegisteredState
 import kotlin.reflect.KClass
 
 /**
@@ -148,7 +149,7 @@ class RegistrationViewModel(
       is RegistrationRoute.PinCreate,
       is RegistrationRoute.PinEntryForSvrRestore,
       is RegistrationRoute.RemoteRestore -> true
-      is RegistrationRoute.ArchiveRestoreSelection -> !this.isPreRegistration
+      is RegistrationRoute.ArchiveRestoreSelection -> this.registeredState != RegisteredState.NotRegistered
       else -> false
     }
   }

feature/registration/src/main/java/org/signal/registration/screens/devicetransfer/complete/DeviceTransferCompleteViewModel.kt

@@ -43,7 +43,8 @@ class DeviceTransferCompleteViewModel(
     when (event) {
       DeviceTransferCompleteScreenEvents.ContinueClicked -> {
         repository.setRestoreDecision(RestoreDecision.COMPLETED)
-        repository.finishRegistrationOrCreateProfile(parentEventEmitter)
+        repository.restoreAccountRecord()
+        parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
       }
       DeviceTransferCompleteScreenEvents.ConsumeOneTimeEvent -> {
         stateEmitter(state.copy(oneTimeEvent = null))

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

@@ -120,7 +120,8 @@ class LocalBackupRestoreViewModel(
       parentEventEmitter.navigateBack()
     } else {
       repository.setRestoreDecision(RestoreDecision.COMPLETED)
-      repository.finishRegistrationOrCreateProfile(parentEventEmitter)
+      repository.restoreAccountRecord()
+      parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
     }
   }
 

feature/registration/src/main/java/org/signal/registration/screens/pincreation/PinCreationViewModel.kt

@@ -125,7 +125,8 @@ class PinCreationViewModel(
       is RequestResult.Success -> {
         Log.i(TAG, "[PinSubmitted] Successfully backed up master key to SVR.")
         repository.setRestoreDecision(RestoreDecision.NEW_ACCOUNT)
-        repository.finishRegistrationOrCreateProfile(parentEventEmitter)
+        repository.restoreAccountRecord()
+        parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
         state
       }
 

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

@@ -67,7 +67,7 @@ class PinEntryForRegistrationLockViewModel(
         stateEmitter(applyPinEntered(localState, event, parentEventEmitter))
       }
       is PinEntryScreenEvents.Skip -> {
-        handleSkip()
+        throw NotImplementedError("Skip is not a valid action during registration lock PIN entry")
       }
       is PinEntryScreenEvents.CreateNewPin,
       is PinEntryScreenEvents.ContactSupport -> Unit
@@ -142,9 +142,11 @@ class PinEntryForRegistrationLockViewModel(
         Log.i(TAG, "[PinEntered] Successfully registered!")
         val (response, keyMaterial) = registerResult.result
         parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool, response.storageCapable))
+        repository.enqueueSvrResetGuessCountJob()
+        repository.restoreAccountRecord()
         when {
-          response.reregistration -> parentEventEmitter.navigateTo(RegistrationRoute.ArchiveRestoreSelection.forPostRegister())
-          else -> repository.finishRegistrationOrCreateProfile(parentEventEmitter)
+          response.reregistration -> parentEventEmitter.navigateTo(RegistrationRoute.ArchiveRestoreSelection.forPostRegisterWithPinKnown())
+          else -> parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
         }
         state
       }
@@ -191,12 +193,6 @@ class PinEntryForRegistrationLockViewModel(
     }
   }
 
-  private fun handleSkip() {
-    // Registration lock is enforced server-side, so there's no way to register without the PIN. The skip option is
-    // never shown in this mode, so reaching here indicates a bug.
-    throw NotImplementedError("Skip is not a valid action during registration lock PIN entry")
-  }
-
   class Factory(
     private val repository: RegistrationRepository,
     private val parentState: StateFlow<RegistrationFlowState>,

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

@@ -151,7 +151,8 @@ class PinEntryForSmsBypassViewModel(
     return when (val result = repository.registerAccountWithRecoveryPassword(e164, recoveryPassword, registrationLock, skipDeviceTransfer = true)) {
       is RequestResult.Success -> {
         repository.enqueueSvrResetGuessCountJob()
-        repository.finishRegistrationOrCreateProfile(parentEventEmitter)
+        repository.restoreAccountRecord()
+        parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
         state
       }
       is RequestResult.RetryableNetworkError -> {

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

@@ -98,7 +98,7 @@ class PinEntryForSvrRestoreViewModel(
         result.result
       }
       is RequestResult.NonSuccess<NetworkController.GetSvrCredentialsError> -> {
-        when (val error = result.error) {
+        when (result.error) {
           NetworkController.GetSvrCredentialsError.NoServiceCredentialsAvailable -> {
             Log.w(TAG, "[PinEntered] No service credentials available when restoring from SVR. This should not happen. Resetting.")
             parentEventEmitter(RegistrationFlowEvent.ResetState)
@@ -125,7 +125,8 @@ class PinEntryForSvrRestoreViewModel(
         repository.enqueueSvrResetGuessCountJob()
         repository.setRestoreDecision(RestoreDecision.COMPLETED)
         parentEventEmitter(RegistrationFlowEvent.MasterKeyRestoredFromSvr(result.result.masterKey))
-        repository.finishRegistrationOrCreateProfile(parentEventEmitter)
+        repository.restoreAccountRecord()
+        parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
         state
       }
       is RequestResult.NonSuccess -> {

feature/registration/src/main/java/org/signal/registration/screens/remotebackuprestore/RemoteBackupRestoreViewModel.kt

@@ -122,7 +122,8 @@ class RemoteBackupRestoreViewModel(
             parentEventEmitter(RegistrationFlowEvent.UserSuppliedAepVerified(aep))
             repository.persistRemoteBackupRestoredState(progress.restoredSvrPin, progress.restoredProfileKey)
             repository.setRestoreDecision(RestoreDecision.COMPLETED)
-            repository.finishRegistrationOrCreateProfile(parentEventEmitter)
+            repository.restoreAccountRecord()
+            parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
           }
           is RemoteBackupRestoreProgress.NetworkError -> {
             Log.w(TAG, "[restoreBackup] Remote restore failed with network error.", progress.cause)

feature/registration/src/main/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionViewModel.kt

@@ -34,7 +34,7 @@ import org.signal.registration.screens.util.navigateTo
  */
 class ArchiveRestoreSelectionViewModel(
   private val restoreOptions: List<ArchiveRestoreOption>,
-  private val isPreRegistration: Boolean,
+  private val registeredState: RegisteredState,
   private val repository: RegistrationRepository,
   private val parentState: StateFlow<RegistrationFlowState>,
   private val parentEventEmitter: (RegistrationFlowEvent) -> Unit
@@ -74,21 +74,27 @@ class ArchiveRestoreSelectionViewModel(
         when (event.option) {
           ArchiveRestoreOption.SignalSecureBackup -> {
             notifyOldDevice(state.restoreMethodToken, NetworkController.RestoreMethod.REMOTE_BACKUP)
-            if (isPreRegistration) {
-              parentEventEmitter(RegistrationFlowEvent.PendingRestoreOptionSelected(PendingRestoreOption.RemoteBackup))
-              parentEventEmitter.navigateTo(RegistrationRoute.PhoneNumberEntry)
-            } else {
-              parentEventEmitter.navigateTo(RegistrationRoute.EnterAepForRemoteBackupPostRegistration)
+            when (registeredState) {
+              RegisteredState.NotRegistered -> {
+                parentEventEmitter(RegistrationFlowEvent.PendingRestoreOptionSelected(PendingRestoreOption.RemoteBackup))
+                parentEventEmitter.navigateTo(RegistrationRoute.PhoneNumberEntry)
+              }
+              else -> {
+                parentEventEmitter.navigateTo(RegistrationRoute.EnterAepForRemoteBackupPostRegistration)
+              }
             }
             state
           }
           ArchiveRestoreOption.LocalBackup -> {
             notifyOldDevice(state.restoreMethodToken, NetworkController.RestoreMethod.LOCAL_BACKUP)
-            if (isPreRegistration) {
-              parentEventEmitter(RegistrationFlowEvent.PendingRestoreOptionSelected(PendingRestoreOption.LocalBackup))
-              parentEventEmitter.navigateTo(RegistrationRoute.PhoneNumberEntry)
-            } else {
-              parentEventEmitter.navigateTo(RegistrationRoute.LocalBackupRestore(isPreRegistration = false))
+            when (registeredState) {
+              RegisteredState.NotRegistered -> {
+                parentEventEmitter(RegistrationFlowEvent.PendingRestoreOptionSelected(PendingRestoreOption.LocalBackup))
+                parentEventEmitter.navigateTo(RegistrationRoute.PhoneNumberEntry)
+              }
+              else -> {
+                parentEventEmitter.navigateTo(RegistrationRoute.LocalBackupRestore(isPreRegistration = false))
+              }
             }
             state
           }
@@ -103,20 +109,29 @@ class ArchiveRestoreSelectionViewModel(
         }
       }
       is ArchiveRestoreSelectionScreenEvents.ConfirmSkip -> {
-        if (isPreRegistration) {
-          parentEventEmitter.navigateTo(RegistrationRoute.PhoneNumberEntry)
-          state.copy(showSkipWarningDialog = false)
-        } else {
-          notifyOldDevice(state.restoreMethodToken, NetworkController.RestoreMethod.DECLINE)
-          repository.setRestoreDecision(RestoreDecision.SKIPPED)
-          if (state.storageCapable) {
-            Log.i(TAG, "[ConfirmSkip] Account is storage capable. Navigating to PIN entry to restore the existing PIN.")
-            parentEventEmitter.navigateTo(RegistrationRoute.PinEntryForSvrRestore)
-          } else {
-            Log.i(TAG, "[ConfirmSkip] Account is not storage capable. Navigating to PIN creation.")
-            parentEventEmitter.navigateTo(RegistrationRoute.PinCreate)
+        when (registeredState) {
+          RegisteredState.NotRegistered -> {
+            parentEventEmitter.navigateTo(RegistrationRoute.PhoneNumberEntry)
+            state.copy(showSkipWarningDialog = false)
+          }
+          RegisteredState.RegisteredAndPinUnknown -> {
+            notifyOldDevice(state.restoreMethodToken, NetworkController.RestoreMethod.DECLINE)
+            repository.setRestoreDecision(RestoreDecision.SKIPPED)
+            if (state.storageCapable) {
+              Log.i(TAG, "[ConfirmSkip] Account is storage capable. Navigating to PIN entry to restore the existing PIN.")
+              parentEventEmitter.navigateTo(RegistrationRoute.PinEntryForSvrRestore)
+            } else {
+              Log.i(TAG, "[ConfirmSkip] Account is not storage capable. Navigating to PIN creation.")
+              parentEventEmitter.navigateTo(RegistrationRoute.PinCreate)
+            }
+            state.copy(showSkipWarningDialog = false)
+          }
+          RegisteredState.RegisteredAndPinKnown -> {
+            notifyOldDevice(state.restoreMethodToken, NetworkController.RestoreMethod.DECLINE)
+            repository.setRestoreDecision(RestoreDecision.SKIPPED)
+            parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
+            state.copy(showSkipWarningDialog = false)
           }
-          state.copy(showSkipWarningDialog = false)
         }
       }
       is ArchiveRestoreSelectionScreenEvents.DismissSkipWarning -> {
@@ -145,13 +160,13 @@ class ArchiveRestoreSelectionViewModel(
 
   class Factory(
     private val restoreOptions: List<ArchiveRestoreOption>,
-    private val isPreRegistration: Boolean,
+    private val registeredState: RegisteredState,
     private val repository: RegistrationRepository,
     private val parentState: StateFlow<RegistrationFlowState>,
     private val parentEventEmitter: (RegistrationFlowEvent) -> Unit
   ) : ViewModelProvider.Factory {
     override fun <T : ViewModel> create(modelClass: Class<T>): T {
-      return ArchiveRestoreSelectionViewModel(restoreOptions, isPreRegistration, repository, parentState, parentEventEmitter) as T
+      return ArchiveRestoreSelectionViewModel(restoreOptions, registeredState, repository, parentState, parentEventEmitter) as T
     }
   }
 }

feature/registration/src/main/java/org/signal/registration/screens/restoreselection/RegisteredState.kt

@@ -0,0 +1,13 @@
+/*
+ * Copyright 2026 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.signal.registration.screens.restoreselection
+
+/**
+ * Just an enum for letting the view model know what the registered status is so it can navigate appropriately.
+ */
+enum class RegisteredState {
+  NotRegistered, RegisteredAndPinUnknown, RegisteredAndPinKnown
+}

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

@@ -355,7 +355,7 @@ class VerificationCodeViewModel(
         parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool, response.storageCapable))
 
         when {
-          response.reregistration -> parentEventEmitter.navigateTo(RegistrationRoute.ArchiveRestoreSelection.forPostRegister())
+          response.reregistration -> parentEventEmitter.navigateTo(RegistrationRoute.ArchiveRestoreSelection.forPostRegisterWithPinUnknown())
           response.storageCapable -> parentEventEmitter.navigateTo(RegistrationRoute.PinEntryForSvrRestore)
           else -> parentEventEmitter.navigateTo(RegistrationRoute.PinCreate)
         }

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

@@ -475,7 +475,7 @@ class RegistrationViewModelTest {
       )
     )
 
-    val postRegisterSelection = RegistrationRoute.ArchiveRestoreSelection.forPostRegister()
+    val postRegisterSelection = RegistrationRoute.ArchiveRestoreSelection.forPostRegisterWithPinKnown()
 
     val result = viewModel.applyEvent(
       initialState,

feature/registration/src/test/java/org/signal/registration/screens/devicetransfer/complete/DeviceTransferCompleteViewModelTest.kt

@@ -6,6 +6,7 @@
 package org.signal.registration.screens.devicetransfer.complete
 
 import assertk.assertThat
+import assertk.assertions.contains
 import assertk.assertions.hasSize
 import assertk.assertions.isEmpty
 import assertk.assertions.isNull
@@ -54,7 +55,7 @@ class DeviceTransferCompleteViewModelTest {
   }
 
   @Test
-  fun `ContinueClicked hands off to finishRegistrationOrCreateProfile`() = runTest {
+  fun `ContinueClicked restores account record and completes registration`() = runTest {
     viewModel.applyEvent(
       DeviceTransferCompleteState(),
       DeviceTransferCompleteScreenEvents.ContinueClicked,
@@ -64,7 +65,8 @@ class DeviceTransferCompleteViewModelTest {
     )
 
     coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) }
-    coVerify { mockRepository.finishRegistrationOrCreateProfile(parentEventEmitter, any()) }
+    coVerify { mockRepository.restoreAccountRecord(any()) }
+    assertThat(emittedEvents).contains(RegistrationFlowEvent.RegistrationComplete)
   }
 
   @Test
@@ -79,7 +81,7 @@ class DeviceTransferCompleteViewModelTest {
 
     coVerifyOrder {
       mockRepository.setRestoreDecision(RestoreDecision.COMPLETED)
-      mockRepository.finishRegistrationOrCreateProfile(parentEventEmitter, any())
+      mockRepository.restoreAccountRecord(any())
     }
   }
 
@@ -109,7 +111,7 @@ class DeviceTransferCompleteViewModelTest {
     assertThat(emittedStates).hasSize(1)
     assertThat(emittedStates.last().oneTimeEvent).isNull()
     coVerify(exactly = 0) { mockRepository.setRestoreDecision(any()) }
-    coVerify(exactly = 0) { mockRepository.finishRegistrationOrCreateProfile(any(), any()) }
+    coVerify(exactly = 0) { mockRepository.restoreAccountRecord(any()) }
   }
 
   @Test
@@ -118,6 +120,7 @@ class DeviceTransferCompleteViewModelTest {
     testDispatcher.scheduler.advanceUntilIdle()
 
     coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) }
-    coVerify { mockRepository.finishRegistrationOrCreateProfile(parentEventEmitter, any()) }
+    coVerify { mockRepository.restoreAccountRecord(any()) }
+    assertThat(emittedEvents).contains(RegistrationFlowEvent.RegistrationComplete)
   }
 }

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

@@ -7,6 +7,7 @@ package org.signal.registration.screens.localbackuprestore
 
 import android.net.Uri
 import assertk.assertThat
+import assertk.assertions.contains
 import assertk.assertions.hasSize
 import assertk.assertions.isEmpty
 import assertk.assertions.isEqualTo
@@ -266,6 +267,7 @@ class LocalBackupRestoreViewModelTest {
     viewModel.applyEvent(initialState, LocalBackupRestoreEvents.PassphraseSubmitted("passphrase"), stateEmitter)
 
     coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) }
-    coVerify { mockRepository.finishRegistrationOrCreateProfile(parentEventEmitter, any()) }
+    coVerify { mockRepository.restoreAccountRecord(any()) }
+    assertThat(emittedParentEvents).contains(RegistrationFlowEvent.RegistrationComplete)
   }
 }

feature/registration/src/test/java/org/signal/registration/screens/pincreation/PinCreationViewModelTest.kt

@@ -6,6 +6,7 @@
 package org.signal.registration.screens.pincreation
 
 import assertk.assertThat
+import assertk.assertions.contains
 import assertk.assertions.hasSize
 import assertk.assertions.isEqualTo
 import assertk.assertions.isFalse
@@ -130,7 +131,7 @@ class PinCreationViewModelTest {
   // ==================== PinSubmitted Success Tests ====================
 
   @Test
-  fun `matching confirmation PIN with valid AEP and successful SVR backup hands off to finishRegistrationOrCreateProfile`() = runTest(testDispatcher) {
+  fun `matching confirmation PIN with valid AEP and successful SVR backup completes registration`() = runTest(testDispatcher) {
     val states = collectStates()
     val aep = AccountEntropyPool.generate()
     val confirmState = PinCreationState(accountEntropyPool = aep, isConfirmEnabled = true, firstPin = "123456")
@@ -141,7 +142,8 @@ class PinCreationViewModelTest {
     viewModel.applyEvent(confirmState, PinCreationScreenEvents.PinSubmitted("123456"))
 
     coVerify { mockRepository.setRestoreDecision(RestoreDecision.NEW_ACCOUNT) }
-    coVerify { mockRepository.finishRegistrationOrCreateProfile(parentEventEmitter, any()) }
+    coVerify { mockRepository.restoreAccountRecord(any()) }
+    assertThat(emittedParentEvents).contains(RegistrationFlowEvent.RegistrationComplete)
     assertThat(states.last().loading).isTrue()
   }
 

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

@@ -82,10 +82,36 @@ class PinEntryForRegistrationLockViewModelTest {
 
     viewModel.applyEvent(initialState, PinEntryScreenEvents.PinEntered("123456"), parentEventEmitter, stateEmitter)
 
-    assertThat(emittedParentEvents).hasSize(2)
+    assertThat(emittedParentEvents).hasSize(3)
     assertThat(emittedParentEvents[0]).isInstanceOf<RegistrationFlowEvent.MasterKeyRestoredFromSvr>()
     assertThat(emittedParentEvents[1]).isInstanceOf<RegistrationFlowEvent.Registered>()
-    coVerify { mockRepository.finishRegistrationOrCreateProfile(parentEventEmitter, any()) }
+    assertThat(emittedParentEvents[2]).isEqualTo(RegistrationFlowEvent.RegistrationComplete)
+    coVerify { mockRepository.restoreAccountRecord(any()) }
+    assertThat(emittedStates.last().loading).isEqualTo(true)
+  }
+
+  @Test
+  fun `PinEntered with correct PIN on re-registration navigates to post-register restore selection`() = runTest {
+    val masterKey = mockk<MasterKey>(relaxed = true)
+    val keyMaterial = mockk<KeyMaterial>(relaxed = true)
+    val registerResponse = createRegisterAccountResponse(reregistration = true)
+    val initialState = PinEntryState(mode = PinEntryState.Mode.RegistrationLock)
+
+    coEvery { mockRepository.restoreMasterKeyFromSvr(any(), any(), forRegistrationLock = true) } returns
+      RequestResult.Success(NetworkController.MasterKeyResponse(masterKey))
+    coEvery { mockRepository.registerAccountWithSession(any(), any(), any(), any()) } returns
+      RequestResult.Success(registerResponse to keyMaterial)
+
+    viewModel.applyEvent(initialState, PinEntryScreenEvents.PinEntered("123456"), parentEventEmitter, stateEmitter)
+
+    assertThat(emittedParentEvents).hasSize(3)
+    assertThat(emittedParentEvents[0]).isInstanceOf<RegistrationFlowEvent.MasterKeyRestoredFromSvr>()
+    assertThat(emittedParentEvents[1]).isInstanceOf<RegistrationFlowEvent.Registered>()
+    assertThat(emittedParentEvents[2])
+      .isInstanceOf<RegistrationFlowEvent.NavigateToScreen>()
+      .prop(RegistrationFlowEvent.NavigateToScreen::route)
+      .isEqualTo(RegistrationRoute.ArchiveRestoreSelection.forPostRegisterWithPinKnown())
+    coVerify { mockRepository.restoreAccountRecord(any()) }
     assertThat(emittedStates.last().loading).isEqualTo(true)
   }
 
@@ -427,7 +453,8 @@ class PinEntryForRegistrationLockViewModelTest {
     aci: String = "test-aci",
     pni: String = "test-pni",
     e164: String = "+15551234567",
-    storageCapable: Boolean = true
+    storageCapable: Boolean = true,
+    reregistration: Boolean = false
   ) = NetworkController.RegisterAccountResponse(
     aci = aci,
     pni = pni,
@@ -436,6 +463,6 @@ class PinEntryForRegistrationLockViewModelTest {
     usernameLinkHandle = null,
     storageCapable = storageCapable,
     entitlements = null,
-    reregistration = false
+    reregistration = reregistration
   )
 }

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

@@ -65,7 +65,7 @@ class PinEntryForSmsBypassViewModelTest {
   // ==================== PinEntered - Restore Master Key Tests ====================
 
   @Test
-  fun `PinEntered with correct PIN restores master key and hands off to finishRegistrationOrCreateProfile`() = runTest {
+  fun `PinEntered with correct PIN restores master key and completes registration`() = runTest {
     val masterKey = mockk<MasterKey>(relaxed = true)
     val initialState = PinEntryState(mode = PinEntryState.Mode.SmsBypass, e164 = "+15551234567")
 
@@ -76,9 +76,10 @@ class PinEntryForSmsBypassViewModelTest {
 
     viewModel.applyEvent(initialState, PinEntryScreenEvents.PinEntered("123456"), parentEventEmitter, stateEmitter)
 
-    assertThat(emittedParentEvents).hasSize(1)
+    assertThat(emittedParentEvents).hasSize(2)
     assertThat(emittedParentEvents[0]).isInstanceOf<RegistrationFlowEvent.MasterKeyRestoredFromSvr>()
-    coVerify { mockRepository.finishRegistrationOrCreateProfile(parentEventEmitter, any()) }
+    assertThat(emittedParentEvents[1]).isEqualTo(RegistrationFlowEvent.RegistrationComplete)
+    coVerify { mockRepository.restoreAccountRecord(any()) }
     assertThat(emittedStates.last().loading).isEqualTo(true)
   }
 
@@ -290,9 +291,10 @@ class PinEntryForSmsBypassViewModelTest {
 
     viewModel.applyEvent(initialState, PinEntryScreenEvents.PinEntered("123456"), parentEventEmitter, stateEmitter)
 
-    assertThat(emittedParentEvents).hasSize(1)
+    assertThat(emittedParentEvents).hasSize(2)
     assertThat(emittedParentEvents[0]).isInstanceOf<RegistrationFlowEvent.MasterKeyRestoredFromSvr>()
-    coVerify { mockRepository.finishRegistrationOrCreateProfile(parentEventEmitter, any()) }
+    assertThat(emittedParentEvents[1]).isEqualTo(RegistrationFlowEvent.RegistrationComplete)
+    coVerify { mockRepository.restoreAccountRecord(any()) }
     assertThat(emittedStates.last().loading).isEqualTo(true)
   }
 

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

@@ -59,7 +59,7 @@ class PinEntryForSvrRestoreViewModelTest {
   // ==================== PinEntered Success Tests ====================
 
   @Test
-  fun `PinEntered with correct PIN restores master key and hands off to finishRegistrationOrCreateProfile`() = runTest {
+  fun `PinEntered with correct PIN restores master key and completes registration`() = runTest {
     val masterKey = mockk<MasterKey>(relaxed = true)
     val svrCredentials = NetworkController.SvrCredentials(
       username = "test-username",
@@ -74,10 +74,11 @@ class PinEntryForSvrRestoreViewModelTest {
 
     viewModel.applyEvent(initialState, PinEntryScreenEvents.PinEntered("123456"), parentEventEmitter, stateEmitter)
 
-    assertThat(emittedParentEvents).hasSize(1)
+    assertThat(emittedParentEvents).hasSize(2)
     assertThat(emittedParentEvents[0]).isInstanceOf<RegistrationFlowEvent.MasterKeyRestoredFromSvr>()
+    assertThat(emittedParentEvents[1]).isEqualTo(RegistrationFlowEvent.RegistrationComplete)
     coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) }
-    coVerify { mockRepository.finishRegistrationOrCreateProfile(parentEventEmitter, any()) }
+    coVerify { mockRepository.restoreAccountRecord(any()) }
     assertThat(emittedStates.last().loading).isEqualTo(true)
   }
 

feature/registration/src/test/java/org/signal/registration/screens/remotebackuprestore/RemoteBackupRestoreViewModelTest.kt

@@ -347,7 +347,7 @@ class RemoteBackupRestoreViewModelTest {
   }
 
   @Test
-  fun `Complete progress emits UserSuppliedAepVerified and hands off to finishRegistrationOrCreateProfile`() = runTest(testDispatcher) {
+  fun `Complete progress emits UserSuppliedAepVerified and completes registration`() = runTest(testDispatcher) {
     every { mockRepository.restoreRemoteBackup(any()) } returns flowOf(
       RemoteBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null)
     )
@@ -361,10 +361,11 @@ class RemoteBackupRestoreViewModelTest {
       stateEmitter
     )
 
-    assertThat(emittedParentEvents).hasSize(1)
+    assertThat(emittedParentEvents).hasSize(2)
     assertThat(emittedParentEvents[0]).isInstanceOf<RegistrationFlowEvent.UserSuppliedAepVerified>()
+    assertThat(emittedParentEvents[1]).isEqualTo(RegistrationFlowEvent.RegistrationComplete)
     coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) }
-    coVerify { mockRepository.finishRegistrationOrCreateProfile(parentEventEmitter, any()) }
+    coVerify { mockRepository.restoreAccountRecord(any()) }
   }
 
   @Test
@@ -462,6 +463,6 @@ class RemoteBackupRestoreViewModelTest {
     viewModel.applyEvent(RemoteBackupRestoreState(aep = aep), RemoteBackupRestoreScreenEvents.BackupRestoreBackup, stateEmitter)
 
     assertThat(states.last().restoreState).isEqualTo(RemoteBackupRestoreState.RestoreState.Restored)
-    coVerify { mockRepository.finishRegistrationOrCreateProfile(parentEventEmitter, any()) }
+    coVerify { mockRepository.restoreAccountRecord(any()) }
   }
 }

feature/registration/src/test/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionViewModelTest.kt

@@ -48,11 +48,11 @@ class ArchiveRestoreSelectionViewModelTest {
       ArchiveRestoreOption.LocalBackup,
       ArchiveRestoreOption.DeviceTransfer
     ),
-    isPreRegistration: Boolean = false
+    registeredState: RegisteredState = RegisteredState.RegisteredAndPinUnknown
   ): ArchiveRestoreSelectionViewModel {
     return ArchiveRestoreSelectionViewModel(
       restoreOptions = restoreOptions,
-      isPreRegistration = isPreRegistration,
+      registeredState = registeredState,
       repository = mockRepository,
       parentState = MutableStateFlow(RegistrationFlowState()),
       parentEventEmitter = parentEventEmitter
@@ -63,7 +63,7 @@ class ArchiveRestoreSelectionViewModelTest {
 
   @Test
   fun `SignalSecureBackup pre-registration emits PendingRestoreOptionSelected and navigates to PhoneNumberEntry`() = runTest {
-    val viewModel = createViewModel(isPreRegistration = true)
+    val viewModel = createViewModel(registeredState = RegisteredState.NotRegistered)
     val initialState = ArchiveRestoreSelectionState()
 
     viewModel.applyEvent(
@@ -85,7 +85,7 @@ class ArchiveRestoreSelectionViewModelTest {
 
   @Test
   fun `SignalSecureBackup post-registration navigates to EnterAepForRemoteBackupPostRegistration`() = runTest {
-    val viewModel = createViewModel(isPreRegistration = false)
+    val viewModel = createViewModel(registeredState = RegisteredState.RegisteredAndPinUnknown)
     val initialState = ArchiveRestoreSelectionState()
 
     viewModel.applyEvent(
@@ -103,7 +103,7 @@ class ArchiveRestoreSelectionViewModelTest {
 
   @Test
   fun `LocalBackup pre-registration emits PendingRestoreOptionSelected and navigates to PhoneNumberEntry`() = runTest {
-    val viewModel = createViewModel(isPreRegistration = true)
+    val viewModel = createViewModel(registeredState = RegisteredState.NotRegistered)
     val initialState = ArchiveRestoreSelectionState()
 
     viewModel.applyEvent(
@@ -125,7 +125,7 @@ class ArchiveRestoreSelectionViewModelTest {
 
   @Test
   fun `LocalBackup post-registration navigates to LocalBackupRestore`() = runTest {
-    val viewModel = createViewModel(isPreRegistration = false)
+    val viewModel = createViewModel(registeredState = RegisteredState.RegisteredAndPinUnknown)
     val initialState = ArchiveRestoreSelectionState()
 
     viewModel.applyEvent(
@@ -143,7 +143,7 @@ class ArchiveRestoreSelectionViewModelTest {
 
   @Test
   fun `DeviceTransfer navigates to DeviceTransferInstructions`() = runTest {
-    val viewModel = createViewModel(isPreRegistration = false)
+    val viewModel = createViewModel(registeredState = RegisteredState.RegisteredAndPinUnknown)
     val initialState = ArchiveRestoreSelectionState()
 
     viewModel.applyEvent(
@@ -161,7 +161,7 @@ class ArchiveRestoreSelectionViewModelTest {
 
   @Test
   fun `None option sets showSkipWarningDialog to true`() = runTest {
-    val viewModel = createViewModel(isPreRegistration = false)
+    val viewModel = createViewModel(registeredState = RegisteredState.RegisteredAndPinUnknown)
     val initialState = ArchiveRestoreSelectionState()
 
     viewModel.applyEvent(
@@ -178,7 +178,7 @@ class ArchiveRestoreSelectionViewModelTest {
 
   @Test
   fun `ConfirmSkip pre-registration navigates to PhoneNumberEntry and clears dialog without recording a skip`() = runTest {
-    val viewModel = createViewModel(isPreRegistration = true)
+    val viewModel = createViewModel(registeredState = RegisteredState.NotRegistered)
     val initialState = ArchiveRestoreSelectionState(showSkipWarningDialog = true)
 
     viewModel.applyEvent(initialState, ArchiveRestoreSelectionScreenEvents.ConfirmSkip, stateEmitter)
@@ -194,7 +194,7 @@ class ArchiveRestoreSelectionViewModelTest {
 
   @Test
   fun `ConfirmSkip post-registration when not storage capable navigates to PinCreate and clears dialog`() = runTest {
-    val viewModel = createViewModel(isPreRegistration = false)
+    val viewModel = createViewModel(registeredState = RegisteredState.RegisteredAndPinUnknown)
     val initialState = ArchiveRestoreSelectionState(showSkipWarningDialog = true, storageCapable = false)
 
     viewModel.applyEvent(initialState, ArchiveRestoreSelectionScreenEvents.ConfirmSkip, stateEmitter)
@@ -210,7 +210,7 @@ class ArchiveRestoreSelectionViewModelTest {
 
   @Test
   fun `ConfirmSkip post-registration when storage capable navigates to PinEntryForSvrRestore and clears dialog`() = runTest {
-    val viewModel = createViewModel(isPreRegistration = false)
+    val viewModel = createViewModel(registeredState = RegisteredState.RegisteredAndPinUnknown)
     val initialState = ArchiveRestoreSelectionState(showSkipWarningDialog = true, storageCapable = true)
 
     viewModel.applyEvent(initialState, ArchiveRestoreSelectionScreenEvents.ConfirmSkip, stateEmitter)
@@ -224,11 +224,24 @@ class ArchiveRestoreSelectionViewModelTest {
     assertThat(emittedStates.last().showSkipWarningDialog).isFalse()
   }
 
+  @Test
+  fun `ConfirmSkip post-registration when PIN is known records skip and completes registration`() = runTest {
+    val viewModel = createViewModel(registeredState = RegisteredState.RegisteredAndPinKnown)
+    val initialState = ArchiveRestoreSelectionState(showSkipWarningDialog = true)
+
+    viewModel.applyEvent(initialState, ArchiveRestoreSelectionScreenEvents.ConfirmSkip, stateEmitter)
+
+    coVerify { mockRepository.setRestoreDecision(RestoreDecision.SKIPPED) }
+    assertThat(emittedParentEvents).hasSize(1)
+    assertThat(emittedParentEvents.first()).isEqualTo(RegistrationFlowEvent.RegistrationComplete)
+    assertThat(emittedStates.last().showSkipWarningDialog).isFalse()
+  }
+
   // ==================== DismissSkipWarning Tests ====================
 
   @Test
   fun `DismissSkipWarning sets showSkipWarningDialog to false`() = runTest {
-    val viewModel = createViewModel(isPreRegistration = false)
+    val viewModel = createViewModel(registeredState = RegisteredState.RegisteredAndPinUnknown)
     val initialState = ArchiveRestoreSelectionState(showSkipWarningDialog = true)
 
     viewModel.applyEvent(

lib/libsignal-service/src/main/java/org/whispersystems/signalservice/api/svr/Svr2Socket.kt

@@ -151,6 +151,8 @@ internal class Svr2Socket(
     override fun onFailure(webSocket: WebSocket, t: Throwable, response: OkHttpResponse?) {
       val exception = if (t.message?.contains("404") == true) {
         NonSuccessfulResponseCodeException(404)
+      } else if (t.message?.contains("429") == true) {
+        NonSuccessfulResponseCodeException(429)
       } else {
         IOException(t)
       }