Add initial support for backup and restore of message and media to staging.

Co-authored-by: Cody Henthorne <cody@signal.org>
2026-04-26 03:40:56 +01:00 · 2024-04-12 11:57:34 -04:00
parent 8617a074ad
commit 689eacd618
71 changed files with 3198 additions and 744 deletions
--- a/app/src/main/java/org/thoughtcrime/securesms/backup/RestoreState.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/RestoreState.kt
@@ -0,0 +1,36 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.backup
+
+import org.signal.core.util.LongSerializer
+
+enum class RestoreState(val id: Int, val inProgress: Boolean) {
+  FAILED(-1, false),
+  NONE(0, false),
+  PENDING(1, true),
+  RESTORING_DB(2, true),
+  RESTORING_MEDIA(3, true);
+
+  companion object {
+    val serializer: LongSerializer<RestoreState> = Serializer()
+  }
+
+  class Serializer : LongSerializer<RestoreState> {
+    override fun serialize(data: RestoreState): Long {
+      return data.id.toLong()
+    }
+
+    override fun deserialize(data: Long): RestoreState {
+      return when (data.toInt()) {
+        FAILED.id -> FAILED
+        PENDING.id -> PENDING
+        RESTORING_DB.id -> RESTORING_DB
+        RESTORING_MEDIA.id -> RESTORING_MEDIA
+        else -> NONE
+      }
+    }
+  }
+}
--- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/BackupRepository.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/BackupRepository.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 Signal Messenger, LLC
+ * Copyright 2024 Signal Messenger, LLC
 * SPDX-License-Identifier: AGPL-3.0-only
 */

@@ -14,6 +14,7 @@ import org.signal.libsignal.messagebackup.MessageBackup.ValidationResult
 import org.signal.libsignal.messagebackup.MessageBackupKey
 import org.signal.libsignal.protocol.ServiceId.Aci
 import org.signal.libsignal.zkgroup.profiles.ProfileKey
+import org.thoughtcrime.securesms.attachments.AttachmentId
 import org.thoughtcrime.securesms.attachments.DatabaseAttachment
 import org.thoughtcrime.securesms.backup.v2.database.ChatItemImportInserter
 import org.thoughtcrime.securesms.backup.v2.database.clearAllDataForBackupRestore
@@ -37,17 +38,20 @@ import org.thoughtcrime.securesms.recipients.RecipientId
 import org.whispersystems.signalservice.api.NetworkResult
 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.ArchiveServiceCredential
-import org.whispersystems.signalservice.api.archive.BatchArchiveMediaResponse
 import org.whispersystems.signalservice.api.archive.DeleteArchivedMediaRequest
+import org.whispersystems.signalservice.api.archive.GetArchiveCdnCredentialsResponse
 import org.whispersystems.signalservice.api.backup.BackupKey
+import org.whispersystems.signalservice.api.backup.MediaName
 import org.whispersystems.signalservice.api.crypto.AttachmentCipherStreamUtil
+import org.whispersystems.signalservice.api.messages.SignalServiceAttachment.ProgressListener
 import org.whispersystems.signalservice.api.push.ServiceId.ACI
 import org.whispersystems.signalservice.api.push.ServiceId.PNI
 import org.whispersystems.signalservice.internal.crypto.PaddingInputStream
 import java.io.ByteArrayOutputStream
+import java.io.File
 import java.io.InputStream
