Fixes for restoring a backup after completing registration.

app/src/main/java/org/thoughtcrime/securesms/registration/RegistrationUtil.java

@@ -28,6 +28,7 @@ public final class RegistrationUtil {
     {
       Log.i(TAG, "Marking registration completed.", new Throwable());
       SignalStore.registration().setRegistrationComplete();
+      SignalStore.registration().clearLocalRegistrationMetadata();
 
       if (SignalStore.phoneNumberPrivacy().getPhoneNumberDiscoverabilityMode() == PhoneNumberDiscoverabilityMode.UNDECIDED) {
         Log.w(TAG, "Phone number discoverability mode is still UNDECIDED. Setting to DISCOVERABLE.");

app/src/main/java/org/thoughtcrime/securesms/registration/data/LocalRegistrationMetadataUtil.kt

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.registration.data
+
+import okio.ByteString.Companion.toByteString
+import org.signal.libsignal.protocol.IdentityKey
+import org.signal.libsignal.protocol.state.KyberPreKeyRecord
+import org.signal.libsignal.protocol.state.SignedPreKeyRecord
+import org.thoughtcrime.securesms.database.model.databaseprotos.LocalRegistrationMetadata
+import org.thoughtcrime.securesms.registration.RegistrationData
+import org.whispersystems.signalservice.api.account.PreKeyCollection
+
+/**
+ * Takes the two sources of registration data ([RegistrationData], [RegistrationRepository.AccountRegistrationResult])
+ * and combines them into a proto-backed class [LocalRegistrationMetadata] so they can be serialized & stored.
+ */
+object LocalRegistrationMetadataUtil {
+  fun createLocalRegistrationMetadata(registrationData: RegistrationData, remoteResult: RegistrationRepository.AccountRegistrationResult, reglockEnabled: Boolean): LocalRegistrationMetadata {
+    return LocalRegistrationMetadata.Builder().apply {
+      aciIdentityKey = remoteResult.aciPreKeyCollection.identityKey.serialize().toByteString()
+      aciSignedPreKey = remoteResult.aciPreKeyCollection.signedPreKey.serialize().toByteString()
+      aciLastRestoreKyberPreKey = remoteResult.aciPreKeyCollection.signedPreKey.serialize().toByteString()
+      pniIdentityKey = remoteResult.pniPreKeyCollection.identityKey.serialize().toByteString()
+      pniSignedPreKey = remoteResult.pniPreKeyCollection.signedPreKey.serialize().toByteString()
+      pniLastRestoreKyberPreKey = remoteResult.pniPreKeyCollection.signedPreKey.serialize().toByteString()
+      aci = remoteResult.uuid
+      pni = remoteResult.pni
+      hasPin = remoteResult.storageCapable
+      remoteResult.pin?.let {
+        pin = it
+      }
+      remoteResult.masterKey?.serialize()?.toByteString()?.let {
+        masterKey = it
+      }
+      e164 = registrationData.e164
+      fcmEnabled = registrationData.isFcm
+      profileKey = registrationData.profileKey.serialize().toByteString()
+      servicePassword = registrationData.password
+      this.reglockEnabled = reglockEnabled
+    }.build()
+  }
+
+  fun LocalRegistrationMetadata.getAciPreKeyCollection(): PreKeyCollection {
+    return PreKeyCollection(
+      IdentityKey(aciIdentityKey.toByteArray()),
+      SignedPreKeyRecord(aciSignedPreKey.toByteArray()),
+      KyberPreKeyRecord(aciLastRestoreKyberPreKey.toByteArray())
+    )
+  }
+
+  fun LocalRegistrationMetadata.getPniPreKeyCollection(): PreKeyCollection {
+    return PreKeyCollection(
+      IdentityKey(pniIdentityKey.toByteArray()),
+      SignedPreKeyRecord(pniSignedPreKey.toByteArray()),
+      KyberPreKeyRecord(pniLastRestoreKyberPreKey.toByteArray())
+    )
+  }
+}

app/src/main/java/org/thoughtcrime/securesms/registration/data/RegistrationRepository.kt

@@ -28,6 +28,7 @@ import org.thoughtcrime.securesms.crypto.storage.PreKeyMetadataStore
 import org.thoughtcrime.securesms.crypto.storage.SignalServiceAccountDataStoreImpl
 import org.thoughtcrime.securesms.database.IdentityTable
 import org.thoughtcrime.securesms.database.SignalDatabase
+import org.thoughtcrime.securesms.database.model.databaseprotos.LocalRegistrationMetadata
 import org.thoughtcrime.securesms.dependencies.AppDependencies
 import org.thoughtcrime.securesms.gcm.FcmUtil
 import org.thoughtcrime.securesms.jobs.DirectoryRefreshJob
@@ -45,6 +46,8 @@ import org.thoughtcrime.securesms.recipients.RecipientId
 import org.thoughtcrime.securesms.registration.PushChallengeRequest
 import org.thoughtcrime.securesms.registration.RegistrationData
 import org.thoughtcrime.securesms.registration.VerifyAccountRepository
+import org.thoughtcrime.securesms.registration.data.LocalRegistrationMetadataUtil.getAciPreKeyCollection
+import org.thoughtcrime.securesms.registration.data.LocalRegistrationMetadataUtil.getPniPreKeyCollection
 import org.thoughtcrime.securesms.registration.data.network.BackupAuthCheckResult
 import org.thoughtcrime.securesms.registration.data.network.RegisterAccountResult
 import org.thoughtcrime.securesms.registration.data.network.RegistrationSessionCheckResult
@@ -159,14 +162,14 @@ object RegistrationRepository {
    * Takes a server response from a successful registration and persists the relevant data.
    */
   @JvmStatic
-  suspend fun registerAccountLocally(context: Context, registrationData: RegistrationData, response: AccountRegistrationResult, reglockEnabled: Boolean) =
+  suspend fun registerAccountLocally(context: Context, data: LocalRegistrationMetadata) =
     withContext(Dispatchers.IO) {
       Log.v(TAG, "registerAccountLocally()")
-      val aciPreKeyCollection: PreKeyCollection = response.aciPreKeyCollection
-      val pniPreKeyCollection: PreKeyCollection = response.pniPreKeyCollection
-      val aci: ACI = ACI.parseOrThrow(response.uuid)
-      val pni: PNI = PNI.parseOrThrow(response.pni)
-      val hasPin: Boolean = response.storageCapable
+      val aciPreKeyCollection = data.getAciPreKeyCollection()
+      val pniPreKeyCollection = data.getPniPreKeyCollection()
+      val aci: ACI = ACI.parseOrThrow(data.aci)
+      val pni: PNI = PNI.parseOrThrow(data.pni)
+      val hasPin: Boolean = data.hasPin
 
       SignalStore.account.setAci(aci)
       SignalStore.account.setPni(pni)
@@ -187,30 +190,31 @@ object RegistrationRepository {
       storeSignedAndLastResortPreKeys(pniProtocolStore, pniMetadataStore, pniPreKeyCollection)
 
       val recipientTable = SignalDatabase.recipients
-      val selfId = Recipient.trustedPush(aci, pni, registrationData.e164).id
+      val selfId = Recipient.trustedPush(aci, pni, data.e164).id
 
       recipientTable.setProfileSharing(selfId, true)
       recipientTable.markRegisteredOrThrow(selfId, aci)
-      recipientTable.linkIdsForSelf(aci, pni, registrationData.e164)
-      recipientTable.setProfileKey(selfId, registrationData.profileKey)
+      recipientTable.linkIdsForSelf(aci, pni, data.e164)
+      recipientTable.setProfileKey(selfId, ProfileKey(data.profileKey.toByteArray()))
 
       AppDependencies.recipientCache.clearSelf()
 
-      SignalStore.account.setE164(registrationData.e164)
-      SignalStore.account.fcmToken = registrationData.fcmToken
-      SignalStore.account.fcmEnabled = registrationData.isFcm
+      SignalStore.account.setE164(data.e164)
+      SignalStore.account.fcmToken = data.fcmToken
+      SignalStore.account.fcmEnabled = data.fcmEnabled
 
       val now = System.currentTimeMillis()
       saveOwnIdentityKey(selfId, aci, aciProtocolStore, now)
       saveOwnIdentityKey(selfId, pni, pniProtocolStore, now)
 
-      SignalStore.account.setServicePassword(registrationData.password)
+      SignalStore.account.setServicePassword(data.servicePassword)
       SignalStore.account.setRegistered(true)
       TextSecurePreferences.setPromptedPushRegistration(context, true)
       TextSecurePreferences.setUnauthorizedReceived(context, false)
       NotificationManagerCompat.from(context).cancel(NotificationIds.UNREGISTERED_NOTIFICATION_ID)
 
-      SvrRepository.onRegistrationComplete(response.masterKey, response.pin, hasPin, reglockEnabled)
+      val masterKey = if (data.masterKey != null) MasterKey(data.masterKey.toByteArray()) else null
+      SvrRepository.onRegistrationComplete(masterKey, data.pin, hasPin, data.reglockEnabled)
 
       AppDependencies.resetNetwork()
       AppDependencies.incomingMessageObserver

app/src/main/java/org/thoughtcrime/securesms/registration/ui/RegistrationViewModel.kt

@@ -35,6 +35,7 @@ import org.thoughtcrime.securesms.pin.SvrRepository
 import org.thoughtcrime.securesms.pin.SvrWrongPinException
 import org.thoughtcrime.securesms.registration.RegistrationData
 import org.thoughtcrime.securesms.registration.RegistrationUtil
+import org.thoughtcrime.securesms.registration.data.LocalRegistrationMetadataUtil
 import org.thoughtcrime.securesms.registration.data.RegistrationRepository
 import org.thoughtcrime.securesms.registration.data.network.BackupAuthCheckResult
 import org.thoughtcrime.securesms.registration.data.network.Challenge
@@ -815,7 +816,11 @@ class RegistrationViewModel : ViewModel() {
 
   private suspend fun onSuccessfulRegistration(context: Context, registrationData: RegistrationData, remoteResult: RegistrationRepository.AccountRegistrationResult, reglockEnabled: Boolean) {
     Log.v(TAG, "onSuccessfulRegistration()")
-    RegistrationRepository.registerAccountLocally(context, registrationData, remoteResult, reglockEnabled)
+    val metadata = LocalRegistrationMetadataUtil.createLocalRegistrationMetadata(registrationData, remoteResult, reglockEnabled)
+    if (RemoteConfig.restoreAfterRegistration) {
+      SignalStore.registration.localRegistrationMetadata = metadata
+    }
+    RegistrationRepository.registerAccountLocally(context, metadata)
 
     if (reglockEnabled) {
       SignalStore.onboarding.clearAll()

app/src/main/java/org/thoughtcrime/securesms/registration/ui/restore/RemoteRestoreActivity.kt

@@ -59,6 +59,7 @@ import org.thoughtcrime.securesms.backup.v2.ui.subscription.RemoteRestoreViewMod
 import org.thoughtcrime.securesms.conversation.v2.registerForLifecycle
 import org.thoughtcrime.securesms.dependencies.AppDependencies
 import org.thoughtcrime.securesms.jobs.ProfileUploadJob
+import org.thoughtcrime.securesms.keyvalue.SignalStore
 import org.thoughtcrime.securesms.profiles.AvatarHelper
 import org.thoughtcrime.securesms.profiles.edit.CreateProfileActivity
 import org.thoughtcrime.securesms.recipients.Recipient
@@ -101,6 +102,7 @@ class RemoteRestoreActivity : BaseActivity() {
           )
           if (state.importState == RemoteRestoreViewModel.ImportState.RESTORED) {
             SideEffect {
+              SignalStore.registration.markRestoreCompleted()
               RegistrationUtil.maybeMarkRegistrationComplete()
               AppDependencies.jobManager.add(ProfileUploadJob())
               startActivity(MainActivity.clearTop(this))