From f375750cf89bdccef43b6cfb628d48bef75a199a Mon Sep 17 00:00:00 2001 From: Greyson Parrelli Date: Tue, 23 Jun 2026 11:05:10 -0400 Subject: [PATCH] Fix storageCapable bug in regV5. --- .../signal/registration/PersistedFlowState.kt | 9 ++++-- .../registration/RegistrationFlowEvent.kt | 11 +++++-- .../registration/RegistrationFlowState.kt | 5 +++- .../registration/RegistrationViewModel.kt | 2 +- ...ForRemoteBackupPreRegistrationViewModel.kt | 4 +-- .../phonenumber/PhoneNumberEntryViewModel.kt | 4 +-- .../PinEntryForRegistrationLockViewModel.kt | 2 +- .../quickrestore/QuickRestoreQrViewModel.kt | 2 +- .../ArchiveRestoreSelectionState.kt | 6 ++-- .../ArchiveRestoreSelectionViewModel.kt | 10 +++++-- .../VerificationCodeViewModel.kt | 2 +- .../registration/PersistedFlowStateTest.kt | 6 +++- .../registration/RegistrationViewModelTest.kt | 5 ++-- .../ArchiveRestoreSelectionViewModelTest.kt | 29 +++++++++++++++++-- 14 files changed, 73 insertions(+), 24 deletions(-) diff --git a/feature/registration/src/main/java/org/signal/registration/PersistedFlowState.kt b/feature/registration/src/main/java/org/signal/registration/PersistedFlowState.kt index 59f74bf77a..21ac82b7bf 100644 --- a/feature/registration/src/main/java/org/signal/registration/PersistedFlowState.kt +++ b/feature/registration/src/main/java/org/signal/registration/PersistedFlowState.kt @@ -24,7 +24,8 @@ data class PersistedFlowState( val doNotAttemptRecoveryPassword: Boolean, val pendingRestoreOption: PendingRestoreOption? = null, val restoredAepValue: String? = null, - val restoreMethodToken: String? = null + val restoreMethodToken: String? = null, + val storageCapable: Boolean = false ) /** @@ -38,7 +39,8 @@ fun RegistrationFlowState.toPersistedFlowState(): PersistedFlowState { doNotAttemptRecoveryPassword = doNotAttemptRecoveryPassword, pendingRestoreOption = pendingRestoreOption, restoredAepValue = unverifiedRestoredAep?.value, - restoreMethodToken = restoreMethodToken + restoreMethodToken = restoreMethodToken, + storageCapable = storageCapable ) } @@ -64,6 +66,7 @@ fun PersistedFlowState.toRegistrationFlowState( doNotAttemptRecoveryPassword = doNotAttemptRecoveryPassword, pendingRestoreOption = pendingRestoreOption, unverifiedRestoredAep = restoredAepValue?.let { AccountEntropyPool(it) }, - restoreMethodToken = restoreMethodToken + restoreMethodToken = restoreMethodToken, + storageCapable = storageCapable ) } diff --git a/feature/registration/src/main/java/org/signal/registration/RegistrationFlowEvent.kt b/feature/registration/src/main/java/org/signal/registration/RegistrationFlowEvent.kt index e99ebec16d..2dfd73a818 100644 --- a/feature/registration/src/main/java/org/signal/registration/RegistrationFlowEvent.kt +++ b/feature/registration/src/main/java/org/signal/registration/RegistrationFlowEvent.kt @@ -25,9 +25,14 @@ sealed interface RegistrationFlowEvent { /** The e164 associated with this registration attempt has been updated. */ data class E164Chosen(val e164: String) : RegistrationFlowEvent - /** The user has successfully registered. */ - data class Registered(val accountEntropyPool: AccountEntropyPool) : RegistrationFlowEvent { - override fun toString(): String = "Registered(accountEntropyPool=${accountEntropyPool.displayValue.censor()})" + /** + * The user has successfully registered. + * + * @param storageCapable Whether the server reports that this account already has SVR/PIN data, as returned in the + * registration response. Used later (e.g. when skipping a restore) to decide between PIN entry and PIN creation. + */ + data class Registered(val accountEntropyPool: AccountEntropyPool, val storageCapable: Boolean) : RegistrationFlowEvent { + override fun toString(): String = "Registered(accountEntropyPool=${accountEntropyPool.displayValue.censor()}, storageCapable=$storageCapable)" } /** The master key has been restored from SVR. */ diff --git a/feature/registration/src/main/java/org/signal/registration/RegistrationFlowState.kt b/feature/registration/src/main/java/org/signal/registration/RegistrationFlowState.kt index 6981a2a9d4..286d49b930 100644 --- a/feature/registration/src/main/java/org/signal/registration/RegistrationFlowState.kt +++ b/feature/registration/src/main/java/org/signal/registration/RegistrationFlowState.kt @@ -30,6 +30,9 @@ data class RegistrationFlowState( /** The AEP we generated as part of this registration. */ val accountEntropyPool: AccountEntropyPool? = null, + /** Whether the server reported that this account already has SVR/PIN data, captured from the registration response. */ + val storageCapable: Boolean = false, + /** The master key we restored from SVR. Needed for initial storage service restore, but afterwards we'll generate a new one. */ val temporaryMasterKey: MasterKey? = null, @@ -55,6 +58,6 @@ data class RegistrationFlowState( val isRestoringNavigationState: Boolean = true ) : Parcelable { override fun toString(): String { - return "RegistrationFlowState(backStack=${backStack.joinToString()}, sessionMetadata=${sessionMetadata.let { "present" }}, sessionE164=$sessionE164, accountEntropyPool=${accountEntropyPool?.displayValue?.censor()}, temporaryMasterKey=${temporaryMasterKey?.toString()?.censor()}, preExistingRegistrationData=${preExistingRegistrationData?.let { "present" }}, doNotAttemptRecoveryPassword=$doNotAttemptRecoveryPassword, pendingRestoreOption=$pendingRestoreOption, unverifiedRestoredAep=${unverifiedRestoredAep?.displayValue?.censor()}, restoreMethodToken=${restoreMethodToken?.censor()}, isRestoringNavigation=$isRestoringNavigationState)" + return "RegistrationFlowState(backStack=${backStack.joinToString()}, sessionMetadata=${sessionMetadata.let { "present" }}, sessionE164=$sessionE164, accountEntropyPool=${accountEntropyPool?.displayValue?.censor()}, storageCapable=$storageCapable, temporaryMasterKey=${temporaryMasterKey?.toString()?.censor()}, preExistingRegistrationData=${preExistingRegistrationData?.let { "present" }}, doNotAttemptRecoveryPassword=$doNotAttemptRecoveryPassword, pendingRestoreOption=$pendingRestoreOption, unverifiedRestoredAep=${unverifiedRestoredAep?.displayValue?.censor()}, restoreMethodToken=${restoreMethodToken?.censor()}, isRestoringNavigation=$isRestoringNavigationState)" } } diff --git a/feature/registration/src/main/java/org/signal/registration/RegistrationViewModel.kt b/feature/registration/src/main/java/org/signal/registration/RegistrationViewModel.kt index c70850bbcc..c4760f1d2e 100644 --- a/feature/registration/src/main/java/org/signal/registration/RegistrationViewModel.kt +++ b/feature/registration/src/main/java/org/signal/registration/RegistrationViewModel.kt @@ -70,7 +70,7 @@ class RegistrationViewModel(private val repository: RegistrationRepository, save is RegistrationFlowEvent.ResetState -> RegistrationFlowState(isRestoringNavigationState = false) is RegistrationFlowEvent.SessionUpdated -> state.copy(sessionMetadata = event.session) is RegistrationFlowEvent.E164Chosen -> state.copy(sessionE164 = event.e164) - is RegistrationFlowEvent.Registered -> state.copy(accountEntropyPool = event.accountEntropyPool) + is RegistrationFlowEvent.Registered -> state.copy(accountEntropyPool = event.accountEntropyPool, storageCapable = event.storageCapable) is RegistrationFlowEvent.MasterKeyRestoredFromSvr -> state.copy(temporaryMasterKey = event.masterKey) is RegistrationFlowEvent.NavigateToScreen -> applyNavigationToScreenEvent(state, event) is RegistrationFlowEvent.NavigateBack -> { 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 fe5caeb9cf..10c453db16 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 @@ -71,10 +71,10 @@ class EnterAepForRemoteBackupPreRegistrationViewModel( when (val result = repository.registerAccountWithRecoveryPassword(e164, recoveryPassword, existingAccountEntropyPool = aep)) { is RequestResult.Success -> { Log.i(TAG, "[Submit] Successfully registered using RRP from user-supplied AEP.") - val (_, keyMaterial) = result.result + val (response, keyMaterial) = result.result stateEmitter(inputState.copy(isRegistering = false)) - parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool)) + parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool, response.storageCapable)) parentEventEmitter.navigateTo(RegistrationRoute.RemoteRestore(aep)) } is RequestResult.NonSuccess -> { 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 3ab02a0407..fb3b598407 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 @@ -234,7 +234,7 @@ class PhoneNumberEntryViewModel( Log.i(TAG, "[Register] Successfully re-registered using RRP from pre-existing data.") val (response, keyMaterial) = registerResult.result - parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool)) + parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool, response.storageCapable)) if (response.storageCapable) { parentEventEmitter.navigateTo(RegistrationRoute.PinEntryForSvrRestore) @@ -320,7 +320,7 @@ class PhoneNumberEntryViewModel( Log.i(TAG, "[LocalRestore] Successfully registered using RRP from restored AEP.") val (response, keyMaterial) = result.result - parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool)) + parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool, response.storageCapable)) if (response.storageCapable) { parentEventEmitter.navigateTo(RegistrationRoute.PinEntryForSvrRestore) 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 150ff4b7d5..9389d7004c 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 @@ -134,7 +134,7 @@ class PinEntryForRegistrationLockViewModel( is RequestResult.Success -> { Log.i(TAG, "[PinEntered] Successfully registered!") val (response, keyMaterial) = registerResult.result - parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool)) + parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool, response.storageCapable)) when { response.reregistration -> parentEventEmitter.navigateTo(RegistrationRoute.ArchiveRestoreSelection.forPostRegister()) else -> repository.finishRegistrationOrCreateProfile(parentEventEmitter) diff --git a/feature/registration/src/main/java/org/signal/registration/screens/quickrestore/QuickRestoreQrViewModel.kt b/feature/registration/src/main/java/org/signal/registration/screens/quickrestore/QuickRestoreQrViewModel.kt index 42c7c3ab8f..3bf74a3f0a 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/quickrestore/QuickRestoreQrViewModel.kt +++ b/feature/registration/src/main/java/org/signal/registration/screens/quickrestore/QuickRestoreQrViewModel.kt @@ -112,7 +112,7 @@ class QuickRestoreQrViewModel( is RequestResult.Success -> { val (response, keyMaterial) = registerResult.result Log.i(TAG, "[Register] Success! reregistration: ${response.reregistration}") - parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool)) + parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool, response.storageCapable)) parentEventEmitter.navigateTo(RegistrationRoute.ArchiveRestoreSelection.forQuickRestore(hasRemoteBackup = message.tier != null)) } is RequestResult.NonSuccess -> { diff --git a/feature/registration/src/main/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionState.kt b/feature/registration/src/main/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionState.kt index d1d15b26e5..83a95ceb0b 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionState.kt +++ b/feature/registration/src/main/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionState.kt @@ -11,9 +11,11 @@ data class ArchiveRestoreSelectionState( val restoreOptions: List = emptyList(), val showSkipWarningDialog: Boolean = false, /** Token that, if present, indicates that the user did a quick restore, and we should hit a network endpoint to indicate our restore selection. */ - val restoreMethodToken: String? = null + val restoreMethodToken: String? = null, + /** Whether the account already has SVR/PIN data on the server. Determines whether skipping restore leads to PIN entry or PIN creation. */ + val storageCapable: Boolean = false ) { - override fun toString(): String = "ArchiveRestoreSelectionState(restoreOptions=$restoreOptions, showSkipWarningDialog=$showSkipWarningDialog, restoreMethodToken=${restoreMethodToken?.censor()})" + override fun toString(): String = "ArchiveRestoreSelectionState(restoreOptions=$restoreOptions, showSkipWarningDialog=$showSkipWarningDialog, restoreMethodToken=${restoreMethodToken?.censor()}, storageCapable=$storageCapable)" val showSkipButton: Boolean get() = ArchiveRestoreOption.None !in restoreOptions } diff --git a/feature/registration/src/main/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionViewModel.kt b/feature/registration/src/main/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionViewModel.kt index fa609c08ea..5c2ecffbd1 100644 --- a/feature/registration/src/main/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionViewModel.kt +++ b/feature/registration/src/main/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionViewModel.kt @@ -64,7 +64,7 @@ class ArchiveRestoreSelectionViewModel( @VisibleForTesting fun applyParentState(state: ArchiveRestoreSelectionState, parentState: RegistrationFlowState): ArchiveRestoreSelectionState { - return state.copy(restoreMethodToken = parentState.restoreMethodToken) + return state.copy(restoreMethodToken = parentState.restoreMethodToken, storageCapable = parentState.storageCapable) } @VisibleForTesting @@ -108,7 +108,13 @@ class ArchiveRestoreSelectionViewModel( is ArchiveRestoreSelectionScreenEvents.ConfirmSkip -> { notifyOldDevice(state.restoreMethodToken, NetworkController.RestoreMethod.DECLINE) repository.setRestoreDecision(RestoreDecision.SKIPPED) - parentEventEmitter.navigateTo(RegistrationRoute.PinCreate) + if (state.storageCapable) { + Log.i(TAG, "[ConfirmSkip] Account is storage capable. Navigating to PIN entry to restore the existing PIN.") + parentEventEmitter.navigateTo(RegistrationRoute.PinEntryForSvrRestore) + } else { + Log.i(TAG, "[ConfirmSkip] Account is not storage capable. Navigating to PIN creation.") + parentEventEmitter.navigateTo(RegistrationRoute.PinCreate) + } state.copy(showSkipWarningDialog = false) } is ArchiveRestoreSelectionScreenEvents.DismissSkipWarning -> { 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 e8ffb19e2c..b3cc9a0fdd 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 @@ -173,7 +173,7 @@ class VerificationCodeViewModel( is RequestResult.Success -> { val (response, keyMaterial) = registerResult.result - parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool)) + parentEventEmitter(RegistrationFlowEvent.Registered(keyMaterial.accountEntropyPool, response.storageCapable)) when { response.reregistration -> parentEventEmitter.navigateTo(RegistrationRoute.ArchiveRestoreSelection.forPostRegister()) diff --git a/feature/registration/src/test/java/org/signal/registration/PersistedFlowStateTest.kt b/feature/registration/src/test/java/org/signal/registration/PersistedFlowStateTest.kt index b6ffee6706..1fe0374895 100644 --- a/feature/registration/src/test/java/org/signal/registration/PersistedFlowStateTest.kt +++ b/feature/registration/src/test/java/org/signal/registration/PersistedFlowStateTest.kt @@ -174,6 +174,7 @@ class PersistedFlowStateTest { sessionMetadata = session, sessionE164 = "+15551234567", accountEntropyPool = AccountEntropyPool.generate(), + storageCapable = true, temporaryMasterKey = MasterKey(ByteArray(32)), doNotAttemptRecoveryPassword = true ) @@ -184,6 +185,7 @@ class PersistedFlowStateTest { assertThat(persisted.sessionMetadata).isEqualTo(session) assertThat(persisted.sessionE164).isEqualTo("+15551234567") assertThat(persisted.doNotAttemptRecoveryPassword).isEqualTo(true) + assertThat(persisted.storageCapable).isEqualTo(true) } @Test @@ -202,7 +204,8 @@ class PersistedFlowStateTest { backStack = listOf(RegistrationRoute.Welcome, RegistrationRoute.PinCreate), sessionMetadata = session, sessionE164 = "+15551234567", - doNotAttemptRecoveryPassword = true + doNotAttemptRecoveryPassword = true, + storageCapable = true ) val aep = AccountEntropyPool.generate() @@ -221,5 +224,6 @@ class PersistedFlowStateTest { assertThat(flowState.temporaryMasterKey).isEqualTo(masterKey) assertThat(flowState.preExistingRegistrationData).isNull() assertThat(flowState.doNotAttemptRecoveryPassword).isEqualTo(true) + assertThat(flowState.storageCapable).isEqualTo(true) } } diff --git a/feature/registration/src/test/java/org/signal/registration/RegistrationViewModelTest.kt b/feature/registration/src/test/java/org/signal/registration/RegistrationViewModelTest.kt index 0ba5e104ca..1935815f35 100644 --- a/feature/registration/src/test/java/org/signal/registration/RegistrationViewModelTest.kt +++ b/feature/registration/src/test/java/org/signal/registration/RegistrationViewModelTest.kt @@ -236,7 +236,7 @@ class RegistrationViewModelTest { val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle()) advanceUntilIdle() - viewModel.onEvent(RegistrationFlowEvent.Registered(AccountEntropyPool.generate())) + viewModel.onEvent(RegistrationFlowEvent.Registered(AccountEntropyPool.generate(), storageCapable = false)) advanceUntilIdle() coVerify(exactly = 0) { mockRepository.saveFlowState(any()) } @@ -380,10 +380,11 @@ class RegistrationViewModelTest { val result = viewModel.applyEvent( RegistrationFlowState(), - RegistrationFlowEvent.Registered(aep) + RegistrationFlowEvent.Registered(aep, storageCapable = true) ) assertThat(result.accountEntropyPool).isEqualTo(aep) + assertThat(result.storageCapable).isTrue() } @Test diff --git a/feature/registration/src/test/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionViewModelTest.kt b/feature/registration/src/test/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionViewModelTest.kt index d6d9693035..b47b74b957 100644 --- a/feature/registration/src/test/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionViewModelTest.kt +++ b/feature/registration/src/test/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionViewModelTest.kt @@ -190,9 +190,9 @@ class ArchiveRestoreSelectionViewModelTest { // ==================== ConfirmSkip Tests ==================== @Test - fun `ConfirmSkip navigates to PinCreate and clears dialog`() = runTest { + fun `ConfirmSkip when not storage capable navigates to PinCreate and clears dialog`() = runTest { val viewModel = createViewModel(isPreRegistration = false) - val initialState = ArchiveRestoreSelectionState(showSkipWarningDialog = true) + val initialState = ArchiveRestoreSelectionState(showSkipWarningDialog = true, storageCapable = false) viewModel.applyEvent(initialState, ArchiveRestoreSelectionScreenEvents.ConfirmSkip, stateEmitter) @@ -205,6 +205,22 @@ class ArchiveRestoreSelectionViewModelTest { assertThat(emittedStates.last().showSkipWarningDialog).isFalse() } + @Test + fun `ConfirmSkip when storage capable navigates to PinEntryForSvrRestore and clears dialog`() = runTest { + val viewModel = createViewModel(isPreRegistration = false) + val initialState = ArchiveRestoreSelectionState(showSkipWarningDialog = true, storageCapable = true) + + viewModel.applyEvent(initialState, ArchiveRestoreSelectionScreenEvents.ConfirmSkip, stateEmitter) + + coVerify { mockRepository.setRestoreDecision(RestoreDecision.SKIPPED) } + assertThat(emittedParentEvents).hasSize(1) + assertThat(emittedParentEvents.first()) + .isInstanceOf() + .prop(RegistrationFlowEvent.NavigateToScreen::route) + .isEqualTo(RegistrationRoute.PinEntryForSvrRestore) + assertThat(emittedStates.last().showSkipWarningDialog).isFalse() + } + // ==================== DismissSkipWarning Tests ==================== @Test @@ -249,4 +265,13 @@ class ArchiveRestoreSelectionViewModelTest { assertThat(viewModel.state.value.showSkipButton).isTrue() } + + @Test + fun `applyParentState copies storageCapable from parent`() = runTest { + val viewModel = createViewModel() + + val result = viewModel.applyParentState(ArchiveRestoreSelectionState(), RegistrationFlowState(storageCapable = true)) + + assertThat(result.storageCapable).isTrue() + } }