Ensure we are restoring media per spec with full media and thumbnail rules.

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

@@ -6,24 +6,30 @@ import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.FlakyTest
 import androidx.test.platform.app.InstrumentationRegistry
 import assertk.assertThat
+import assertk.assertions.hasSize
 import assertk.assertions.isEqualTo
 import assertk.assertions.isNotEqualTo
+import assertk.assertions.isTrue
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertNotEquals
 import org.junit.Before
 import org.junit.Ignore
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.signal.core.util.Base64.decodeBase64OrThrow
 import org.signal.core.util.copyTo
 import org.signal.core.util.stream.NullOutputStream
+import org.thoughtcrime.securesms.attachments.ArchivedAttachment
 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.mms.IncomingMessage
 import org.thoughtcrime.securesms.mms.MediaStream
 import org.thoughtcrime.securesms.mms.SentMediaQuality
 import org.thoughtcrime.securesms.providers.BlobProvider
+import org.thoughtcrime.securesms.testing.SignalActivityRule
 import org.thoughtcrime.securesms.util.MediaUtil
 import org.whispersystems.signalservice.api.crypto.AttachmentCipherOutputStream
 import org.whispersystems.signalservice.api.crypto.NoCipherOutputStream
@@ -32,10 +38,19 @@ import org.whispersystems.signalservice.api.messages.SignalServiceAttachmentRemo
 import java.io.ByteArrayOutputStream
 import java.io.File
 import java.util.Optional
