Fix pin restore from remote backup in regV5.

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

@@ -722,6 +722,7 @@ private fun EntryProviderScope<NavKey>.navigationEntries(
       factory = RemoteBackupRestoreViewModel.Factory(
         aep = key.aep,
         repository = registrationRepository,
+        parentState = registrationViewModel.state,
         parentEventEmitter = registrationViewModel::onEvent
       )
     )

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

@@ -814,6 +814,14 @@ class RegistrationRepository(val context: Context, val networkController: Networ
     storageController.readInProgressRegistrationData().lastUpdatedMillis.takeIf { it > 0 }
   }
 
+  /**
+   * True if a PIN is already known for this registration -- restored from a backup or provided by the old device
+   * during a quick restore -- and persisted in the in-progress registration data.
+   */
+  suspend fun hasKnownPin(): Boolean = withContext(Dispatchers.IO) {
+    storageController.readInProgressRegistrationData().pin.isNotEmpty()
+  }
+
   /**
    * Clears any persisted flow state JSON from the in-progress registration data.
    */

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

@@ -22,15 +22,19 @@ 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.RestoreDecision
 import org.signal.registration.screens.EventDrivenViewModel
 import org.signal.registration.screens.util.navigateBack
+import org.signal.registration.screens.util.navigateTo
 import kotlin.coroutines.CoroutineContext
 
 class RemoteBackupRestoreViewModel(
   private val aep: AccountEntropyPool,
   private val repository: RegistrationRepository,
+  private val parentState: StateFlow<RegistrationFlowState>,
   private val parentEventEmitter: (RegistrationFlowEvent) -> Unit,
   private val ioDispatcher: CoroutineContext = Dispatchers.IO
 ) : EventDrivenViewModel<RemoteBackupRestoreScreenEvents>(TAG) {
@@ -122,8 +126,21 @@ class RemoteBackupRestoreViewModel(
             )
             repository.persistRestoredBackupState(progress.restoredSvrPin, progress.restoredProfileKey)
             repository.setRestoreDecision(RestoreDecision.COMPLETED)
-            repository.restoreAccountRecord()
-            parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
+
+            when {
+              repository.hasKnownPin() -> {
+                repository.restoreAccountRecord()
+                parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
+              }
+              parentState.value.storageCapable -> {
+                Log.i(TAG, "[restoreBackup] No PIN is known and the account is storage capable. Navigating to PIN entry to restore the existing PIN.")
+                parentEventEmitter.navigateTo(RegistrationRoute.PinEntryForSvrRestore)
+              }
+              else -> {
+                Log.i(TAG, "[restoreBackup] No PIN is known and the account is not storage capable. Navigating to PIN creation.")
+                parentEventEmitter.navigateTo(RegistrationRoute.PinCreate)
+              }
+            }
           }
           is RemoteBackupRestoreProgress.NetworkError -> {
             Log.w(TAG, "[restoreBackup] Remote restore failed with network error.", progress.cause)
@@ -236,10 +253,11 @@ class RemoteBackupRestoreViewModel(
   class Factory(
     private val aep: AccountEntropyPool,
     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 RemoteBackupRestoreViewModel(aep, repository, parentEventEmitter) as T
+      return RemoteBackupRestoreViewModel(aep, repository, parentState, parentEventEmitter) as T
     }
   }
 }

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

@@ -53,6 +53,7 @@ import org.signal.registration.NetworkController.SvrCredentials
 import org.signal.registration.fakes.FakeNetworkController
 import org.signal.registration.fakes.FakeStorageController
 import org.signal.registration.fakes.SystemOutLogger
+import org.signal.registration.screens.remotebackuprestore.RemoteBackupRestoreProgress
 import org.signal.registration.screens.util.MockMultiplePermissionsState
 import org.signal.registration.screens.util.MockPermissionsState
 import org.signal.registration.test.TestTags
