Ensure we always use the user-entered AEP in regV5.

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

@@ -412,29 +412,23 @@ class AppRegistrationStorageController(private val context: Context) : StorageCo
           is Result.Success -> {
             AppDependencies.jobManager.add(LocalBackupRestoreMediaJob.create(rootUri))
 
-            // Only adopt the entered recovery key as the account's AEP if the backup actually belongs to this account.
-            // Otherwise we'd overwrite the account's real AEP with a foreign backup's key. Messages are still imported.
-            val actualBackupId = LocalArchiver.getBackupId(snapshotFileSystem, messageBackupKey)
-            val expectedBackupId = SignalStore.account.accountEntropyPool.deriveMessageBackupKey().deriveBackupId(selfAci)
-            if (actualBackupId?.value?.contentEquals(expectedBackupId.value) == true) {
-              Log.i(TAG, "V2 local backup belongs to current account; adopting entered recovery key.")
-              SignalStore.account.restoreAccountEntropyPool(aep)
-              updateInProgressRegistrationData { this.accountEntropyPool = aep.value }
+            // The entered recovery key decrypted the backup the user chose to restore, so it always becomes the
+            // account's AEP -- even if the backup was made by a different account.
+            SignalStore.account.restoreAccountEntropyPool(aep)
+            updateInProgressRegistrationData { this.accountEntropyPool = aep.value }
 
-              // Re-enable new-style local backups pointing at the restored location, so the user keeps getting backups.
-              // Skip it if the folder is the SignalBackups directory itself, since it can't be reused as a destination.
-              val archiveFileSystem = ArchiveFileSystem.openForRestore(context, rootUri)
-              if (archiveFileSystem != null && !archiveFileSystem.isRootedAtSignalBackups) {
-                SignalStore.backup.newLocalBackupsDirectory = rootUri.toString()
-                SignalStore.backup.newLocalBackupsEnabled = true
-                LocalBackupListener.setNextBackupTimeToIntervalFromNow(context)
-                LocalBackupListener.schedule(context)
-              } else {
-                Log.w(TAG, "V2 local backup directory can't be reused as a destination; not re-enabling local backups.")
-              }
+            // Re-enable new-style local backups pointing at the restored location, so the user keeps getting backups.
+            // Skip it if the folder is the SignalBackups directory itself, since it can't be reused as a destination.
+            val archiveFileSystem = ArchiveFileSystem.openForRestore(context, rootUri)
+            if (archiveFileSystem != null && !archiveFileSystem.isRootedAtSignalBackups) {
+              SignalStore.backup.newLocalBackupsDirectory = rootUri.toString()
+              SignalStore.backup.newLocalBackupsEnabled = true
+              LocalBackupListener.setNextBackupTimeToIntervalFromNow(context)
+              LocalBackupListener.schedule(context)
             } else {
-              Log.w(TAG, "V2 local backup does not belong to current account; keeping existing recovery key.")
+              Log.w(TAG, "V2 local backup directory can't be reused as a destination; not re-enabling local backups.")
             }
+
             trySend(readRestoredLocalBackupState())
             Log.d(TAG, "V2 restore complete.")
           }

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

@@ -154,6 +154,12 @@ class LocalBackupRestoreViewModel(
       resultBus.sendResult(resultKey, LocalBackupRestoreResult.Success(state.aep ?: progress.restoredAccountEntropyPool))
       parentEventEmitter.navigateBack()
     } else {
+      // The entered key decrypted the restored backup, so it becomes the canonical AEP -- even if it didn't match the
+      // AEP the server had associated with the account. Downstream screens (e.g. PIN creation) must use it too.
+      if (backupType == LocalBackupInfo.BackupType.V2 && state.aep != null) {
+        parentEventEmitter(RegistrationFlowEvent.UserSuppliedAepVerified(state.aep))
+      }
+
       repository.setRestoreDecision(RestoreDecision.COMPLETED)
 
       if (progress.restoredSvrPin != null) {
@@ -231,7 +237,7 @@ class LocalBackupRestoreViewModel(
             storageCapable = currentState.storageCapable
           )
           is LocalBackupRestoreProgress.Complete -> {
-            onRestoreComplete(_state.value.copy(aep = currentState.aep, v1Passphrase = currentState.v1Passphrase, storageCapable = currentState.storageCapable), progress, backup.type)
+            onRestoreComplete(_state.value.copy(aep = aep, v1Passphrase = currentState.v1Passphrase, storageCapable = currentState.storageCapable), progress, backup.type)
             _state.value
           }
           is LocalBackupRestoreProgress.IncorrectCredential -> {

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

@@ -734,8 +734,11 @@ class RegistrationEndToEndTest {
     assert(committed != null) { "Expected registration data to be committed" }
     assert(committed!!.accountData?.e164 == E164) { "Expected committed e164 $E164 but was ${committed.accountData?.e164}" }
     assert(committed.pin == PIN) { "Expected committed pin $PIN but was ${committed.pin}" }
-    assert(committed.accountEntropyPool.isNotEmpty() && committed.accountEntropyPool != aep.value) {
-      "Expected a fresh AEP to be committed rather than the foreign backup's AEP"
+    assert(committed.accountEntropyPool == aep.value) {
+      "Expected the entered recovery key to be adopted as the committed AEP"
+    }
+    assert(networkController.lastSetPinRequest?.masterKey == aep.deriveMasterKey()) {
+      "Expected the new PIN to be backed up to SVR with the master key derived from the adopted AEP"
     }
     assert(storageController.restoreDecision == RestoreDecision.COMPLETED) { "Expected COMPLETED restore decision but was ${storageController.restoreDecision}" }
   }

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

@@ -234,6 +234,50 @@ class LocalBackupRestoreViewModelTest {
       .isEqualTo(RegistrationRoute.PinCreate)
   }
 
+  // ==================== AEP adoption Tests ====================
+
+  @Test
+  fun `V2 restore promotes the entered recovery key to the canonical AEP`() = runTest(testDispatcher) {
+    val viewModel = createViewModel(isPreRegistration = false, storageCapable = false, knownAep = AccountEntropyPool(VALID_AEP))
+    val backupInfo = LocalBackupInfo(
+      type = LocalBackupInfo.BackupType.V2,
+      date = LocalDateTime.now(),
+      name = "signal-backup",
+      uri = mockk()
+    )
+    val initialState = LocalBackupRestoreState(backupInfo = backupInfo, selectedFolderUri = mockk())
+
+    every { mockRepository.restoreV2Backup(any(), any(), any()) } returns flowOf(
+      LocalBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null)
+    )
+
+    viewModel.applyEvent(initialState, LocalBackupRestoreEvents.PassphraseSubmitted(VALID_AEP), stateEmitter)
+
+    val verifiedEvents = emittedParentEvents.filterIsInstance<RegistrationFlowEvent.UserSuppliedAepVerified>()
+    assertThat(verifiedEvents).hasSize(1)
+    assertThat(verifiedEvents.single().aep.value).isEqualTo(VALID_AEP)
+  }
+
+  @Test
+  fun `post-registration V1 restore does not promote any AEP`() = 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, selectedFolderUri = mockk())
+
+    every { mockRepository.restoreV1Backup(any(), any(), any()) } returns flowOf(
+      LocalBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null)
+    )
+
+    viewModel.applyEvent(initialState, LocalBackupRestoreEvents.PassphraseSubmitted("passphrase"), stateEmitter)
+
+    assertThat(emittedParentEvents.filterIsInstance<RegistrationFlowEvent.UserSuppliedAepVerified>()).isEmpty()
+  }
+
   // ==================== RestoreBackup with no backup Tests ====================
 
   @Test