+import java.util.UUID
+import kotlin.random.Random
+import kotlin.time.Duration
+import kotlin.time.Duration.Companion.days
+import kotlin.time.Duration.Companion.milliseconds
+import kotlin.time.Duration.Companion.seconds
 
 @RunWith(AndroidJUnit4::class)
 class AttachmentTableTest {
 
+  @get:Rule
+  val harness = SignalActivityRule(othersCount = 10)
+
   @Before
   fun setUp() {
     SignalDatabase.attachments.deleteAllAttachments()
@@ -195,6 +210,57 @@ class AttachmentTableTest {
     assertThat(SignalDatabase.attachments.getAttachment(attachmentId)!!.archiveTransferState).isEqualTo(AttachmentTable.ArchiveTransferState.NONE)
   }
 
+  @Test
+  fun given10NewerAnd10OlderAttachments_whenIGetEachBatch_thenIExpectProperBucketing() {
+    val now = System.currentTimeMillis().milliseconds
+    val attachments = (0 until 20).map {
+      createArchivedAttachment()
+    }
+
+    val newMessages = attachments.take(10).mapIndexed { index, attachment ->
+      createIncomingMessage(serverTime = now - index.seconds, attachment = attachment)
+    }
+
+    val twoMonthsAgo = now - 60.days
+    val oldMessages = attachments.drop(10).mapIndexed { index, attachment ->
+      createIncomingMessage(serverTime = twoMonthsAgo - index.seconds, attachment = attachment)
+    }
+
+    (newMessages + oldMessages).forEach {
+      SignalDatabase.messages.insertMessageInbox(it)
+    }
+
+    val firstAttachmentsToDownload = SignalDatabase.attachments.getLast30DaysOfRestorableAttachments(500)
+    val nextAttachmentsToDownload = SignalDatabase.attachments.getOlderRestorableAttachments(500)
+
+    assertThat(firstAttachmentsToDownload).hasSize(10)
+    val resultNewMessages = SignalDatabase.messages.getMessages(firstAttachmentsToDownload.map { it.mmsId })
+    resultNewMessages.forEach {
+      assertThat(it.serverTimestamp.milliseconds >= now - 30.days).isTrue()
+    }
+
+    assertThat(nextAttachmentsToDownload).hasSize(10)
+    val resultOldMessages = SignalDatabase.messages.getMessages(nextAttachmentsToDownload.map { it.mmsId })
+    resultOldMessages.forEach {
+      assertThat(it.serverTimestamp.milliseconds < now - 30.days).isTrue()
+    }
+  }
+
+  private fun createIncomingMessage(
+    serverTime: Duration,
+    attachment: Attachment
+  ): IncomingMessage {
+    return IncomingMessage(
+      type = MessageType.NORMAL,
+      from = harness.others[0],
+      body = null,
+      sentTimeMillis = serverTime.inWholeMilliseconds,
+      serverTimeMillis = serverTime.inWholeMilliseconds,
+      receivedTimeMillis = serverTime.inWholeMilliseconds,
+      attachments = listOf(attachment)
+    )
+  }
+
   private fun createAttachmentPointer(key: ByteArray, digest: ByteArray, size: Int): Attachment {
     return PointerAttachment.forPointer(
       pointer = Optional.of(
@@ -223,6 +289,32 @@ class AttachmentTableTest {
     ).get()
   }
 
+  private fun createArchivedAttachment(): Attachment {
+    return ArchivedAttachment(
+      contentType = "image/jpeg",
+      size = 1024,
+      cdn = 3,
+      uploadTimestamp = 0,
+      key = Random.nextBytes(8),
+      cdnKey = "password",
+      archiveCdn = 3,
+      plaintextHash = Random.nextBytes(8),
+      incrementalMac = Random.nextBytes(8),
+      incrementalMacChunkSize = 8,
+      width = 100,
+      height = 100,
+      caption = null,
+      blurHash = null,
+      voiceNote = false,
+      borderless = false,
+      stickerLocator = null,
+      gif = false,
+      quote = false,
+      uuid = UUID.randomUUID(),
+      fileName = null
+    )
+  }
+
   private fun createAttachment(id: Long, uri: Uri, transformProperties: AttachmentTable.TransformProperties): UriAttachment {
     return UriAttachmentBuilder.build(
       id,

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

@@ -119,6 +119,7 @@ import java.util.Optional
 import java.util.UUID
 import kotlin.time.Duration
 import kotlin.time.Duration.Companion.days
+import kotlin.time.Duration.Companion.milliseconds
 
 class AttachmentTable(
   context: Context,
@@ -510,6 +511,54 @@ class AttachmentTable(
       }
   }
 
+  /**
+   * Grabs the last 30 days worth of restorable attachments with respect to the message's server timestamp,
+   * up to the given batch size.
+   */
+  fun getLast30DaysOfRestorableAttachments(batchSize: Int): List<RestorableAttachment> {
+    val thirtyDaysAgo = System.currentTimeMillis().milliseconds - 30.days
+    return readableDatabase
+      .select("$TABLE_NAME.$ID", MESSAGE_ID, DATA_SIZE, DATA_HASH_END, REMOTE_KEY)
+      .from("$TABLE_NAME INNER JOIN ${MessageTable.TABLE_NAME} ON ${MessageTable.TABLE_NAME}.${MessageTable.ID} = $TABLE_NAME.$MESSAGE_ID")
+      .where("$TRANSFER_STATE = ? AND ${MessageTable.TABLE_NAME}.${MessageTable.DATE_RECEIVED} >= ?", TRANSFER_NEEDS_RESTORE, thirtyDaysAgo.inWholeMilliseconds)
+      .limit(batchSize)
+      .orderBy("$TABLE_NAME.$ID DESC")
+      .run()
+      .readToList {
+        RestorableAttachment(
+          attachmentId = AttachmentId(it.requireLong(ID)),
+          mmsId = it.requireLong(MESSAGE_ID),
+          size = it.requireLong(DATA_SIZE),
+          plaintextHash = it.requireBlob(DATA_HASH_END),
+          remoteKey = it.requireBlob(REMOTE_KEY)
+        )
+      }
+  }
+
+  /**
+   * Grabs attachments outside of the last 30 days with respect to the message's server timestamp,
+   * up to the given batch size.
+   */
+  fun getOlderRestorableAttachments(batchSize: Int): List<RestorableAttachment> {
+    val thirtyDaysAgo = System.currentTimeMillis().milliseconds - 30.days
+    return readableDatabase
+      .select("$TABLE_NAME.$ID", MESSAGE_ID, DATA_SIZE, DATA_HASH_END, REMOTE_KEY)
+      .from("$TABLE_NAME INNER JOIN ${MessageTable.TABLE_NAME} ON ${MessageTable.TABLE_NAME}.${MessageTable.ID} = $TABLE_NAME.$MESSAGE_ID")
+      .where("$TRANSFER_STATE = ? AND ${MessageTable.TABLE_NAME}.${MessageTable.DATE_RECEIVED} < ?", TRANSFER_NEEDS_RESTORE, thirtyDaysAgo.inWholeMilliseconds)
+      .limit(batchSize)
+      .orderBy("$TABLE_NAME.$ID DESC")
+      .run()
+      .readToList {
+        RestorableAttachment(
+          attachmentId = AttachmentId(it.requireLong(ID)),
+          mmsId = it.requireLong(MESSAGE_ID),
+          size = it.requireLong(DATA_SIZE),
+          plaintextHash = it.requireBlob(DATA_HASH_END),
+          remoteKey = it.requireBlob(REMOTE_KEY)
+        )
+      }
+  }
+
   fun getRestorableAttachments(batchSize: Int): List<RestorableAttachment> {
     return readableDatabase
       .select(ID, MESSAGE_ID, DATA_SIZE, DATA_HASH_END, REMOTE_KEY)

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

@@ -61,7 +61,16 @@ class BackupRestoreMediaJob private constructor(parameters: Parameters) : BaseJo
       val restoreThumbnailOnlyAttachmentsIds: MutableList<AttachmentId> = mutableListOf()
       val notRestorable: MutableList<AttachmentId> = mutableListOf()
 
-      val attachmentBatch = SignalDatabase.attachments.getRestorableAttachments(batchSize)
+      val last30DaysAttachments = SignalDatabase.attachments.getLast30DaysOfRestorableAttachments(batchSize)
+      val remainingSize = batchSize - last30DaysAttachments.size
+
+      val remaining = if (remainingSize > 0) {
+        SignalDatabase.attachments.getOlderRestorableAttachments(batchSize = remainingSize)
+      } else {
+        listOf()
+      }
+
+      val attachmentBatch = last30DaysAttachments + remaining
       val messageIds = attachmentBatch.map { it.mmsId }.toSet()
       val messageMap = SignalDatabase.messages.getMessages(messageIds).associate { it.id to (it as MmsMessageRecord) }
 
@@ -75,18 +84,18 @@ class BackupRestoreMediaJob private constructor(parameters: Parameters) : BaseJo
           continue
         }
 
-        restoreThumbnailJobs += RestoreAttachmentThumbnailJob(
-          messageId = attachment.mmsId,
-          attachmentId = attachment.attachmentId,
-          highPriority = false
-        )
-
         if (isWallpaper || shouldRestoreFullSize(message!!, restoreTime, SignalStore.backup.optimizeStorage)) {
           restoreFullAttachmentJobs += RestoreAttachmentJob.forInitialRestore(
             messageId = attachment.mmsId,
             attachmentId = attachment.attachmentId
           )
         } else {
+          restoreThumbnailJobs += RestoreAttachmentThumbnailJob(
+            messageId = attachment.mmsId,
+            attachmentId = attachment.attachmentId,
+            highPriority = false
+          )
+
           restoreThumbnailOnlyAttachmentsIds += attachment.attachmentId
         }
       }

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

@@ -203,6 +203,15 @@ class RestoreAttachmentJob private constructor(
       Log.w(TAG, format(this, "onFailure() messageId: $messageId  attachmentId: $attachmentId"))
 
       markFailed(attachmentId)
+
+      Log.w(TAG, "onFailure(): Attempting to fall back on attachment thumbnail.")
+      val restoreThumbnailAttachmentJob = RestoreAttachmentThumbnailJob(
+        messageId = messageId,
+        attachmentId = attachmentId,
+        highPriority = manual
+      )
+
+      AppDependencies.jobManager.add(restoreThumbnailAttachmentJob)
     }
   }