From 7a07a96fbaf86429ce80fead848745b689aaf3ba Mon Sep 17 00:00:00 2001 From: Greyson Parrelli Date: Wed, 15 Jul 2026 13:00:08 -0400 Subject: [PATCH] Fix local backup v2 restore pre-registration in regV5. --- .../backup/v2/local/LocalArchiver.kt | 9 +- .../v2/AppRegistrationStorageController.kt | 10 + .../dependencies/DemoStorageController.kt | 2 + .../registration/RegistrationNavigation.kt | 32 +- .../registration/RegistrationRepository.kt | 5 + .../signal/registration/StorageController.kt | 7 + .../screens/aepentry/EnterAepEvents.kt | 6 + .../EnterAepForLocalBackupViewModel.kt | 173 ++++++++++- ...orRemoteBackupPostRegistrationViewModel.kt | 4 + ...ForRemoteBackupPreRegistrationViewModel.kt | 4 + .../screens/aepentry/EnterAepScreen.kt | 20 ++ .../screens/aepentry/EnterAepState.kt | 4 +- .../LocalBackupRestoreEvents.kt | 3 + .../LocalBackupRestoreResult.kt | 7 + .../LocalBackupRestoreViewModel.kt | 21 +- .../phonenumber/PhoneNumberEntryViewModel.kt | 11 +- .../PinEntryForRegistrationLockViewModel.kt | 20 ++ .../VerificationCodeViewModel.kt | 20 ++ .../src/main/res/values/strings.xml | 6 + .../registration/RegistrationEndToEndTest.kt | 154 +++++++++- .../fakes/FakeStorageController.kt | 11 +- .../EnterAepForLocalBackupViewModelTest.kt | 276 ++++++++++++++---- .../LocalBackupRestoreViewModelTest.kt | 84 +++++- ...inEntryForRegistrationLockViewModelTest.kt | 56 ++++ 24 files changed, 855 insertions(+), 90 deletions(-) diff --git a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/local/LocalArchiver.kt b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/local/LocalArchiver.kt index 9a0358010d..690332793d 100644 --- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/local/LocalArchiver.kt +++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/local/LocalArchiver.kt @@ -26,6 +26,7 @@ import org.signal.core.util.toJson import org.signal.libsignal.crypto.Aes256Ctr32 import org.thoughtcrime.securesms.backup.LocalExportProgress import org.thoughtcrime.securesms.backup.v2.BackupRepository +import org.thoughtcrime.securesms.backup.v2.ImportResult import org.thoughtcrime.securesms.database.AttachmentTable import org.thoughtcrime.securesms.database.SignalDatabase import org.thoughtcrime.securesms.keyvalue.SignalStore @@ -269,13 +270,18 @@ object LocalArchiver { val mainStreamLength = snapshotFileSystem.mainLength() ?: return ArchiveResult.failure(RestoreFailure.MainStream) - BackupRepository.importLocal( + val importResult = BackupRepository.importLocal( mainStreamFactory = { snapshotFileSystem.mainInputStream()!! }, mainStreamLength = mainStreamLength, selfData = selfData, backupId = backupId, messageBackupKey = messageBackupKey ) + + if (importResult is ImportResult.Failure) { + Log.w(TAG, "Local backup import failed") + return RestoreResult.failure(RestoreFailure.ImportFailed) + } } finally { metadataStream?.close() } @@ -350,6 +356,7 @@ object LocalArchiver { data object MainStream : RestoreFailure data object Cancelled : RestoreFailure data object BackupIdMissing : RestoreFailure + data object ImportFailed : RestoreFailure data class VersionMismatch(val backupVersion: Int, val supportedVersion: Int) : RestoreFailure } 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 60978ef8f8..ccc659be96 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 @@ -353,6 +353,16 @@ class AppRegistrationStorageController(private val context: Context) : StorageCo } } + override suspend fun verifyLocalBackupKey(backupUri: Uri, aep: AccountEntropyPool): Boolean = withContext(Dispatchers.IO) { + val backupDir = DocumentFile.fromTreeUri(context, backupUri) + if (backupDir == null || !backupDir.canRead()) { + Log.w(TAG, "[verifyLocalBackupKey] Could not open backup directory.") + return@withContext false + } + + LocalArchiver.canDecryptMainArchive(SnapshotFileSystem(context, backupDir), aep.deriveMessageBackupKey()) + } + override fun restoreLocalBackupV2(rootUri: Uri, backupUri: Uri, aep: AccountEntropyPool): Flow = callbackFlow { Log.d(TAG, "Starting V2 local backup restore from backup=$backupUri, root=$rootUri") 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 40fa0f25ba..36aa4419be 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 @@ -318,6 +318,8 @@ class DemoStorageController(private val context: Context) : StorageController { Log.d(TAG, "Simulated V1 restore complete.") }.flowOn(Dispatchers.IO) + override suspend fun verifyLocalBackupKey(backupUri: Uri, aep: AccountEntropyPool): Boolean = true + override fun restoreLocalBackupV2(rootUri: Uri, backupUri: Uri, aep: AccountEntropyPool): Flow = flow { Log.d(TAG, "Starting simulated V2 local backup restore from backup=$backupUri, root=$rootUri") diff --git a/feature/registration/src/main/java/org/signal/registration/RegistrationNavigation.kt b/feature/registration/src/main/java/org/signal/registration/RegistrationNavigation.kt index d6027b7d35..ae069ab277 100644 --- a/feature/registration/src/main/java/org/signal/registration/RegistrationNavigation.kt +++ b/feature/registration/src/main/java/org/signal/registration/RegistrationNavigation.kt @@ -46,6 +46,7 @@ import org.signal.core.util.serialization.AccountEntropyPoolSerializer import org.signal.registration.screens.accountlocked.AccountLockedScreen import org.signal.registration.screens.accountlocked.AccountLockedScreenEvents import org.signal.registration.screens.accountlocked.AccountLockedState +import org.signal.registration.screens.aepentry.EnterAepForLocalBackupResult import org.signal.registration.screens.aepentry.EnterAepForLocalBackupViewModel import org.signal.registration.screens.aepentry.EnterAepForRemoteBackupPostRegistrationViewModel import org.signal.registration.screens.aepentry.EnterAepForRemoteBackupPreRegistrationViewModel @@ -231,8 +232,19 @@ sealed interface RegistrationRoute : NavKey, Parcelable { @Serializable data object EnterLocalBackupV1Passphrase : RegistrationRoute + /** + * Recovery key entry for a local V2 backup. + * + * When [isPreRegistration] is true (pre-registration manual restore), submitting the key first verifies it can + * decrypt the backup at [backupUri], then registers the account via the recovery password derived from it. A backup + * belonging to a different account is surfaced to the user, who can choose to restore it after verifying over SMS. + * When false (already registered), the key is simply handed back to the restore screen to decrypt the backup. + */ @Serializable - data object EnterAepForLocalBackup : RegistrationRoute + data class EnterAepForLocalBackup( + val isPreRegistration: Boolean = false, + val backupUri: String? = null + ) : RegistrationRoute @Serializable data class EnterAepForRemoteBackupPreRegistration(val e164: String) : RegistrationRoute @@ -275,6 +287,7 @@ sealed interface RegistrationRoute : NavKey, Parcelable { private const val CAPTCHA_RESULT = "captcha_token" private const val COUNTRY_CODE_RESULT = "country_code_result" private const val BACKUP_CREDENTIAL_RESULT = "backup_credential_result" +private const val AEP_FOR_LOCAL_BACKUP_RESULT = "aep_for_local_backup_result" private const val LOCAL_BACKUP_RESTORE_RESULT = "local_backup_restore_result" private const val PHONE_NUMBER_DISCOVERABILITY_RESULT = "phone_number_discoverability_result" private const val PIN_LEARN_MORE_URL = "https://support.signal.org/hc/articles/360007059792" @@ -349,7 +362,7 @@ fun RegistrationNavHost( transitionSpec = { TransitionSpecs.HorizontalSlide.transitionSpec }, popTransitionSpec = { when { - initialState.key == RegistrationRoute.EnterAepForLocalBackup.toString() || initialState.key == RegistrationRoute.EnterAepForRemoteBackupPreRegistration.toString() -> { + initialState.key.toString().startsWith("EnterAepForLocalBackup") || initialState.key == RegistrationRoute.EnterAepForRemoteBackupPreRegistration.toString() -> { TransitionSpecs.HorizontalSlide.transitionSpec } @@ -755,6 +768,13 @@ private fun EntryProviderScope.navigationEntries( } } + ResultEffect(registrationViewModel.resultBus, AEP_FOR_LOCAL_BACKUP_RESULT) { result -> + when (result) { + is EnterAepForLocalBackupResult.RestoreReady -> viewModel.onEvent(LocalBackupRestoreEvents.PassphraseSubmitted(result.key)) + is EnterAepForLocalBackupResult.RegistrationDeferredToSms -> viewModel.onEvent(LocalBackupRestoreEvents.RegistrationDeferredToSms) + } + } + LocalBackupRestoreScreen( state = state, onEvent = { viewModel.onEvent(it) } @@ -777,13 +797,17 @@ private fun EntryProviderScope.navigationEntries( // TODO I think we can re-use the screen but attach different viewmodels to progress forward rather than do for-result flows? // -- Enter AEP - entry { + entry { key -> val context = LocalContext.current val viewModel: EnterAepForLocalBackupViewModel = viewModel( factory = EnterAepForLocalBackupViewModel.Factory( + isPreRegistration = key.isPreRegistration, + backupUri = key.backupUri, + repository = registrationRepository, + parentState = registrationViewModel.state, parentEventEmitter = registrationViewModel::onEvent, resultBus = registrationViewModel.resultBus, - resultKey = BACKUP_CREDENTIAL_RESULT, + resultKey = AEP_FOR_LOCAL_BACKUP_RESULT, isPasswordManagerAvailable = RegistrationCredentialManager.isSupported(context) ) ) 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 576a15d512..e9b8eb9a55 100644 --- a/feature/registration/src/main/java/org/signal/registration/RegistrationRepository.kt +++ b/feature/registration/src/main/java/org/signal/registration/RegistrationRepository.kt @@ -906,6 +906,11 @@ class RegistrationRepository(val context: Context, val networkController: Networ return storageController.restoreLocalBackupV2(rootUri, backupUri, aep) } + /** Verifies that [aep] can decrypt the V2 local backup at [backupUri] without restoring anything. */ + suspend fun verifyLocalBackupKey(backupUri: Uri, aep: AccountEntropyPool): Boolean = withContext(Dispatchers.IO) { + storageController.verifyLocalBackupKey(backupUri, aep) + } + suspend fun scanLocalBackupFolder(folderUri: Uri): List = withContext(Dispatchers.IO) { storageController.scanLocalBackupFolder(folderUri) } diff --git a/feature/registration/src/main/java/org/signal/registration/StorageController.kt b/feature/registration/src/main/java/org/signal/registration/StorageController.kt index 17d2ed58f3..6a03afc47b 100644 --- a/feature/registration/src/main/java/org/signal/registration/StorageController.kt +++ b/feature/registration/src/main/java/org/signal/registration/StorageController.kt @@ -125,6 +125,13 @@ interface StorageController { */ fun restoreLocalBackupV2(rootUri: Uri, backupUri: Uri, aep: AccountEntropyPool): Flow + /** + * Verifies that [aep] can decrypt the V2 (folder-based) backup at [backupUri], without restoring anything. + * Used to distinguish a mistyped recovery key from a key that belongs to a different account before + * attempting recovery-password registration. + */ + suspend fun verifyLocalBackupKey(backupUri: Uri, aep: AccountEntropyPool): Boolean + /** * Begins restoring from a remote (server-hosted) backup. * diff --git a/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepEvents.kt b/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepEvents.kt index 422c8206fb..3228793fbe 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepEvents.kt +++ b/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepEvents.kt @@ -21,4 +21,10 @@ sealed class EnterAepEvents { /** Dismiss a registration error dialog. */ data object DismissError : EnterAepEvents() + + /** User confirmed restoring a backup that belongs to a different account, deferring the restore until after SMS verification. */ + data object ConfirmDifferentAccountRestore : EnterAepEvents() + + /** User dismissed the different-account warning dialog without restoring. */ + data object DismissDifferentAccountDialog : EnterAepEvents() } diff --git a/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForLocalBackupViewModel.kt b/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForLocalBackupViewModel.kt index 1c3b6214e4..e064c1770c 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForLocalBackupViewModel.kt +++ b/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForLocalBackupViewModel.kt @@ -5,6 +5,8 @@ package org.signal.registration.screens.aepentry +import androidx.annotation.VisibleForTesting +import androidx.core.net.toUri import androidx.lifecycle.ViewModel import androidx.lifecycle.ViewModelProvider import androidx.lifecycle.viewModelScope @@ -13,14 +15,35 @@ import kotlinx.coroutines.flow.StateFlow import kotlinx.coroutines.flow.asStateFlow import kotlinx.coroutines.flow.launchIn import kotlinx.coroutines.flow.onEach -import kotlinx.coroutines.flow.update +import org.signal.core.models.AccountEntropyPool import org.signal.core.ui.compose.EventDrivenViewModel import org.signal.core.ui.navigation.ResultEventBus import org.signal.core.util.logging.Log +import org.signal.libsignal.net.RequestResult +import org.signal.registration.NetworkController import org.signal.registration.RegistrationFlowEvent +import org.signal.registration.RegistrationFlowState +import org.signal.registration.RegistrationRepository +import org.signal.registration.RegistrationRoute import org.signal.registration.screens.util.navigateBack +import org.signal.registration.screens.util.navigateTo +/** + * Recovery key entry for a local V2 backup restore. + * + * When [isPreRegistration] is true (pre-registration manual restore), submitting the key first verifies it can + * decrypt the selected backup, then attempts recovery-password registration with it -- the restore itself only ever + * runs against a registered account. If the server rejects the recovery password, the key is valid but the backup + * belongs to a different account: the user is warned and can choose to restore it anyway, which defers the import + * until after they verify their number over SMS. + * + * When [isPreRegistration] is false (already registered), the key is simply handed back to the restore screen. + */ class EnterAepForLocalBackupViewModel( + private val isPreRegistration: Boolean, + private val backupUri: String?, + private val repository: RegistrationRepository, + private val parentState: StateFlow, private val parentEventEmitter: (RegistrationFlowEvent) -> Unit, private val resultBus: ResultEventBus, private val resultKey: String, @@ -41,33 +64,167 @@ class EnterAepForLocalBackupViewModel( } override suspend fun processEvent(event: EnterAepEvents) { + applyEvent(_state.value, event) { _state.value = it } + } + + @VisibleForTesting + suspend fun applyEvent(inputState: EnterAepState, event: EnterAepEvents, stateEmitter: (EnterAepState) -> Unit) { when (event) { is EnterAepEvents.BackupKeyChanged -> { - _state.update { EnterAepScreenEventHandler.applyEvent(it, event) } + stateEmitter(EnterAepScreenEventHandler.applyEvent(inputState, event)) } is EnterAepEvents.Submit -> { - if (_state.value.isBackupKeyValid) { - resultBus.sendResult(resultKey, _state.value.backupKey) - parentEventEmitter.navigateBack() - } + applySubmit(inputState, stateEmitter) } is EnterAepEvents.Cancel -> { parentEventEmitter.navigateBack() } is EnterAepEvents.DismissError -> { - _state.update { EnterAepScreenEventHandler.applyEvent(it, event) } + stateEmitter(EnterAepScreenEventHandler.applyEvent(inputState, event)) + } + is EnterAepEvents.ConfirmDifferentAccountRestore -> { + applyConfirmDifferentAccountRestore(inputState, stateEmitter) + } + is EnterAepEvents.DismissDifferentAccountDialog -> { + stateEmitter(inputState.copy(showDifferentAccountDialog = false)) } } } + private suspend fun applySubmit(inputState: EnterAepState, stateEmitter: (EnterAepState) -> Unit) { + check(inputState.isBackupKeyValid) { "AEP is not valid, should not have gotten here." } + + if (!isPreRegistration) { + resultBus.sendResult(resultKey, EnterAepForLocalBackupResult.RestoreReady(inputState.backupKey)) + parentEventEmitter.navigateBack() + return + } + + val aep = AccountEntropyPool(inputState.backupKey) + + stateEmitter(inputState.copy(isRegistering = true)) + + // Confirm the key actually decrypts the backup before going to the server, so a recovery-password rejection can + // only mean the backup belongs to a different account rather than a mistyped key. + if (!repository.verifyLocalBackupKey(checkNotNull(backupUri).toUri(), aep)) { + Log.w(TAG, "[Submit] Entered key cannot decrypt the selected backup.") + stateEmitter(inputState.copy(isRegistering = false, aepValidationError = AepValidationError.Incorrect)) + return + } + + parentEventEmitter(RegistrationFlowEvent.UserSuppliedAepSubmitted(aep)) + + Log.i(TAG, "[Submit] Attempting registration with RRP derived from user-supplied AEP.") + + attemptToRegister(inputState, aep, provideRegistrationLock = false, stateEmitter) + } + + private suspend fun attemptToRegister(inputState: EnterAepState, aep: AccountEntropyPool, provideRegistrationLock: Boolean, stateEmitter: (EnterAepState) -> Unit) { + val e164 = checkNotNull(parentState.value.sessionE164) { "No e164 present in the flow state, should not have gotten here." } + val masterKey = aep.deriveMasterKey() + val recoveryPassword = masterKey.deriveRegistrationRecoveryPassword() + val registrationLock = masterKey.deriveRegistrationLock().takeIf { provideRegistrationLock } + + when (val result = repository.registerAccountWithRecoveryPassword(e164, recoveryPassword, registrationLock, existingAccountEntropyPool = aep)) { + is RequestResult.Success -> { + Log.i(TAG, "[Submit] Successfully registered using RRP from user-supplied AEP. Proceeding with the restore.") + val (response, keyMaterial) = result.result + + stateEmitter(inputState.copy(isRegistering = false)) + parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool, response.storageCapable)) + resultBus.sendResult(resultKey, EnterAepForLocalBackupResult.RestoreReady(inputState.backupKey)) + parentEventEmitter.navigateBack() + } + is RequestResult.NonSuccess -> { + when (val error = result.error) { + is NetworkController.RegisterAccountError.RegistrationRecoveryPasswordIncorrect -> { + Log.w(TAG, "[Submit] RRP incorrect, but the key decrypts the backup. The backup belongs to a different account. Message: ${error.message}") + stateEmitter(inputState.copy(isRegistering = false, showDifferentAccountDialog = true)) + } + is NetworkController.RegisterAccountError.InvalidRequest -> { + Log.w(TAG, "[Submit] Invalid request. Message: ${error.message}") + stateEmitter( + inputState.copy( + isRegistering = false, + registrationError = RegistrationError.UnknownError + ) + ) + } + is NetworkController.RegisterAccountError.RegistrationLock -> { + if (provideRegistrationLock) { + Log.w(TAG, "[Submit] Still registration locked after providing the reglock token derived from the AEP. Falling back to PIN entry.") + stateEmitter(inputState.copy(isRegistering = false)) + parentEventEmitter.navigateTo( + RegistrationRoute.PinEntryForRegistrationLock( + timeRemaining = error.data.timeRemaining, + svrCredentials = error.data.svr2Credentials + ) + ) + } else { + Log.w(TAG, "[Submit] Registration locked. Retrying with the reglock token derived from the AEP.") + attemptToRegister(inputState, aep, provideRegistrationLock = true, stateEmitter) + } + } + is NetworkController.RegisterAccountError.RateLimited -> { + Log.w(TAG, "[Submit] Rate limited (retryAfter: ${error.retryAfter}).") + stateEmitter(inputState.copy(isRegistering = false, registrationError = RegistrationError.RateLimited)) + } + is NetworkController.RegisterAccountError.SessionNotFoundOrNotVerified -> { + error("[Submit] Session not found or not verified. This should not happen with RRP-based registration.") + } + is NetworkController.RegisterAccountError.DeviceTransferPossible -> { + error("[Submit] Device transfer possible. This should not happen with RRP-based registration.") + } + } + } + is RequestResult.RetryableNetworkError -> { + Log.w(TAG, "[Submit] Network error.", result.networkError) + stateEmitter(inputState.copy(isRegistering = false, registrationError = RegistrationError.NetworkError)) + } + is RequestResult.ApplicationError -> { + Log.w(TAG, "[Submit] Application error.", result.cause) + stateEmitter(inputState.copy(isRegistering = false, registrationError = RegistrationError.UnknownError)) + } + } + } + + /** + * The user chose to restore the different-account backup anyway. Force the session/SMS path (the recovery password + * won't work), then hand control back so the flow verifies the number over SMS. The restore is resumed once + * registration completes, keyed off the still-set [RegistrationFlowState.pendingRestoreOption] and the entered + * [RegistrationFlowState.unverifiedRestoredAep]. + */ + private fun applyConfirmDifferentAccountRestore(inputState: EnterAepState, stateEmitter: (EnterAepState) -> Unit) { + Log.i(TAG, "[ConfirmDifferentAccountRestore] Deferring the restore until after SMS verification.") + + stateEmitter(inputState.copy(showDifferentAccountDialog = false)) + + parentEventEmitter(RegistrationFlowEvent.RecoveryPasswordInvalid) + resultBus.sendResult(resultKey, EnterAepForLocalBackupResult.RegistrationDeferredToSms) + parentEventEmitter.navigateBack() + } + class Factory( + private val isPreRegistration: Boolean, + private val backupUri: String?, + private val repository: RegistrationRepository, + private val parentState: StateFlow, private val parentEventEmitter: (RegistrationFlowEvent) -> Unit, private val resultBus: ResultEventBus, private val resultKey: String, private val isPasswordManagerAvailable: Boolean = false ) : ViewModelProvider.Factory { override fun create(modelClass: Class): T { - return EnterAepForLocalBackupViewModel(parentEventEmitter, resultBus, resultKey, isPasswordManagerAvailable) as T + return EnterAepForLocalBackupViewModel(isPreRegistration, backupUri, repository, parentState, parentEventEmitter, resultBus, resultKey, isPasswordManagerAvailable) as T } } } + +/** Result sent back to the local backup restore screen from [EnterAepForLocalBackupViewModel]. */ +sealed interface EnterAepForLocalBackupResult { + /** The account is registered (either it already was, or RRP registration just succeeded) and the backup can be restored with [key]. */ + data class RestoreReady(val key: String) : EnterAepForLocalBackupResult + + /** The backup belongs to a different account and the user chose to restore it anyway. Registration must happen over SMS first. */ + data object RegistrationDeferredToSms : EnterAepForLocalBackupResult +} diff --git a/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForRemoteBackupPostRegistrationViewModel.kt b/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForRemoteBackupPostRegistrationViewModel.kt index df641edd3b..5d7cde2a33 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForRemoteBackupPostRegistrationViewModel.kt +++ b/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForRemoteBackupPostRegistrationViewModel.kt @@ -53,6 +53,10 @@ class EnterAepForRemoteBackupPostRegistrationViewModel( is EnterAepEvents.DismissError -> { stateEmitter(EnterAepScreenEventHandler.applyEvent(inputState, event)) } + is EnterAepEvents.ConfirmDifferentAccountRestore, + is EnterAepEvents.DismissDifferentAccountDialog -> { + error("Different-account handling only exists for local backup restores.") + } } } diff --git a/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForRemoteBackupPreRegistrationViewModel.kt b/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForRemoteBackupPreRegistrationViewModel.kt index def6b46bcb..d4203911e5 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForRemoteBackupPreRegistrationViewModel.kt +++ b/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForRemoteBackupPreRegistrationViewModel.kt @@ -63,6 +63,10 @@ class EnterAepForRemoteBackupPreRegistrationViewModel( is EnterAepEvents.DismissError -> { stateEmitter(EnterAepScreenEventHandler.applyEvent(inputState, event)) } + is EnterAepEvents.ConfirmDifferentAccountRestore, + is EnterAepEvents.DismissDifferentAccountDialog -> { + error("Different-account handling only exists for local backup restores.") + } } } diff --git a/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepScreen.kt b/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepScreen.kt index e0480da802..3f0d653ad1 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepScreen.kt +++ b/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepScreen.kt @@ -74,12 +74,32 @@ fun EnterAepScreen( ) { RegistrationErrorDialog(state.registrationError, onEvent) + if (state.showDifferentAccountDialog) { + DifferentAccountDialog(onEvent) + } + when (val layoutParams = RegistrationScaffold.rememberLayoutParams()) { is RegistrationScaffold.Params.OnePane -> OnePaneLayout(layoutParams, state, onEvent, modifier) is RegistrationScaffold.Params.TwoPane -> TwoPaneLayout(layoutParams, state, onEvent, modifier) } } +/** + * Warns that the entered key decrypts the backup but the backup was created by a different account, offering to + * restore it anyway (which requires verifying the phone number over SMS first). + */ +@Composable +private fun DifferentAccountDialog(onEvent: (EnterAepEvents) -> Unit) { + Dialogs.SimpleAlertDialog( + title = stringResource(R.string.EnterAepScreen__restore_to_new_account), + body = stringResource(R.string.EnterAepScreen__restore_to_new_account_body), + confirm = stringResource(R.string.EnterAepScreen__restore), + dismiss = stringResource(android.R.string.cancel), + onConfirm = { onEvent(EnterAepEvents.ConfirmDifferentAccountRestore) }, + onDismiss = { onEvent(EnterAepEvents.DismissDifferentAccountDialog) } + ) +} + /** * Shows a dismissable dialog for generic registration errors (network/rate-limit/unknown). Incorrect-key errors are * surfaced inline on the text field instead, so they are intentionally not shown here. diff --git a/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepState.kt b/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepState.kt index 1526c27774..088b1654bb 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepState.kt +++ b/feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepState.kt @@ -17,10 +17,12 @@ data class EnterAepState( val chunkLength: Int = 4, val isRegistering: Boolean = false, val registrationError: RegistrationError? = null, + /** The entered key decrypts the backup, but the backup belongs to a different account. Asks whether to restore it anyway after SMS verification. */ + val showDifferentAccountDialog: Boolean = false, /** Whether a password manager / credential provider is available to fill the recovery key. */ val isPasswordManagerAvailable: Boolean = false ) { - override fun toString(): String = "EnterAepState(enteredText=${enteredText.censor()}, backupKey=${backupKey.censor()}, isBackupKeyValid=$isBackupKeyValid, aepValidationError=$aepValidationError, chunkLength=$chunkLength, isRegistering=$isRegistering, registrationError=$registrationError, isPasswordManagerAvailable=$isPasswordManagerAvailable)" + override fun toString(): String = "EnterAepState(enteredText=${enteredText.censor()}, backupKey=${backupKey.censor()}, isBackupKeyValid=$isBackupKeyValid, aepValidationError=$aepValidationError, chunkLength=$chunkLength, isRegistering=$isRegistering, registrationError=$registrationError, showDifferentAccountDialog=$showDifferentAccountDialog, isPasswordManagerAvailable=$isPasswordManagerAvailable)" } sealed interface AepValidationError { diff --git a/feature/registration/src/main/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreEvents.kt b/feature/registration/src/main/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreEvents.kt index 2f2f7727b9..addcecbb1f 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreEvents.kt +++ b/feature/registration/src/main/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreEvents.kt @@ -34,6 +34,9 @@ sealed interface LocalBackupRestoreEvents { override fun toString(): String = "PassphraseSubmitted(credential=${credential.censor()})" } + /** The backup belongs to a different account and the user chose to restore it after verifying over SMS. Hands control back to the phone number screen. */ + data object RegistrationDeferredToSms : LocalBackupRestoreEvents + /** The folder picker was dismissed without selecting a folder. */ data object FolderPickerDismissed : LocalBackupRestoreEvents diff --git a/feature/registration/src/main/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreResult.kt b/feature/registration/src/main/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreResult.kt index cd59242726..3428ff7179 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreResult.kt +++ b/feature/registration/src/main/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreResult.kt @@ -6,6 +6,7 @@ package org.signal.registration.screens.localbackuprestore import org.signal.core.models.AccountEntropyPool +import org.signal.registration.RegistrationFlowState /** * Result communicated back from the pre-registration local backup restore flow @@ -15,6 +16,12 @@ sealed interface LocalBackupRestoreResult { /** The restore completed successfully. Contains the AEP if V2 backup, or the restored AEP for V1. */ data class Success(val aep: AccountEntropyPool?) : LocalBackupRestoreResult + /** + * The backup belongs to a different account and the user chose to restore it anyway. The phone number must be + * verified over SMS, after which the restore recorded in [RegistrationFlowState.pendingLocalBackupRestore] runs. + */ + data object DeferredToSms : LocalBackupRestoreResult + /** The user canceled the restore flow. The pending restore option should be cleared. */ data object Canceled : LocalBackupRestoreResult } 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 1c55084569..c888834cf7 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 @@ -81,6 +81,10 @@ class LocalBackupRestoreViewModel( is LocalBackupRestoreEvents.PassphraseSubmitted -> { applyPassphraseSubmitted(state, event.credential, stateEmitter) } + is LocalBackupRestoreEvents.RegistrationDeferredToSms -> { + resultBus.sendResult(resultKey, LocalBackupRestoreResult.DeferredToSms) + parentEventEmitter.navigateBack() + } is LocalBackupRestoreEvents.ChooseDifferentFolder -> { stateEmitter(LocalBackupRestoreState(launchFolderPicker = true, storageCapable = state.storageCapable)) } @@ -116,7 +120,10 @@ class LocalBackupRestoreViewModel( val credentialRoute = when (backup.type) { LocalBackupInfo.BackupType.V1 -> RegistrationRoute.EnterLocalBackupV1Passphrase - LocalBackupInfo.BackupType.V2 -> RegistrationRoute.EnterAepForLocalBackup + LocalBackupInfo.BackupType.V2 -> RegistrationRoute.EnterAepForLocalBackup( + isPreRegistration = isPreRegistration, + backupUri = backup.uri.toString() + ) } parentEventEmitter.navigateTo(credentialRoute) } @@ -132,10 +139,16 @@ class LocalBackupRestoreViewModel( startRestore(backup, state.selectedFolderUri, credential, aep) } - private suspend fun onRestoreComplete(state: LocalBackupRestoreState, progress: LocalBackupRestoreProgress.Complete) { + private suspend fun onRestoreComplete(state: LocalBackupRestoreState, progress: LocalBackupRestoreProgress.Complete, backupType: LocalBackupInfo.BackupType) { repository.persistRestoredBackupState(progress.restoredSvrPin, progress.restoredProfileKey) - if (isPreRegistration) { + // The restore ran, so clear the pending-restore signal. Otherwise a post-registration screen (verification code or + // reglock PIN) would see it still set and route back here to restore again. + parentEventEmitter(RegistrationFlowEvent.PendingRestoreOptionSelected(null)) + + // V1 backups restore before registration, then the phone number screen registers with the restored data. + // V2 backups only restore against a registered account, so completion always continues the post-registration flow. + if (isPreRegistration && backupType == LocalBackupInfo.BackupType.V1) { repository.persistRestoredIdentityKeys(progress.restoredAciIdentityKey, progress.restoredPniIdentityKey) repository.setRestoreDecision(RestoreDecision.COMPLETED) resultBus.sendResult(resultKey, LocalBackupRestoreResult.Success(state.aep ?: progress.restoredAccountEntropyPool)) @@ -218,7 +231,7 @@ class LocalBackupRestoreViewModel( storageCapable = currentState.storageCapable ) is LocalBackupRestoreProgress.Complete -> { - onRestoreComplete(_state.value.copy(aep = currentState.aep, v1Passphrase = currentState.v1Passphrase, storageCapable = currentState.storageCapable), progress) + onRestoreComplete(_state.value.copy(aep = currentState.aep, v1Passphrase = currentState.v1Passphrase, storageCapable = currentState.storageCapable), progress, backup.type) _state.value } is LocalBackupRestoreProgress.IncorrectCredential -> { diff --git a/feature/registration/src/main/java/org/signal/registration/screens/phonenumber/PhoneNumberEntryViewModel.kt b/feature/registration/src/main/java/org/signal/registration/screens/phonenumber/PhoneNumberEntryViewModel.kt index 3d563fb255..d51aee6bc9 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/phonenumber/PhoneNumberEntryViewModel.kt +++ b/feature/registration/src/main/java/org/signal/registration/screens/phonenumber/PhoneNumberEntryViewModel.kt @@ -142,6 +142,13 @@ class PhoneNumberEntryViewModel( localState = applyLocalBackupRestoreCompleted(localState, event.result.aep, parentEventEmitter) stateEmitter(localState.copy(showSpinner = false)) } + is LocalBackupRestoreResult.DeferredToSms -> { + Log.i(TAG, "[LocalRestore] Backup belongs to a different account. Verifying the number over SMS before restoring.") + var localState = state.copy(showSpinner = true) + stateEmitter(localState) + localState = applySessionBasedRegistration(localState, localState.sessionE164 ?: "+${localState.countryCode}${localState.nationalNumber}", parentEventEmitter) + stateEmitter(localState.copy(showSpinner = false)) + } is LocalBackupRestoreResult.Canceled -> { parentEventEmitter(RegistrationFlowEvent.PendingRestoreOptionSelected(null)) } @@ -344,8 +351,8 @@ class PhoneNumberEntryViewModel( } /** - * Handles the result of a pre-registration local backup restore. - * If an AEP was obtained (V2 backup), attempts RRP-based registration. + * Handles the result of a pre-registration V1 local backup restore (V2 backups register before restoring instead). + * If the restored database contained an AEP, attempts RRP-based registration with it. * Falls back to SVR check and SMS verification if RRP fails or no AEP is available. */ private suspend fun applyLocalBackupRestoreCompleted( diff --git a/feature/registration/src/main/java/org/signal/registration/screens/pinentry/PinEntryForRegistrationLockViewModel.kt b/feature/registration/src/main/java/org/signal/registration/screens/pinentry/PinEntryForRegistrationLockViewModel.kt index 914a2e547b..ad8e81f15a 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/pinentry/PinEntryForRegistrationLockViewModel.kt +++ b/feature/registration/src/main/java/org/signal/registration/screens/pinentry/PinEntryForRegistrationLockViewModel.kt @@ -19,6 +19,7 @@ import org.signal.core.ui.compose.EventDrivenViewModel import org.signal.core.util.logging.Log import org.signal.libsignal.net.RequestResult import org.signal.registration.NetworkController +import org.signal.registration.PendingRestoreOption import org.signal.registration.RegistrationFlowEvent import org.signal.registration.RegistrationFlowState import org.signal.registration.RegistrationRepository @@ -163,7 +164,12 @@ class PinEntryForRegistrationLockViewModel( parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool, response.storageCapable)) repository.enqueueSvrResetGuessCountJob() repository.restoreAccountRecord() + val pendingRestore = pendingRestoreNavigation() when { + pendingRestore != null -> { + Log.i(TAG, "[PinEntered] A restore was pending behind the registration lock. Resuming it now.") + parentEventEmitter.navigateTo(pendingRestore) + } response.reregistration && parentState.value.pendingRestoreOption == null -> parentEventEmitter.navigateTo(RegistrationRoute.ArchiveRestoreSelection.forPostRegisterWithPinKnown()) else -> parentEventEmitter(RegistrationFlowEvent.RegistrationComplete) } @@ -212,6 +218,20 @@ class PinEntryForRegistrationLockViewModel( } } + /** + * If the user pre-selected a restore (see [RegistrationFlowState.pendingRestoreOption]) and it hasn't run yet, + * returns the restore screen to resume it now that the account is registered; otherwise null. Used to pick a + * restore back up when it was blocked behind a registration lock that had to be cleared with the PIN first. + */ + private fun pendingRestoreNavigation(): RegistrationRoute? { + val aep = parentState.value.unverifiedRestoredAep ?: return null + return when (parentState.value.pendingRestoreOption) { + PendingRestoreOption.LocalBackup -> RegistrationRoute.LocalBackupRestore(isPreRegistration = false, aep = aep) + PendingRestoreOption.RemoteBackup -> RegistrationRoute.RemoteRestore(aep) + null -> null + } + } + class Factory( private val repository: RegistrationRepository, private val parentState: StateFlow, diff --git a/feature/registration/src/main/java/org/signal/registration/screens/verificationcode/VerificationCodeViewModel.kt b/feature/registration/src/main/java/org/signal/registration/screens/verificationcode/VerificationCodeViewModel.kt index b7d324f0dd..af53c2bf8c 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/verificationcode/VerificationCodeViewModel.kt +++ b/feature/registration/src/main/java/org/signal/registration/screens/verificationcode/VerificationCodeViewModel.kt @@ -31,6 +31,7 @@ import org.signal.core.ui.compose.EventDrivenViewModel import org.signal.core.util.logging.Log import org.signal.libsignal.net.RequestResult import org.signal.registration.NetworkController +import org.signal.registration.PendingRestoreOption import org.signal.registration.RegistrationFlowEvent import org.signal.registration.RegistrationFlowState import org.signal.registration.RegistrationRepository @@ -364,7 +365,12 @@ class VerificationCodeViewModel( parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool, response.storageCapable)) + val pendingRestore = pendingRestoreNavigation() when { + pendingRestore != null -> { + Log.i(TAG, "[Register] A restore was deferred until after SMS verification. Resuming it now.") + parentEventEmitter.navigateTo(pendingRestore) + } response.reregistration && parentState.value.pendingRestoreOption == null -> parentEventEmitter.navigateTo(RegistrationRoute.ArchiveRestoreSelection.forPostRegisterWithPinUnknown()) response.storageCapable -> parentEventEmitter.navigateTo(RegistrationRoute.PinEntryForSvrRestore) else -> parentEventEmitter.navigateTo(RegistrationRoute.PinCreate) @@ -415,6 +421,20 @@ class VerificationCodeViewModel( } } + /** + * If the user pre-selected a restore (see [RegistrationFlowState.pendingRestoreOption]) and it hasn't run yet, + * returns the restore screen to resume it now that the account is registered; otherwise null. Used to pick a + * restore back up after it was deferred to SMS verification (e.g. a local backup that belongs to a different account). + */ + private fun pendingRestoreNavigation(): RegistrationRoute? { + val aep = parentState.value.unverifiedRestoredAep ?: return null + return when (parentState.value.pendingRestoreOption) { + PendingRestoreOption.LocalBackup -> RegistrationRoute.LocalBackupRestore(isPreRegistration = false, aep = aep) + PendingRestoreOption.RemoteBackup -> RegistrationRoute.RemoteRestore(aep) + null -> null + } + } + private suspend fun applyResendCode( state: VerificationCodeState, transport: NetworkController.VerificationCodeTransport diff --git a/feature/registration/src/main/res/values/strings.xml b/feature/registration/src/main/res/values/strings.xml index 5d834531a7..76cd8ec543 100644 --- a/feature/registration/src/main/res/values/strings.xml +++ b/feature/registration/src/main/res/values/strings.xml @@ -228,6 +228,12 @@ Incorrect recovery key Fill from password manager + + Restore to new account? + + The backup you are restoring was created using a different Signal account. You can restore this backup to a new account, but you will no longer be a member of any groups that were restored. + + Restore Preparing restore… diff --git a/feature/registration/src/test/java/org/signal/registration/RegistrationEndToEndTest.kt b/feature/registration/src/test/java/org/signal/registration/RegistrationEndToEndTest.kt index 13048eeab7..7deda5a889 100644 --- a/feature/registration/src/test/java/org/signal/registration/RegistrationEndToEndTest.kt +++ b/feature/registration/src/test/java/org/signal/registration/RegistrationEndToEndTest.kt @@ -233,7 +233,7 @@ class RegistrationEndToEndTest { } @Test - fun `restoring a remote backup for a reglocked account whose reglock is not derived from the aep falls back to pin entry and registers without a session`() { + fun `restoring a remote backup for a reglocked account whose reglock is not derived from the aep falls back to pin entry then resumes the restore`() { val aep = AccountEntropyPool.generate() val svrMasterKey = MasterKey(ByteArray(32) { it.toByte() }) @@ -261,6 +261,11 @@ class RegistrationEndToEndTest { } } + // The backup contains the user's PIN, so no PIN screens are needed after the restore + storageController.onRestoreRemoteBackup = { + flowOf(RemoteBackupRestoreProgress.Complete(restoredSvrPin = PIN, restoredProfileKey = null)) + } + var registrationComplete = false launchRegistrationFlow(onRegistrationComplete = { registrationComplete = true }) @@ -274,6 +279,9 @@ class RegistrationEndToEndTest { composeTestRule.onNodeWithTag(TestTags.PIN_ENTRY_INPUT).performTextInput(PIN) composeTestRule.onNodeWithTag(TestTags.PIN_ENTRY_CONTINUE_BUTTON).performClick() + // With the reglock cleared, the remote restore the user chose resumes rather than being dropped + startRemoteRestore() + waitFor("registration to complete") { registrationComplete } assert(networkController.lastRestoreMasterKeyRequest?.pin == PIN) { "Expected master key restore with pin $PIN but was ${networkController.lastRestoreMasterKeyRequest}" } @@ -284,6 +292,7 @@ class RegistrationEndToEndTest { val committed = storageController.committedData assert(committed != null) { "Expected registration data to be committed" } assert(committed!!.accountData?.e164 == E164) { "Expected committed e164 $E164 but was ${committed.accountData?.e164}" } + assert(storageController.restoreDecision == RestoreDecision.COMPLETED) { "Expected COMPLETED restore decision (the remote backup was restored) but was ${storageController.restoreDecision}" } } @Test @@ -670,11 +679,13 @@ class RegistrationEndToEndTest { } @Test - fun `restoring a local backup whose aep the server rejects falls back to sms verification`() { + fun `restoring a local backup for a different account warns the user, verifies over sms, then imports the backup`() { val aep = AccountEntropyPool.generate() // The AEP decrypts the backup fine, but it doesn't belong to the account for this phone number + val registerRequests = mutableListOf() networkController.onRegisterAccount = { request -> + registerRequests += request if (request.recoveryPassword != null) { RequestResult.NonSuccess(RegisterAccountError.RegistrationRecoveryPasswordIncorrect("wrong recovery password")) } else { @@ -682,6 +693,14 @@ class RegistrationEndToEndTest { } } + // The import only ever runs against a registered account, so it must come after SMS verification + var registerAttemptsWhenRestoreRan = -1 + val defaultLocalRestore = storageController.onRestoreLocalBackupV2 + storageController.onRestoreLocalBackupV2 = { backupUri, restoreAep -> + registerAttemptsWhenRestoreRan = registerRequests.size + defaultLocalRestore(backupUri, restoreAep) + } + var registrationComplete = false launchRegistrationFlow(folderPickerResult = backupFolderUri, onRegistrationComplete = { registrationComplete = true }) @@ -690,19 +709,146 @@ class RegistrationEndToEndTest { enterPhoneNumber() restoreLocalBackup(aep) - // The backup was restored locally, but recovery-password registration was rejected, so the flow - // falls back to verifying the phone number over SMS + // The key decrypts the backup but the recovery password is rejected, meaning the backup belongs to a different + // account. The user is warned and confirms restoring it to this account anyway. + waitForTag(Dialogs.TEST_TAG_ALERT_DIALOG_CONFIRM_BUTTON) + composeTestRule.onNodeWithTag(Dialogs.TEST_TAG_ALERT_DIALOG_CONFIRM_BUTTON).performClick() + + // Confirming requires verifying the phone number over SMS. Afterwards the restore resumes on the standard + // post-registration local restore screen, where the folder is re-selected (the entered key is reused). submitVerificationCode(VERIFICATION_CODE) + restoreFoundLocalBackup() createPin(PIN) waitFor("registration to complete") { registrationComplete } + assert(registerRequests.first().recoveryPassword == aep.deriveMasterKey().deriveRegistrationRecoveryPassword()) { + "Expected the first registration attempt to use the recovery password derived from the backup's AEP" + } assert(networkController.lastRegisterAccountRequest?.sessionId != null) { "Expected the final registration to use a verified session" } + assert(registerAttemptsWhenRestoreRan == registerRequests.size) { + "Expected the backup to be imported only after registration completed, but the import ran after $registerAttemptsWhenRestoreRan of ${registerRequests.size} registration attempts" + } val committed = storageController.committedData 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(storageController.restoreDecision == RestoreDecision.COMPLETED) { "Expected COMPLETED restore decision but was ${storageController.restoreDecision}" } + } + + @Test + fun `declining to restore a local backup for a different account returns to the recovery key entry screen`() { + val aep = AccountEntropyPool.generate() + + networkController.onRegisterAccount = { request -> + if (request.recoveryPassword != null) { + RequestResult.NonSuccess(RegisterAccountError.RegistrationRecoveryPasswordIncorrect("wrong recovery password")) + } else { + RequestResult.Success(networkController.registerAccountResponse(request.e164)) + } + } + + launchRegistrationFlow(folderPickerResult = backupFolderUri) + + startManualRestore() + chooseRestoreOption(TestTags.ARCHIVE_RESTORE_SELECTION_FROM_BACKUP_FOLDER) + enterPhoneNumber() + restoreLocalBackup(aep) + + // The user is warned that the backup belongs to a different account and declines + waitForTag(Dialogs.TEST_TAG_ALERT_DIALOG_DISMISS_BUTTON) + composeTestRule.onNodeWithTag(Dialogs.TEST_TAG_ALERT_DIALOG_DISMISS_BUTTON).performClick() + + // The user stays on the recovery key entry screen and nothing was restored or registered + waitForTag(TestTags.ENTER_AEP_SCREEN) + assert(storageController.committedData == null) { "Expected no registration data to be committed" } + assert(storageController.restoreDecision == null) { "Expected no restore decision to have been made" } + } + + @Test + fun `restoring a local backup for a reglocked account whose reglock is not derived from the aep verifies the pin then restores`() { + val aep = AccountEntropyPool.generate() + val svrMasterKey = MasterKey(ByteArray(32) { it.toByte() }) + + // The AEP is valid for the account (its recovery password is accepted), but the account's reglock is governed by a + // separate master key held in SVR, so the reglock proof derived from the AEP is rejected. + networkController.onRegisterAccount = { request -> + when { + request.registrationLock == svrMasterKey.deriveRegistrationLock() -> RequestResult.Success(networkController.registerAccountResponse(request.e164)) + else -> RequestResult.NonSuccess( + RegisterAccountError.RegistrationLock( + RegistrationLockResponse( + timeRemaining = 14.days.inWholeMilliseconds, + svr2Credentials = SvrCredentials(username = "svr-user", password = "svr-pass") + ) + ) + ) + } + } + + networkController.onRestoreMasterKeyFromSvr = { request -> + if (request.pin == PIN) { + RequestResult.Success(MasterKeyResponse(svrMasterKey)) + } else { + RequestResult.NonSuccess(RestoreMasterKeyError.WrongPin(triesRemaining = 3)) + } + } + + // The backup contains the user's PIN, so the flow finishes without any PIN screens after the restore + storageController.onRestoreLocalBackupV2 = { _, _ -> + flowOf(LocalBackupRestoreProgress.Complete(restoredSvrPin = PIN, restoredProfileKey = null)) + } + + var registrationComplete = false + launchRegistrationFlow(folderPickerResult = backupFolderUri, onRegistrationComplete = { registrationComplete = true }) + + startManualRestore() + chooseRestoreOption(TestTags.ARCHIVE_RESTORE_SELECTION_FROM_BACKUP_FOLDER) + enterPhoneNumber() + restoreLocalBackup(aep) + + // The AEP-derived reglock proof was rejected, so the user must prove their PIN to clear the registration lock + waitForTag(TestTags.PIN_ENTRY_SCREEN) + composeTestRule.onNodeWithTag(TestTags.PIN_ENTRY_INPUT).performTextInput(PIN) + composeTestRule.onNodeWithTag(TestTags.PIN_ENTRY_CONTINUE_BUTTON).performClick() + + // With the reglock cleared, the pending backup restore resumes on the post-registration local restore screen + restoreFoundLocalBackup() + + waitFor("registration to complete") { registrationComplete } + + assert(networkController.lastRestoreMasterKeyRequest?.pin == PIN) { "Expected the PIN to be used to restore the master key from SVR" } + assert(networkController.lastRegisterAccountRequest?.registrationLock == svrMasterKey.deriveRegistrationLock()) { "Expected registration with the reglock proof derived from the SVR master key" } + + val committed = storageController.committedData + assert(committed != null) { "Expected registration data to be committed" } + assert(committed!!.accountData?.e164 == E164) { "Expected committed e164 $E164 but was ${committed.accountData?.e164}" } + assert(storageController.restoreDecision == RestoreDecision.COMPLETED) { "Expected COMPLETED restore decision (the backup was restored) but was ${storageController.restoreDecision}" } + } + + @Test + fun `entering a recovery key that cannot decrypt the local backup shows an inline error without registering`() { + val aep = AccountEntropyPool.generate() + + storageController.onVerifyLocalBackupKey = { _, _ -> false } + + launchRegistrationFlow(folderPickerResult = backupFolderUri) + + startManualRestore() + chooseRestoreOption(TestTags.ARCHIVE_RESTORE_SELECTION_FROM_BACKUP_FOLDER) + enterPhoneNumber() + restoreLocalBackup(aep) + + // The key can't decrypt the backup, so submission is rejected inline before any registration attempt + waitFor("the incorrect key to be rejected") { + composeTestRule.onAllNodesWithTag(TestTags.ENTER_AEP_NEXT_BUTTON).fetchSemanticsNodes().firstOrNull() + ?.config?.getOrNull(SemanticsProperties.Disabled) != null + } + assert(networkController.lastRegisterAccountRequest == null) { "Expected no registration attempt for a key that cannot decrypt the backup" } } @Test diff --git a/feature/registration/src/test/java/org/signal/registration/fakes/FakeStorageController.kt b/feature/registration/src/test/java/org/signal/registration/fakes/FakeStorageController.kt index 16a4e6d41c..e8d25762da 100644 --- a/feature/registration/src/test/java/org/signal/registration/fakes/FakeStorageController.kt +++ b/feature/registration/src/test/java/org/signal/registration/fakes/FakeStorageController.kt @@ -64,6 +64,8 @@ class FakeStorageController : StorageController { flowOf(LocalBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null)) } + var onVerifyLocalBackupKey: suspend (backupUri: Uri, aep: AccountEntropyPool) -> Boolean = { _, _ -> true } + var onRestoreRemoteBackup: (aep: AccountEntropyPool) -> Flow = { flowOf(RemoteBackupRestoreProgress.Complete(restoredSvrPin = null, restoredProfileKey = null)) } @@ -92,7 +94,10 @@ class FakeStorageController : StorageController { } override suspend fun setRestoreDecision(decision: RestoreDecision) { - restoreDecision = decision + // Mirrors the real controller: only the first decision sticks, later ones are ignored + if (restoreDecision == null) { + restoreDecision = decision + } } override fun restoreLocalBackupV1(rootUri: Uri, backupUri: Uri, passphrase: String): Flow { @@ -103,6 +108,10 @@ class FakeStorageController : StorageController { return onRestoreLocalBackupV2(backupUri, aep) } + override suspend fun verifyLocalBackupKey(backupUri: Uri, aep: AccountEntropyPool): Boolean { + return onVerifyLocalBackupKey(backupUri, aep) + } + override fun restoreRemoteBackup(aep: AccountEntropyPool): Flow { return onRestoreRemoteBackup(aep) } diff --git a/feature/registration/src/test/java/org/signal/registration/screens/aepentry/EnterAepForLocalBackupViewModelTest.kt b/feature/registration/src/test/java/org/signal/registration/screens/aepentry/EnterAepForLocalBackupViewModelTest.kt index 092787cc54..570d212487 100644 --- a/feature/registration/src/test/java/org/signal/registration/screens/aepentry/EnterAepForLocalBackupViewModelTest.kt +++ b/feature/registration/src/test/java/org/signal/registration/screens/aepentry/EnterAepForLocalBackupViewModelTest.kt @@ -5,114 +5,280 @@ package org.signal.registration.screens.aepentry +import android.net.Uri import assertk.assertThat import assertk.assertions.hasSize -import assertk.assertions.isEmpty import assertk.assertions.isEqualTo +import assertk.assertions.isInstanceOf import assertk.assertions.isNull -import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.ExperimentalCoroutinesApi -import kotlinx.coroutines.test.StandardTestDispatcher -import kotlinx.coroutines.test.advanceUntilIdle -import kotlinx.coroutines.test.resetMain +import assertk.assertions.prop +import io.mockk.coEvery +import io.mockk.coVerify +import io.mockk.every +import io.mockk.mockk +import io.mockk.mockkStatic +import io.mockk.unmockkStatic +import kotlinx.coroutines.flow.MutableStateFlow import kotlinx.coroutines.test.runTest -import kotlinx.coroutines.test.setMain import org.junit.After import org.junit.Before import org.junit.Test +import org.signal.core.models.AccountEntropyPool import org.signal.core.ui.navigation.ResultEventBus +import org.signal.libsignal.net.RequestResult +import org.signal.registration.KeyMaterial +import org.signal.registration.NetworkController import org.signal.registration.RegistrationFlowEvent +import org.signal.registration.RegistrationFlowState +import org.signal.registration.RegistrationRepository +import org.signal.registration.RegistrationRoute -@OptIn(ExperimentalCoroutinesApi::class) class EnterAepForLocalBackupViewModelTest { - private lateinit var viewModel: EnterAepForLocalBackupViewModel + private lateinit var mockRepository: RegistrationRepository private lateinit var resultBus: ResultEventBus + private lateinit var parentState: MutableStateFlow private lateinit var emittedParentEvents: MutableList private lateinit var parentEventEmitter: (RegistrationFlowEvent) -> Unit - - private val resultKey = "test-result-key" - - private val testDispatcher = StandardTestDispatcher() + private lateinit var emittedStates: MutableList + private lateinit var stateEmitter: (EnterAepState) -> Unit @Before fun setup() { - Dispatchers.setMain(testDispatcher) + mockkStatic(Uri::class) + every { Uri.parse(any()) } answers { mockk(relaxed = true) } + + mockRepository = mockk(relaxed = true) resultBus = ResultEventBus() + parentState = MutableStateFlow(RegistrationFlowState(sessionE164 = E164)) emittedParentEvents = mutableListOf() parentEventEmitter = { event -> emittedParentEvents.add(event) } - viewModel = EnterAepForLocalBackupViewModel( - parentEventEmitter = parentEventEmitter, - resultBus = resultBus, - resultKey = resultKey - ) + emittedStates = mutableListOf() + stateEmitter = { state -> emittedStates.add(state) } } @After fun tearDown() { - Dispatchers.resetMain() + unmockkStatic(Uri::class) } - // ==================== BackupKeyChanged Tests ==================== - - @Test - fun `BackupKeyChanged updates backup key in state`() = runTest { - val testKey = VALID_AEP - - viewModel.onEvent(EnterAepEvents.BackupKeyChanged(testKey)) - advanceUntilIdle() - - assertThat(viewModel.state.value.backupKey).isEqualTo(testKey) + private fun createViewModel(isPreRegistration: Boolean = true): EnterAepForLocalBackupViewModel { + return EnterAepForLocalBackupViewModel( + isPreRegistration = isPreRegistration, + backupUri = BACKUP_URI, + repository = mockRepository, + parentState = parentState, + parentEventEmitter = parentEventEmitter, + resultBus = resultBus, + resultKey = RESULT_KEY + ) } - // ==================== Submit Tests ==================== + private fun latestResult(): EnterAepForLocalBackupResult? { + return resultBus.channelMap[RESULT_KEY]?.tryReceive()?.getOrNull() as EnterAepForLocalBackupResult? + } + + // ==================== Already-registered mode ==================== @Test - fun `Submit with valid key sends result via resultBus and emits NavigateBack`() = runTest { - viewModel.onEvent(EnterAepEvents.BackupKeyChanged(VALID_AEP)) - viewModel.onEvent(EnterAepEvents.Submit) - advanceUntilIdle() + fun `Submit without requiring registration hands the key back and navigates back`() = runTest { + val viewModel = createViewModel(isPreRegistration = false) + val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true) - val result = resultBus.channelMap[resultKey]?.tryReceive()?.getOrNull() - assertThat(result).isEqualTo(VALID_AEP) + viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter) + + assertThat(latestResult()).isEqualTo(EnterAepForLocalBackupResult.RestoreReady(VALID_AEP)) assertThat(emittedParentEvents).hasSize(1) assertThat(emittedParentEvents.first()).isEqualTo(RegistrationFlowEvent.NavigateBack) + coVerify(exactly = 0) { mockRepository.registerAccountWithRecoveryPassword(any(), any(), any(), any(), any(), any()) } + } + + // ==================== Register-first mode ==================== + + @Test + fun `Submit with a key that cannot decrypt the backup shows an inline error without registering`() = runTest { + val viewModel = createViewModel() + val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true) + + coEvery { mockRepository.verifyLocalBackupKey(any(), any()) } returns false + + viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter) + + assertThat(emittedStates.last().aepValidationError).isEqualTo(AepValidationError.Incorrect) + assertThat(emittedStates.last().isRegistering).isEqualTo(false) + coVerify(exactly = 0) { mockRepository.registerAccountWithRecoveryPassword(any(), any(), any(), any(), any(), any()) } } @Test - fun `Submit with invalid key does not send result or navigate`() = runTest { - viewModel.onEvent(EnterAepEvents.BackupKeyChanged("too-short")) - viewModel.onEvent(EnterAepEvents.Submit) - advanceUntilIdle() + fun `Submit with successful registration emits UserSuppliedAepSubmitted and Registered, then hands the key back`() = runTest { + val viewModel = createViewModel() + val aep = AccountEntropyPool(VALID_AEP) + val mockKeyMaterial = mockk(relaxed = true) { + io.mockk.every { accountEntropyPool } returns aep + } + val mockResponse = mockk(relaxed = true) + val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true) - assertThat(resultBus.channelMap[resultKey]).isNull() - assertThat(emittedParentEvents).isEmpty() + coEvery { mockRepository.verifyLocalBackupKey(any(), any()) } returns true + coEvery { mockRepository.registerAccountWithRecoveryPassword(any(), any(), any(), any(), any(), any()) } returns + RequestResult.Success(mockResponse to mockKeyMaterial) + + viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter) + + assertThat(emittedParentEvents).hasSize(3) + assertThat(emittedParentEvents[0]).isInstanceOf() + assertThat(emittedParentEvents[1]).isInstanceOf() + assertThat(emittedParentEvents[2]).isEqualTo(RegistrationFlowEvent.NavigateBack) + assertThat(latestResult()).isEqualTo(EnterAepForLocalBackupResult.RestoreReady(VALID_AEP)) + } + + @Test + fun `Submit with RegistrationRecoveryPasswordIncorrect shows the different account dialog`() = runTest { + val viewModel = createViewModel() + val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true) + + coEvery { mockRepository.verifyLocalBackupKey(any(), any()) } returns true + coEvery { mockRepository.registerAccountWithRecoveryPassword(any(), any(), any(), any(), any(), any()) } returns + RequestResult.NonSuccess( + NetworkController.RegisterAccountError.RegistrationRecoveryPasswordIncorrect("Incorrect") + ) + + viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter) + + assertThat(emittedStates.last().showDifferentAccountDialog).isEqualTo(true) + assertThat(emittedStates.last().isRegistering).isEqualTo(false) + assertThat(latestResult()).isNull() + } + + @Test + fun `ConfirmDifferentAccountRestore forces the session path and hands control back for SMS verification`() = runTest { + val viewModel = createViewModel() + val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true, showDifferentAccountDialog = true) + + viewModel.applyEvent(initialState, EnterAepEvents.ConfirmDifferentAccountRestore, stateEmitter) + + assertThat(emittedStates.last().showDifferentAccountDialog).isEqualTo(false) + assertThat(emittedParentEvents).hasSize(2) + assertThat(emittedParentEvents[0]).isEqualTo(RegistrationFlowEvent.RecoveryPasswordInvalid) + assertThat(emittedParentEvents[1]).isEqualTo(RegistrationFlowEvent.NavigateBack) + assertThat(latestResult()).isEqualTo(EnterAepForLocalBackupResult.RegistrationDeferredToSms) + } + + @Test + fun `DismissDifferentAccountDialog clears the dialog`() = runTest { + val viewModel = createViewModel() + val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true, showDifferentAccountDialog = true) + + viewModel.applyEvent(initialState, EnterAepEvents.DismissDifferentAccountDialog, stateEmitter) + + assertThat(emittedStates.last().showDifferentAccountDialog).isEqualTo(false) + assertThat(emittedParentEvents).hasSize(0) + } + + @Test + fun `Submit with RegistrationLock retries with the reglock token derived from the AEP`() = runTest { + val viewModel = createViewModel() + val aep = AccountEntropyPool(VALID_AEP) + val mockKeyMaterial = mockk(relaxed = true) { + io.mockk.every { accountEntropyPool } returns aep + } + val mockResponse = mockk(relaxed = true) + val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true) + val registrationLockData = NetworkController.RegistrationLockResponse( + timeRemaining = 86400000L, + svr2Credentials = NetworkController.SvrCredentials(username = "test-username", password = "test-password") + ) + + coEvery { mockRepository.verifyLocalBackupKey(any(), any()) } returns true + coEvery { mockRepository.registerAccountWithRecoveryPassword(any(), any(), registrationLock = any(), any(), any(), any()) } returns + RequestResult.Success(mockResponse to mockKeyMaterial) + coEvery { mockRepository.registerAccountWithRecoveryPassword(any(), any(), registrationLock = null, any(), any(), any()) } returns + RequestResult.NonSuccess( + NetworkController.RegisterAccountError.RegistrationLock(registrationLockData) + ) + + viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter) + + coVerify { + mockRepository.registerAccountWithRecoveryPassword(any(), any(), registrationLock = aep.deriveMasterKey().deriveRegistrationLock(), any(), any(), any()) + } + assertThat(latestResult()).isEqualTo(EnterAepForLocalBackupResult.RestoreReady(VALID_AEP)) + } + + @Test + fun `Submit with RegistrationLock when already providing the reglock token navigates to PinEntryForRegistrationLock`() = runTest { + val viewModel = createViewModel() + val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true) + val registrationLockData = NetworkController.RegistrationLockResponse( + timeRemaining = 86400000L, + svr2Credentials = NetworkController.SvrCredentials(username = "test-username", password = "test-password") + ) + + coEvery { mockRepository.verifyLocalBackupKey(any(), any()) } returns true + coEvery { mockRepository.registerAccountWithRecoveryPassword(any(), any(), any(), any(), any(), any()) } returns + RequestResult.NonSuccess( + NetworkController.RegisterAccountError.RegistrationLock(registrationLockData) + ) + + viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter) + + assertThat(emittedStates.last().isRegistering).isEqualTo(false) + assertThat(emittedParentEvents.last()) + .isInstanceOf() + .prop(RegistrationFlowEvent.NavigateToScreen::route) + .isInstanceOf() + } + + @Test + fun `Submit with RateLimited sets registrationError to RateLimited`() = runTest { + val viewModel = createViewModel() + val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true) + + coEvery { mockRepository.verifyLocalBackupKey(any(), any()) } returns true + coEvery { mockRepository.registerAccountWithRecoveryPassword(any(), any(), any(), any(), any(), any()) } returns + RequestResult.NonSuccess( + NetworkController.RegisterAccountError.RateLimited(kotlin.time.Duration.parse("1m")) + ) + + viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter) + + assertThat(emittedStates.last().registrationError).isEqualTo(RegistrationError.RateLimited) + assertThat(emittedStates.last().isRegistering).isEqualTo(false) + } + + @Test + fun `Submit with RetryableNetworkError sets registrationError to NetworkError`() = runTest { + val viewModel = createViewModel() + val initialState = EnterAepState(backupKey = VALID_AEP, isBackupKeyValid = true) + + coEvery { mockRepository.verifyLocalBackupKey(any(), any()) } returns true + coEvery { mockRepository.registerAccountWithRecoveryPassword(any(), any(), any(), any(), any(), any()) } returns + RequestResult.RetryableNetworkError(java.io.IOException("Network error")) + + viewModel.applyEvent(initialState, EnterAepEvents.Submit, stateEmitter) + + assertThat(emittedStates.last().registrationError).isEqualTo(RegistrationError.NetworkError) + assertThat(emittedStates.last().isRegistering).isEqualTo(false) } // ==================== Cancel Tests ==================== @Test fun `Cancel emits NavigateBack`() = runTest { - viewModel.onEvent(EnterAepEvents.Cancel) - advanceUntilIdle() + val viewModel = createViewModel() + + viewModel.applyEvent(EnterAepState(), EnterAepEvents.Cancel, stateEmitter) assertThat(emittedParentEvents).hasSize(1) assertThat(emittedParentEvents.first()).isEqualTo(RegistrationFlowEvent.NavigateBack) } - // ==================== DismissError Tests ==================== - - @Test - fun `DismissError clears registrationError from state`() = runTest { - viewModel.onEvent(EnterAepEvents.DismissError) - advanceUntilIdle() - - assertThat(viewModel.state.value.registrationError).isNull() - } - // ==================== Constants ==================== companion object { private const val VALID_AEP = "uy38jh2778hjjhj8lk19ga61s672jsj089r023s6a57809bap92j2yh5t326vv7t" + private const val E164 = "+15551234567" + private const val BACKUP_URI = "content://test/backups/signal-backup-2026-01-01-12-00-00" + private const val RESULT_KEY = "test-result-key" } } 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 5b2b78973c..d408b04e3b 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 @@ -72,14 +72,20 @@ class LocalBackupRestoreViewModelTest { Dispatchers.resetMain() } - private fun createViewModel(isPreRegistration: Boolean, storageCapable: Boolean = true): LocalBackupRestoreViewModel { + private fun createViewModel( + isPreRegistration: Boolean, + storageCapable: Boolean = true, + knownAep: AccountEntropyPool? = null, + parentState: RegistrationFlowState = RegistrationFlowState(storageCapable = storageCapable) + ): LocalBackupRestoreViewModel { return LocalBackupRestoreViewModel( repository = mockRepository, - parentState = flowOf(RegistrationFlowState(storageCapable = storageCapable)), + parentState = flowOf(parentState), parentEventEmitter = parentEventEmitter, isPreRegistration = isPreRegistration, resultBus = resultBus, - resultKey = resultKey + resultKey = resultKey, + knownAep = knownAep ) } @@ -149,7 +155,7 @@ class LocalBackupRestoreViewModelTest { // ==================== RestoreBackup with V2 Tests ==================== @Test - fun `RestoreBackup with V2 backup navigates to EnterAepForLocalBackup`() = runTest { + fun `RestoreBackup with V2 backup post-registration navigates to EnterAepForLocalBackup without requiring registration`() = runTest { val viewModel = createViewModel(isPreRegistration = false) val backupInfo = LocalBackupInfo( type = LocalBackupInfo.BackupType.V2, @@ -162,10 +168,70 @@ class LocalBackupRestoreViewModelTest { viewModel.applyEvent(initialState, LocalBackupRestoreEvents.RestoreBackup, stateEmitter) assertThat(emittedParentEvents).hasSize(1) - assertThat(emittedParentEvents.first()) + val route = assertThat(emittedParentEvents.first()) .isInstanceOf() .prop(RegistrationFlowEvent.NavigateToScreen::route) - .isEqualTo(RegistrationRoute.EnterAepForLocalBackup) + route.isInstanceOf().prop(RegistrationRoute.EnterAepForLocalBackup::isPreRegistration).isEqualTo(false) + } + + @Test + fun `RestoreBackup with V2 backup pre-registration navigates to EnterAepForLocalBackup requiring registration`() = runTest { + val viewModel = createViewModel(isPreRegistration = true) + val backupInfo = LocalBackupInfo( + type = LocalBackupInfo.BackupType.V2, + date = LocalDateTime.now(), + name = "backup.bin", + uri = mockk(relaxed = true) + ) + val initialState = LocalBackupRestoreState(backupInfo = backupInfo, selectedFolderUri = mockk(relaxed = true)) + + viewModel.applyEvent(initialState, LocalBackupRestoreEvents.RestoreBackup, stateEmitter) + + assertThat(emittedParentEvents).hasSize(1) + val route = assertThat(emittedParentEvents.first()) + .isInstanceOf() + .prop(RegistrationFlowEvent.NavigateToScreen::route) + route.isInstanceOf().prop(RegistrationRoute.EnterAepForLocalBackup::isPreRegistration).isEqualTo(true) + } + + // ==================== Deferred restore Tests ==================== + + @Test + fun `RegistrationDeferredToSms forwards the deferral to the phone number screen and navigates back`() = runTest { + val viewModel = createViewModel(isPreRegistration = true) + val initialState = LocalBackupRestoreState() + + viewModel.applyEvent(initialState, LocalBackupRestoreEvents.RegistrationDeferredToSms, stateEmitter) + + val result = resultBus.channelMap[resultKey]?.tryReceive()?.getOrNull() + assertThat(result).isNotNull().isEqualTo(LocalBackupRestoreResult.DeferredToSms) + assertThat(emittedParentEvents).hasSize(1) + assertThat(emittedParentEvents.first()).isEqualTo(RegistrationFlowEvent.NavigateBack) + } + + @Test + fun `a completed restore clears the pending restore option so it is not resumed again post-registration`() = 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) + + coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) } + assertThat(emittedParentEvents).contains(RegistrationFlowEvent.PendingRestoreOptionSelected(null)) + assertThat(emittedParentEvents.last()) + .isInstanceOf() + .prop(RegistrationFlowEvent.NavigateToScreen::route) + .isEqualTo(RegistrationRoute.PinCreate) } // ==================== RestoreBackup with no backup Tests ==================== @@ -312,8 +378,7 @@ class LocalBackupRestoreViewModelTest { coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) } coVerify(exactly = 0) { mockRepository.restoreAccountRecord(any()) } - assertThat(emittedParentEvents).hasSize(1) - assertThat(emittedParentEvents.first()) + assertThat(emittedParentEvents.last()) .isInstanceOf() .prop(RegistrationFlowEvent.NavigateToScreen::route) .isEqualTo(RegistrationRoute.PinEntryForSvrRestore) @@ -340,8 +405,7 @@ class LocalBackupRestoreViewModelTest { coVerify { mockRepository.setRestoreDecision(RestoreDecision.COMPLETED) } coVerify(exactly = 0) { mockRepository.restoreAccountRecord(any()) } - assertThat(emittedParentEvents).hasSize(1) - assertThat(emittedParentEvents.first()) + assertThat(emittedParentEvents.last()) .isInstanceOf() .prop(RegistrationFlowEvent.NavigateToScreen::route) .isEqualTo(RegistrationRoute.PinCreate) diff --git a/feature/registration/src/test/java/org/signal/registration/screens/pinentry/PinEntryForRegistrationLockViewModelTest.kt b/feature/registration/src/test/java/org/signal/registration/screens/pinentry/PinEntryForRegistrationLockViewModelTest.kt index 9c7a48cf22..ebc422d66d 100644 --- a/feature/registration/src/test/java/org/signal/registration/screens/pinentry/PinEntryForRegistrationLockViewModelTest.kt +++ b/feature/registration/src/test/java/org/signal/registration/screens/pinentry/PinEntryForRegistrationLockViewModelTest.kt @@ -19,10 +19,12 @@ import kotlinx.coroutines.flow.MutableStateFlow import kotlinx.coroutines.test.runTest import org.junit.Before import org.junit.Test +import org.signal.core.models.AccountEntropyPool import org.signal.core.models.MasterKey import org.signal.libsignal.net.RequestResult import org.signal.registration.KeyMaterial import org.signal.registration.NetworkController +import org.signal.registration.PendingRestoreOption import org.signal.registration.RegistrationFlowEvent import org.signal.registration.RegistrationFlowState import org.signal.registration.RegistrationRepository @@ -116,6 +118,60 @@ class PinEntryForRegistrationLockViewModelTest { assertThat(emittedStates.last().loading).isEqualTo(true) } + @Test + fun `PinEntered resumes a pending local backup restore after clearing the registration lock`() = runTest { + val masterKey = mockk(relaxed = true) + val keyMaterial = mockk(relaxed = true) + val restoreAep = AccountEntropyPool.generate() + val initialState = PinEntryState(mode = PinEntryState.Mode.RegistrationLock) + + parentState.value = parentState.value.copy( + pendingRestoreOption = PendingRestoreOption.LocalBackup, + unverifiedRestoredAep = restoreAep + ) + + coEvery { mockRepository.restoreMasterKeyFromSvr(any(), any(), forRegistrationLock = true) } returns + RequestResult.Success(NetworkController.MasterKeyResponse(masterKey)) + coEvery { mockRepository.registerAccountWithSession(any(), any(), any(), any()) } returns + RequestResult.Success(createRegisterAccountResponse() to keyMaterial) + + viewModel.applyEvent(initialState, PinEntryScreenEvents.PinEntered("123456"), parentEventEmitter, stateEmitter) + + assertThat(emittedParentEvents.last()) + .isInstanceOf() + .prop(RegistrationFlowEvent.NavigateToScreen::route) + .isInstanceOf() + .prop(RegistrationRoute.LocalBackupRestore::aep) + .isEqualTo(restoreAep) + } + + @Test + fun `PinEntered resumes a pending remote backup restore after clearing the registration lock`() = runTest { + val masterKey = mockk(relaxed = true) + val keyMaterial = mockk(relaxed = true) + val restoreAep = AccountEntropyPool.generate() + val initialState = PinEntryState(mode = PinEntryState.Mode.RegistrationLock) + + parentState.value = parentState.value.copy( + pendingRestoreOption = PendingRestoreOption.RemoteBackup, + unverifiedRestoredAep = restoreAep + ) + + coEvery { mockRepository.restoreMasterKeyFromSvr(any(), any(), forRegistrationLock = true) } returns + RequestResult.Success(NetworkController.MasterKeyResponse(masterKey)) + coEvery { mockRepository.registerAccountWithSession(any(), any(), any(), any()) } returns + RequestResult.Success(createRegisterAccountResponse() to keyMaterial) + + viewModel.applyEvent(initialState, PinEntryScreenEvents.PinEntered("123456"), parentEventEmitter, stateEmitter) + + assertThat(emittedParentEvents.last()) + .isInstanceOf() + .prop(RegistrationFlowEvent.NavigateToScreen::route) + .isInstanceOf() + .prop(RegistrationRoute.RemoteRestore::aep) + .isEqualTo(restoreAep) + } + @Test fun `PinEntered with wrong PIN returns state with tries remaining`() = runTest { val triesRemaining = 3