From bf338a68357d062226146ee981bf2f0d04b9f825 Mon Sep 17 00:00:00 2001 From: Greyson Parrelli Date: Tue, 24 Sep 2024 16:39:59 -0400 Subject: [PATCH] Keep remote fields in sync when deduping downloads. --- .../database/AttachmentTableTest_deduping.kt | 84 +++++++++++++++++++ .../securesms/database/AttachmentTable.kt | 9 +- 2 files changed, 92 insertions(+), 1 deletion(-) diff --git a/app/src/androidTest/java/org/thoughtcrime/securesms/database/AttachmentTableTest_deduping.kt b/app/src/androidTest/java/org/thoughtcrime/securesms/database/AttachmentTableTest_deduping.kt index 17912fd2de..771c871b20 100644 --- a/app/src/androidTest/java/org/thoughtcrime/securesms/database/AttachmentTableTest_deduping.kt +++ b/app/src/androidTest/java/org/thoughtcrime/securesms/database/AttachmentTableTest_deduping.kt @@ -12,9 +12,12 @@ import org.junit.Before import org.junit.Test import org.junit.runner.RunWith import org.signal.core.util.Base64 +import org.signal.core.util.readFully +import org.signal.core.util.stream.LimitedInputStream import org.signal.core.util.update import org.thoughtcrime.securesms.attachments.AttachmentId import org.thoughtcrime.securesms.attachments.Cdn +import org.thoughtcrime.securesms.attachments.PointerAttachment import org.thoughtcrime.securesms.backup.v2.BackupRepository.getMediaName import org.thoughtcrime.securesms.database.AttachmentTable.TransformProperties import org.thoughtcrime.securesms.keyvalue.SignalStore @@ -30,6 +33,7 @@ import org.whispersystems.signalservice.api.attachment.AttachmentUploadResult import org.whispersystems.signalservice.api.backup.MediaId import org.whispersystems.signalservice.api.messages.SignalServiceAttachmentRemoteId import org.whispersystems.signalservice.api.push.ServiceId +import org.whispersystems.signalservice.internal.crypto.PaddingInputStream import java.io.File import java.util.UUID import kotlin.random.Random @@ -378,6 +382,37 @@ class AttachmentTableTest_deduping { } } + @Test + fun downloads() { + // Normal attachment download that dupes with an existing attachment + test { + val id1 = insertWithData(DATA_A) + upload(id1) + + val id2 = insertUndownloadedPlaceholder() + download(id2, DATA_A) + + assertDataFilesAreTheSame(id1, id2) + assertDataHashEndMatches(id1, id2) + assertRemoteFieldsMatch(id1, id2) + assertArchiveFieldsMatch(id1, id2) + } + + // Attachment download that dupes with an existing attachment, but has bad padding + test { + val id1 = insertWithData(DATA_A) + upload(id1) + + val id2 = insertUndownloadedPlaceholder() + download(id2, DATA_A, properPadding = false) + + assertDataFilesAreTheSame(id1, id2) + assertDataHashEndMatches(id1, id2) + assertRemoteFieldsMatch(id1, id2) + assertArchiveFieldsMatch(id1, id2) + } + } + /** * Various deletion scenarios to ensure that duped files don't deleted while there's still references. */ @@ -614,6 +649,39 @@ class AttachmentTableTest_deduping { } private class TestContext { + fun insertUndownloadedPlaceholder(): AttachmentId { + return SignalDatabase.attachments.insertAttachmentsForMessage( + mmsId = 1, + attachments = listOf( + PointerAttachment( + contentType = "image/jpeg", + transferState = AttachmentTable.TRANSFER_PROGRESS_PENDING, + size = 100, + fileName = null, + cdn = Cdn.CDN_3, + location = "somelocation", + key = Base64.encodeWithPadding(Util.getSecretBytes(64)), + iv = null, + digest = Util.getSecretBytes(64), + incrementalDigest = null, + incrementalMacChunkSize = 0, + fastPreflightId = null, + voiceNote = false, + borderless = false, + videoGif = false, + width = 100, + height = 100, + uploadTimestamp = System.currentTimeMillis(), + caption = null, + stickerLocator = null, + blurHash = null, + uuid = UUID.randomUUID() + ) + ), + quoteAttachment = emptyList() + ).values.first() + } + fun insertWithData(data: ByteArray, transformProperties: TransformProperties = TransformProperties.empty()): AttachmentId { val uri = BlobProvider.getInstance().forData(data).createForSingleSessionInMemory() @@ -675,6 +743,22 @@ class AttachmentTableTest_deduping { ) } + fun download(attachmentId: AttachmentId, data: ByteArray, properPadding: Boolean = true) { + val paddedData = if (properPadding) { + PaddingInputStream(data.inputStream(), data.size.toLong()).readFully() + } else { + val badPadding = ByteArray(16) { 42 } + data + badPadding + } + + SignalDatabase.attachments.finalizeAttachmentAfterDownload( + mmsId = 1, + attachmentId = attachmentId, + inputStream = LimitedInputStream(paddedData.inputStream(), data.size.toLong()), + iv = Util.getSecretBytes(16) + ) + } + fun delete(attachmentId: AttachmentId) { SignalDatabase.attachments.deleteAttachment(attachmentId) } 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 98002b8856..da99a9e687 100644 --- a/app/src/main/java/org/thoughtcrime/securesms/database/AttachmentTable.kt +++ b/app/src/main/java/org/thoughtcrime/securesms/database/AttachmentTable.kt @@ -1140,8 +1140,13 @@ class AttachmentTable( values.put(TRANSFER_FILE, null as String?) values.put(TRANSFORM_PROPERTIES, TransformProperties.forSkipTransform().serialize()) values.put(ARCHIVE_TRANSFER_FILE, null as String?) + values.put(REMOTE_LOCATION, existingPlaceholder.remoteLocation) + values.put(CDN_NUMBER, existingPlaceholder.cdn.serialize()) + values.put(REMOTE_KEY, existingPlaceholder.remoteKey!!) values.put(REMOTE_IV, iv) values.put(REMOTE_DIGEST, digest) + values.put(REMOTE_INCREMENTAL_DIGEST, existingPlaceholder.incrementalDigest) + values.put(REMOTE_INCREMENTAL_DIGEST_CHUNK_SIZE, existingPlaceholder.incrementalMacChunkSize) if (digestChanged) { values.put(UPLOAD_TIMESTAMP, 0) @@ -1151,9 +1156,11 @@ class AttachmentTable( values.put(OFFLOAD_RESTORED_AT, offloadRestoredAt.inWholeMilliseconds) } + val dataFilePath = hashMatch?.file?.absolutePath ?: fileWriteResult.file.absolutePath + db.update(TABLE_NAME) .values(values) - .where("$ID = ?", attachmentId.id) + .where("$ID = ? OR $DATA_FILE = ?", attachmentId.id, dataFilePath) .run() Log.i(TAG, "[finalizeAttachmentAfterDownload] Finalized downloaded data for $attachmentId. (MessageId: $mmsId, $attachmentId)")