Add post-registration restore for backups v2 as well as error messaging.

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

@@ -5,17 +5,29 @@
 
 package org.thoughtcrime.securesms.registration.ui.restore
 
+import android.net.Uri
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
 import androidx.lifecycle.ViewModel
+import androidx.lifecycle.viewModelScope
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.update
+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.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() {
 
@@ -30,6 +42,8 @@ class EnterBackupKeyViewModel : ViewModel() {
     )
   )
 
+  private val verifyGeneration = AtomicInteger(0)
+
   var backupKey by mutableStateOf("")
     private set
 
@@ -49,6 +63,56 @@ class EnterBackupKeyViewModel : ViewModel() {
     }
   }
 
+  fun verifyLocalBackupKey(selectedTimestamp: Long) {
+    if (!state.value.backupKeyValid) {
+      return
+    }
+
+    val generation = verifyGeneration.incrementAndGet()
+    store.update { it.copy(backupKeyValid = false) }
+
+    viewModelScope.launch(Dispatchers.IO) {
+      val result = verifyKey(selectedTimestamp)
+
+      if (verifyGeneration.get() == generation) {
+        if (result) {
+          store.update { it.copy(backupKeyValid = true) }
+        } else {
+          store.update { it.copy(aepValidationError = AccountEntropyPoolVerification.AEPValidationError.Incorrect) }
+        }
+      }
+    }
+  }
+
+  private fun verifyKey(selectedTimestamp: Long): Boolean {
+    try {
+      val aep = AccountEntropyPool.parseOrNull(backupKey) ?: return false
+
+      val dirUri = SignalStore.backup.newLocalBackupsDirectory ?: return false
+      val archiveFileSystem = ArchiveFileSystem.fromUri(AppDependencies.application, Uri.parse(dirUri)) ?: return false
+      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)
+    } catch (e: Exception) {
+      Log.w(TAG, "Failed to verify local backup key", e)
+      return false
+    }
+  }
+
   fun registering() {
     store.update { it.copy(isRegistering = true) }
   }

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