+import java.io.OutputStream
 import kotlin.time.Duration.Companion.milliseconds

 object BackupRepository {
@@ -55,10 +59,8 @@ object BackupRepository {
  private val TAG = Log.tag(BackupRepository::class.java)
  private const val VERSION = 1L

-  fun export(plaintext: Boolean = false): ByteArray {
+  fun export(outputStream: OutputStream, append: (ByteArray) -> Unit, plaintext: Boolean = false) {
    val eventTimer = EventTimer()
-
-    val outputStream = ByteArrayOutputStream()
    val writer: BackupExportWriter = if (plaintext) {
      PlainTextBackupWriter(outputStream)
    } else {
@@ -66,11 +68,11 @@ object BackupRepository {
        key = SignalStore.svr().getOrCreateMasterKey().deriveBackupKey(),
        aci = SignalStore.account().aci!!,
        outputStream = outputStream,
-        append = { mac -> outputStream.write(mac) }
+        append = append
      )
    }

-    val exportState = ExportState(System.currentTimeMillis())
+    val exportState = ExportState(backupTime = System.currentTimeMillis(), allowMediaBackup = true)

    writer.use {
      writer.write(
@@ -110,7 +112,11 @@ object BackupRepository {
    }

    Log.d(TAG, "export() ${eventTimer.stop().summary}")
+  }

+  fun export(plaintext: Boolean = false): ByteArray {
+    val outputStream = ByteArrayOutputStream()
+    export(outputStream = outputStream, append = { mac -> outputStream.write(mac) }, plaintext = plaintext)
    return outputStream.toByteArray()
  }

@@ -124,11 +130,13 @@ object BackupRepository {
  fun import(length: Long, inputStreamFactory: () -> InputStream, selfData: SelfData, plaintext: Boolean = false) {
    val eventTimer = EventTimer()

+    val backupKey = SignalStore.svr().getOrCreateMasterKey().deriveBackupKey()
+
    val frameReader = if (plaintext) {
      PlainTextBackupReader(inputStreamFactory())
    } else {
      EncryptedBackupReader(
-        key = SignalStore.svr().getOrCreateMasterKey().deriveBackupKey(),
+        key = backupKey,
        aci = selfData.aci,
        streamLength = length,
        dataStream = inputStreamFactory
@@ -160,7 +168,7 @@ object BackupRepository {
      SignalDatabase.recipients.setProfileSharing(selfId, true)

      eventTimer.emit("setup")
-      val backupState = BackupState()
+      val backupState = BackupState(backupKey)
      val chatItemInserter: ChatItemImportInserter = ChatItemBackupProcessor.beginImport(backupState)

      for (frame in frameReader) {
@@ -281,6 +289,24 @@ object BackupRepository {
      .also { Log.i(TAG, "OverallResult: $it") } is NetworkResult.Success
  }

+  fun downloadBackupFile(destination: File, listener: ProgressListener? = null): Boolean {
+    val api = ApplicationDependencies.getSignalServiceAccountManager().archiveApi
+    val backupKey = SignalStore.svr().getOrCreateMasterKey().deriveBackupKey()
+
+    return api
+      .triggerBackupIdReservation(backupKey)
+      .then { getAuthCredential() }
+      .then { credential ->
+        api.getBackupInfo(backupKey, credential)
+      }
+      .then { info -> getCdnReadCredentials().map { it.headers to info } }
+      .map { pair ->
+        val (cdnCredentials, info) = pair
+        val messageReceiver = ApplicationDependencies.getSignalServiceMessageReceiver()
+        messageReceiver.retrieveBackup(info.cdn!!, cdnCredentials, "backups/${info.backupDir}/${info.backupName}", destination, listener)
+      } is NetworkResult.Success
+  }
+
  /**
   * Returns an object with details about the remote backup state.
   */
@@ -296,7 +322,7 @@ object BackupRepository {
      }
  }

-  fun archiveMedia(attachment: DatabaseAttachment): NetworkResult<ArchiveMediaResponse> {
+  fun archiveMedia(attachment: DatabaseAttachment): NetworkResult<Unit> {
    val api = ApplicationDependencies.getSignalServiceAccountManager().archiveApi
    val backupKey = SignalStore.svr().getOrCreateMasterKey().deriveBackupKey()

@@ -304,16 +330,23 @@ object BackupRepository {
      .triggerBackupIdReservation(backupKey)
      .then { getAuthCredential() }
      .then { credential ->
-        api.archiveAttachmentMedia(
-          backupKey = backupKey,
-          serviceCredential = credential,
-          item = attachment.toArchiveMediaRequest(backupKey)
-        )
+        val mediaName = attachment.getMediaName()
+        val request = attachment.toArchiveMediaRequest(mediaName, backupKey)
+        api
+          .archiveAttachmentMedia(
+            backupKey = backupKey,
+            serviceCredential = credential,
+            item = request
+          )
+          .map { Triple(mediaName, request.mediaId, it) }
      }
-      .also { Log.i(TAG, "backupMediaResult: $it") }
+      .map { (mediaName, mediaId, response) ->
+        SignalDatabase.attachments.setArchiveData(attachmentId = attachment.attachmentId, archiveCdn = response.cdn, archiveMediaName = mediaName.name, archiveMediaId = mediaId)
+      }
+      .also { Log.i(TAG, "archiveMediaResult: $it") }
  }

-  fun archiveMedia(attachments: List<DatabaseAttachment>): NetworkResult<BatchArchiveMediaResponse> {
+  fun archiveMedia(databaseAttachments: List<DatabaseAttachment>): NetworkResult<BatchArchiveMediaResult> {
    val api = ApplicationDependencies.getSignalServiceAccountManager().archiveApi
    val backupKey = SignalStore.svr().getOrCreateMasterKey().deriveBackupKey()

@@ -321,24 +354,55 @@ object BackupRepository {
      .triggerBackupIdReservation(backupKey)
      .then { getAuthCredential() }
      .then { credential ->
-        api.archiveAttachmentMedia(
-          backupKey = backupKey,
-          serviceCredential = credential,
-          items = attachments.map { it.toArchiveMediaRequest(backupKey) }
-        )
+        val requests = mutableListOf<ArchiveMediaRequest>()
+        val mediaIdToAttachmentId = mutableMapOf<String, AttachmentId>()
+        val attachmentIdToMediaName = mutableMapOf<AttachmentId, String>()
+
+        databaseAttachments.forEach {
+          val mediaName = it.getMediaName()
+          val request = it.toArchiveMediaRequest(mediaName, backupKey)
+          requests += request
+          mediaIdToAttachmentId[request.mediaId] = it.attachmentId
+          attachmentIdToMediaName[it.attachmentId] = mediaName.name
+        }
+
+        api
+          .archiveAttachmentMedia(
+            backupKey = backupKey,
+            serviceCredential = credential,
+            items = requests
+          )
+          .map { BatchArchiveMediaResult(it, mediaIdToAttachmentId, attachmentIdToMediaName) }
      }
-      .also { Log.i(TAG, "backupMediaResult: $it") }
+      .map { result ->
+        result
+          .successfulResponses
+          .forEach {
+            val attachmentId = result.mediaIdToAttachmentId(it.mediaId)
+            val mediaName = result.attachmentIdToMediaName(attachmentId)
+            SignalDatabase.attachments.setArchiveData(attachmentId = attachmentId, archiveCdn = it.cdn!!, archiveMediaName = mediaName, archiveMediaId = it.mediaId)
+          }
+        result
+      }
+      .also { Log.i(TAG, "archiveMediaResult: $it") }
  }

  fun deleteArchivedMedia(attachments: List<DatabaseAttachment>): NetworkResult<Unit> {
    val api = ApplicationDependencies.getSignalServiceAccountManager().archiveApi
    val backupKey = SignalStore.svr().getOrCreateMasterKey().deriveBackupKey()

-    val mediaToDelete = attachments.map {
-      DeleteArchivedMediaRequest.ArchivedMediaObject(
-        cdn = 3, // TODO [cody] store and reuse backup cdn returned from copy/move call
-        mediaId = backupKey.deriveMediaId(Base64.decode(it.dataHash!!)).toString()
-      )
+    val mediaToDelete = attachments
+      .filter { it.archiveMediaId != null }
+      .map {
+        DeleteArchivedMediaRequest.ArchivedMediaObject(
+          cdn = it.archiveCdn,
+          mediaId = it.archiveMediaId!!
+        )
+      }
+
+    if (mediaToDelete.isEmpty()) {
+      Log.i(TAG, "No media to delete, quick success")
+      return NetworkResult.Success(Unit)
    }

    return getAuthCredential()
@@ -349,7 +413,101 @@ object BackupRepository {
          mediaToDelete = mediaToDelete
        )
      }
-      .also { Log.i(TAG, "deleteBackupMediaResult: $it") }
+      .map {
+        SignalDatabase.attachments.clearArchiveData(attachments.map { it.attachmentId })
+      }
+      .also { Log.i(TAG, "deleteArchivedMediaResult: $it") }
+  }
+
+  fun debugDeleteAllArchivedMedia(): NetworkResult<Unit> {
+    val api = ApplicationDependencies.getSignalServiceAccountManager().archiveApi
+    val backupKey = SignalStore.svr().getOrCreateMasterKey().deriveBackupKey()
+
+    return debugGetArchivedMediaState()
+      .then { archivedMedia ->
+        val mediaToDelete = archivedMedia
+          .map {
+            DeleteArchivedMediaRequest.ArchivedMediaObject(
+              cdn = it.cdn,
+              mediaId = it.mediaId
+            )
+          }
+
+        if (mediaToDelete.isEmpty()) {
+          Log.i(TAG, "No media to delete, quick success")
+          NetworkResult.Success(Unit)
+        } else {
+          getAuthCredential()
+            .then { credential ->
+              api.deleteArchivedMedia(
+                backupKey = backupKey,
+                serviceCredential = credential,
+                mediaToDelete = mediaToDelete
+              )
+            }
+        }
+      }
+      .map {
+        SignalDatabase.attachments.clearAllArchiveData()
+      }
+      .also { Log.i(TAG, "debugDeleteAllArchivedMediaResult: $it") }
+  }
+
+  /**
+   * Retrieve credentials for reading from the backup cdn.
+   */
+  fun getCdnReadCredentials(): NetworkResult<GetArchiveCdnCredentialsResponse> {
+    val cached = SignalStore.backup().cdnReadCredentials
+    if (cached != null) {
+      return NetworkResult.Success(cached)
+    }
+
+    val api = ApplicationDependencies.getSignalServiceAccountManager().archiveApi
+    val backupKey = SignalStore.svr().getOrCreateMasterKey().deriveBackupKey()
+
+    return getAuthCredential()
+      .then { credential ->
+        api.getCdnReadCredentials(
+          backupKey = backupKey,
+          serviceCredential = credential
+        )
+      }
+      .also {
+        if (it is NetworkResult.Success) {
+          SignalStore.backup().cdnReadCredentials = it.result
+        }
+      }
+      .also { Log.i(TAG, "getCdnReadCredentialsResult: $it") }
+  }
+
+  /**
+   * Retrieves backupDir and mediaDir, preferring cached value if available.
+   *
+   * These will only ever change if the backup expires.
+   */
+  fun getCdnBackupDirectories(): NetworkResult<BackupDirectories> {
+    val cachedBackupDirectory = SignalStore.backup().cachedBackupDirectory
+    val cachedBackupMediaDirectory = SignalStore.backup().cachedBackupMediaDirectory
+
+    if (cachedBackupDirectory != null && cachedBackupMediaDirectory != null) {
+      return NetworkResult.Success(BackupDirectories(cachedBackupDirectory, cachedBackupMediaDirectory))
+    }
+
+    val api = ApplicationDependencies.getSignalServiceAccountManager().archiveApi
+    val backupKey = SignalStore.svr().getOrCreateMasterKey().deriveBackupKey()
+
+    return getAuthCredential()
+      .then { credential ->
+        api.getBackupInfo(backupKey, credential).map {
+          BackupDirectories(it.backupDir!!, it.mediaDir!!)
+        }
+      }
+      .also {
+        if (it is NetworkResult.Success) {
+          SignalStore.backup().cachedBackupDirectory = it.result.backupDir
+          SignalStore.backup().cachedBackupMediaDirectory = it.result.mediaDir
+        }
+      }
  }

  /**
@@ -380,15 +538,20 @@ object BackupRepository {
    val profileKey: ProfileKey
  )

-  private fun DatabaseAttachment.toArchiveMediaRequest(backupKey: BackupKey): ArchiveMediaRequest {
-    val mediaSecrets = backupKey.deriveMediaSecrets(Base64.decode(dataHash!!))
+  fun DatabaseAttachment.getMediaName(): MediaName {
+    return MediaName.fromDigest(remoteDigest!!)
+  }
+
+  private fun DatabaseAttachment.toArchiveMediaRequest(mediaName: MediaName, backupKey: BackupKey): ArchiveMediaRequest {
+    val mediaSecrets = backupKey.deriveMediaSecrets(mediaName)
+
    return ArchiveMediaRequest(
      sourceAttachment = ArchiveMediaRequest.SourceAttachment(
-        cdn = cdnNumber,
+        cdn = cdn.cdnNumber,
        key = remoteLocation!!
      ),
      objectLength = AttachmentCipherStreamUtil.getCiphertextLength(PaddingInputStream.getPaddedSize(size)).toInt(),
-      mediaId = mediaSecrets.id.toString(),
+      mediaId = mediaSecrets.id.encode(),
      hmacKey = Base64.encodeWithPadding(mediaSecrets.macKey),
      encryptionKey = Base64.encodeWithPadding(mediaSecrets.cipherKey),
      iv = Base64.encodeWithPadding(mediaSecrets.iv)
@@ -396,12 +559,14 @@ object BackupRepository {
  }
 }

-class ExportState(val backupTime: Long) {
+data class BackupDirectories(val backupDir: String, val mediaDir: String)
+
+class ExportState(val backupTime: Long, val allowMediaBackup: Boolean) {
  val recipientIds = HashSet<Long>()
  val threadIds = HashSet<Long>()
 }

-class BackupState {
+class BackupState(val backupKey: BackupKey) {
  val backupToLocalRecipientId = HashMap<Long, RecipientId>()
  val chatIdToLocalThreadId = HashMap<Long, Long>()
  val chatIdToLocalRecipientId = HashMap<Long, RecipientId>()
--- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/BackupRestoreManager.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/BackupRestoreManager.kt
@@ -0,0 +1,46 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.backup.v2
+
+import org.signal.core.util.concurrent.SignalExecutors
+import org.thoughtcrime.securesms.attachments.AttachmentId
+import org.thoughtcrime.securesms.attachments.DatabaseAttachment
+import org.thoughtcrime.securesms.database.AttachmentTable
+import org.thoughtcrime.securesms.database.model.MessageRecord
+import org.thoughtcrime.securesms.database.model.MmsMessageRecord
+import org.thoughtcrime.securesms.jobs.RestoreAttachmentJob
+
+/**
+ * Responsible for managing logic around restore prioritization
+ */
+object BackupRestoreManager {
+
+  private val reprioritizedAttachments: HashSet<AttachmentId> = HashSet()
+
+  /**
+   * Raise priority of all attachments for the included message records.
+   *
+   * This is so we can make certain attachments get downloaded more quickly
+   */
+  fun prioritizeAttachmentsIfNeeded(messageRecords: List<MessageRecord>) {
+    SignalExecutors.BOUNDED.execute {
+      synchronized(this) {
+        val restoringAttachments: List<AttachmentId> = messageRecords
+          .mapNotNull { (it as? MmsMessageRecord?)?.slideDeck?.slides }
+          .flatten()
+          .mapNotNull { it.asAttachment() as? DatabaseAttachment }
+          .filter { it.transferState == AttachmentTable.TRANSFER_RESTORE_IN_PROGRESS && !reprioritizedAttachments.contains(it.attachmentId) }
+          .map { it.attachmentId }
+
+        reprioritizedAttachments += restoringAttachments
+
+        if (restoringAttachments.isNotEmpty()) {
+          RestoreAttachmentJob.modifyPriorities(restoringAttachments.toSet(), 1)
+        }
+      }
+    }
+  }
+}
--- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/BatchArchiveMediaResult.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/BatchArchiveMediaResult.kt
@@ -0,0 +1,39 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.backup.v2
+
+import org.thoughtcrime.securesms.attachments.AttachmentId
+import org.whispersystems.signalservice.api.archive.BatchArchiveMediaResponse
+
+/**
+ * Result of attempting to batch copy multiple attachments at once with helpers for
+ * processing the collection of mini-responses.
+ */
+data class BatchArchiveMediaResult(
+  private val response: BatchArchiveMediaResponse,
+  private val mediaIdToAttachmentId: Map<String, AttachmentId>,
+  private val attachmentIdToMediaName: Map<AttachmentId, String>
+) {
+  val successfulResponses: Sequence<BatchArchiveMediaResponse.BatchArchiveMediaItemResponse>
+    get() = response
+      .responses
+      .asSequence()
+      .filter { it.status == 200 }
+
+  val sourceNotFoundResponses: Sequence<BatchArchiveMediaResponse.BatchArchiveMediaItemResponse>
+    get() = response
+      .responses
+      .asSequence()
+      .filter { it.status == 410 }
+
+  fun mediaIdToAttachmentId(mediaId: String): AttachmentId {
+    return mediaIdToAttachmentId[mediaId]!!
+  }
+
+  fun attachmentIdToMediaName(attachmentId: AttachmentId): String {
+    return attachmentIdToMediaName[attachmentId]!!
+  }
+}
--- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/database/ChatItemExportIterator.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/database/ChatItemExportIterator.kt
@@ -17,7 +17,9 @@ import org.signal.core.util.requireBoolean
 import org.signal.core.util.requireInt
 import org.signal.core.util.requireLong
 import org.signal.core.util.requireString
+import org.thoughtcrime.securesms.attachments.Cdn
 import org.thoughtcrime.securesms.attachments.DatabaseAttachment
+import org.thoughtcrime.securesms.backup.v2.BackupRepository.getMediaName
 import org.thoughtcrime.securesms.backup.v2.proto.CallChatUpdate
 import org.thoughtcrime.securesms.backup.v2.proto.ChatItem
 import org.thoughtcrime.securesms.backup.v2.proto.ChatUpdateMessage
@@ -36,6 +38,7 @@ import org.thoughtcrime.securesms.backup.v2.proto.SimpleChatUpdate
 import org.thoughtcrime.securesms.backup.v2.proto.StandardMessage
 import org.thoughtcrime.securesms.backup.v2.proto.Text
 import org.thoughtcrime.securesms.backup.v2.proto.ThreadMergeChatUpdate
+import org.thoughtcrime.securesms.database.AttachmentTable
 import org.thoughtcrime.securesms.database.GroupReceiptTable
 import org.thoughtcrime.securesms.database.MessageTable
 import org.thoughtcrime.securesms.database.MessageTypes
@@ -73,7 +76,7 @@ import org.thoughtcrime.securesms.backup.v2.proto.BodyRange as BackupBodyRange
 *
 * All of this complexity is hidden from the user -- they just get a normal iterator interface.
 */
-class ChatItemExportIterator(private val cursor: Cursor, private val batchSize: Int) : Iterator<ChatItem>, Closeable {
+class ChatItemExportIterator(private val cursor: Cursor, private val batchSize: Int, private val archiveMedia: Boolean) : Iterator<ChatItem>, Closeable {

  companion object {
    private val TAG = Log.tag(ChatItemExportIterator::class.java)
@@ -139,6 +142,7 @@ class ChatItemExportIterator(private val cursor: Cursor, private val batchSize:
          builder.expiresInMs = null
        }
        MessageTypes.isProfileChange(record.type) -> {
+          if (record.body == null) continue
          builder.updateMessage = ChatUpdateMessage(
            profileChange = try {
              val decoded: ByteArray = Base64.decode(record.body!!)
@@ -354,24 +358,46 @@ class ChatItemExportIterator(private val cursor: Cursor, private val batchSize:
  }

  private fun DatabaseAttachment.toBackupAttachment(): MessageAttachment {
+    val builder = FilePointer.Builder()
+    builder.contentType = contentType
+    builder.incrementalMac = incrementalDigest?.toByteString()
+    builder.incrementalMacChunkSize = incrementalMacChunkSize
+    builder.fileName = fileName
+    builder.width = width
+    builder.height = height
+    builder.caption = caption
+    builder.blurHash = blurHash?.hash
+
+    if (remoteKey.isNullOrBlank() || remoteDigest == null || size == 0L) {
+      builder.invalidAttachmentLocator = FilePointer.InvalidAttachmentLocator()
+    } else {
+      if (archiveMedia) {
+        builder.backupLocator = FilePointer.BackupLocator(
+          mediaName = archiveMediaName ?: this.getMediaName().toString(),
+          cdnNumber = if (archiveMediaName != null) archiveCdn else Cdn.CDN_3.cdnNumber, // TODO (clark): Update when new proto with optional cdn is landed
+          key = decode(remoteKey).toByteString(),
+          size = this.size.toInt(),
+          digest = remoteDigest.toByteString()
+        )
+      } else {
+        if (remoteLocation.isNullOrBlank()) {
+          builder.invalidAttachmentLocator = FilePointer.InvalidAttachmentLocator()
+        } else {
+          builder.attachmentLocator = FilePointer.AttachmentLocator(
+            cdnKey = this.remoteLocation,
+            cdnNumber = this.cdn.cdnNumber,
+            uploadTimestamp = this.uploadTimestamp,
+            key = decode(remoteKey).toByteString(),
+            size = this.size.toInt(),
+            digest = remoteDigest.toByteString()
+          )
+        }
+      }
+    }
    return MessageAttachment(
-      pointer = FilePointer(
-        attachmentLocator = FilePointer.AttachmentLocator(
-          cdnKey = this.remoteLocation ?: "",
-          cdnNumber = this.cdnNumber,
-          uploadTimestamp = this.uploadTimestamp
-        ),
-        key = if (remoteKey != null) decode(remoteKey).toByteString() else null,
-        contentType = this.contentType,
-        size = this.size.toInt(),
-        incrementalMac = this.incrementalDigest?.toByteString(),
-        incrementalMacChunkSize = this.incrementalMacChunkSize,
-        fileName = this.fileName,
-        width = this.width,
-        height = this.height,
-        caption = this.caption,
-        blurHash = this.blurHash?.hash
-      )
+      pointer = builder.build(),
+      wasDownloaded = this.transferState == AttachmentTable.TRANSFER_PROGRESS_DONE || this.transferState == AttachmentTable.TRANSFER_NEEDS_RESTORE,
+      flag = if (voiceNote) MessageAttachment.Flag.VOICE_MESSAGE else if (videoGif) MessageAttachment.Flag.GIF else if (borderless) MessageAttachment.Flag.BORDERLESS else MessageAttachment.Flag.NONE
    )
  }

--- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/database/ChatItemImportInserter.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/database/ChatItemImportInserter.kt
@@ -13,8 +13,11 @@ import org.signal.core.util.logging.Log
 import org.signal.core.util.orNull
 import org.signal.core.util.requireLong
 import org.signal.core.util.toInt
+import org.thoughtcrime.securesms.attachments.ArchivedAttachment
 import org.thoughtcrime.securesms.attachments.Attachment
+import org.thoughtcrime.securesms.attachments.Cdn
 import org.thoughtcrime.securesms.attachments.PointerAttachment
+import org.thoughtcrime.securesms.attachments.TombstoneAttachment
 import org.thoughtcrime.securesms.backup.v2.BackupState
 import org.thoughtcrime.securesms.backup.v2.proto.BodyRange
 import org.thoughtcrime.securesms.backup.v2.proto.ChatItem
@@ -26,6 +29,7 @@ import org.thoughtcrime.securesms.backup.v2.proto.Reaction
 import org.thoughtcrime.securesms.backup.v2.proto.SendStatus
 import org.thoughtcrime.securesms.backup.v2.proto.SimpleChatUpdate
 import org.thoughtcrime.securesms.backup.v2.proto.StandardMessage
+import org.thoughtcrime.securesms.database.AttachmentTable
 import org.thoughtcrime.securesms.database.CallTable
 import org.thoughtcrime.securesms.database.GroupReceiptTable
 import org.thoughtcrime.securesms.database.MessageTable
@@ -48,11 +52,12 @@ import org.thoughtcrime.securesms.mms.QuoteModel
 import org.thoughtcrime.securesms.recipients.Recipient
 import org.thoughtcrime.securesms.recipients.RecipientId
 import org.thoughtcrime.securesms.util.JsonUtils
+import org.whispersystems.signalservice.api.backup.MediaName
 import org.whispersystems.signalservice.api.messages.SignalServiceAttachmentPointer
 import org.whispersystems.signalservice.api.messages.SignalServiceAttachmentRemoteId
-import org.whispersystems.signalservice.api.messages.SignalServiceDataMessage
 import org.whispersystems.signalservice.api.push.ServiceId
 import org.whispersystems.signalservice.api.util.UuidUtil
+import org.whispersystems.signalservice.internal.push.DataMessage
 import java.util.Optional

 /**
@@ -570,12 +575,12 @@ class ChatItemImportInserter(
        pointer.attachmentLocator.cdnNumber,
        SignalServiceAttachmentRemoteId.from(pointer.attachmentLocator.cdnKey),
        contentType,
-        pointer.key?.toByteArray(),
-        Optional.ofNullable(pointer.size),
+        pointer.attachmentLocator.key.toByteArray(),
+        Optional.ofNullable(pointer.attachmentLocator.size),
        Optional.empty(),
        pointer.width ?: 0,
        pointer.height ?: 0,
-        Optional.empty(),
+        Optional.ofNullable(pointer.attachmentLocator.digest.toByteArray()),
        Optional.ofNullable(pointer.incrementalMac?.toByteArray()),
        pointer.incrementalMacChunkSize ?: 0,
        Optional.ofNullable(fileName),
@@ -586,14 +591,51 @@ class ChatItemImportInserter(
        Optional.ofNullable(pointer.blurHash),
        pointer.attachmentLocator.uploadTimestamp
      )
-      return PointerAttachment.forPointer(Optional.of(signalAttachmentPointer)).orNull()
+      return PointerAttachment.forPointer(
+        pointer = Optional.of(signalAttachmentPointer),
+        transferState = if (wasDownloaded) AttachmentTable.TRANSFER_NEEDS_RESTORE else AttachmentTable.TRANSFER_PROGRESS_PENDING
+      ).orNull()
+    } else if (pointer.invalidAttachmentLocator != null) {
+      return TombstoneAttachment(
+        contentType = contentType,
+        incrementalMac = pointer.incrementalMac?.toByteArray(),
+        incrementalMacChunkSize = pointer.incrementalMacChunkSize,
+        width = pointer.width,
+        height = pointer.height,
+        caption = pointer.caption,
+        blurHash = pointer.blurHash,
+        voiceNote = flag == MessageAttachment.Flag.VOICE_MESSAGE,
+        borderless = flag == MessageAttachment.Flag.BORDERLESS,
+        gif = flag == MessageAttachment.Flag.GIF,
+        quote = false
+      )
+    } else if (pointer.backupLocator != null) {
+      return ArchivedAttachment(
+        contentType = contentType,
+        size = pointer.backupLocator.size.toLong(),
+        cdn = Cdn.fromCdnNumber(pointer.backupLocator.cdnNumber),
+        cdnKey = pointer.backupLocator.key.toByteArray(),
+        archiveMediaName = pointer.backupLocator.mediaName,
+        archiveMediaId = backupState.backupKey.deriveMediaId(MediaName(pointer.backupLocator.mediaName)).encode(),
+        digest = pointer.backupLocator.digest.toByteArray(),
+        incrementalMac = pointer.incrementalMac?.toByteArray(),
+        incrementalMacChunkSize = pointer.incrementalMacChunkSize,
+        width = pointer.width,
+        height = pointer.height,
+        caption = pointer.caption,
+        blurHash = pointer.blurHash,
+        voiceNote = flag == MessageAttachment.Flag.VOICE_MESSAGE,
+        borderless = flag == MessageAttachment.Flag.BORDERLESS,
+        gif = flag == MessageAttachment.Flag.GIF,
+        quote = false
+      )
    }
    return null
  }

  private fun Quote.QuotedAttachment.toLocalAttachment(): Attachment? {
    return thumbnail?.toLocalAttachment(this.contentType, this.fileName)
-      ?: if (this.contentType == null) null else PointerAttachment.forPointer(SignalServiceDataMessage.Quote.QuotedAttachment(contentType = this.contentType!!, fileName = this.fileName, thumbnail = null)).orNull()
+      ?: if (this.contentType == null) null else PointerAttachment.forPointer(quotedAttachment = DataMessage.Quote.QuotedAttachment(contentType = this.contentType, fileName = this.fileName, thumbnail = null)).orNull()
  }

  private class MessageInsert(val contentValues: ContentValues, val followUp: ((Long) -> Unit)?)
--- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/database/MessageTableBackupExtensions.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/database/MessageTableBackupExtensions.kt
@@ -16,7 +16,7 @@ import java.util.concurrent.TimeUnit
 private val TAG = Log.tag(MessageTable::class.java)
 private const val BASE_TYPE = "base_type"

-fun MessageTable.getMessagesForBackup(backupTime: Long): ChatItemExportIterator {
+fun MessageTable.getMessagesForBackup(backupTime: Long, archiveMedia: Boolean): ChatItemExportIterator {
  val cursor = readableDatabase
    .select(
      MessageTable.ID,
@@ -64,7 +64,7 @@ fun MessageTable.getMessagesForBackup(backupTime: Long): ChatItemExportIterator
    .orderBy("${MessageTable.DATE_RECEIVED} ASC")
    .run()

-  return ChatItemExportIterator(cursor, 100)
+  return ChatItemExportIterator(cursor, 100, archiveMedia)
 }

 fun MessageTable.createChatItemInserter(backupState: BackupState): ChatItemImportInserter {
--- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/processor/ChatItemBackupProcessor.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/processor/ChatItemBackupProcessor.kt
@@ -19,7 +19,7 @@ object ChatItemBackupProcessor {
  val TAG = Log.tag(ChatItemBackupProcessor::class.java)

  fun export(exportState: ExportState, emitter: BackupFrameEmitter) {
-    SignalDatabase.messages.getMessagesForBackup(exportState.backupTime).use { chatItems ->
+    SignalDatabase.messages.getMessagesForBackup(exportState.backupTime, exportState.allowMediaBackup).use { chatItems ->
      for (chatItem in chatItems) {
        if (exportState.threadIds.contains(chatItem.chatId)) {
          emitter.emit(Frame(chatItem = chatItem))
--- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/ui/MessageBackupsTestRestoreActivity.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/ui/MessageBackupsTestRestoreActivity.kt
@@ -26,6 +26,7 @@ import androidx.compose.material3.Surface
 import androidx.compose.material3.Switch
 import androidx.compose.material3.Text
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.SideEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
@@ -80,6 +81,15 @@ class MessageBackupsTestRestoreActivity : BaseActivity() {
            .fillMaxSize()
            .padding(16.dp)
        ) {
+          Buttons.LargePrimary(
+            onClick = this@MessageBackupsTestRestoreActivity::restoreFromServer,
+            enabled = !state.importState.inProgress
+          ) {
+            Text("Restore")
+          }
+
+          Spacer(modifier = Modifier.height(8.dp))
+
          Row(
            verticalAlignment = Alignment.CenterVertically
          ) {
@@ -120,9 +130,20 @@ class MessageBackupsTestRestoreActivity : BaseActivity() {
          }
        }
      }
+      if (state.importState == MessageBackupsTestRestoreViewModel.ImportState.RESTORED) {
+        SideEffect {
+          RegistrationUtil.maybeMarkRegistrationComplete()
+          ApplicationDependencies.getJobManager().add(ProfileUploadJob())
+          startActivity(MainActivity.clearTop(this))
+        }
+      }
    }
  }

+  private fun restoreFromServer() {
+    viewModel.restore()
+  }
+
  private fun continueRegistration() {
    if (Recipient.self().profileName.isEmpty || !AvatarHelper.hasAvatar(this, Recipient.self().id)) {
      val main = MainActivity.clearTop(this)
--- a/app/src/main/java/org/thoughtcrime/securesms/backup/v2/ui/MessageBackupsTestRestoreViewModel.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/backup/v2/ui/MessageBackupsTestRestoreViewModel.kt
@@ -15,8 +15,12 @@ import io.reactivex.rxjava3.disposables.CompositeDisposable
 import io.reactivex.rxjava3.kotlin.plusAssign
 import io.reactivex.rxjava3.kotlin.subscribeBy
 import io.reactivex.rxjava3.schedulers.Schedulers
+import org.signal.core.util.orNull
 import org.signal.libsignal.zkgroup.profiles.ProfileKey
 import org.thoughtcrime.securesms.backup.v2.BackupRepository
+import org.thoughtcrime.securesms.dependencies.ApplicationDependencies
+import org.thoughtcrime.securesms.jobmanager.JobTracker
+import org.thoughtcrime.securesms.jobs.BackupRestoreJob
 import org.thoughtcrime.securesms.recipients.Recipient
 import java.io.InputStream

@@ -40,6 +44,19 @@ class MessageBackupsTestRestoreViewModel : ViewModel() {
      }
  }

+  fun restore() {
+    _state.value = _state.value.copy(importState = ImportState.IN_PROGRESS)
+    disposables += Single.fromCallable {
+      val jobState = ApplicationDependencies.getJobManager().runSynchronously(BackupRestoreJob(), 120_000)
+      jobState.orNull() == JobTracker.JobState.SUCCESS
+    }
+      .subscribeOn(Schedulers.io())
+      .observeOn(AndroidSchedulers.mainThread())
+      .subscribeBy {
+        _state.value = _state.value.copy(importState = ImportState.RESTORED)
+      }
+  }
+
  fun onPlaintextToggled() {
    _state.value = _state.value.copy(plaintext = !_state.value.plaintext)
  }
@@ -54,6 +71,6 @@ class MessageBackupsTestRestoreViewModel : ViewModel() {
  )

  enum class ImportState(val inProgress: Boolean = false) {
-    NONE, IN_PROGRESS(true)
+    NONE, IN_PROGRESS(true), RESTORED
  }
 }