Properly restore keys from local backups in regV5.

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

@@ -70,6 +70,7 @@ import org.whispersystems.signalservice.api.link.TransferArchiveResponse
 import java.io.File
 import java.io.IOException
 import java.time.LocalDateTime
+import kotlin.jvm.optionals.getOrNull
 
 /**
  * Implementation of [StorageController] that bridges to the app's existing storage infrastructure.
@@ -311,7 +312,7 @@ class AppRegistrationStorageController(private val context: Context) : StorageCo
 
       SignalDatabase.runPostBackupRestoreTasks(database)
 
-      emit(LocalBackupRestoreProgress.Complete)
+      emit(readRestoredLocalBackupState(includeIdentityKeys = true))
       Log.d(TAG, "V1 restore complete.")
     } catch (e: FullBackupImporter.DatabaseDowngradeException) {
       Log.w(TAG, "V1 restore failed: database downgrade", e)
@@ -369,7 +370,7 @@ class AppRegistrationStorageController(private val context: Context) : StorageCo
           } else {
             Log.w(TAG, "V2 local backup does not belong to current account; keeping existing recovery key.")
           }
-          emit(LocalBackupRestoreProgress.Complete)
+          emit(readRestoredLocalBackupState())
           Log.d(TAG, "V2 restore complete.")
         }
         is Result.Failure -> {
@@ -383,6 +384,24 @@ class AppRegistrationStorageController(private val context: Context) : StorageCo
     }
   }.flowOn(Dispatchers.IO)
 
+  private fun readRestoredLocalBackupState(includeIdentityKeys: Boolean = false): LocalBackupRestoreProgress.Complete {
+    val restoredPin = SignalStore.svr.pin?.takeIf { it.isNotBlank() }
+    val restoredProfileKey = SignalStore.account.aci
+      ?.let { SignalDatabase.recipients.getByAci(it).getOrNull() }
+      ?.let { SignalDatabase.recipients.getRecord(it).profileKey }
+      ?.let { ProfileKey(it) }
+
+    val restoredAciIdentityKey = if (includeIdentityKeys && SignalStore.account.hasAciIdentityKey()) SignalStore.account.aciIdentityKey else null
+    val restoredPniIdentityKey = if (includeIdentityKeys && SignalStore.account.hasPniIdentityKey()) SignalStore.account.pniIdentityKey else null
+
+    return LocalBackupRestoreProgress.Complete(
+      restoredSvrPin = restoredPin,
+      restoredProfileKey = restoredProfileKey,
+      restoredAciIdentityKey = restoredAciIdentityKey,
+      restoredPniIdentityKey = restoredPniIdentityKey
+    )
+  }
+
   override suspend fun scanLocalBackupFolder(folderUri: Uri): List<LocalBackupInfo> = withContext(Dispatchers.IO) {
     val folder = DocumentFile.fromTreeUri(context, folderUri) ?: return@withContext emptyList()
     val children = folder.listFiles()

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

@@ -312,7 +312,7 @@ class DemoStorageController(private val context: Context) : StorageController {
       delay(SIMULATED_STAGE_DELAY_MS)
     }
 
-    emit(LocalBackupRestoreProgress.Complete)
+    emit(LocalBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null))
     Log.d(TAG, "Simulated V1 restore complete.")
   }.flowOn(Dispatchers.IO)
 
@@ -330,7 +330,7 @@ class DemoStorageController(private val context: Context) : StorageController {
       delay(SIMULATED_STAGE_DELAY_MS)
     }
 
-    emit(LocalBackupRestoreProgress.Complete)
+    emit(LocalBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null))
     Log.d(TAG, "Simulated V2 restore complete.")
   }.flowOn(Dispatchers.IO)
 

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

@@ -549,10 +549,16 @@ class RegistrationRepository(val context: Context, val networkController: Networ
 
     Log.i(TAG, "[registerAccount] Starting registration for $e164. sessionId: ${sessionId != null}, recoveryPassword: ${recoveryPassword != null}, registrationLock: ${registrationLock != null}, skipDeviceTransfer: $skipDeviceTransfer, existingAep: ${existingAccountEntropyPool != null}")
 
+    val inProgressData = storageController.readInProgressRegistrationData()
+    val resumedAciIdentityKeyPair = inProgressData.aciIdentityKeyPair.takeIf { it.size > 0 }?.let { IdentityKeyPair(it.toByteArray()) }
+    val resumedPniIdentityKeyPair = inProgressData.pniIdentityKeyPair.takeIf { it.size > 0 }?.let { IdentityKeyPair(it.toByteArray()) }
+    val resumedProfileKey = inProgressData.profileKey.takeIf { it.size > 0 }?.let { ProfileKey(it.toByteArray()) }
+
     val keyMaterial = generateKeyMaterial(
       existingAccountEntropyPool = existingAccountEntropyPool,
-      existingAciIdentityKeyPair = existingAciIdentityKeyPair,
-      existingPniIdentityKeyPair = existingPniIdentityKeyPair
+      existingAciIdentityKeyPair = existingAciIdentityKeyPair ?: resumedAciIdentityKeyPair,
+      existingPniIdentityKeyPair = existingPniIdentityKeyPair ?: resumedPniIdentityKeyPair,
+      profileKey = resumedProfileKey
     )
 
     storageController.updateInProgressRegistrationData {
@@ -706,16 +712,27 @@ class RegistrationRepository(val context: Context, val networkController: Networ
   }
 
   /**
-   * Persist any data in our scratch storage that was restored as part of a remote backup so that we don't accidentally overwrite it
-   * when we commit it.
+   * Persist any data in our scratch storage that was restored as part of a backup (remote or local) so that we don't
+   * accidentally overwrite it when we commit it.
    */
