From 8a5826555ce0ca9b7ffd47e9e484163fa9b4c94f Mon Sep 17 00:00:00 2001 From: Greyson Parrelli Date: Wed, 8 Jul 2026 15:14:41 -0400 Subject: [PATCH] Properly restore keys from local backups in regV5. --- .../v2/AppRegistrationStorageController.kt | 23 ++++++++++-- .../dependencies/DemoStorageController.kt | 4 +-- .../registration/RegistrationRepository.kt | 27 +++++++++++--- .../LocalBackupRestoreViewModel.kt | 7 ++-- .../RemoteBackupRestoreViewModel.kt | 2 +- .../LocalBackupRestoreViewModelTest.kt | 35 ++++++++++++++++++- .../archive/LocalBackupRestoreProgress.kt | 21 +++++++++-- 7 files changed, 104 insertions(+), 15 deletions(-) diff --git a/app/src/main/java/org/thoughtcrime/securesms/registration/v2/AppRegistrationStorageController.kt b/app/src/main/java/org/thoughtcrime/securesms/registration/v2/AppRegistrationStorageController.kt index 4a6fa26971..c8c6904ec4 100644 --- a/app/src/main/java/org/thoughtcrime/securesms/registration/v2/AppRegistrationStorageController.kt +++ b/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 = withContext(Dispatchers.IO) { val folder = DocumentFile.fromTreeUri(context, folderUri) ?: return@withContext emptyList() val children = folder.listFiles() diff --git a/demo/registration/src/main/java/org/signal/registration/sample/dependencies/DemoStorageController.kt b/demo/registration/src/main/java/org/signal/registration/sample/dependencies/DemoStorageController.kt index 22af077aa3..0ca5817dae 100644 --- a/demo/registration/src/main/java/org/signal/registration/sample/dependencies/DemoStorageController.kt +++ b/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) diff --git a/feature/registration/src/main/java/org/signal/registration/RegistrationRepository.kt b/feature/registration/src/main/java/org/signal/registration/RegistrationRepository.kt index d5ea55590e..880ab62102 100644 --- a/feature/registration/src/main/java/org/signal/registration/RegistrationRepository.kt +++ b/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 diff --git a/feature/registration/src/main/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreViewModel.kt b/feature/registration/src/main/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreViewModel.kt index 20d5679e63..75aa081338 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreViewModel.kt +++ b/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 -> { diff --git a/feature/registration/src/main/java/org/signal/registration/screens/remotebackuprestore/RemoteBackupRestoreViewModel.kt b/feature/registration/src/main/java/org/signal/registration/screens/remotebackuprestore/RemoteBackupRestoreViewModel.kt index f293b190fd..b86b99baa3 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/remotebackuprestore/RemoteBackupRestoreViewModel.kt +++ b/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) diff --git a/feature/registration/src/test/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreViewModelTest.kt b/feature/registration/src/test/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreViewModelTest.kt index 82cf5c6bcf..f9bcd3fd3d 100644 --- a/feature/registration/src/test/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreViewModelTest.kt +++ b/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" } diff --git a/lib/archive/src/main/java/org/signal/archive/LocalBackupRestoreProgress.kt b/lib/archive/src/main/java/org/signal/archive/LocalBackupRestoreProgress.kt index 6ac6f7676d..16f7d2f9cb 100644 --- a/lib/archive/src/main/java/org/signal/archive/LocalBackupRestoreProgress.kt +++ b/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