Fix local backup restore AEP handling and conditional re-enable.

2026-04-19 08:09:12 +01:00 · 2026-03-19 12:50:00 -03:00
parent c7a6c7ad9e
commit 78d3db319c
16 changed files with 528 additions and 58 deletions
--- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/BackupRepository.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/BackupRepository.kt
@@ -21,6 +21,7 @@ import org.greenrobot.eventbus.EventBus
 import org.signal.core.models.AccountEntropyPool
 import org.signal.core.models.ServiceId.ACI
 import org.signal.core.models.ServiceId.PNI
+import org.signal.core.models.backup.BackupId
 import org.signal.core.models.backup.MediaName
 import org.signal.core.models.backup.MediaRootBackupKey
 import org.signal.core.models.backup.MessageBackupKey
@@ -1088,13 +1089,13 @@ object BackupRepository {
  /**
   * Imports a local backup file that was exported to disk.
   */
-  fun importLocal(mainStreamFactory: () -> InputStream, mainStreamLength: Long, selfData: SelfData): ImportResult {
-    val backupKey = SignalStore.backup.messageBackupKey
+  fun importLocal(mainStreamFactory: () -> InputStream, mainStreamLength: Long, selfData: SelfData, backupId: BackupId, messageBackupKey: MessageBackupKey): ImportResult {
+    val backupKey = messageBackupKey

    val frameReader = try {
      EncryptedBackupReader.createForLocalOrLinking(
        key = backupKey,
-        aci = selfData.aci,
+        backupId = backupId,
        length = mainStreamLength,
        dataStream = mainStreamFactory
      )
--- 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
@@ -8,6 +8,7 @@ package org.thoughtcrime.securesms.backup.v2.local
 import okio.ByteString.Companion.toByteString
 import org.signal.core.models.backup.BackupId
 import org.signal.core.models.backup.MediaName
+import org.signal.core.models.backup.MessageBackupKey
 import org.signal.core.util.Stopwatch
 import org.signal.core.util.StreamUtil
 import org.signal.core.util.Util
@@ -17,6 +18,7 @@ import org.thoughtcrime.securesms.attachments.AttachmentId
 import org.thoughtcrime.securesms.backup.v2.BackupRepository
 import org.thoughtcrime.securesms.backup.v2.local.proto.FilesFrame
 import org.thoughtcrime.securesms.backup.v2.local.proto.Metadata
+import org.thoughtcrime.securesms.backup.v2.stream.EncryptedBackupReader
 import org.thoughtcrime.securesms.database.AttachmentTable
 import org.thoughtcrime.securesms.keyvalue.SignalStore
 import org.thoughtcrime.securesms.keyvalue.protos.LocalBackupCreationProgress
@@ -152,7 +154,7 @@ object LocalArchiver {
  /**
   * Import archive data from a folder on the system. Does not restore attachments.
   */
-  fun import(snapshotFileSystem: SnapshotFileSystem, selfData: BackupRepository.SelfData): RestoreResult {
+  fun import(snapshotFileSystem: SnapshotFileSystem, selfData: BackupRepository.SelfData, messageBackupKey: MessageBackupKey): RestoreResult {
    var metadataStream: InputStream? = null

    try {
@@ -169,19 +171,16 @@ object LocalArchiver {
        return RestoreResult.failure(RestoreFailure.BackupIdMissing)
      }

-      val backupId = decryptBackupId(metadata.backupId)
-
-      if (!backupId.value.contentEquals(SignalStore.backup.messageBackupKey.deriveBackupId(SignalStore.account.requireAci()).value)) {
-        Log.w(TAG, "Local backup metadata backup id does not match derived backup id, likely from another account")
-        return RestoreResult.failure(RestoreFailure.BackupIdMismatch)
-      }
+      val backupId = decryptBackupId(metadata.backupId, messageBackupKey)

      val mainStreamLength = snapshotFileSystem.mainLength() ?: return ArchiveResult.failure(RestoreFailure.MainStream)

      BackupRepository.importLocal(
        mainStreamFactory = { snapshotFileSystem.mainInputStream()!! },
        mainStreamLength = mainStreamLength,
-        selfData = selfData
+        selfData = selfData,
+        backupId = backupId,
+        messageBackupKey = messageBackupKey
      )
    } finally {
      metadataStream?.close()
@@ -190,8 +189,41 @@ object LocalArchiver {
    return RestoreResult.success(RestoreSuccess.FullSuccess)
  }

-  private fun decryptBackupId(encryptedBackupId: Metadata.EncryptedBackupId): BackupId {
-    val metadataKey = SignalStore.backup.messageBackupKey.deriveLocalBackupMetadataKey()
+  /**
+   * Verifies that the provided [messageBackupKey] can decrypt and authenticate the snapshot's main archive.
+   */
+  fun canDecryptMainArchive(snapshotFileSystem: SnapshotFileSystem, messageBackupKey: MessageBackupKey): Boolean {
+    return try {
+      val backupId = getBackupId(snapshotFileSystem, messageBackupKey) ?: return false
+      val mainStreamLength = snapshotFileSystem.mainLength() ?: return false
+
+      EncryptedBackupReader.createForLocalOrLinking(
+        key = messageBackupKey,
+        backupId = backupId,
+        length = mainStreamLength,
+        dataStream = { snapshotFileSystem.mainInputStream() ?: error("Missing main archive stream") }
+      ).use { reader ->
+        reader.getHeader() != null
+      }
+    } catch (e: Exception) {
+      Log.w(TAG, "Unable to verify local backup archive", e)
+      false
+    }
+  }
+
+  fun getBackupId(snapshotFileSystem: SnapshotFileSystem, messageBackupKey: MessageBackupKey): BackupId? {
+    return try {
+      val metadata = snapshotFileSystem.metadataInputStream()?.use { Metadata.ADAPTER.decode(it) } ?: return null
+      val encryptedBackupId = metadata.backupId ?: return null
+      decryptBackupId(encryptedBackupId, messageBackupKey)
+    } catch (e: Exception) {
+      Log.w(TAG, "Unable to decrypt local backup id", e)
+      null
+    }
+  }
+
+  private fun decryptBackupId(encryptedBackupId: Metadata.EncryptedBackupId, messageBackupKey: MessageBackupKey): BackupId {
+    val metadataKey = messageBackupKey.deriveLocalBackupMetadataKey()
    val iv = encryptedBackupId.iv.toByteArray()
    val backupIdCipher = encryptedBackupId.encryptedId.toByteArray()

@@ -227,7 +259,6 @@ object LocalArchiver {
    data object MainStream : RestoreFailure
    data object Cancelled : RestoreFailure
    data object BackupIdMissing : RestoreFailure
-    data object BackupIdMismatch : RestoreFailure
    data class VersionMismatch(val backupVersion: Int, val supportedVersion: Int) : RestoreFailure
  }

--- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/stream/EncryptedBackupReader.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/stream/EncryptedBackupReader.kt
@@ -8,6 +8,7 @@ package org.thoughtcrime.securesms.backup.v2.stream
 import androidx.annotation.VisibleForTesting
 import com.google.common.io.CountingInputStream
 import org.signal.core.models.ServiceId.ACI
+import org.signal.core.models.backup.BackupId
 import org.signal.core.models.backup.MessageBackupKey
 import org.signal.core.util.readFully
 import org.signal.core.util.readNBytesOrThrow
@@ -66,9 +67,23 @@ class EncryptedBackupReader private constructor(
      forwardSecrecyToken: BackupForwardSecrecyToken,
      length: Long,
      dataStream: () -> InputStream
+    ): EncryptedBackupReader {
+      return createForSignalBackup(key, key.deriveBackupId(aci), forwardSecrecyToken, length, dataStream)
+    }
+
+    /**
+     * Create a reader for a backup from the archive CDN, using a [BackupId] directly
+     * instead of deriving it from an ACI.
+     */
+    fun createForSignalBackup(
+      key: MessageBackupKey,
+      backupId: BackupId,
+      forwardSecrecyToken: BackupForwardSecrecyToken,
+      length: Long,
+      dataStream: () -> InputStream
    ): EncryptedBackupReader {
      return EncryptedBackupReader(
-        keyMaterial = key.deriveBackupSecrets(aci, forwardSecrecyToken),
+        keyMaterial = key.deriveBackupSecrets(backupId, forwardSecrecyToken),
        length = length,
        dataStream = dataStream
      )
@@ -79,8 +94,16 @@ class EncryptedBackupReader private constructor(
     * The key difference is that we don't require forward secrecy data.
     */
    fun createForLocalOrLinking(key: MessageBackupKey, aci: ACI, length: Long, dataStream: () -> InputStream): EncryptedBackupReader {
+      return createForLocalOrLinking(key, key.deriveBackupId(aci), length, dataStream)
+    }
+
+    /**
+     * Create a reader for a local backup or for a transfer to a linked device, using a [BackupId] directly
+     * instead of deriving it from an ACI.
+     */
+    fun createForLocalOrLinking(key: MessageBackupKey, backupId: BackupId, length: Long, dataStream: () -> InputStream): EncryptedBackupReader {
      return EncryptedBackupReader(
-        keyMaterial = key.deriveBackupSecrets(aci, forwardSecrecyToken = null),
+        keyMaterial = key.deriveBackupSecrets(backupId, forwardSecrecyToken = null),
        length = length,
        dataStream = dataStream
      )
--- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/stream/EncryptedBackupWriter.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/stream/EncryptedBackupWriter.kt
@@ -6,6 +6,7 @@
 package org.thoughtcrime.securesms.backup.v2.stream

 import org.signal.core.models.ServiceId.ACI
+import org.signal.core.models.backup.BackupId
 import org.signal.core.models.backup.MessageBackupKey
 import org.signal.core.util.Util
 import org.signal.core.util.stream.MacOutputStream
@@ -13,7 +14,6 @@ import org.signal.core.util.writeVarInt32
 import org.signal.libsignal.messagebackup.BackupForwardSecrecyToken
 import org.thoughtcrime.securesms.backup.v2.proto.BackupInfo
 import org.thoughtcrime.securesms.backup.v2.proto.Frame
-import org.thoughtcrime.securesms.backup.v2.stream.EncryptedBackupReader.Companion.createForSignalBackup
 import java.io.IOException
 import java.io.OutputStream
 import javax.crypto.Cipher
@@ -29,8 +29,7 @@ import javax.crypto.spec.SecretKeySpec
 * to the end of the [outputStream].
 */
 class EncryptedBackupWriter private constructor(
-  key: MessageBackupKey,
-  aci: ACI,
+  keyMaterial: MessageBackupKey.BackupKeyMaterial,
  forwardSecrecyToken: BackupForwardSecrecyToken?,
  forwardSecrecyMetadata: ByteArray?,
  private val outputStream: OutputStream,
@@ -54,10 +53,24 @@ class EncryptedBackupWriter private constructor(
      forwardSecrecyMetadata: ByteArray,
      outputStream: OutputStream,
      append: (ByteArray) -> Unit
+    ): EncryptedBackupWriter {
+      return createForSignalBackup(key, key.deriveBackupId(aci), forwardSecrecyToken, forwardSecrecyMetadata, outputStream, append)
+    }
+
+    /**
+     * Create a writer for a backup from the archive CDN, using a [BackupId] directly
+     * instead of deriving it from an ACI.
+     */
+    fun createForSignalBackup(
+      key: MessageBackupKey,
+      backupId: BackupId,
+      forwardSecrecyToken: BackupForwardSecrecyToken,
+      forwardSecrecyMetadata: ByteArray,
+      outputStream: OutputStream,
+      append: (ByteArray) -> Unit
    ): EncryptedBackupWriter {
      return EncryptedBackupWriter(
-        key = key,
-        aci = aci,
+        keyMaterial = key.deriveBackupSecrets(backupId, forwardSecrecyToken),
        forwardSecrecyToken = forwardSecrecyToken,
        forwardSecrecyMetadata = forwardSecrecyMetadata,
        outputStream = outputStream,
@@ -74,10 +87,22 @@ class EncryptedBackupWriter private constructor(
      aci: ACI,
      outputStream: OutputStream,
      append: (ByteArray) -> Unit
+    ): EncryptedBackupWriter {
+      return createForLocalOrLinking(key, key.deriveBackupId(aci), outputStream, append)
+    }
+
+    /**
+     * Create a writer for a local backup or for a transfer to a linked device, using a [BackupId] directly
+     * instead of deriving it from an ACI.
+     */
+    fun createForLocalOrLinking(
+      key: MessageBackupKey,
+      backupId: BackupId,
+      outputStream: OutputStream,
+      append: (ByteArray) -> Unit
    ): EncryptedBackupWriter {
      return EncryptedBackupWriter(
-        key = key,
-        aci = aci,
+        keyMaterial = key.deriveBackupSecrets(backupId, forwardSecrecyToken = null),
        forwardSecrecyToken = null,
        forwardSecrecyMetadata = null,
        outputStream = outputStream,
@@ -99,8 +124,6 @@ class EncryptedBackupWriter private constructor(
      outputStream.flush()
    }

-    val keyMaterial = key.deriveBackupSecrets(aci, forwardSecrecyToken)
-
    val iv: ByteArray = Util.getSecretBytes(16)
    outputStream.write(iv)
    outputStream.flush()
--- a/app/src/main/java/org/thoughtcrime/securesms/keyvalue/BackupValues.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/keyvalue/BackupValues.kt
@@ -107,6 +107,8 @@ class BackupValues(store: KeyValueStore) : SignalStoreValues(store) {
    private const val KEY_NEW_LOCAL_BACKUPS_SELECTED_SNAPSHOT_TIMESTAMP = "backup.new_local_backups_selected_snapshot_timestamp"
    private const val KEY_NEW_LOCAL_BACKUPS_CREATION_PROGRESS = "backup.new_local_backups_creation_progress"

+    private const val KEY_LOCAL_RESTORE_ACCOUNT_ENTROPY_POOL = "backup.local_restore_account_entropy_pool"
+
    private const val KEY_UPLOAD_BANNER_VISIBLE = "backup.upload_banner_visible"

    private val cachedCdnCredentialsExpiresIn: Duration = 12.hours
@@ -501,6 +503,12 @@ class BackupValues(store: KeyValueStore) : SignalStoreValues(store) {
   */
  var newLocalBackupsSelectedSnapshotTimestamp: Long by longValue(KEY_NEW_LOCAL_BACKUPS_SELECTED_SNAPSHOT_TIMESTAMP, -1L)

+  /**
+   * Temporary storage for the AEP used to decrypt a local backup file. This is kept separate from
+   * the account AEP because the local backup may belong to a different account (e.g., after ACI change).
+   */
+  var localRestoreAccountEntropyPool: String? by stringValue(KEY_LOCAL_RESTORE_ACCOUNT_ENTROPY_POOL, null as String?)
+
  /**
   * When we are told by the server that we are out of storage space, we should show
   * UX treatment to make the user aware of this.
--- a/app/src/main/java/org/thoughtcrime/securesms/registration/ui/RegistrationViewModel.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/registration/ui/RegistrationViewModel.kt
@@ -649,6 +649,11 @@ class RegistrationViewModel : ViewModel() {
    }
  }

+  /** Clears the recovery password from state, e.g. when a backup-key-based registration attempt fails and the stale password must not be retried. */
+  fun clearRecoveryPassword() {
+    setRecoveryPassword(null)
+  }
+
  private fun setRecoveryPassword(recoveryPassword: String?) {
    store.update {
      it.copy(recoveryPassword = recoveryPassword)
@@ -930,7 +935,7 @@ class RegistrationViewModel : ViewModel() {
      Log.w(TAG, "Unable to start auth websocket", e)
    }

-    if (!remoteResult.reRegistration && SignalStore.registration.restoreDecisionState.isDecisionPending) {
+    if (!remoteResult.reRegistration && SignalStore.registration.restoreDecisionState.isDecisionPending && SignalStore.backup.localRestoreAccountEntropyPool == null) {
      Log.v(TAG, "Not re-registration, and still pending restore decision, likely an account with no data to restore, skipping post register restore")
      SignalStore.registration.restoreDecisionState = RestoreDecisionState.NewAccount
    }
--- a/app/src/main/java/org/thoughtcrime/securesms/registration/ui/restore/EnterBackupKeyViewModel.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/registration/ui/restore/EnterBackupKeyViewModel.kt
@@ -19,15 +19,12 @@ import kotlinx.coroutines.launch
 import org.signal.core.models.AccountEntropyPool
 import org.signal.core.util.logging.Log
 import org.thoughtcrime.securesms.backup.v2.local.ArchiveFileSystem
+import org.thoughtcrime.securesms.backup.v2.local.LocalArchiver
 import org.thoughtcrime.securesms.backup.v2.local.SnapshotFileSystem
-import org.thoughtcrime.securesms.backup.v2.local.proto.Metadata
 import org.thoughtcrime.securesms.dependencies.AppDependencies
 import org.thoughtcrime.securesms.keyvalue.SignalStore
 import org.thoughtcrime.securesms.registration.data.network.RegisterAccountResult
 import java.util.concurrent.atomic.AtomicInteger
-import javax.crypto.Cipher
-import javax.crypto.spec.IvParameterSpec
-import javax.crypto.spec.SecretKeySpec

 class EnterBackupKeyViewModel : ViewModel() {

@@ -93,20 +90,7 @@ class EnterBackupKeyViewModel : ViewModel() {
      val snapshot = archiveFileSystem.listSnapshots().firstOrNull { it.timestamp == selectedTimestamp } ?: return false

      val snapshotFs = SnapshotFileSystem(AppDependencies.application, snapshot.file)
-      val metadata = snapshotFs.metadataInputStream()?.use { Metadata.ADAPTER.decode(it) } ?: return false
-      val encryptedBackupId = metadata.backupId ?: return false
-
-      val messageBackupKey = aep.deriveMessageBackupKey()
-      val metadataKey = messageBackupKey.deriveLocalBackupMetadataKey()
-      val iv = encryptedBackupId.iv.toByteArray()
-      val backupIdCipher = encryptedBackupId.encryptedId.toByteArray()
-
-      val cipher = Cipher.getInstance("AES/CTR/NoPadding")
-      cipher.init(Cipher.DECRYPT_MODE, SecretKeySpec(metadataKey, "AES"), IvParameterSpec(iv))
-      val decryptedBackupId = cipher.doFinal(backupIdCipher)
-
-      val expectedBackupId = messageBackupKey.deriveBackupId(SignalStore.account.requireAci())
-      return decryptedBackupId.contentEquals(expectedBackupId.value)
+      return LocalArchiver.canDecryptMainArchive(snapshotFs, aep.deriveMessageBackupKey())
    } catch (e: Exception) {
      Log.w(TAG, "Failed to verify local backup key", e)
      return false
@@ -117,6 +101,11 @@ class EnterBackupKeyViewModel : ViewModel() {
    store.update { it.copy(isRegistering = true) }
  }

+  /** Resets [EnterBackupKeyState.isRegistering] without triggering a registration error dialog. Use when navigation away from this screen is handling the error itself. */
+  fun cancelRegistering() {
+    store.update { it.copy(isRegistering = false) }
+  }
+
  fun handleRegistrationFailure(registerAccountResult: RegisterAccountResult) {
    store.update {
      if (it.isRegistering) {
--- a/app/src/main/java/org/thoughtcrime/securesms/registration/ui/restore/local/RestoreLocalBackupActivity.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/registration/ui/restore/local/RestoreLocalBackupActivity.kt
@@ -50,6 +50,7 @@ import org.thoughtcrime.securesms.R
 import org.thoughtcrime.securesms.components.contactsupport.ContactSupportCallbacks
 import org.thoughtcrime.securesms.components.contactsupport.ContactSupportDialog
 import org.thoughtcrime.securesms.components.contactsupport.ContactSupportViewModel
+import org.thoughtcrime.securesms.components.settings.app.AppSettingsActivity
 import org.thoughtcrime.securesms.restore.RestoreActivity
 import kotlin.math.max

@@ -84,9 +85,11 @@ class RestoreLocalBackupActivity : BaseActivity() {
      LaunchedEffect(state.restorePhase) {
        when (state.restorePhase) {
          RestorePhase.COMPLETE -> {
-            startActivity(MainActivity.clearTop(this@RestoreLocalBackupActivity))
-            if (finishActivity) {
-              finishAffinity()
+            if (!state.showLocalBackupsDisabledDialog) {
+              startActivity(MainActivity.clearTop(this@RestoreLocalBackupActivity))
+              if (finishActivity) {
+                finishAffinity()
+              }
            }
          }

@@ -105,6 +108,25 @@ class RestoreLocalBackupActivity : BaseActivity() {
        RestoreLocalBackupScreen(
          state = state,
          onContactSupportClick = contactSupportViewModel::showContactSupport,
+          onLocalBackupsDisabledDialogConfirm = {
+            viewModel.dismissLocalBackupsDisabledDialog()
+            startActivities(
+              arrayOf(
+                MainActivity.clearTop(this@RestoreLocalBackupActivity),
+                AppSettingsActivity.backups(this@RestoreLocalBackupActivity)
+              )
+            )
+            if (finishActivity) {
+              finishAffinity()
+            }
+          },
+          onLocalBackupsDisabledDialogDeny = {
+            viewModel.dismissLocalBackupsDisabledDialog()
+            startActivity(MainActivity.clearTop(this@RestoreLocalBackupActivity))
+            if (finishActivity) {
+              finishAffinity()
+            }
+          },
          onFailureDialogConfirm = {
            if (finishActivity) {
              viewModel.resetRestoreState()
@@ -126,6 +148,8 @@ class RestoreLocalBackupActivity : BaseActivity() {
 private fun RestoreLocalBackupScreen(
  state: RestoreLocalBackupScreenState,
  onFailureDialogConfirm: () -> Unit,
+  onLocalBackupsDisabledDialogConfirm: () -> Unit,
+  onLocalBackupsDisabledDialogDeny: () -> Unit,
  onContactSupportClick: () -> Unit,
  contactSupportState: ContactSupportViewModel.ContactSupportState<Unit>,
  contactSupportCallbacks: ContactSupportCallbacks
@@ -241,6 +265,17 @@ private fun RestoreLocalBackupScreen(
        )
      }
    }
+
+    if (state.showLocalBackupsDisabledDialog) {
+      Dialogs.SimpleAlertDialog(
+        title = stringResource(R.string.RestoreLocalBackupActivity__turn_on_on_device_backups),
+        body = stringResource(R.string.RestoreLocalBackupActivity__to_continue_using_on_device_backups),
+        confirm = stringResource(R.string.RestoreLocalBackupActivity__turn_on_now),
+        dismiss = stringResource(R.string.RestoreLocalBackupActivity__not_now),
+        onConfirm = onLocalBackupsDisabledDialogConfirm,
+        onDeny = onLocalBackupsDisabledDialogDeny
+      )
+    }
  }
 }

@@ -251,6 +286,8 @@ private fun RestoreLocalBackupScreenPreview() {
    RestoreLocalBackupScreen(
      state = RestoreLocalBackupScreenState(),
      onFailureDialogConfirm = {},
+      onLocalBackupsDisabledDialogConfirm = {},
+      onLocalBackupsDisabledDialogDeny = {},
      onContactSupportClick = {},
      contactSupportState = ContactSupportViewModel.ContactSupportState(),
      contactSupportCallbacks = ContactSupportCallbacks.Empty
--- a/app/src/main/java/org/thoughtcrime/securesms/registration/ui/restore/local/RestoreLocalBackupActivityViewModel.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/registration/ui/restore/local/RestoreLocalBackupActivityViewModel.kt
@@ -16,6 +16,7 @@ import kotlinx.coroutines.launch
 import org.greenrobot.eventbus.EventBus
 import org.greenrobot.eventbus.Subscribe
 import org.greenrobot.eventbus.ThreadMode
+import org.signal.core.models.AccountEntropyPool
 import org.signal.core.util.ByteSize
 import org.signal.core.util.Result
 import org.signal.core.util.bytes
@@ -28,7 +29,6 @@ import org.thoughtcrime.securesms.backup.v2.local.LocalArchiver
 import org.thoughtcrime.securesms.backup.v2.local.SnapshotFileSystem
 import org.thoughtcrime.securesms.database.model.databaseprotos.RestoreDecisionState
 import org.thoughtcrime.securesms.dependencies.AppDependencies
-import org.thoughtcrime.securesms.jobs.LocalBackupJob
 import org.thoughtcrime.securesms.jobs.RestoreLocalAttachmentJob
 import org.thoughtcrime.securesms.keyvalue.Completed
 import org.thoughtcrime.securesms.keyvalue.SignalStore
@@ -110,23 +110,50 @@ class RestoreLocalBackupActivityViewModel : ViewModel() {
        return@launch
      }

+      val localAep = SignalStore.backup.localRestoreAccountEntropyPool
+      if (localAep == null) {
+        Log.w(TAG, "No local restore AEP set")
+        internalState.update { it.copy(restorePhase = RestorePhase.FAILED) }
+        return@launch
+      }
+      val localAepPool = AccountEntropyPool(localAep)
+      val messageBackupKey = localAepPool.deriveMessageBackupKey()
+
      val snapshotFileSystem = SnapshotFileSystem(AppDependencies.application, snapshotInfo.file)
-      val result = LocalArchiver.import(snapshotFileSystem, selfData)
+      val result = LocalArchiver.import(snapshotFileSystem, selfData, messageBackupKey)

      if (result is Result.Success) {
        Log.i(TAG, "Local backup import succeeded")
        val mediaNameToFileInfo = archiveFileSystem.filesFileSystem.allFiles()
        RestoreLocalAttachmentJob.enqueueRestoreLocalAttachmentsJobs(mediaNameToFileInfo)

+        val actualBackupId = LocalArchiver.getBackupId(snapshotFileSystem, messageBackupKey)
+        val expectedBackupId = SignalStore.account.accountEntropyPool
+          .deriveMessageBackupKey()
+          .deriveBackupId(self.aci.get())
+
+        SignalStore.backup.localRestoreAccountEntropyPool = null
        SignalStore.registration.restoreDecisionState = RestoreDecisionState.Completed
        SignalStore.backup.backupSecretRestoreRequired = false
        SignalStore.backup.newLocalBackupsSelectedSnapshotTimestamp = -1L
-        SignalStore.backup.newLocalBackupsEnabled = true
-        LocalBackupJob.enqueueArchive(false)
+
+        val backupIdMatchesCurrentAccount = actualBackupId?.value?.contentEquals(expectedBackupId.value) == true
+        if (backupIdMatchesCurrentAccount) {
+          SignalStore.account.restoreAccountEntropyPool(localAepPool)
+          SignalStore.backup.newLocalBackupsEnabled = true
+        } else {
+          Log.w(TAG, "Local backup does not match current account, not re-enabling local backups")
+        }
+
        StorageServiceRestore.restore()
        RegistrationUtil.maybeMarkRegistrationComplete()

-        internalState.update { it.copy(restorePhase = RestorePhase.COMPLETE) }
+        internalState.update {
+          it.copy(
+            restorePhase = RestorePhase.COMPLETE,
+            showLocalBackupsDisabledDialog = !backupIdMatchesCurrentAccount
+          )
+        }
      } else {
        Log.w(TAG, "Local backup import failed")
        internalState.update { it.copy(restorePhase = RestorePhase.FAILED) }
@@ -134,6 +161,10 @@ class RestoreLocalBackupActivityViewModel : ViewModel() {
    }
  }

+  fun dismissLocalBackupsDisabledDialog() {
+    internalState.update { it.copy(showLocalBackupsDisabledDialog = false) }
+  }
+
  fun resetRestoreState() {
    SignalStore.registration.restoreDecisionState = RestoreDecisionState(decisionState = RestoreDecisionState.State.START)
  }
@@ -143,7 +174,8 @@ data class RestoreLocalBackupScreenState(
  val restorePhase: RestorePhase = RestorePhase.RESTORING,
  val bytesRead: ByteSize = 0L.bytes,
  val totalBytes: ByteSize = 0L.bytes,
-  val progress: Float = 0f
+  val progress: Float = 0f,
+  val showLocalBackupsDisabledDialog: Boolean = false
 )

 enum class RestorePhase {
--- a/app/src/main/java/org/thoughtcrime/securesms/registration/ui/restore/local/RestoreLocalBackupFragment.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/registration/ui/restore/local/RestoreLocalBackupFragment.kt
@@ -34,6 +34,7 @@ import org.signal.core.ui.compose.ComposeFragment
 import org.signal.core.ui.compose.theme.SignalTheme
 import org.signal.core.util.logging.Log
 import org.thoughtcrime.securesms.keyvalue.SignalStore
+import org.thoughtcrime.securesms.registration.data.network.RegisterAccountResult
 import org.thoughtcrime.securesms.registration.ui.RegistrationViewModel
 import org.thoughtcrime.securesms.registration.ui.phonenumber.EnterPhoneNumberMode
 import org.thoughtcrime.securesms.registration.ui.restore.EnterBackupKeyViewModel
@@ -79,7 +80,18 @@ class RestoreLocalBackupFragment : ComposeFragment() {
          .filterNotNull()
          .collect {
            sharedViewModel.registerAccountErrorShown()
-            enterBackupKeyViewModel.handleRegistrationFailure(it)
+            if (it is RegisterAccountResult.IncorrectRecoveryPassword) {
+              SignalStore.account.resetAccountEntropyPool()
+              SignalStore.account.resetAciAndPniIdentityKeysAfterFailedRestore()
+              sharedViewModel.clearRecoveryPassword()
+              enterBackupKeyViewModel.cancelRegistering()
+              sharedViewModel.intendToRestore(hasOldDevice = false, fromRemote = false, fromLocalV2 = true)
+              findNavController().safeNavigate(
+                RestoreLocalBackupFragmentDirections.goToEnterPhoneNumber(EnterPhoneNumberMode.RESTART_AFTER_COLLECTION)
+              )
+            } else {
+              enterBackupKeyViewModel.handleRegistrationFailure(it)
+            }
          }
      }
    }
@@ -137,6 +149,8 @@ class RestoreLocalBackupFragment : ComposeFragment() {
    override fun submitBackupKey() {
      enterBackupKeyViewModel.registering()

+      SignalStore.backup.localRestoreAccountEntropyPool = enterBackupKeyViewModel.backupKey
+
      val selectedTimestamp = restoreLocalBackupViewModel.state.value.selectedBackup?.timestamp ?: -1L
      SignalStore.backup.newLocalBackupsSelectedSnapshotTimestamp = selectedTimestamp

@@ -152,6 +166,8 @@ class RestoreLocalBackupFragment : ComposeFragment() {

    override fun onBackupKeyChanged(key: String) {
      enterBackupKeyViewModel.updateBackupKey(key)
+      val timestamp = restoreLocalBackupViewModel.state.value.selectedBackup?.timestamp ?: return
+      enterBackupKeyViewModel.verifyLocalBackupKey(timestamp)
    }

    override fun clearRegistrationError() {
--- a/app/src/main/java/org/thoughtcrime/securesms/restore/local/PostRegistrationRestoreLocalBackupFragment.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/restore/local/PostRegistrationRestoreLocalBackupFragment.kt
@@ -100,8 +100,9 @@ class PostRegistrationRestoreLocalBackupFragment : ComposeFragment() {
    }

    override fun submitBackupKey() {
-      val aep = AccountEntropyPool.parseOrNull(enterBackupKeyViewModel.backupKey) ?: return
-      SignalStore.account.restoreAccountEntropyPool(aep)
+      AccountEntropyPool.parseOrNull(enterBackupKeyViewModel.backupKey) ?: return
+
+      SignalStore.backup.localRestoreAccountEntropyPool = enterBackupKeyViewModel.backupKey

      val selectedTimestamp = restoreLocalBackupViewModel.state.value.selectedBackup?.timestamp ?: -1L
      SignalStore.backup.newLocalBackupsSelectedSnapshotTimestamp = selectedTimestamp