Fix local backup v1 restores post-registration in regV5.

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

@@ -726,6 +726,7 @@ private fun EntryProviderScope<NavKey>.navigationEntries(
     val viewModel: LocalBackupRestoreViewModel = viewModel(
       factory = LocalBackupRestoreViewModel.Factory(
         repository = registrationRepository,
+        parentState = registrationViewModel.state,
         parentEventEmitter = registrationViewModel::onEvent,
         isPreRegistration = key.isPreRegistration,
         resultBus = registrationViewModel.resultBus,

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

@@ -7,31 +7,36 @@ package org.signal.registration.screens.localbackuprestore
 
 import android.net.Uri
 import org.signal.core.util.censor
+import org.signal.registration.RegistrationFlowState
+
+sealed interface LocalBackupRestoreEvents {
+
+  /** Parent state changed. */
+  data class ParentStateChanged(val state: RegistrationFlowState) : LocalBackupRestoreEvents
 
-sealed class LocalBackupRestoreEvents {
   /** User tapped the button to pick a backup folder. */
-  data object PickBackupFolder : LocalBackupRestoreEvents()
+  data object PickBackupFolder : LocalBackupRestoreEvents
 
   /** User selected a backup folder via the folder picker. */
-  data class BackupFolderSelected(val uri: Uri) : LocalBackupRestoreEvents()
+  data class BackupFolderSelected(val uri: Uri) : LocalBackupRestoreEvents
 
   /** User wants to restore the found backup (navigates to credential entry). */
-  data object RestoreBackup : LocalBackupRestoreEvents()
+  data object RestoreBackup : LocalBackupRestoreEvents
 
   /** User wants to choose a different folder. */
-  data object ChooseDifferentFolder : LocalBackupRestoreEvents()
+  data object ChooseDifferentFolder : LocalBackupRestoreEvents
 
   /** User selected a specific backup from the backup picker. */
-  data class BackupSelected(val backup: LocalBackupInfo) : LocalBackupRestoreEvents()
+  data class BackupSelected(val backup: LocalBackupInfo) : LocalBackupRestoreEvents
 
   /** A credential (passphrase or AEP) was received from the credential entry screen. */
-  data class PassphraseSubmitted(val credential: String) : LocalBackupRestoreEvents() {
+  data class PassphraseSubmitted(val credential: String) : LocalBackupRestoreEvents {
     override fun toString(): String = "PassphraseSubmitted(credential=${credential.censor()})"
   }
 
   /** The folder picker was dismissed without selecting a folder. */
-  data object FolderPickerDismissed : LocalBackupRestoreEvents()
+  data object FolderPickerDismissed : LocalBackupRestoreEvents
 
   /** User wants to cancel. */
-  data object Cancel : LocalBackupRestoreEvents()
+  data object Cancel : LocalBackupRestoreEvents
 }

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

@@ -18,7 +18,8 @@ data class LocalBackupRestoreState(
   val errorMessage: String? = null,
   val launchFolderPicker: Boolean = false,
   val aep: AccountEntropyPool? = null,
-  val v1Passphrase: String? = null
+  val v1Passphrase: String? = null,
+  val storageCapable: Boolean = true
 ) {
 
   override fun toString(): String = "LocalBackupRestoreState(restorePhase=$restorePhase, backupInfo=$backupInfo, allBackups=$allBackups, selectedFolderUri=$selectedFolderUri, progressFraction=$progressFraction, errorMessage=$errorMessage, launchFolderPicker=$launchFolderPicker, aep=${aep?.displayValue?.censor()}, v1Passphrase=${v1Passphrase?.censor()})"

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

@@ -11,6 +11,7 @@ import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.Job
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.launchIn
@@ -21,6 +22,7 @@ import org.signal.core.models.AccountEntropyPool
 import org.signal.core.ui.navigation.ResultEventBus
 import org.signal.core.util.logging.Log
 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
@@ -30,6 +32,7 @@ import org.signal.registration.screens.util.navigateTo
 
 class LocalBackupRestoreViewModel(
   private val repository: RegistrationRepository,
+  parentState: Flow<RegistrationFlowState>,
   private val parentEventEmitter: (RegistrationFlowEvent) -> Unit,
   private val isPreRegistration: Boolean,
   private val resultBus: ResultEventBus,
@@ -49,6 +52,10 @@ class LocalBackupRestoreViewModel(
     _state
       .onEach { Log.d(TAG, "[State] $it") }
       .launchIn(viewModelScope)
+
+    parentState
+      .onEach { onEvent(LocalBackupRestoreEvents.ParentStateChanged(it)) }
+      .launchIn(viewModelScope)
   }
 
   override suspend fun processEvent(event: LocalBackupRestoreEvents) {
@@ -58,6 +65,9 @@ class LocalBackupRestoreViewModel(
   @VisibleForTesting
   suspend fun applyEvent(state: LocalBackupRestoreState, event: LocalBackupRestoreEvents, stateEmitter: (LocalBackupRestoreState) -> Unit) {
     when (event) {
+      is LocalBackupRestoreEvents.ParentStateChanged -> {
+        stateEmitter(state.copy(storageCapable = event.state.storageCapable))
+      }
       is LocalBackupRestoreEvents.PickBackupFolder -> {
         stateEmitter(state.copy(launchFolderPicker = true))
       }
@@ -123,8 +133,15 @@ class LocalBackupRestoreViewModel(
       parentEventEmitter.navigateBack()
     } else {
       repository.setRestoreDecision(RestoreDecision.COMPLETED)
-      repository.restoreAccountRecord()
-      parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
+
+      if (progress.restoredSvrPin != null) {
+        repository.restoreAccountRecord()
+        parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
+      } else if (state.storageCapable) {
+        parentEventEmitter.navigateTo(RegistrationRoute.PinEntryForSvrRestore)
+      } else {
+        parentEventEmitter.navigateTo(RegistrationRoute.PinCreate)
+      }
     }
   }
 
@@ -210,13 +227,14 @@ class LocalBackupRestoreViewModel(
 
   class Factory(
     private val repository: RegistrationRepository,
+    private val parentState: Flow<RegistrationFlowState>,
     private val parentEventEmitter: (RegistrationFlowEvent) -> Unit,
     private val isPreRegistration: Boolean,
     private val resultBus: ResultEventBus,
     private val resultKey: String
   ) : ViewModelProvider.Factory {
     override fun <T : ViewModel> create(modelClass: Class<T>): T {
-      return LocalBackupRestoreViewModel(repository, parentEventEmitter, isPreRegistration, resultBus, resultKey) as T
+      return LocalBackupRestoreViewModel(repository, parentState, parentEventEmitter, isPreRegistration, resultBus, resultKey) as T
     }
   }
 }

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

@@ -35,6 +35,7 @@ import org.signal.core.ui.navigation.ResultEventBus
 import org.signal.libsignal.protocol.IdentityKeyPair
 import org.signal.libsignal.zkgroup.profiles.ProfileKey
 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
@@ -70,9 +71,10 @@ class LocalBackupRestoreViewModelTest {
     Dispatchers.resetMain()
   }
 
-  private fun createViewModel(isPreRegistration: Boolean): LocalBackupRestoreViewModel {
+  private fun createViewModel(isPreRegistration: Boolean, storageCapable: Boolean = true): LocalBackupRestoreViewModel {
     return LocalBackupRestoreViewModel(
       repository = mockRepository,
+      parentState = flowOf(RegistrationFlowState(storageCapable = storageCapable)),
       parentEventEmitter = parentEventEmitter,
       isPreRegistration = isPreRegistration,
       resultBus = resultBus,
@@ -108,6 +110,19 @@ class LocalBackupRestoreViewModelTest {
     assertThat(emittedStates.last().selectedFolderUri).isEqualTo(folderUri)
   }
 
+  // ==================== ParentStateChanged Tests ====================
+
+  @Test
+  fun `ParentStateChanged copies storageCapable from parent state`() = runTest {
+    val viewModel = createViewModel(isPreRegistration = false)
+    val initialState = LocalBackupRestoreState(storageCapable = true)
+
+    viewModel.applyEvent(initialState, LocalBackupRestoreEvents.ParentStateChanged(RegistrationFlowState(storageCapable = false)), stateEmitter)
+
+    assertThat(emittedStates).hasSize(1)
+    assertThat(emittedStates.last().storageCapable).isEqualTo(false)
+  }
+
   // ==================== RestoreBackup with V1 Tests ====================
 
   @Test
@@ -255,7 +270,7 @@ class LocalBackupRestoreViewModelTest {
   // ==================== Restore Completion Tests ====================
 
   @Test
-  fun `successful V1 restore records COMPLETED restore decision and finishes registration`() = runTest(testDispatcher) {
+  fun `V1 restore that recovers a PIN records COMPLETED restore decision and finishes registration`() = runTest(testDispatcher) {
     val viewModel = createViewModel(isPreRegistration = false)
     val backupInfo = LocalBackupInfo(
       type = LocalBackupInfo.BackupType.V1,
@@ -265,14 +280,62 @@ class LocalBackupRestoreViewModelTest {
     )
     val initialState = LocalBackupRestoreState(backupInfo = backupInfo)
 
+    every { mockRepository.restoreV1Backup(any(), any()) } returns flowOf(LocalBackupRestoreProgress.Complete(restoredSvrPin = "1234", restoredProfileKey = null))
+
+    viewModel.applyEvent(initialState, LocalBackupRestoreEvents.PassphraseSubmitted("passphrase"), stateEmitter)
+
+    coVerify { mockRepository.persistRestoredBackupState("1234", null) }
+    coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) }
+    coVerify { mockRepository.restoreAccountRecord(any()) }
+    assertThat(emittedParentEvents).contains(RegistrationFlowEvent.RegistrationComplete)
+  }
+
+  @Test
+  fun `V1 restore without a PIN when storage capable navigates to PinEntryForSvrRestore`() = runTest(testDispatcher) {
+    val viewModel = createViewModel(isPreRegistration = false, storageCapable = true)
+    val backupInfo = LocalBackupInfo(
+      type = LocalBackupInfo.BackupType.V1,
+      date = LocalDateTime.now(),
+      name = "backup.backup",
+      uri = mockk()
+    )
+    val initialState = LocalBackupRestoreState(backupInfo = backupInfo)
+
     every { mockRepository.restoreV1Backup(any(), any()) } returns flowOf(LocalBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null))
 
     viewModel.applyEvent(initialState, LocalBackupRestoreEvents.PassphraseSubmitted("passphrase"), stateEmitter)
 
-    coVerify { mockRepository.persistRestoredBackupState(null, null) }
     coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) }
-    coVerify { mockRepository.restoreAccountRecord(any()) }
-    assertThat(emittedParentEvents).contains(RegistrationFlowEvent.RegistrationComplete)
+    coVerify(exactly = 0) { mockRepository.restoreAccountRecord(any()) }
+    assertThat(emittedParentEvents).hasSize(1)
+    assertThat(emittedParentEvents.first())
+      .isInstanceOf<RegistrationFlowEvent.NavigateToScreen>()
+      .prop(RegistrationFlowEvent.NavigateToScreen::route)
+      .isEqualTo(RegistrationRoute.PinEntryForSvrRestore)
+  }
+
+  @Test
+  fun `V1 restore without a PIN when not storage capable navigates to PinCreate`() = runTest(testDispatcher) {
+    val viewModel = createViewModel(isPreRegistration = false, storageCapable = false)
+    val backupInfo = LocalBackupInfo(
+      type = LocalBackupInfo.BackupType.V1,
+      date = LocalDateTime.now(),
+      name = "backup.backup",
+      uri = mockk()
+    )
+    val initialState = LocalBackupRestoreState(backupInfo = backupInfo)
+
+    every { mockRepository.restoreV1Backup(any(), any()) } returns flowOf(LocalBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null))
+
+    viewModel.applyEvent(initialState, LocalBackupRestoreEvents.PassphraseSubmitted("passphrase"), stateEmitter)
+
+    coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) }
+    coVerify(exactly = 0) { mockRepository.restoreAccountRecord(any()) }
+    assertThat(emittedParentEvents).hasSize(1)
+    assertThat(emittedParentEvents.first())
+      .isInstanceOf<RegistrationFlowEvent.NavigateToScreen>()
+      .prop(RegistrationFlowEvent.NavigateToScreen::route)
+      .isEqualTo(RegistrationRoute.PinCreate)
   }
 
   // ==================== Incorrect Credential Tests ====================