@@ -240,6 +241,11 @@ class RegistrationEndToEndTest {
   fun `restoring a remote backup before registering completes registration`() {
     val aep = AccountEntropyPool.generate()
 
+    // 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 })
 
@@ -259,6 +265,7 @@ class RegistrationEndToEndTest {
     assert(committed != null) { "Expected registration data to be committed" }
     assert(committed!!.e164 == E164) { "Expected committed e164 $E164 but was ${committed.e164}" }
     assert(committed.accountEntropyPool == aep.value) { "Expected the committed AEP to be the one the user entered" }
+    assert(committed.pin == PIN) { "Expected the pin from the restored backup but was ${committed.pin}" }
     assert(storageController.restoreDecision == RestoreDecision.COMPLETED) { "Expected COMPLETED restore decision but was ${storageController.restoreDecision}" }
   }
 
@@ -298,6 +305,11 @@ class RegistrationEndToEndTest {
       RequestResult.Success(networkController.registerAccountResponse(request.e164, reregistration = true))
     }
 
+    // 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 })
 
@@ -314,9 +326,41 @@ class RegistrationEndToEndTest {
     val committed = storageController.committedData
     assert(committed != null) { "Expected registration data to be committed" }
     assert(committed!!.accountEntropyPool == aep.value) { "Expected the committed AEP to be the one the user entered" }
+    assert(committed.pin == PIN) { "Expected the pin from the restored backup but was ${committed.pin}" }
     assert(storageController.restoreDecision == RestoreDecision.COMPLETED) { "Expected COMPLETED restore decision but was ${storageController.restoreDecision}" }
   }
 
+  @Test
+  fun `restoring a remote backup without a pin after registering requires creating a pin`() {
+    val aep = AccountEntropyPool.generate()
+
+    networkController.onRegisterAccount = { request ->
+      RequestResult.Success(networkController.registerAccountResponse(request.e164, reregistration = true))
+    }
+
+    var registrationComplete = false
+    launchRegistrationFlow(onRegistrationComplete = { registrationComplete = true })
+
+    submitPhoneNumber()
+    submitVerificationCode(VERIFICATION_CODE)
+
+    // The user is re-registering, so they're offered a restore
+    chooseRestoreOption(TestTags.ARCHIVE_RESTORE_SELECTION_FROM_SIGNAL_BACKUPS)
+    enterAep(aep)
+    startRemoteRestore()
+
+    // The restored backup had no PIN and the account is not storage capable, so the user must create a PIN
+    createPin(PIN)
+
+    waitFor("registration to complete") { registrationComplete }
+
+    val committed = storageController.committedData
+    assert(committed != null) { "Expected registration data to be committed" }
+    assert(committed!!.e164 == E164) { "Expected committed e164 $E164 but was ${committed.e164}" }
+    assert(committed.pin == PIN) { "Expected committed pin $PIN but was ${committed.pin}" }
+    assert(networkController.lastSetPinRequest?.pin == PIN) { "Expected pin $PIN on SVR but was ${networkController.lastSetPinRequest?.pin}" }
+  }
+
   @Test
   fun `restoring a local backup after registering completes registration`() {
     val aep = AccountEntropyPool.generate()
@@ -392,6 +436,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 })
 
@@ -462,11 +511,12 @@ class RegistrationEndToEndTest {
   fun `quick restore with a remote backup completes registration`() {
     val aep = AccountEntropyPool.generate()
 
-    // The old device scans the QR code as soon as it is shown and sends its provisioning data
+    // The old device scans the QR code as soon as it is shown and sends its provisioning data, including the PIN,
+    // so no PIN screens are needed after the restore
     networkController.onStartProvisioning = {
       flowOf(
         ProvisioningEvent.QrCodeReady("https://signal.test/qr"),
-        ProvisioningEvent.MessageReceived(networkController.provisioningMessage(aep = aep, e164 = E164))
+        ProvisioningEvent.MessageReceived(networkController.provisioningMessage(aep = aep, e164 = E164, pin = PIN))
       )
     }
 
@@ -490,6 +540,7 @@ class RegistrationEndToEndTest {
     assert(committed != null) { "Expected registration data to be committed" }
     assert(committed!!.e164 == E164) { "Expected committed e164 $E164 but was ${committed.e164}" }
     assert(committed.accountEntropyPool == aep.value) { "Expected the committed AEP to be the provisioned one" }
+    assert(committed.pin == PIN) { "Expected the provisioned pin $PIN but was ${committed.pin}" }
     assert(storageController.restoreDecision == RestoreDecision.COMPLETED) { "Expected COMPLETED restore decision but was ${storageController.restoreDecision}" }
   }
 

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

@@ -15,6 +15,7 @@ import io.mockk.every
 import io.mockk.mockk
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.emptyFlow
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.launch
@@ -30,7 +31,9 @@ import org.signal.core.models.AccountEntropyPool
 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.RestoreDecision
 import java.io.IOException
 import kotlin.time.Duration.Companion.seconds
@@ -66,10 +69,11 @@ class RemoteBackupRestoreViewModelTest {
     Dispatchers.resetMain()
   }
 
-  private fun createViewModel(): RemoteBackupRestoreViewModel {
+  private fun createViewModel(storageCapable: Boolean = false): RemoteBackupRestoreViewModel {
     return RemoteBackupRestoreViewModel(
       aep = aep,
       repository = mockRepository,
+      parentState = MutableStateFlow(RegistrationFlowState(storageCapable = storageCapable)),
       parentEventEmitter = parentEventEmitter,
       ioDispatcher = testDispatcher
     )
@@ -346,10 +350,11 @@ class RemoteBackupRestoreViewModelTest {
   }
 
   @Test
-  fun `Complete progress completes registration`() = runTest(testDispatcher) {
+  fun `Complete progress with a known pin completes registration`() = runTest(testDispatcher) {
     every { mockRepository.restoreRemoteBackup(any()) } returns flowOf(
-      RemoteBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null)
+      RemoteBackupRestoreProgress.Complete(restoredSvrPin = "1234", restoredProfileKey = null)
     )
+    coEvery { mockRepository.hasKnownPin() } returns true
 
     val viewModel = createViewModel()
     val initialState = RemoteBackupRestoreState(aep = aep)
@@ -362,10 +367,51 @@ class RemoteBackupRestoreViewModelTest {
 
     assertThat(emittedParentEvents).hasSize(1)
     assertThat(emittedParentEvents[0]).isEqualTo(RegistrationFlowEvent.RegistrationComplete)
+    coVerify { mockRepository.persistRestoredBackupState("1234", null) }
     coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) }
     coVerify { mockRepository.restoreAccountRecord(any()) }
   }
 
+  @Test
+  fun `Complete progress without a known pin navigates to pin creation when not storage capable`() = runTest(testDispatcher) {
+    every { mockRepository.restoreRemoteBackup(any()) } returns flowOf(
+      RemoteBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null)
+    )
+    coEvery { mockRepository.hasKnownPin() } returns false
+
+    val viewModel = createViewModel(storageCapable = false)
+
+    viewModel.applyEvent(
+      RemoteBackupRestoreState(aep = aep),
+      RemoteBackupRestoreScreenEvents.BackupRestoreBackup,
+      stateEmitter
+    )
+
+    assertThat(emittedParentEvents).hasSize(1)
+    assertThat(emittedParentEvents[0]).isEqualTo(RegistrationFlowEvent.NavigateToScreen(RegistrationRoute.PinCreate))
+    coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) }
+  }
+
+  @Test
+  fun `Complete progress without a known pin navigates to SVR pin entry when storage capable`() = runTest(testDispatcher) {
+    every { mockRepository.restoreRemoteBackup(any()) } returns flowOf(
+      RemoteBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null)
+    )
+    coEvery { mockRepository.hasKnownPin() } returns false
+
+    val viewModel = createViewModel(storageCapable = true)
+
+    viewModel.applyEvent(
+      RemoteBackupRestoreState(aep = aep),
+      RemoteBackupRestoreScreenEvents.BackupRestoreBackup,
+      stateEmitter
+    )
+
+    assertThat(emittedParentEvents).hasSize(1)
+    assertThat(emittedParentEvents[0]).isEqualTo(RegistrationFlowEvent.NavigateToScreen(RegistrationRoute.PinEntryForSvrRestore))
+    coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) }
+  }
+
   @Test
   fun `successful backup info emits UserSuppliedAepVerified`() = runTest(testDispatcher) {
     coEvery { mockRepository.getRemoteBackupInfo(any()) } returns RequestResult.Success(backupInfo())
@@ -464,8 +510,9 @@ class RemoteBackupRestoreViewModelTest {
       RemoteBackupRestoreProgress.Downloading(bytesDownloaded = 10, totalBytes = 100),
       RemoteBackupRestoreProgress.Restoring(bytesRead = 60, totalBytes = 100),
       RemoteBackupRestoreProgress.Finalizing,
-      RemoteBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null)
+      RemoteBackupRestoreProgress.Complete(restoredSvrPin = "1234", restoredProfileKey = null)
     )
+    coEvery { mockRepository.hasKnownPin() } returns true
     val viewModel = createViewModel()
     val states = collectStatesOf(viewModel)