Refactor how archive service access is managed during restore.

2026-03-03 15:58:40 +00:00 · 2024-11-06 10:30:54 -05:00
parent c878da30ae
commit 743e2aaa82
7 changed files with 174 additions and 184 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
@@ -72,7 +72,8 @@ import org.whispersystems.signalservice.api.StatusCodeErrorAction
 import org.whispersystems.signalservice.api.archive.ArchiveGetMediaItemsResponse
 import org.whispersystems.signalservice.api.archive.ArchiveMediaRequest
 import org.whispersystems.signalservice.api.archive.ArchiveMediaResponse
-import org.whispersystems.signalservice.api.archive.ArchiveServiceCredentialPair
+import org.whispersystems.signalservice.api.archive.ArchiveServiceAccess
+import org.whispersystems.signalservice.api.archive.ArchiveServiceAccessPair
 import org.whispersystems.signalservice.api.archive.DeleteArchivedMediaRequest
 import org.whispersystems.signalservice.api.archive.GetArchiveCdnCredentialsResponse
 import org.whispersystems.signalservice.api.backup.MediaName
@@ -723,28 +724,24 @@ object BackupRepository {
  }

  fun listRemoteMediaObjects(limit: Int, cursor: String? = null): NetworkResult<ArchiveGetMediaItemsResponse> {
-    val backupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+    return initBackupAndFetchAuth()
      .then { credential ->
-        SignalNetwork.archive.getArchiveMediaItemsPage(mediaRootBackupKey, SignalStore.account.requireAci(), credential.mediaCredential, limit, cursor)
+        SignalNetwork.archive.getArchiveMediaItemsPage(SignalStore.account.requireAci(), credential.mediaBackupAccess, limit, cursor)
+      }.runOnStatusCodeError {
+        SignalStore.backup.mediaCredentials.clearAll()
      }
  }

  fun getRemoteBackupUsedSpace(): NetworkResult<Long?> {
-    val backupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+    return initBackupAndFetchAuth()
      .then { credential ->
-        SignalNetwork.archive.getBackupInfo(mediaRootBackupKey, SignalStore.account.requireAci(), credential.mediaCredential)
+        SignalNetwork.archive.getBackupInfo(SignalStore.account.requireAci(), credential.mediaBackupAccess)
          .map { it.usedSpace }
      }
  }

  /**
-   * If backups are enabled, sync with the network. Otherwise, return a 404.a
+   * If backups are enabled, sync with the network. Otherwise, return a 404.
   * Used in instrumentation tests.
   */
  fun getBackupTier(): NetworkResult<MessageBackupTier> {
@@ -761,12 +758,9 @@ object BackupRepository {
   * to be the case.
   */
  private fun getBackupTier(aci: ACI): NetworkResult<MessageBackupTier> {
-    val backupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+    return initBackupAndFetchAuth()
      .map { credential ->
-        val zkCredential = SignalNetwork.archive.getZkCredential(backupKey, aci, credential.messageCredential)
+        val zkCredential = SignalNetwork.archive.getZkCredential(aci, credential.messageBackupAccess)
        if (zkCredential.backupLevel == BackupLevel.PAID) {
          MessageBackupTier.PAID
        } else {
@@ -779,17 +773,14 @@ object BackupRepository {
   * Returns an object with details about the remote backup state.
   */
  fun getRemoteBackupState(): NetworkResult<BackupMetadata> {
-    val backupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+    return initBackupAndFetchAuth()
      .then { credential ->
-        SignalNetwork.archive.getBackupInfo(mediaRootBackupKey, SignalStore.account.requireAci(), credential.mediaCredential)
+        SignalNetwork.archive.getBackupInfo(SignalStore.account.requireAci(), credential.mediaBackupAccess)
          .map { it to credential }
      }
      .then { pair ->
        val (mediaBackupInfo, credential) = pair
-        SignalNetwork.archive.debugGetUploadedMediaItemMetadata(mediaRootBackupKey, SignalStore.account.requireAci(), credential.mediaCredential)
+        SignalNetwork.archive.debugGetUploadedMediaItemMetadata(SignalStore.account.requireAci(), credential.mediaBackupAccess)
          .also { Log.i(TAG, "MediaItemMetadataResult: $it") }
          .map { mediaObjects ->
            BackupMetadata(
@@ -806,12 +797,9 @@ object BackupRepository {
   * @return True if successful, otherwise false.
   */
  fun uploadBackupFile(backupStream: InputStream, backupStreamLength: Long): NetworkResult<Unit> {
-    val backupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+    return initBackupAndFetchAuth()
      .then { credential ->
-        SignalNetwork.archive.getMessageBackupUploadForm(backupKey, SignalStore.account.requireAci(), credential.messageCredential)
+        SignalNetwork.archive.getMessageBackupUploadForm(SignalStore.account.requireAci(), credential.messageBackupAccess)
          .also { Log.i(TAG, "UploadFormResult: $it") }
      }
      .then { form ->
@@ -826,29 +814,23 @@ object BackupRepository {
      }
  }

-  fun downloadBackupFile(destination: File, listener: ProgressListener? = null): Boolean {
-    val backupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+  fun downloadBackupFile(destination: File, listener: ProgressListener? = null): NetworkResult<Unit> {
+    return initBackupAndFetchAuth()
      .then { credential ->
-        SignalNetwork.archive.getBackupInfo(backupKey, SignalStore.account.requireAci(), credential.messageCredential)
+        SignalNetwork.archive.getBackupInfo(SignalStore.account.requireAci(), credential.messageBackupAccess)
      }
      .then { info -> getCdnReadCredentials(CredentialType.MESSAGE, info.cdn ?: Cdn.CDN_3.cdnNumber).map { it.headers to info } }
      .map { pair ->
        val (cdnCredentials, info) = pair
        val messageReceiver = AppDependencies.signalServiceMessageReceiver
        messageReceiver.retrieveBackup(info.cdn!!, cdnCredentials, "backups/${info.backupDir}/${info.backupName}", destination, listener)
-      } is NetworkResult.Success
+      }
  }

  fun getBackupFileLastModified(): NetworkResult<ZonedDateTime?> {
-    val backupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+    return initBackupAndFetchAuth()
      .then { credential ->
-        SignalNetwork.archive.getBackupInfo(backupKey, SignalStore.account.requireAci(), credential.messageCredential)
+        SignalNetwork.archive.getBackupInfo(SignalStore.account.requireAci(), credential.messageBackupAccess)
      }
      .then { info -> getCdnReadCredentials(CredentialType.MESSAGE, info.cdn ?: Cdn.CDN_3.cdnNumber).map { it.headers to info } }
      .then { pair ->
@@ -864,12 +846,9 @@ object BackupRepository {
   * Returns an object with details about the remote backup state.
   */
  private fun debugGetArchivedMediaState(): NetworkResult<List<ArchiveGetMediaItemsResponse.StoredMediaObject>> {
-    val backupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+    return initBackupAndFetchAuth()
      .then { credential ->
-        SignalNetwork.archive.debugGetUploadedMediaItemMetadata(mediaRootBackupKey, SignalStore.account.requireAci(), credential.mediaCredential)
+        SignalNetwork.archive.debugGetUploadedMediaItemMetadata(SignalStore.account.requireAci(), credential.mediaBackupAccess)
      }
  }

@@ -880,12 +859,9 @@ object BackupRepository {
   * It's important to note that in order to get this to the archive cdn, you still need to use [copyAttachmentToArchive].
   */
  fun getAttachmentUploadForm(): NetworkResult<AttachmentUploadForm> {
-    val backupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+    return initBackupAndFetchAuth()
      .then { credential ->
-        SignalNetwork.archive.getMediaUploadForm(mediaRootBackupKey, SignalStore.account.requireAci(), credential.mediaCredential)
+        SignalNetwork.archive.getMediaUploadForm(SignalStore.account.requireAci(), credential.mediaBackupAccess)
      }
  }

@@ -893,16 +869,13 @@ object BackupRepository {
   * Copies a thumbnail that has been uploaded to the transit cdn to the archive cdn.
   */
  fun copyThumbnailToArchive(thumbnailAttachment: Attachment, parentAttachment: DatabaseAttachment): NetworkResult<ArchiveMediaResponse> {
-    val backupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-    val request = thumbnailAttachment.toArchiveMediaRequest(parentAttachment.getThumbnailMediaName(), mediaRootBackupKey)
-
-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+    return initBackupAndFetchAuth()
      .then { credential ->
+        val request = thumbnailAttachment.toArchiveMediaRequest(parentAttachment.getThumbnailMediaName(), credential.mediaBackupAccess.backupKey)
+
        SignalNetwork.archive.copyAttachmentToArchive(
-          mediaRootBackupKey = mediaRootBackupKey,
          aci = SignalStore.account.requireAci(),
-          serviceCredential = credential.mediaCredential,
+          archiveServiceAccess = credential.mediaBackupAccess,
          item = request
        )
      }
@@ -912,34 +885,28 @@ object BackupRepository {
   * Copies an attachment that has been uploaded to the transit cdn to the archive cdn.
   */
  fun copyAttachmentToArchive(attachment: DatabaseAttachment): NetworkResult<Unit> {
-    val backupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+    return initBackupAndFetchAuth()
      .then { credential ->
        val mediaName = attachment.getMediaName()
-        val request = attachment.toArchiveMediaRequest(mediaName, mediaRootBackupKey)
+        val request = attachment.toArchiveMediaRequest(mediaName, credential.mediaBackupAccess.backupKey)
        SignalNetwork.archive
          .copyAttachmentToArchive(
-            mediaRootBackupKey = mediaRootBackupKey,
            aci = SignalStore.account.requireAci(),
-            serviceCredential = credential.mediaCredential,
+            archiveServiceAccess = credential.mediaBackupAccess,
            item = request
          )
-          .map { Triple(mediaName, request.mediaId, it) }
+          .map { credential to Triple(mediaName, request.mediaId, it) }
      }
-      .map { (mediaName, mediaId, response) ->
-        val thumbnailId = mediaRootBackupKey.deriveMediaId(attachment.getThumbnailMediaName()).encode()
+      .map { (credential, triple) ->
+        val (mediaName, mediaId, response) = triple
+        val thumbnailId = credential.mediaBackupAccess.backupKey.deriveMediaId(attachment.getThumbnailMediaName()).encode()
        SignalDatabase.attachments.setArchiveData(attachmentId = attachment.attachmentId, archiveCdn = response.cdn, archiveMediaName = mediaName.name, archiveMediaId = mediaId, archiveThumbnailMediaId = thumbnailId)
      }
      .also { Log.i(TAG, "archiveMediaResult: $it") }
  }

  fun copyAttachmentToArchive(databaseAttachments: List<DatabaseAttachment>): NetworkResult<BatchArchiveMediaResult> {
-    val backupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+    return initBackupAndFetchAuth()
      .then { credential ->
        val requests = mutableListOf<ArchiveMediaRequest>()
        val mediaIdToAttachmentId = mutableMapOf<String, AttachmentId>()
@@ -947,7 +914,7 @@ object BackupRepository {

        databaseAttachments.forEach {
          val mediaName = it.getMediaName()
-          val request = it.toArchiveMediaRequest(mediaName, mediaRootBackupKey)
+          val request = it.toArchiveMediaRequest(mediaName, credential.mediaBackupAccess.backupKey)
          requests += request
          mediaIdToAttachmentId[request.mediaId] = it.attachmentId
          attachmentIdToMediaName[it.attachmentId] = mediaName.name
@@ -955,20 +922,19 @@ object BackupRepository {

        SignalNetwork.archive
          .copyAttachmentToArchive(
-            mediaRootBackupKey = mediaRootBackupKey,
            aci = SignalStore.account.requireAci(),
-            serviceCredential = credential.mediaCredential,
+            archiveServiceAccess = credential.mediaBackupAccess,
            items = requests
          )
-          .map { BatchArchiveMediaResult(it, mediaIdToAttachmentId, attachmentIdToMediaName) }
+          .map { credential to BatchArchiveMediaResult(it, mediaIdToAttachmentId, attachmentIdToMediaName) }
      }
-      .map { result ->
+      .map { (credential, result) ->
        result
          .successfulResponses
          .forEach {
            val attachmentId = result.mediaIdToAttachmentId(it.mediaId)
            val mediaName = result.attachmentIdToMediaName(attachmentId)
-            val thumbnailId = mediaRootBackupKey.deriveMediaId(MediaName.forThumbnailFromMediaName(mediaName = mediaName)).encode()
+            val thumbnailId = credential.mediaBackupAccess.backupKey.deriveMediaId(MediaName.forThumbnailFromMediaName(mediaName = mediaName)).encode()
            SignalDatabase.attachments.setArchiveData(attachmentId = attachmentId, archiveCdn = it.cdn!!, archiveMediaName = mediaName, archiveMediaId = it.mediaId, thumbnailId)
          }
        result
@@ -977,9 +943,6 @@ object BackupRepository {
  }

  fun deleteArchivedMedia(attachments: List<DatabaseAttachment>): NetworkResult<Unit> {
-    val backupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
    val mediaToDelete = attachments
      .filter { it.archiveMediaId != null }
      .map {
@@ -994,12 +957,11 @@ object BackupRepository {
      return NetworkResult.Success(Unit)
    }

-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+    return initBackupAndFetchAuth()
      .then { credential ->
        SignalNetwork.archive.deleteArchivedMedia(
-          mediaRootBackupKey = mediaRootBackupKey,
          aci = SignalStore.account.requireAci(),
-          serviceCredential = credential.mediaCredential,
+          archiveServiceAccess = credential.mediaBackupAccess,
          mediaToDelete = mediaToDelete
        )
      }
@@ -1010,9 +972,6 @@ object BackupRepository {
  }

  fun deleteAbandonedMediaObjects(mediaObjects: Collection<ArchivedMediaObject>): NetworkResult<Unit> {
-    val backupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
    val mediaToDelete = mediaObjects
      .map {
        DeleteArchivedMediaRequest.ArchivedMediaObject(
@@ -1026,12 +985,11 @@ object BackupRepository {
      return NetworkResult.Success(Unit)
    }

-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+    return initBackupAndFetchAuth()
      .then { credential ->
        SignalNetwork.archive.deleteArchivedMedia(
-          mediaRootBackupKey = mediaRootBackupKey,
          aci = SignalStore.account.requireAci(),
-          serviceCredential = credential.mediaCredential,
+          archiveServiceAccess = credential.mediaBackupAccess,
          mediaToDelete = mediaToDelete
        )
      }
@@ -1039,8 +997,6 @@ object BackupRepository {
  }

  fun debugDeleteAllArchivedMedia(): NetworkResult<Unit> {
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
    return debugGetArchivedMediaState()
      .then { archivedMedia ->
        val mediaToDelete = archivedMedia
@@ -1055,12 +1011,11 @@ object BackupRepository {
          Log.i(TAG, "No media to delete, quick success")
          NetworkResult.Success(Unit)
        } else {
-          getAuthCredentialPair()
+          getArchiveServiceAccessPair()
            .then { credential ->
              SignalNetwork.archive.deleteArchivedMedia(
-                mediaRootBackupKey = mediaRootBackupKey,
                aci = SignalStore.account.requireAci(),
-                serviceCredential = credential.mediaCredential,
+                archiveServiceAccess = credential.mediaBackupAccess,
                mediaToDelete = mediaToDelete
              )
            }
@@ -1086,24 +1041,17 @@ object BackupRepository {
      return NetworkResult.Success(cached)
    }

-    val messageBackupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
-    val credentialBackupKey = when (credentialType) {
-      CredentialType.MESSAGE -> messageBackupKey
-      CredentialType.MEDIA -> mediaRootBackupKey
-    }
-
-    return initBackupAndFetchAuth(messageBackupKey, mediaRootBackupKey)
+    return initBackupAndFetchAuth()
      .then { credential ->
+        val archiveServiceAccess = when (credentialType) {
+          CredentialType.MESSAGE -> credential.messageBackupAccess
+          CredentialType.MEDIA -> credential.mediaBackupAccess
+        }
+
        SignalNetwork.archive.getCdnReadCredentials(
          cdnNumber = cdnNumber,
-          backupKey = credentialBackupKey,
          aci = SignalStore.account.requireAci(),
-          serviceCredential = when (credentialType) {
-            CredentialType.MESSAGE -> credential.messageCredential
-            CredentialType.MEDIA -> credential.mediaCredential
-          }
+          archiveServiceAccess = archiveServiceAccess
        )
      }
      .also {
@@ -1152,12 +1100,9 @@ object BackupRepository {
      return NetworkResult.Success(cachedMediaPath)
    }

-    val backupKey = SignalStore.backup.messageBackupKey
-    val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
-
-    return initBackupAndFetchAuth(backupKey, mediaRootBackupKey)
+    return initBackupAndFetchAuth()
      .then { credential ->
-        SignalNetwork.archive.getBackupInfo(mediaRootBackupKey, SignalStore.account.requireAci(), credential.mediaCredential).map {
+        SignalNetwork.archive.getBackupInfo(SignalStore.account.requireAci(), credential.mediaBackupAccess).map {
          SignalStore.backup.usedBackupMediaSpace = it.usedSpace ?: 0L
          "${it.backupDir!!.urlEncode()}/${it.mediaDir!!.urlEncode()}"
        }
@@ -1221,18 +1166,27 @@ object BackupRepository {
  }

  /**
-   * Ensures that the backupId has been reserved and that your public key has been set, while also returning an auth credential.
-   * Should be the basis of all backup operations.
+   * During normal operation, ensures that the backupId has been reserved and that your public key has been set,
+   * while also returning an archive access data. Should be the basis of all backup operations.
+   *
+   * When called during registration before backups are initialized, will only fetch access data and not initialize backups. This
+   * prevents early initialization with incorrect keys before we have restored them.
   */
-  private fun initBackupAndFetchAuth(messageBackupKey: MessageBackupKey, mediaRootBackupKey: MediaRootBackupKey): NetworkResult<ArchiveServiceCredentialPair> {
+  private fun initBackupAndFetchAuth(): NetworkResult<ArchiveServiceAccessPair> {
    return if (SignalStore.backup.backupsInitialized) {
-      getAuthCredentialPair().runOnStatusCodeError(resetInitializedStateErrorAction)
+      getArchiveServiceAccessPair().runOnStatusCodeError(resetInitializedStateErrorAction)
+    } else if (isPreRestoreDuringRegistration()) {
+      Log.w(TAG, "Requesting/using auth credentials in pre-restore state")
+      getArchiveServiceAccessPair()
    } else {
+      val messageBackupKey = SignalStore.backup.messageBackupKey
+      val mediaRootBackupKey = SignalStore.backup.mediaRootBackupKey
+
      return SignalNetwork.archive
        .triggerBackupIdReservation(messageBackupKey, mediaRootBackupKey, SignalStore.account.requireAci())
-        .then { getAuthCredentialPair() }
-        .then { credential -> SignalNetwork.archive.setPublicKey(messageBackupKey, SignalStore.account.requireAci(), credential.messageCredential).map { credential } }
-        .then { credential -> SignalNetwork.archive.setPublicKey(mediaRootBackupKey, SignalStore.account.requireAci(), credential.mediaCredential).map { credential } }
+        .then { getArchiveServiceAccessPair() }
+        .then { credential -> SignalNetwork.archive.setPublicKey(SignalStore.account.requireAci(), credential.messageBackupAccess).map { credential } }
+        .then { credential -> SignalNetwork.archive.setPublicKey(SignalStore.account.requireAci(), credential.mediaBackupAccess).map { credential } }
        .runIfSuccessful { SignalStore.backup.backupsInitialized = true }
        .runOnStatusCodeError(resetInitializedStateErrorAction)
    }
@@ -1241,14 +1195,19 @@ object BackupRepository {
  /**
   * Retrieves an auth credential, preferring a cached value if available.
   */
-  private fun getAuthCredentialPair(): NetworkResult<ArchiveServiceCredentialPair> {
+  private fun getArchiveServiceAccessPair(): NetworkResult<ArchiveServiceAccessPair> {
    val currentTime = System.currentTimeMillis()

    val messageCredential = SignalStore.backup.messageCredentials.byDay.getForCurrentTime(currentTime.milliseconds)
    val mediaCredential = SignalStore.backup.mediaCredentials.byDay.getForCurrentTime(currentTime.milliseconds)

    if (messageCredential != null && mediaCredential != null) {
-      return NetworkResult.Success(ArchiveServiceCredentialPair(messageCredential, mediaCredential))
+      return NetworkResult.Success(
+        ArchiveServiceAccessPair(
+          messageBackupAccess = ArchiveServiceAccess(messageCredential, SignalStore.backup.messageBackupKey),
+          mediaBackupAccess = ArchiveServiceAccess(mediaCredential, SignalStore.backup.mediaRootBackupKey)
+        )
+      )
    }

    Log.w(TAG, "No credentials found for today, need to fetch new ones! This shouldn't happen under normal circumstances. We should ensure the routine fetch is running properly.")
@@ -1260,13 +1219,20 @@ object BackupRepository {
      SignalStore.backup.mediaCredentials.add(result.mediaCredentials)
      SignalStore.backup.mediaCredentials.clearOlderThan(currentTime)

-      ArchiveServiceCredentialPair(
-        messageCredential = SignalStore.backup.messageCredentials.byDay.getForCurrentTime(currentTime.milliseconds)!!,
-        mediaCredential = SignalStore.backup.mediaCredentials.byDay.getForCurrentTime(currentTime.milliseconds)!!
+      ArchiveServiceAccessPair(
+        messageBackupAccess = ArchiveServiceAccess(SignalStore.backup.messageCredentials.byDay.getForCurrentTime(currentTime.milliseconds)!!, SignalStore.backup.messageBackupKey),
+        mediaBackupAccess = ArchiveServiceAccess(SignalStore.backup.mediaCredentials.byDay.getForCurrentTime(currentTime.milliseconds)!!, SignalStore.backup.mediaRootBackupKey)
      )
    }
  }

+  private fun isPreRestoreDuringRegistration(): Boolean {
+    return !SignalStore.registration.isRegistrationComplete &&
+      !SignalStore.registration.hasCompletedRestore() &&
+      !SignalStore.registration.hasSkippedTransferOrRestore() &&
+      RemoteConfig.restoreAfterRegistration
+  }
+
  private fun File.deleteAllFilesWithPrefix(prefix: String) {
    this.listFiles()?.filter { it.name.startsWith(prefix) }?.forEach { it.delete() }
  }
--- a/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/internal/backup/InternalBackupPlaygroundViewModel.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/internal/backup/InternalBackupPlaygroundViewModel.kt
@@ -548,9 +548,13 @@ class InternalBackupPlaygroundViewModel : ViewModel() {
    SignalExecutors.BOUNDED_IO.execute {
      Log.d(TAG, "Downloading file...")
      val tempBackupFile = BlobProvider.getInstance().forNonAutoEncryptingSingleSessionOnDisk(AppDependencies.application)
-      if (!BackupRepository.downloadBackupFile(tempBackupFile)) {
-        Log.e(TAG, "Failed to download backup file")
-        throw IOException()
+
+      when (val result = BackupRepository.downloadBackupFile(tempBackupFile)) {
+        is NetworkResult.Success -> Log.i(TAG, "Download successful")
+        else -> {
+          Log.w(TAG, "Failed to download backup file", result.getCause())
+          throw IOException(result.getCause())
+        }
      }

      val encryptedStream = tempBackupFile.inputStream()
--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/BackupRestoreJob.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/BackupRestoreJob.kt
@@ -20,6 +20,7 @@ import org.thoughtcrime.securesms.net.NotPushRegisteredException
 import org.thoughtcrime.securesms.providers.BlobProvider
 import org.thoughtcrime.securesms.recipients.Recipient
 import org.thoughtcrime.securesms.service.BackupProgressService
+import org.whispersystems.signalservice.api.NetworkResult
 import org.whispersystems.signalservice.api.messages.SignalServiceAttachment.ProgressListener
 import java.io.IOException

@@ -81,9 +82,12 @@ class BackupRestoreJob private constructor(parameters: Parameters) : BaseJob(par
    }

    val tempBackupFile = BlobProvider.getInstance().forNonAutoEncryptingSingleSessionOnDisk(AppDependencies.application)
-    if (!BackupRepository.downloadBackupFile(tempBackupFile, progressListener)) {
-      Log.e(TAG, "Failed to download backup file")
-      throw IOException()
+    when (val result = BackupRepository.downloadBackupFile(tempBackupFile, progressListener)) {
+      is NetworkResult.Success -> Log.i(TAG, "Download successful")
+      else -> {
+        Log.w(TAG, "Failed to download backup file", result.getCause())
+        throw IOException(result.getCause())
+      }
    }

    if (isCanceled) {
--- a/libsignal-service/src/main/java/org/whispersystems/signalservice/api/archive/ArchiveApi.kt
+++ b/libsignal-service/src/main/java/org/whispersystems/signalservice/api/archive/ArchiveApi.kt
@@ -71,10 +71,10 @@ class ArchiveApi(private val pushServiceSocket: PushServiceSocket) {
   * - 403: Forbidden
   * - 429: Rate-limited
   */
-  fun getCdnReadCredentials(cdnNumber: Int, backupKey: BackupKey, aci: ACI, serviceCredential: ArchiveServiceCredential): NetworkResult<GetArchiveCdnCredentialsResponse> {
+  fun getCdnReadCredentials(cdnNumber: Int, aci: ACI, archiveServiceAccess: ArchiveServiceAccess<*>): NetworkResult<GetArchiveCdnCredentialsResponse> {
    return NetworkResult.fromFetch {
-      val zkCredential = getZkCredential(backupKey, aci, serviceCredential)
-      val presentationData = CredentialPresentationData.from(backupKey, aci, zkCredential, backupServerPublicParams)
+      val zkCredential = getZkCredential(aci, archiveServiceAccess)
+      val presentationData = CredentialPresentationData.from(archiveServiceAccess.backupKey, aci, zkCredential, backupServerPublicParams)

      pushServiceSocket.getArchiveCdnReadCredentials(cdnNumber, presentationData.toArchiveCredentialPresentation())
    }
@@ -111,10 +111,10 @@ class ArchiveApi(private val pushServiceSocket: PushServiceSocket) {
   * - 403: Forbidden
   * - 429: Rate-limited
   */
-  fun setPublicKey(backupKey: BackupKey, aci: ACI, serviceCredential: ArchiveServiceCredential): NetworkResult<Unit> {
+  fun setPublicKey(aci: ACI, archiveServiceAccess: ArchiveServiceAccess<*>): NetworkResult<Unit> {
    return NetworkResult.fromFetch {
-      val zkCredential = getZkCredential(backupKey, aci, serviceCredential)
-      val presentationData = CredentialPresentationData.from(backupKey, aci, zkCredential, backupServerPublicParams)
+      val zkCredential = getZkCredential(aci, archiveServiceAccess)
+      val presentationData = CredentialPresentationData.from(archiveServiceAccess.backupKey, aci, zkCredential, backupServerPublicParams)
      pushServiceSocket.setArchivePublicKey(presentationData.publicKey, presentationData.toArchiveCredentialPresentation())
    }
  }
@@ -128,10 +128,10 @@ class ArchiveApi(private val pushServiceSocket: PushServiceSocket) {
   * - 403: Insufficient permissions
   * - 429: Rate-limited
   */
-  fun getMessageBackupUploadForm(messageBackupKey: MessageBackupKey, aci: ACI, serviceCredential: ArchiveServiceCredential): NetworkResult<AttachmentUploadForm> {
+  fun getMessageBackupUploadForm(aci: ACI, archiveServiceAccess: ArchiveServiceAccess<MessageBackupKey>): NetworkResult<AttachmentUploadForm> {
    return NetworkResult.fromFetch {
-      val zkCredential = getZkCredential(messageBackupKey, aci, serviceCredential)
-      val presentationData = CredentialPresentationData.from(messageBackupKey, aci, zkCredential, backupServerPublicParams)
+      val zkCredential = getZkCredential(aci, archiveServiceAccess)
+      val presentationData = CredentialPresentationData.from(archiveServiceAccess.backupKey, aci, zkCredential, backupServerPublicParams)
      pushServiceSocket.getArchiveMessageBackupUploadForm(presentationData.toArchiveCredentialPresentation())
    }
  }
@@ -142,10 +142,10 @@ class ArchiveApi(private val pushServiceSocket: PushServiceSocket) {
   *
   * Will return a [NetworkResult.StatusCodeError] with status code 404 if you haven't uploaded a backup yet.
   */
-  fun getBackupInfo(backupKey: BackupKey, aci: ACI, messageServiceCredential: ArchiveServiceCredential): NetworkResult<ArchiveGetBackupInfoResponse> {
+  fun getBackupInfo(aci: ACI, archiveServiceAccess: ArchiveServiceAccess<*>): NetworkResult<ArchiveGetBackupInfoResponse> {
    return NetworkResult.fromFetch {
-      val zkCredential = getZkCredential(backupKey, aci, messageServiceCredential)
-      val presentationData = CredentialPresentationData.from(backupKey, aci, zkCredential, backupServerPublicParams)
+      val zkCredential = getZkCredential(aci, archiveServiceAccess)
+      val presentationData = CredentialPresentationData.from(archiveServiceAccess.backupKey, aci, zkCredential, backupServerPublicParams)
      pushServiceSocket.getArchiveBackupInfo(presentationData.toArchiveCredentialPresentation())
    }
  }
@@ -153,10 +153,10 @@ class ArchiveApi(private val pushServiceSocket: PushServiceSocket) {
  /**
   * Lists the media objects in the backup
   */
-  fun listMediaObjects(mediaRootBackupKey: MediaRootBackupKey, aci: ACI, serviceCredential: ArchiveServiceCredential, limit: Int, cursor: String? = null): NetworkResult<ArchiveGetMediaItemsResponse> {
+  fun listMediaObjects(aci: ACI, archiveServiceAccess: ArchiveServiceAccess<MediaRootBackupKey>, limit: Int, cursor: String? = null): NetworkResult<ArchiveGetMediaItemsResponse> {
    return NetworkResult.fromFetch {
-      val zkCredential = getZkCredential(mediaRootBackupKey, aci, serviceCredential)
-      val presentationData = CredentialPresentationData.from(mediaRootBackupKey, aci, zkCredential, backupServerPublicParams)
+      val zkCredential = getZkCredential(aci, archiveServiceAccess)
+      val presentationData = CredentialPresentationData.from(archiveServiceAccess.backupKey, aci, zkCredential, backupServerPublicParams)
      pushServiceSocket.getArchiveMediaItemsPage(presentationData.toArchiveCredentialPresentation(), limit, cursor)
    }
  }
@@ -194,10 +194,10 @@ class ArchiveApi(private val pushServiceSocket: PushServiceSocket) {
   * - 403: Forbidden
   * - 429: Rate-limited
   */
-  fun getMediaUploadForm(mediaRootBackupKey: MediaRootBackupKey, aci: ACI, serviceCredential: ArchiveServiceCredential): NetworkResult<AttachmentUploadForm> {
+  fun getMediaUploadForm(aci: ACI, archiveServiceAccess: ArchiveServiceAccess<MediaRootBackupKey>): NetworkResult<AttachmentUploadForm> {
    return NetworkResult.fromFetch {
-      val zkCredential = getZkCredential(mediaRootBackupKey, aci, serviceCredential)
-      val presentationData = CredentialPresentationData.from(mediaRootBackupKey, aci, zkCredential, backupServerPublicParams)
+      val zkCredential = getZkCredential(aci, archiveServiceAccess)
+      val presentationData = CredentialPresentationData.from(archiveServiceAccess.backupKey, aci, zkCredential, backupServerPublicParams)
      pushServiceSocket.getArchiveMediaUploadForm(presentationData.toArchiveCredentialPresentation())
    }
  }
@@ -206,13 +206,13 @@ class ArchiveApi(private val pushServiceSocket: PushServiceSocket) {
   * Retrieves all media items in the user's archive. Note that this could be a very large number of items, making this only suitable for debugging.
   * Use [getArchiveMediaItemsPage] in production.
   */
-  fun debugGetUploadedMediaItemMetadata(mediaRootBackupKey: MediaRootBackupKey, aci: ACI, serviceCredential: ArchiveServiceCredential): NetworkResult<List<StoredMediaObject>> {
+  fun debugGetUploadedMediaItemMetadata(aci: ACI, archiveServiceAccess: ArchiveServiceAccess<MediaRootBackupKey>): NetworkResult<List<StoredMediaObject>> {
    return NetworkResult.fromFetch {
      val mediaObjects: MutableList<StoredMediaObject> = ArrayList()

      var cursor: String? = null
      do {
-        val response: ArchiveGetMediaItemsResponse = getArchiveMediaItemsPage(mediaRootBackupKey, aci, serviceCredential, 512, cursor).successOrThrow()
+        val response: ArchiveGetMediaItemsResponse = getArchiveMediaItemsPage(aci, archiveServiceAccess, 512, cursor).successOrThrow()
        mediaObjects += response.storedMediaObjects
        cursor = response.cursor
      } while (cursor != null)
@@ -226,10 +226,10 @@ class ArchiveApi(private val pushServiceSocket: PushServiceSocket) {
   * @param limit The maximum number of items to return.
   * @param cursor A token that can be read from your previous response, telling the server where to start the next page.
   */
-  fun getArchiveMediaItemsPage(mediaRootBackupKey: MediaRootBackupKey, aci: ACI, mediaServiceCredential: ArchiveServiceCredential, limit: Int, cursor: String?): NetworkResult<ArchiveGetMediaItemsResponse> {
+  fun getArchiveMediaItemsPage(aci: ACI, archiveServiceAccess: ArchiveServiceAccess<MediaRootBackupKey>, limit: Int, cursor: String?): NetworkResult<ArchiveGetMediaItemsResponse> {
    return NetworkResult.fromFetch {
-      val zkCredential = getZkCredential(mediaRootBackupKey, aci, mediaServiceCredential)
-      val presentationData = CredentialPresentationData.from(mediaRootBackupKey, aci, zkCredential, backupServerPublicParams)
+      val zkCredential = getZkCredential(aci, archiveServiceAccess)
+      val presentationData = CredentialPresentationData.from(archiveServiceAccess.backupKey, aci, zkCredential, backupServerPublicParams)

      pushServiceSocket.getArchiveMediaItemsPage(presentationData.toArchiveCredentialPresentation(), limit, cursor)
    }
@@ -247,14 +247,13 @@ class ArchiveApi(private val pushServiceSocket: PushServiceSocket) {
   *   429: Rate-limited
   */
  fun copyAttachmentToArchive(
-    mediaRootBackupKey: MediaRootBackupKey,
    aci: ACI,
-    serviceCredential: ArchiveServiceCredential,
+    archiveServiceAccess: ArchiveServiceAccess<MediaRootBackupKey>,
    item: ArchiveMediaRequest
  ): NetworkResult<ArchiveMediaResponse> {
    return NetworkResult.fromFetch {
-      val zkCredential = getZkCredential(mediaRootBackupKey, aci, serviceCredential)
-      val presentationData = CredentialPresentationData.from(mediaRootBackupKey, aci, zkCredential, backupServerPublicParams)
+      val zkCredential = getZkCredential(aci, archiveServiceAccess)
+      val presentationData = CredentialPresentationData.from(archiveServiceAccess.backupKey, aci, zkCredential, backupServerPublicParams)

      pushServiceSocket.archiveAttachmentMedia(presentationData.toArchiveCredentialPresentation(), item)
    }
@@ -264,14 +263,13 @@ class ArchiveApi(private val pushServiceSocket: PushServiceSocket) {
   * Copy and re-encrypt media from the attachments cdn into the backup cdn.
   */
  fun copyAttachmentToArchive(
-    mediaRootBackupKey: MediaRootBackupKey,
    aci: ACI,
-    serviceCredential: ArchiveServiceCredential,
+    archiveServiceAccess: ArchiveServiceAccess<MediaRootBackupKey>,
    items: List<ArchiveMediaRequest>
  ): NetworkResult<BatchArchiveMediaResponse> {
    return NetworkResult.fromFetch {
-      val zkCredential = getZkCredential(mediaRootBackupKey, aci, serviceCredential)
-      val presentationData = CredentialPresentationData.from(mediaRootBackupKey, aci, zkCredential, backupServerPublicParams)
+      val zkCredential = getZkCredential(aci, archiveServiceAccess)
+      val presentationData = CredentialPresentationData.from(archiveServiceAccess.backupKey, aci, zkCredential, backupServerPublicParams)

      val request = BatchArchiveMediaRequest(items = items)

@@ -290,27 +288,26 @@ class ArchiveApi(private val pushServiceSocket: PushServiceSocket) {
   * - 429: Rate-limited
   */
  fun deleteArchivedMedia(
-    mediaRootBackupKey: MediaRootBackupKey,
    aci: ACI,
-    serviceCredential: ArchiveServiceCredential,
+    archiveServiceAccess: ArchiveServiceAccess<MediaRootBackupKey>,
    mediaToDelete: List<DeleteArchivedMediaRequest.ArchivedMediaObject>
  ): NetworkResult<Unit> {
    return NetworkResult.fromFetch {
-      val zkCredential = getZkCredential(mediaRootBackupKey, aci, serviceCredential)
-      val presentationData = CredentialPresentationData.from(mediaRootBackupKey, aci, zkCredential, backupServerPublicParams)
+      val zkCredential = getZkCredential(aci, archiveServiceAccess)
+      val presentationData = CredentialPresentationData.from(archiveServiceAccess.backupKey, aci, zkCredential, backupServerPublicParams)
      val request = DeleteArchivedMediaRequest(mediaToDelete = mediaToDelete)

      pushServiceSocket.deleteArchivedMedia(presentationData.toArchiveCredentialPresentation(), request)
    }
  }

-  fun getZkCredential(backupKey: BackupKey, aci: ACI, serviceCredential: ArchiveServiceCredential): BackupAuthCredential {
-    val backupAuthResponse = BackupAuthCredentialResponse(serviceCredential.credential)
-    val backupRequestContext = BackupAuthCredentialRequestContext.create(backupKey.value, aci.rawUuid)
+  fun getZkCredential(aci: ACI, archiveServiceAccess: ArchiveServiceAccess<*>): BackupAuthCredential {
+    val backupAuthResponse = BackupAuthCredentialResponse(archiveServiceAccess.credential.credential)
+    val backupRequestContext = BackupAuthCredentialRequestContext.create(archiveServiceAccess.backupKey.value, aci.rawUuid)

    return backupRequestContext.receiveResponse(
      backupAuthResponse,
-      Instant.ofEpochSecond(serviceCredential.redemptionTime),
+      Instant.ofEpochSecond(archiveServiceAccess.credential.redemptionTime),
      backupServerPublicParams
    )
  }
--- a/libsignal-service/src/main/java/org/whispersystems/signalservice/api/archive/ArchiveServiceAccess.kt
+++ b/libsignal-service/src/main/java/org/whispersystems/signalservice/api/archive/ArchiveServiceAccess.kt
@@ -0,0 +1,16 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.signalservice.api.archive
+
+import org.whispersystems.signalservice.api.backup.BackupKey
+
+/**
+ * Key and credential combo needed to perform backup operations on the server.
+ */
+class ArchiveServiceAccess<T : BackupKey>(
+  val credential: ArchiveServiceCredential,
+  val backupKey: T
+)
--- a/libsignal-service/src/main/java/org/whispersystems/signalservice/api/archive/ArchiveServiceAccessPair.kt
+++ b/libsignal-service/src/main/java/org/whispersystems/signalservice/api/archive/ArchiveServiceAccessPair.kt
@@ -0,0 +1,17 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.signalservice.api.archive
+
+import org.whispersystems.signalservice.api.backup.MediaRootBackupKey
+import org.whispersystems.signalservice.api.backup.MessageBackupKey
+
+/**
+ * A convenient container for passing around both a message and media archive service credential.
+ */
+data class ArchiveServiceAccessPair(
+  val messageBackupAccess: ArchiveServiceAccess<MessageBackupKey>,
+  val mediaBackupAccess: ArchiveServiceAccess<MediaRootBackupKey>
+)
--- a/libsignal-service/src/main/java/org/whispersystems/signalservice/api/archive/ArchiveServiceCredentialPair.kt
+++ b/libsignal-service/src/main/java/org/whispersystems/signalservice/api/archive/ArchiveServiceCredentialPair.kt
@@ -1,14 +0,0 @@
-/*
- * Copyright 2024 Signal Messenger, LLC
- * SPDX-License-Identifier: AGPL-3.0-only
- */
-
-package org.whispersystems.signalservice.api.archive
-
-/**
- * A convenient container for passing around both a message and media archive service credential.
- */
-data class ArchiveServiceCredentialPair(
-  val messageCredential: ArchiveServiceCredential,
-  val mediaCredential: ArchiveServiceCredential
-)