@@ -27,23 +27,30 @@ import androidx.compose.runtime.Composable
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableIntStateOf
+import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.layout.onSizeChanged
+import androidx.compose.ui.platform.LocalContext
 import androidx.compose.ui.platform.LocalDensity
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.text.style.TextAlign
 import androidx.compose.ui.unit.dp
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import org.signal.core.ui.compose.DayNightPreviews
+import org.signal.core.ui.compose.Dialogs
 import org.signal.core.ui.compose.Previews
 import org.signal.core.ui.compose.horizontalGutters
 import org.signal.core.ui.compose.theme.SignalTheme
 import org.thoughtcrime.securesms.BaseActivity
 import org.thoughtcrime.securesms.MainActivity
 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.restore.RestoreActivity
 import kotlin.math.max
 
 /**
@@ -62,6 +69,11 @@ class RestoreLocalBackupActivity : BaseActivity() {
   }
 
   private val viewModel: RestoreLocalBackupActivityViewModel by viewModels()
+  private val contactSupportViewModel: ContactSupportViewModel<Unit> by viewModels()
+
+  private val finishActivity by lazy {
+    intent.getBooleanExtra(KEY_FINISH, false)
+  }
 
   override fun onCreate(savedInstanceState: Bundle?) {
     super.onCreate(savedInstanceState)
@@ -73,26 +85,51 @@ class RestoreLocalBackupActivity : BaseActivity() {
         when (state.restorePhase) {
           RestorePhase.COMPLETE -> {
             startActivity(MainActivity.clearTop(this@RestoreLocalBackupActivity))
-            if (intent.getBooleanExtra(KEY_FINISH, false)) {
+            if (finishActivity) {
               finishAffinity()
             }
           }
+
           RestorePhase.FAILED -> {
             Toast.makeText(this@RestoreLocalBackupActivity, getString(R.string.RestoreLocalBackupActivity__backup_restore_failed), Toast.LENGTH_LONG).show()
           }
+
           else -> Unit
         }
       }
 
+      val contactSupportState by contactSupportViewModel.state.collectAsStateWithLifecycle()
+      val context = LocalContext.current
+
       SignalTheme {
-        RestoreLocalBackupScreen(state = state)
+        RestoreLocalBackupScreen(
+          state = state,
+          onContactSupportClick = contactSupportViewModel::showContactSupport,
+          onFailureDialogConfirm = {
+            if (finishActivity) {
+              viewModel.resetRestoreState()
+              startActivity(RestoreActivity.getRestoreIntent(context))
+            }
+
+            // User invocation here should always finish, it just shouldn't route back to RestoreActivity.
+            supportFinishAfterTransition()
+          },
+          contactSupportState = contactSupportState,
+          contactSupportCallbacks = contactSupportViewModel
+        )
       }
     }
   }
 }
 
 @Composable
-private fun RestoreLocalBackupScreen(state: RestoreLocalBackupScreenState) {
+private fun RestoreLocalBackupScreen(
+  state: RestoreLocalBackupScreenState,
+  onFailureDialogConfirm: () -> Unit,
+  onContactSupportClick: () -> Unit,
+  contactSupportState: ContactSupportViewModel.ContactSupportState<Unit>,
+  contactSupportCallbacks: ContactSupportCallbacks
+) {
   val density = LocalDensity.current
   var headerHeightPx by remember { mutableIntStateOf(0) }
   var contentHeightPx by remember { mutableIntStateOf(0) }
@@ -177,6 +214,33 @@ private fun RestoreLocalBackupScreen(state: RestoreLocalBackupScreenState) {
         }
       }
     }
+
+    if (state.restorePhase == RestorePhase.FAILED) {
+      var wasContactSupportShown by remember { mutableStateOf(false) }
+      LaunchedEffect(contactSupportState.show) {
+        if (wasContactSupportShown && !contactSupportState.show) {
+          onFailureDialogConfirm()
+        }
+
+        wasContactSupportShown = contactSupportState.show
+      }
+
+      if (!contactSupportState.show) {
+        Dialogs.SimpleAlertDialog(
+          title = stringResource(R.string.RestoreLocalBackupActivity__cant_restore_backup),
+          body = stringResource(R.string.RestoreLocalBackupActivity__error_occurred_while_restoring),
+          confirm = stringResource(android.R.string.ok),
+          onConfirm = onFailureDialogConfirm,
+          dismiss = stringResource(R.string.RestoreLocalBackupActivity__contact_support),
+          onDeny = onContactSupportClick
+        )
+      } else {
+        ContactSupportDialog(
+          showInProgress = contactSupportState.showAsProgress,
+          callbacks = contactSupportCallbacks
+        )
+      }
+    }
   }
 }
 
@@ -184,6 +248,12 @@ private fun RestoreLocalBackupScreen(state: RestoreLocalBackupScreenState) {
 @Composable
 private fun RestoreLocalBackupScreenPreview() {
   Previews.Preview {
-    RestoreLocalBackupScreen(state = RestoreLocalBackupScreenState())
+    RestoreLocalBackupScreen(
+      state = RestoreLocalBackupScreenState(),
+      onFailureDialogConfirm = {},
+      onContactSupportClick = {},
+      contactSupportState = ContactSupportViewModel.ContactSupportState(),
+      contactSupportCallbacks = ContactSupportCallbacks.Empty
+    )
   }
 }

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

@@ -133,6 +133,10 @@ class RestoreLocalBackupActivityViewModel : ViewModel() {
       }
     }
   }
+
+  fun resetRestoreState() {
+    SignalStore.registration.restoreDecisionState = RestoreDecisionState(decisionState = RestoreDecisionState.State.START)
+  }
 }
 
 data class RestoreLocalBackupScreenState(

app/src/main/java/org/thoughtcrime/securesms/restore/RestoreViewModel.kt

@@ -25,6 +25,7 @@ import org.thoughtcrime.securesms.keyvalue.skippedRestoreChoice
 import org.thoughtcrime.securesms.registration.data.QuickRegistrationRepository
 import org.thoughtcrime.securesms.registration.ui.restore.RestoreMethod
 import org.thoughtcrime.securesms.registration.ui.restore.StorageServiceRestore
+import org.thoughtcrime.securesms.util.Environment
 import org.whispersystems.signalservice.api.provisioning.RestoreMethod as ApiRestoreMethod
 
 /**
@@ -59,7 +60,11 @@ class RestoreViewModel : ViewModel() {
 
   fun getAvailableRestoreMethods(): List<RestoreMethod> {
     if (SignalStore.registration.isOtherDeviceAndroid || SignalStore.registration.restoreDecisionState.skippedRestoreChoice) {
-      val methods = mutableListOf(RestoreMethod.FROM_LOCAL_BACKUP_V1)
+      val methods = if (Environment.Backups.isNewFormatSupportedForLocalBackup()) {
+        mutableListOf(RestoreMethod.FROM_LOCAL_BACKUP_V2)
+      } else {
+        mutableListOf(RestoreMethod.FROM_LOCAL_BACKUP_V1)
+      }
 
       if (SignalStore.registration.isOtherDeviceAndroid && SignalStore.registration.restoreDecisionState.includeDeviceToDeviceTransfer) {
         methods.add(0, RestoreMethod.FROM_OLD_DEVICE)

app/src/main/java/org/thoughtcrime/securesms/restore/local/PostRegistrationRestoreLocalBackupFragment.kt

@@ -0,0 +1,132 @@
+/*
+ * Copyright 2026 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.restore.local
+
+import android.content.Context
+import android.net.Uri
+import androidx.activity.compose.LocalActivity
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.remember
+import androidx.fragment.app.FragmentActivity
+import androidx.fragment.app.activityViewModels
+import androidx.fragment.app.viewModels
+import androidx.lifecycle.compose.collectAsStateWithLifecycle
+import androidx.lifecycle.lifecycleScope
+import androidx.navigation.fragment.findNavController
+import androidx.navigationevent.compose.LocalNavigationEventDispatcherOwner
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
+import org.signal.core.models.AccountEntropyPool
+import org.signal.core.ui.compose.ComposeFragment
+import org.signal.core.ui.compose.theme.SignalTheme
+import org.thoughtcrime.securesms.MainActivity
+import org.thoughtcrime.securesms.keyvalue.SignalStore
+import org.thoughtcrime.securesms.registration.ui.restore.EnterBackupKeyViewModel
+import org.thoughtcrime.securesms.registration.ui.restore.local.RestoreLocalBackupActivity
+import org.thoughtcrime.securesms.registration.ui.restore.local.RestoreLocalBackupCallback
+import org.thoughtcrime.securesms.registration.ui.restore.local.RestoreLocalBackupNavDisplay
+import org.thoughtcrime.securesms.registration.ui.restore.local.RestoreLocalBackupViewModel
+import org.thoughtcrime.securesms.registration.ui.restore.local.SelectableBackup
+import org.thoughtcrime.securesms.restore.RestoreViewModel
+import org.thoughtcrime.securesms.util.CommunicationActions
+import org.thoughtcrime.securesms.util.navigation.safeNavigate
+
+/**
+ * Post Registration restore fragment for V2 backups.
+ */
+class PostRegistrationRestoreLocalBackupFragment : ComposeFragment() {
+
+  companion object {
+    private const val LEARN_MORE_URL = "https://support.signal.org/hc/articles/360007059752"
+  }
+
+  private val sharedViewModel: RestoreViewModel by activityViewModels()
+  private val restoreLocalBackupViewModel by viewModels<RestoreLocalBackupViewModel>()
+  private val enterBackupKeyViewModel by viewModels<EnterBackupKeyViewModel>()
+
+  @Composable
+  override fun FragmentContent() {
+    val state by restoreLocalBackupViewModel.state.collectAsStateWithLifecycle()
+    val enterBackupKeyState by enterBackupKeyViewModel.state.collectAsStateWithLifecycle()
+
+    SignalTheme {
+      val activity = LocalActivity.current as FragmentActivity
+      CompositionLocalProvider(LocalNavigationEventDispatcherOwner provides activity) {
+        RestoreLocalBackupNavDisplay(
+          state = state,
+          callback = remember { Callbacks() },
+          isRegistrationInProgress = false,
+          enterBackupKeyState = enterBackupKeyState,
+          backupKey = enterBackupKeyViewModel.backupKey
+        )
+      }
+    }
+  }
+
+  private inner class Callbacks : RestoreLocalBackupCallback {
+    override fun setSelectedBackup(backup: SelectableBackup) {
+      restoreLocalBackupViewModel.setSelectedBackup(backup)
+    }
+
+    override fun setSelectedBackupDirectory(context: Context, uri: Uri): Boolean {
+      return restoreLocalBackupViewModel.setSelectedBackupDirectory(context, uri)
+    }
+
+    override fun displaySkipRestoreWarning() {
+      restoreLocalBackupViewModel.displaySkipRestoreWarning()
+    }
+
+    override fun clearDialog() {
+      restoreLocalBackupViewModel.clearDialog()
+    }
+
+    override fun skipRestore() {
+      sharedViewModel.skipRestore()
+
+      viewLifecycleOwner.lifecycleScope.launch {
+        sharedViewModel.performStorageServiceAccountRestoreIfNeeded()
+
+        withContext(Dispatchers.Main) {
+          startActivity(MainActivity.clearTop(requireContext()))
+          activity?.finish()
+        }
+      }
+    }
+
+    override fun submitBackupKey() {
+      val aep = AccountEntropyPool.parseOrNull(enterBackupKeyViewModel.backupKey) ?: return
+      SignalStore.account.restoreAccountEntropyPool(aep)
+
+      val selectedTimestamp = restoreLocalBackupViewModel.state.value.selectedBackup?.timestamp ?: -1L
+      SignalStore.backup.newLocalBackupsSelectedSnapshotTimestamp = selectedTimestamp
+
+      startActivity(RestoreLocalBackupActivity.getIntent(requireContext()))
+      requireActivity().supportFinishAfterTransition()
+    }
+
+    override fun routeToLegacyBackupRestoration(uri: Uri) {
+      sharedViewModel.setBackupFileUri(uri)
+      findNavController().safeNavigate(PostRegistrationRestoreLocalBackupFragmentDirections.restoreLocalV1Backup())
+    }
+
+    override fun onBackupKeyChanged(key: String) {
+      enterBackupKeyViewModel.updateBackupKey(key)
+      val timestamp = restoreLocalBackupViewModel.state.value.selectedBackup?.timestamp ?: return
+      enterBackupKeyViewModel.verifyLocalBackupKey(timestamp)
+    }
+
+    override fun clearRegistrationError() {
+      enterBackupKeyViewModel.clearRegistrationError()
+    }
+
+    override fun onBackupKeyHelp() {
+      CommunicationActions.openBrowserLink(requireContext(), LEARN_MORE_URL)
+    }
+  }
+}

