Be more lenient with attachment restore conditions.

app/src/androidTest/java/org/thoughtcrime/securesms/database/AttachmentTableTest.kt

@@ -19,6 +19,7 @@ import org.junit.Ignore
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.signal.core.models.backup.MediaName
 import org.signal.core.util.Base64
 import org.signal.core.util.Base64.decodeBase64OrThrow
 import org.signal.core.util.copyTo
@@ -28,6 +29,8 @@ import org.thoughtcrime.securesms.attachments.Attachment
 import org.thoughtcrime.securesms.attachments.AttachmentId
 import org.thoughtcrime.securesms.attachments.PointerAttachment
 import org.thoughtcrime.securesms.attachments.UriAttachment
+import org.thoughtcrime.securesms.backup.v2.ArchivedMediaObject
+import org.thoughtcrime.securesms.keyvalue.SignalStore
 import org.thoughtcrime.securesms.mms.IncomingMessage
 import org.thoughtcrime.securesms.mms.MediaStream
 import org.thoughtcrime.securesms.mms.SentMediaQuality
@@ -414,6 +417,126 @@ class AttachmentTableTest {
     assertThat(dbAttachment2.archiveTransferState).isEqualTo(AttachmentTable.ArchiveTransferState.FINISHED)
   }
 
+  @Test
+  fun givenAttachmentsWithMatchingMediaId_whenISetArchiveFinishedForMatchingMediaObjects_thenIExpectThoseAttachmentsToBeMarkedFinished() {
+    // GIVEN
+    val data = byteArrayOf(1, 2, 3, 4, 5)
+
+    val attachment1 = createAttachmentPointer("remote-key-1".toByteArray(), data.size)
+    val attachment2 = createAttachmentPointer("remote-key-2".toByteArray(), data.size)
+    val attachment3 = createAttachmentPointer("remote-key-3".toByteArray(), data.size)
+
+    // Insert messages with attachments
+    val message1Result = SignalDatabase.messages.insertMessageInbox(createIncomingMessage(serverTime = 0.days, attachment = attachment1)).get()
+    val attachment1Id = message1Result.insertedAttachments!![attachment1]!!
+    SignalDatabase.attachments.setTransferState(message1Result.messageId, attachment1Id, AttachmentTable.TRANSFER_PROGRESS_STARTED)
+    SignalDatabase.attachments.finalizeAttachmentAfterDownload(message1Result.messageId, attachment1Id, ByteArrayInputStream(data))
+
+    val message2Result = SignalDatabase.messages.insertMessageInbox(createIncomingMessage(serverTime = 1.days, attachment = attachment2)).get()
+    val attachment2Id = message2Result.insertedAttachments!![attachment2]!!
+    SignalDatabase.attachments.setTransferState(message2Result.messageId, attachment2Id, AttachmentTable.TRANSFER_PROGRESS_STARTED)
+    SignalDatabase.attachments.finalizeAttachmentAfterDownload(message2Result.messageId, attachment2Id, ByteArrayInputStream(data))
+
+    val message3Result = SignalDatabase.messages.insertMessageInbox(createIncomingMessage(serverTime = 2.days, attachment = attachment3)).get()
+    val attachment3Id = message3Result.insertedAttachments!![attachment3]!!
+    SignalDatabase.attachments.setTransferState(message3Result.messageId, attachment3Id, AttachmentTable.TRANSFER_PROGRESS_STARTED)
+    SignalDatabase.attachments.finalizeAttachmentAfterDownload(message3Result.messageId, attachment3Id, ByteArrayInputStream(data))
+
+    // Ensure attachments are in NONE state
+    assertThat(SignalDatabase.attachments.getAttachment(attachment1Id)!!.archiveTransferState).isEqualTo(AttachmentTable.ArchiveTransferState.NONE)
+    assertThat(SignalDatabase.attachments.getAttachment(attachment2Id)!!.archiveTransferState).isEqualTo(AttachmentTable.ArchiveTransferState.NONE)
+    assertThat(SignalDatabase.attachments.getAttachment(attachment3Id)!!.archiveTransferState).isEqualTo(AttachmentTable.ArchiveTransferState.NONE)
+
+    // Build media objects only for attachments 1 and 2 (not 3)
+    val dbAttachment1 = SignalDatabase.attachments.getAttachment(attachment1Id)!!
+    val dbAttachment2 = SignalDatabase.attachments.getAttachment(attachment2Id)!!
+
+    val mediaId1 = MediaName.fromPlaintextHashAndRemoteKey(
+      dbAttachment1.dataHash!!.decodeBase64OrThrow(),
+      dbAttachment1.remoteKey!!.decodeBase64OrThrow()
+    ).toMediaId(SignalStore.backup.mediaRootBackupKey).encode()
+
+    val mediaId2 = MediaName.fromPlaintextHashAndRemoteKey(
+      dbAttachment2.dataHash!!.decodeBase64OrThrow(),
+      dbAttachment2.remoteKey!!.decodeBase64OrThrow()
+    ).toMediaId(SignalStore.backup.mediaRootBackupKey).encode()
+
+    val archivedMediaObjects = setOf(
+      ArchivedMediaObject(mediaId = mediaId1, cdn = 5),
+      ArchivedMediaObject(mediaId = mediaId2, cdn = 6)
+    )
+
+    // WHEN
+    val updatedCount = SignalDatabase.attachments.setArchiveFinishedForMatchingMediaObjects(archivedMediaObjects)
+
+    // THEN
+    assertThat(updatedCount).isEqualTo(2)
+
+    val resultAttachment1 = SignalDatabase.attachments.getAttachment(attachment1Id)!!
+    val resultAttachment2 = SignalDatabase.attachments.getAttachment(attachment2Id)!!
+    val resultAttachment3 = SignalDatabase.attachments.getAttachment(attachment3Id)!!
+
+    // Attachments 1 and 2 should be FINISHED with their respective CDNs
+    assertThat(resultAttachment1.archiveTransferState).isEqualTo(AttachmentTable.ArchiveTransferState.FINISHED)
+    assertThat(resultAttachment1.archiveCdn).isEqualTo(5)
+
+    assertThat(resultAttachment2.archiveTransferState).isEqualTo(AttachmentTable.ArchiveTransferState.FINISHED)
+    assertThat(resultAttachment2.archiveCdn).isEqualTo(6)
+
+    // Attachment 3 should still be NONE (not in the set)
+    assertThat(resultAttachment3.archiveTransferState).isEqualTo(AttachmentTable.ArchiveTransferState.NONE)
+  }
+
+  @Test
+  fun givenEmptyMediaObjectsSet_whenISetArchiveFinishedForMatchingMediaObjects_thenIExpectZeroUpdates() {
+    // GIVEN
+    val data = byteArrayOf(1, 2, 3, 4, 5)
+    val attachment = createAttachmentPointer("remote-key-1".toByteArray(), data.size)
+
+    val messageResult = SignalDatabase.messages.insertMessageInbox(createIncomingMessage(serverTime = 0.days, attachment = attachment)).get()
+    val attachmentId = messageResult.insertedAttachments!![attachment]!!
+    SignalDatabase.attachments.setTransferState(messageResult.messageId, attachmentId, AttachmentTable.TRANSFER_PROGRESS_STARTED)
+    SignalDatabase.attachments.finalizeAttachmentAfterDownload(messageResult.messageId, attachmentId, ByteArrayInputStream(data))
+
+    // WHEN
+    val updatedCount = SignalDatabase.attachments.setArchiveFinishedForMatchingMediaObjects(emptySet())
+
+    // THEN
+    assertThat(updatedCount).isEqualTo(0)
+    assertThat(SignalDatabase.attachments.getAttachment(attachmentId)!!.archiveTransferState).isEqualTo(AttachmentTable.ArchiveTransferState.NONE)
+  }
+
+  @Test
+  fun givenAlreadyFinishedAttachment_whenISetArchiveFinishedForMatchingMediaObjects_thenIExpectNoUpdate() {
+    // GIVEN
+    val data = byteArrayOf(1, 2, 3, 4, 5)
+    val attachment = createAttachmentPointer("remote-key-1".toByteArray(), data.size)
+
+    val messageResult = SignalDatabase.messages.insertMessageInbox(createIncomingMessage(serverTime = 0.days, attachment = attachment)).get()
+    val attachmentId = messageResult.insertedAttachments!![attachment]!!
+    SignalDatabase.attachments.setTransferState(messageResult.messageId, attachmentId, AttachmentTable.TRANSFER_PROGRESS_STARTED)
+    SignalDatabase.attachments.finalizeAttachmentAfterDownload(messageResult.messageId, attachmentId, ByteArrayInputStream(data))
+
+    // Mark as already FINISHED
+    SignalDatabase.attachments.setArchiveTransferState(attachmentId, AttachmentTable.ArchiveTransferState.FINISHED)
+
+    val dbAttachment = SignalDatabase.attachments.getAttachment(attachmentId)!!
+    val mediaId = MediaName.fromPlaintextHashAndRemoteKey(
+      dbAttachment.dataHash!!.decodeBase64OrThrow(),
+      dbAttachment.remoteKey!!.decodeBase64OrThrow()
+    ).toMediaId(SignalStore.backup.mediaRootBackupKey).encode()
+
+    val archivedMediaObjects = setOf(
+      ArchivedMediaObject(mediaId = mediaId, cdn = 5)
+    )
+
+    // WHEN
+    val updatedCount = SignalDatabase.attachments.setArchiveFinishedForMatchingMediaObjects(archivedMediaObjects)
+
+    // THEN - should not update since already FINISHED
+    assertThat(updatedCount).isEqualTo(0)
+  }
+
   private fun createIncomingMessage(
     serverTime: Duration,
     attachment: Attachment,

app/src/main/java/org/thoughtcrime/securesms/database/AttachmentTable.kt

@@ -2393,6 +2393,81 @@ class AttachmentTable(
       .run()
   }
 
+  /**
+   * For attachments that match the given media objects (by computing mediaId from plaintextHash + remoteKey),
+   * update their archive transfer state to FINISHED and set the archive CDN.
+   *
+   * This is used during reconciliation to fix attachments that were restored from a backup but didn't
+   * have the correct archive state because the archive upload hadn't completed when the backup was made.
+   *
+   * @return the number of unique (plaintextHash, remoteKey) pairs that were updated
+   */
+  fun setArchiveFinishedForMatchingMediaObjects(objects: Set<ArchivedMediaObject>): Int {
+    if (objects.isEmpty()) {
+      return 0
+    }
+
+    val objectsByMediaId: Map<String, ArchivedMediaObject> = objects.associateBy { it.mediaId }
+
+    // Collect updates grouped by CDN: Map<cdn, List<Pair<plaintextHash, remoteKey>>>
+    val updatesByCdn: MutableMap<Int, MutableList<Pair<String, String>>> = mutableMapOf()
+
+    readableDatabase
+      .select(DATA_HASH_END, REMOTE_KEY)
+      .from(TABLE_NAME)
+      .where("$REMOTE_KEY NOT NULL AND $DATA_HASH_END NOT NULL AND $ARCHIVE_TRANSFER_STATE != ?", ArchiveTransferState.FINISHED.value)
+      .groupBy("$DATA_HASH_END, $REMOTE_KEY")
+      .run()
+      .forEach { cursor ->
+        val remoteKeyStr = cursor.requireNonNullString(REMOTE_KEY)
+        val plaintextHashStr = cursor.requireNonNullString(DATA_HASH_END)
+        val remoteKey = Base64.decode(remoteKeyStr)
+        val plaintextHash = Base64.decode(plaintextHashStr)
+
+        val mediaId = MediaName.fromPlaintextHashAndRemoteKey(plaintextHash, remoteKey)
+          .toMediaId(SignalStore.backup.mediaRootBackupKey)
+          .encode()
+
+        val matchingObject = objectsByMediaId[mediaId]
+        if (matchingObject != null) {
+          updatesByCdn.getOrPut(matchingObject.cdn) { mutableListOf() }
+            .add(plaintextHashStr to remoteKeyStr)
+        }
+      }
+
+    if (updatesByCdn.isEmpty()) {
+      return 0
+    }
+
+    var updatedCount = 0
+
+    writableDatabase.withinTransaction { db ->
+      for ((cdn, pairs) in updatesByCdn) {
+        // Batch updates - each pair uses 2 query args, so chunk accordingly
+        for (batch in pairs.chunked(500)) {
+          val whereClause = batch.joinToString(" OR ") { "($DATA_HASH_END = ? AND $REMOTE_KEY = ?)" }
+          val whereArgs = batch.flatMap { listOf(it.first, it.second) }.toTypedArray()
+
+          db.update(TABLE_NAME)
+            .values(
+              ARCHIVE_TRANSFER_STATE to ArchiveTransferState.FINISHED.value,
+              ARCHIVE_CDN to cdn
+            )
+            .where(whereClause, *whereArgs)
+            .run()
+
+          updatedCount += batch.size
+        }
+      }
+    }
+
+    if (updatedCount > 0) {
+      AppDependencies.databaseObserver.notifyAttachmentUpdatedObservers()
+    }
+
+    return updatedCount
+  }
+
   fun clearArchiveData(attachmentId: AttachmentId) {
     writableDatabase
       .update(TABLE_NAME)

app/src/main/java/org/thoughtcrime/securesms/jobs/ArchiveAttachmentReconciliationJob.kt

@@ -352,6 +352,9 @@ class ArchiveAttachmentReconciliationJob private constructor(
    * Deletes attachments from the archive CDN, after verifying that they also can't be found anywhere in [org.thoughtcrime.securesms.database.AttachmentTable]
    * either. Checking the attachment table is very expensive and independent of query size, which is why we batch the lookups.
    *
+   * Also fixes archive transfer state for attachments that ARE found locally but may have incorrect state
+   * (e.g., restored from a backup before archive upload completed).
+   *
    * @return A non-successful [Result] in the case of failure, otherwise null for success.
    */
   private fun validateAndDeleteFromRemote(deletes: Set<ArchivedMediaObject>): Result? {
@@ -360,6 +363,17 @@ class ArchiveAttachmentReconciliationJob private constructor(
     Log.d(TAG, "Found that ${validatedDeletes.size}/${deletes.size} requested remote deletes were valid based on current attachment table state.", true)
     stopwatch.split("validate")
 
+    // Fix archive state for attachments that are found locally but weren't in the latest snapshot.
+    // This can happen when restoring from a backup that was made before archive upload completed. The files would be uploaded, but no CDN info would be in the backup.
+    val foundLocally = deletes - validatedDeletes
+    if (foundLocally.isNotEmpty()) {
+      val fixedCount = SignalDatabase.attachments.setArchiveFinishedForMatchingMediaObjects(foundLocally)
+      if (fixedCount > 0) {
+        Log.i(TAG, "Fixed archive transfer state for $fixedCount attachment groups that were found on CDN but had incorrect local state.", true)
+      }
+    }
+    stopwatch.split("fix-state")
+
     if (validatedDeletes.isEmpty()) {
       return null
     }

app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentJob.kt

@@ -359,7 +359,7 @@ class RestoreAttachmentJob private constructor(
       }
 
       if (useArchiveCdn && attachment.archiveTransferState != AttachmentTable.ArchiveTransferState.FINISHED) {
-        throw InvalidAttachmentException("[$attachmentId] Invalid attachment configuration! backsUpMedia: ${SignalStore.backup.backsUpMedia}, forceTransitTier: $forceTransitTier, archiveTransferState: ${attachment.archiveTransferState}")
+        Log.w(TAG, "[$attachmentId] Archive state was not FINISHED, but we backup media and have a dataHash, so we should try anyway. archiveTransferState: ${attachment.archiveTransferState}")
       }
 
       val messageReceiver = AppDependencies.signalServiceMessageReceiver