diff --git a/app/src/androidTest/java/org/thoughtcrime/securesms/database/AttachmentTableTest.kt b/app/src/androidTest/java/org/thoughtcrime/securesms/database/AttachmentTableTest.kt index 9bee3ab029..f3f35f217d 100644 --- a/app/src/androidTest/java/org/thoughtcrime/securesms/database/AttachmentTableTest.kt +++ b/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, diff --git a/app/src/main/java/org/thoughtcrime/securesms/database/AttachmentTable.kt b/app/src/main/java/org/thoughtcrime/securesms/database/AttachmentTable.kt index f7166e9815..0e4df2be03 100644 --- a/app/src/main/java/org/thoughtcrime/securesms/database/AttachmentTable.kt +++ b/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): Int { + if (objects.isEmpty()) { + return 0 + } + + val objectsByMediaId: Map = objects.associateBy { it.mediaId } + + // Collect updates grouped by CDN: Map>> + val updatesByCdn: MutableMap>> = 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) diff --git a/app/src/main/java/org/thoughtcrime/securesms/jobs/ArchiveAttachmentReconciliationJob.kt b/app/src/main/java/org/thoughtcrime/securesms/jobs/ArchiveAttachmentReconciliationJob.kt index 53bdb89f2b..9598f9bb78 100644 --- a/app/src/main/java/org/thoughtcrime/securesms/jobs/ArchiveAttachmentReconciliationJob.kt +++ b/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): 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 } diff --git a/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentJob.kt b/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentJob.kt index 367ba8b414..5a7b16c0a6 100644 --- a/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentJob.kt +++ b/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