app/src/main/java/org/thoughtcrime/securesms/restore/selection/SelectRestoreMethodFragment.kt

@@ -98,7 +98,7 @@ class SelectRestoreMethodFragment : ComposeFragment() {
       }
       RestoreMethod.FROM_OLD_DEVICE -> findNavController().safeNavigate(SelectRestoreMethodFragmentDirections.goToDeviceTransfer())
       RestoreMethod.FROM_LOCAL_BACKUP_V1 -> findNavController().safeNavigate(SelectRestoreMethodFragmentDirections.goToLocalBackupRestore())
-      RestoreMethod.FROM_LOCAL_BACKUP_V2 -> error("Not currently supported")
+      RestoreMethod.FROM_LOCAL_BACKUP_V2 -> findNavController().safeNavigate(SelectRestoreMethodFragmentDirections.goToLocalBackupRestoreV2())
     }
   }
 }

app/src/main/res/navigation/restore.xml

@@ -59,6 +59,14 @@
             app:popEnterAnim="@anim/nav_default_pop_enter_anim"
             app:popExitAnim="@anim/nav_default_pop_exit_anim" />
 
+        <action
+            android:id="@+id/go_to_local_backup_restore_v2"
+            app:destination="@id/choose_local_backup_fragment_v2"
+            app:enterAnim="@anim/nav_default_enter_anim"
+            app:exitAnim="@anim/nav_default_exit_anim"
+            app:popEnterAnim="@anim/nav_default_pop_enter_anim"
+            app:popExitAnim="@anim/nav_default_pop_exit_anim" />
+
         <action
             android:id="@+id/go_to_post_restore_enter_backup_key"
             app:destination="@id/enterBackupKey"
