Restore the AEP in v1 local backups in regV5.

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

@@ -335,7 +335,7 @@ class AppRegistrationStorageController(private val context: Context) : StorageCo
 
         SignalDatabase.runPostBackupRestoreTasks(database)
 
-        trySend(readRestoredLocalBackupState(includeIdentityKeys = true))
+        trySend(readRestoredLocalBackupState(includePreRegistrationKeys = true))
         Log.d(TAG, "V1 restore complete.")
       } catch (e: FullBackupImporter.DatabaseDowngradeException) {
         Log.w(TAG, "V1 restore failed: database downgrade", e)
@@ -414,19 +414,21 @@ class AppRegistrationStorageController(private val context: Context) : StorageCo
     }
   }.flowOn(Dispatchers.IO)
 
-  private fun readRestoredLocalBackupState(includeIdentityKeys: Boolean = false): LocalBackupRestoreProgress.Complete {
+  private fun readRestoredLocalBackupState(includePreRegistrationKeys: 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
+    val restoredAccountEntropyPool = if (includePreRegistrationKeys) SignalStore.account.accountEntropyPoolOrNull else null
+    val restoredAciIdentityKey = if (includePreRegistrationKeys && SignalStore.account.hasAciIdentityKey()) SignalStore.account.aciIdentityKey else null
+    val restoredPniIdentityKey = if (includePreRegistrationKeys && SignalStore.account.hasPniIdentityKey()) SignalStore.account.pniIdentityKey else null
 
     return LocalBackupRestoreProgress.Complete(
       restoredSvrPin = restoredPin,
       restoredProfileKey = restoredProfileKey,
+      restoredAccountEntropyPool = restoredAccountEntropyPool,
       restoredAciIdentityKey = restoredAciIdentityKey,
       restoredPniIdentityKey = restoredPniIdentityKey
     )

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

@@ -130,7 +130,7 @@ class LocalBackupRestoreViewModel(
     if (isPreRegistration) {
       repository.persistRestoredIdentityKeys(progress.restoredAciIdentityKey, progress.restoredPniIdentityKey)
       repository.setRestoreDecision(RestoreDecision.COMPLETED)
-      resultBus.sendResult(resultKey, LocalBackupRestoreResult.Success(state.aep))
+      resultBus.sendResult(resultKey, LocalBackupRestoreResult.Success(state.aep ?: progress.restoredAccountEntropyPool))
       parentEventEmitter.navigateBack()
     } else {
       repository.setRestoreDecision(RestoreDecision.COMPLETED)

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

@@ -31,6 +31,7 @@ import org.junit.After
 import org.junit.Before
 import org.junit.Test
 import org.signal.archive.LocalBackupRestoreProgress
+import org.signal.core.models.AccountEntropyPool
 import org.signal.core.ui.navigation.ResultEventBus
 import org.signal.libsignal.protocol.IdentityKeyPair
 import org.signal.libsignal.zkgroup.profiles.ProfileKey
@@ -414,6 +415,33 @@ class LocalBackupRestoreViewModelTest {
     coVerify { mockRepository.persistRestoredIdentityKeys(aciIdentityKey, pniIdentityKey) }
   }
 
+  @Test
+  fun `pre-registration V1 restore sends restored AEP in the success result`() = 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 restoredAep = AccountEntropyPool(VALID_AEP)
+
+    every { mockRepository.restoreV1Backup(any(), any()) } returns flowOf(
+      LocalBackupRestoreProgress.Complete(
+        restoredSvrPin = null,
+        restoredProfileKey = null,
+        restoredAccountEntropyPool = restoredAep
+      )
+    )
+
+    viewModel.applyEvent(initialState, LocalBackupRestoreEvents.PassphraseSubmitted("passphrase"), stateEmitter)
+
+    val result = resultBus.channelMap[resultKey]?.tryReceive()?.getOrNull()
+    assertThat(result).isNotNull().isEqualTo(LocalBackupRestoreResult.Success(restoredAep))
+  }
+
   companion object {
     private const val VALID_AEP = "uy38jh2778hjjhj8lk19ga61s672jsj089r023s6a57809bap92j2yh5t326vv7t"
   }

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

@@ -5,6 +5,7 @@
 
 package org.signal.archive
 
+import org.signal.core.models.AccountEntropyPool
 import org.signal.libsignal.protocol.IdentityKeyPair
 import org.signal.libsignal.zkgroup.profiles.ProfileKey
 
@@ -32,12 +33,14 @@ sealed interface LocalBackupRestoreProgress {
    * 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.
+   * [restoredAccountEntropyPool], [restoredAciIdentityKey], and [restoredPniIdentityKey] are only populated for V1
+   * backups restored before registration, where we want to preserve the device's existing keys rather than generating
+   * new ones.
    */
   data class Complete(
     val restoredSvrPin: String?,
     val restoredProfileKey: ProfileKey?,
+    val restoredAccountEntropyPool: AccountEntropyPool? = null,
     val restoredAciIdentityKey: IdentityKeyPair? = null,
     val restoredPniIdentityKey: IdentityKeyPair? = null
   ) : LocalBackupRestoreProgress