-  suspend fun persistRemoteBackupRestoredState(restoredPin: String?, restoredProfileKey: ProfileKey?) {
+  suspend fun persistRestoredBackupState(restoredPin: String?, restoredProfileKey: ProfileKey?) {
     storageController.updateInProgressRegistrationData {
       pin = restoredPin ?: pin
       profileKey = restoredProfileKey?.serialize()?.toByteString() ?: profileKey
     }
   }
 
+  /**
+   * Persists identity key pairs restored from a pre-registration local backup into our scratch storage, so that the
+   * upcoming registration reuses the device's existing identity rather than generating a fresh one.
+   */
+  suspend fun persistRestoredIdentityKeys(restoredAciIdentityKey: IdentityKeyPair?, restoredPniIdentityKey: IdentityKeyPair?) {
+    storageController.updateInProgressRegistrationData {
+      aciIdentityKeyPair = restoredAciIdentityKey?.serialize()?.toByteString() ?: aciIdentityKeyPair
+      pniIdentityKeyPair = restoredPniIdentityKey?.serialize()?.toByteString() ?: pniIdentityKeyPair
+    }
+  }
+
   /**
    * Records the terminal restore decision the user reached (new account, skipped a restore, or successfully restored)
    * and commits it. The app translates this into its own restore-decision state so the rest of the app knows what

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

@@ -110,8 +110,11 @@ class LocalBackupRestoreViewModel(
     startRestore(backup, state.selectedFolderUri, credential, aep)
   }
 
-  private suspend fun onRestoreComplete(state: LocalBackupRestoreState) {
+  private suspend fun onRestoreComplete(state: LocalBackupRestoreState, progress: LocalBackupRestoreProgress.Complete) {
+    repository.persistRestoredBackupState(progress.restoredSvrPin, progress.restoredProfileKey)
+
     if (isPreRegistration) {
+      repository.persistRestoredIdentityKeys(progress.restoredAciIdentityKey, progress.restoredPniIdentityKey)
       resultBus.sendResult(resultKey, LocalBackupRestoreResult.Success(state.aep))
       parentEventEmitter.navigateBack()
     } else {
@@ -180,7 +183,7 @@ class LocalBackupRestoreViewModel(
             v1Passphrase = currentState.v1Passphrase
           )
           is LocalBackupRestoreProgress.Complete -> {
-            onRestoreComplete(_localState.value.copy(aep = currentState.aep, v1Passphrase = currentState.v1Passphrase))
+            onRestoreComplete(_localState.value.copy(aep = currentState.aep, v1Passphrase = currentState.v1Passphrase), progress)
             _localState.value
           }
           is LocalBackupRestoreProgress.IncorrectCredential -> {

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

@@ -119,7 +119,7 @@ class RemoteBackupRestoreViewModel(
               restoreState = RemoteBackupRestoreState.RestoreState.Restored,
               restoreProgress = null
             )
-            repository.persistRemoteBackupRestoredState(progress.restoredSvrPin, progress.restoredProfileKey)
+            repository.persistRestoredBackupState(progress.restoredSvrPin, progress.restoredProfileKey)
             repository.setRestoreDecision(RestoreDecision.COMPLETED)
             repository.restoreAccountRecord()
             parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)

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

@@ -32,6 +32,8 @@ import org.junit.Before
 import org.junit.Test
 import org.signal.archive.LocalBackupRestoreProgress
 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.RegistrationRepository
 import org.signal.registration.RegistrationRoute
@@ -263,10 +265,11 @@ class LocalBackupRestoreViewModelTest {
     )
     val initialState = LocalBackupRestoreState(backupInfo = backupInfo)
 
-    every { mockRepository.restoreV1Backup(any(), any()) } returns flowOf(LocalBackupRestoreProgress.Complete)
+    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)
@@ -318,6 +321,36 @@ class LocalBackupRestoreViewModelTest {
     coVerify(exactly = 0) { mockRepository.setRestoreDecision(any()) }
   }
 
+  @Test
+  fun `pre-registration V1 restore persists restored backup state and identity keys`() = runTest(testDispatcher) {
+    val viewModel = createViewModel(isPreRegistration = true)
+    val backupInfo = LocalBackupInfo(
+      type = LocalBackupInfo.BackupType.V1,
+      date = LocalDateTime.now(),
+      name = "backup.backup",
+      uri = mockk()
+    )
+    val initialState = LocalBackupRestoreState(backupInfo = backupInfo)
+
+    val profileKey = ProfileKey(ByteArray(32))
+    val aciIdentityKey = IdentityKeyPair.generate()
+    val pniIdentityKey = IdentityKeyPair.generate()
+
+    every { mockRepository.restoreV1Backup(any(), any()) } returns flowOf(
+      LocalBackupRestoreProgress.Complete(
+        restoredSvrPin = "1234",
+        restoredProfileKey = profileKey,
+        restoredAciIdentityKey = aciIdentityKey,
+        restoredPniIdentityKey = pniIdentityKey
+      )
+    )
+
+    viewModel.applyEvent(initialState, LocalBackupRestoreEvents.PassphraseSubmitted("passphrase"), stateEmitter)
+
+    coVerify { mockRepository.persistRestoredBackupState("1234", profileKey) }
+    coVerify { mockRepository.persistRestoredIdentityKeys(aciIdentityKey, pniIdentityKey) }
+  }
+
   companion object {
     private const val VALID_AEP = "uy38jh2778hjjhj8lk19ga61s672jsj089r023s6a57809bap92j2yh5t326vv7t"
   }

lib/archive/src/main/java/org/signal/archive/LocalBackupRestoreProgress.kt

@@ -5,6 +5,9 @@
 
 package org.signal.archive
 
+import org.signal.libsignal.protocol.IdentityKeyPair
+import org.signal.libsignal.zkgroup.profiles.ProfileKey
+
 /**
  * Represents the progress of a local backup restore operation.
  * Emitted as a flow from the storage controller during restore.
@@ -22,8 +25,22 @@ sealed interface LocalBackupRestoreProgress {
       get() = if (totalBytes > 0) (bytesRead.toFloat() / totalBytes.toFloat()).coerceIn(0f, 1f) else 0f
   }
 
-  /** The restore completed successfully. */
-  data object Complete : LocalBackupRestoreProgress
+  /**
+   * The restore completed successfully.
+   * Provides registration-relevant data that was restored so that it isn't accidentally overridden.
+   *
+   * If any of the args are null, we will assume that they were unavailable in the backup, and will defer to
+   * values generated during registration.
+   *
+   * [restoredAciIdentityKey] and [restoredPniIdentityKey] are only populated for V1 backups restored before
+   * registration, where we want to preserve the device's existing identity rather than generating a new one.
+   */
+  data class Complete(
+    val restoredSvrPin: String?,
+    val restoredProfileKey: ProfileKey?,
+    val restoredAciIdentityKey: IdentityKeyPair? = null,
+    val restoredPniIdentityKey: IdentityKeyPair? = null
+  ) : LocalBackupRestoreProgress
 
   /** The provided passphrase (V1) or recovery key (V2) could not decrypt the backup. */
   data object IncorrectCredential : LocalBackupRestoreProgress