@@ -93,6 +101,17 @@
             app:destination="@id/restore_local_backup_fragment" />
     </fragment>
 
+    <fragment
+        android:id="@+id/choose_local_backup_fragment_v2"
+        android:name="org.thoughtcrime.securesms.restore.local.PostRegistrationRestoreLocalBackupFragment"
+        android:label="choose_local_backup"
+        tools:layout="@layout/fragment_choose_backup">
+
+        <action
+            android:id="@+id/restore_local_v1_backup"
+            app:destination="@id/restore_local_backup_fragment" />
+    </fragment>
+
     <fragment
         android:id="@+id/restore_local_backup_fragment"
         android:name="org.thoughtcrime.securesms.restore.restorelocalbackup.RestoreLocalBackupFragment"

app/src/main/res/values/strings.xml

@@ -8964,6 +8964,15 @@
     <!-- RestoreLocalBackupActivity: Progress text showing bytes read of total with percentage, e.g. "1.2 MB of 5.0 MB (24%)" -->
     <string name="RestoreLocalBackupActivity__s_of_s_d_percent">%1$s of %2$s (%3$d%%)</string>
 
+    <!-- RestoreLocalBackupActivity: Title for the dialog shown when backup restore fails and is not recoverable -->
+    <string name="RestoreLocalBackupActivity__cant_restore_backup">Can\'t restore backup</string>
+
+    <!-- RestoreLocalBackupActivity: Body for the dialog shown when backup restore fails and is not recoverable -->
+    <string name="RestoreLocalBackupActivity__error_occurred_while_restoring">An error occurred while restoring your backup. Your backup is not recoverable. Please contact support for help.</string>
+
+    <!-- RestoreLocalBackupActivity: Label for the button that opens the contact support form -->
+    <string name="RestoreLocalBackupActivity__contact_support">Contact support</string>
+
     <!-- SelectInstructionsSheet: Title for the select backup folder instruction sheet -->
     <string name="SelectInstructionsSheet__select_your_backup_folder">Select your backup folder</string>
     <!-- SelectInstructionsSheet: Instruction to tap the select this folder button -->