Add initial thumbnail restore for message backup.

2026-04-21 17:29:32 +01:00 · 2024-05-15 14:58:29 -04:00
parent 757c0fd2ea
commit b72d586748
25 changed files with 397 additions and 50 deletions
--- a/app/src/main/java/org/thoughtcrime/securesms/database/AttachmentTable.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/database/AttachmentTable.kt
@@ -125,6 +125,8 @@ class AttachmentTable(
    const val DATA_RANDOM = "data_random"
    const val DATA_HASH_START = "data_hash_start"
    const val DATA_HASH_END = "data_hash_end"
+    const val THUMBNAIL_FILE = "thumbnail_file"
+    const val THUMBNAIL_RANDOM = "thumbnail_random"
    const val FILE_NAME = "file_name"
    const val FAST_PREFLIGHT_ID = "fast_preflight_id"
    const val VOICE_NOTE = "voice_note"
@@ -145,6 +147,8 @@ class AttachmentTable(
    const val ARCHIVE_CDN = "archive_cdn"
    const val ARCHIVE_MEDIA_NAME = "archive_media_name"
    const val ARCHIVE_MEDIA_ID = "archive_media_id"
+    const val ARCHIVE_THUMBNAIL_MEDIA_ID = "archive_thumbnail_media_id"
+    const val ARCHIVE_THUMBNAIL_CDN = "archive_thumbnail_cdn"
    const val ARCHIVE_TRANSFER_FILE = "archive_transfer_file"
    const val ARCHIVE_TRANSFER_STATE = "archive_transfer_state"

@@ -159,6 +163,7 @@ class AttachmentTable(
    const val TRANSFER_PROGRESS_PERMANENT_FAILURE = 4
    const val TRANSFER_NEEDS_RESTORE = 5
    const val TRANSFER_RESTORE_IN_PROGRESS = 6
+    const val TRANSFER_RESTORE_OFFLOADED = 7
    const val PREUPLOAD_MESSAGE_ID: Long = -8675309

    private val PROJECTION = arrayOf(
@@ -196,9 +201,11 @@ class AttachmentTable(
      DATA_HASH_START,
      DATA_HASH_END,
      ARCHIVE_CDN,
+      ARCHIVE_THUMBNAIL_CDN,
      ARCHIVE_MEDIA_NAME,
      ARCHIVE_MEDIA_ID,
-      ARCHIVE_TRANSFER_FILE
+      ARCHIVE_TRANSFER_FILE,
+      THUMBNAIL_FILE
    )

    @JvmField
@@ -240,8 +247,12 @@ class AttachmentTable(
        $ARCHIVE_CDN INTEGER DEFAULT 0,
        $ARCHIVE_MEDIA_NAME TEXT DEFAULT NULL,
        $ARCHIVE_MEDIA_ID TEXT DEFAULT NULL,
+        $ARCHIVE_THUMBNAIL_MEDIA_ID TEXT DEFAULT NULL,
+        $ARCHIVE_THUMBNAIL_CDN INTEGER DEFAULT 0,
        $ARCHIVE_TRANSFER_FILE TEXT DEFAULT NULL,
-        $ARCHIVE_TRANSFER_STATE INTEGER DEFAULT ${ArchiveTransferState.NONE.value}
+        $ARCHIVE_TRANSFER_STATE INTEGER DEFAULT ${ArchiveTransferState.NONE.value},
+        $THUMBNAIL_FILE TEXT DEFAULT NULL,
+        $THUMBNAIL_RANDOM BLOB DEFAULT NULL
      )
      """

@@ -273,6 +284,15 @@ class AttachmentTable(
    } ?: throw IOException("No stream for: $attachmentId")
  }

+  @Throws(IOException::class)
+  fun getAttachmentThumbnailStream(attachmentId: AttachmentId, offset: Long): InputStream {
+    return try {
+      getThumbnailStream(attachmentId, offset)
+    } catch (e: FileNotFoundException) {
+      throw IOException("No stream for: $attachmentId", e)
+    } ?: throw IOException("No stream for: $attachmentId")
+  }
+
  /**
   * Returns a [File] for an attachment that has no [DATA_HASH_END] and is in the [TRANSFER_PROGRESS_DONE] state, if present.
   */
@@ -826,6 +846,36 @@ class AttachmentTable(
    }
  }

+  @Throws(IOException::class)
+  fun finalizeAttachmentThumbnailAfterDownload(attachmentId: AttachmentId, archiveMediaId: String, inputStream: InputStream, transferFile: File) {
+    Log.i(TAG, "[finalizeAttachmentThumbnailAfterDownload] Finalizing downloaded data for $attachmentId.")
+    val fileWriteResult: DataFileWriteResult = writeToDataFile(newDataFile(context), inputStream, TransformProperties.empty())
+
+    writableDatabase.withinTransaction { db ->
+      val values = contentValuesOf(
+        THUMBNAIL_FILE to fileWriteResult.file.absolutePath,
+        THUMBNAIL_RANDOM to fileWriteResult.random
+      )
+
+      db.update(TABLE_NAME)
+        .values(values)
+        .where("$ARCHIVE_MEDIA_ID = ?", archiveMediaId)
+        .run()
+
+      db.update(TABLE_NAME)
+        .values(TRANSFER_STATE to TRANSFER_RESTORE_OFFLOADED)
+        .where("$ID = ?", attachmentId.id)
+        .run()
+    }
+
+    notifyConversationListListeners()
+    notifyAttachmentListeners()
+
+    if (!transferFile.delete()) {
+      Log.w(TAG, "Unable to delete transfer file.")
+    }
+  }
+
  /**
   * Needs to be called after an attachment is successfully uploaded. Writes metadata around it's final remote location, as well as calculates
   * it's ending hash, which is critical for backups.
@@ -1158,6 +1208,10 @@ class AttachmentTable(
    return transferFile
  }

+  fun createArchiveThumbnailTransferFile(): File {
+    return newTransferFile()
+  }
+
  fun getDataFileInfo(attachmentId: AttachmentId): DataFileInfo? {
    return readableDatabase
      .select(ID, DATA_FILE, DATA_SIZE, DATA_RANDOM, DATA_HASH_START, DATA_HASH_END, TRANSFORM_PROPERTIES, UPLOAD_TIMESTAMP)
@@ -1173,6 +1227,21 @@ class AttachmentTable(
      }
  }

+  fun getThumbnailFileInfo(attachmentId: AttachmentId): ThumbnailFileInfo? {
+    return readableDatabase
+      .select(ID, THUMBNAIL_FILE, THUMBNAIL_RANDOM)
+      .from(TABLE_NAME)
+      .where("$ID = ?", attachmentId.id)
+      .run()
+      .readToSingleObject { cursor ->
+        if (cursor.isNull(THUMBNAIL_FILE)) {
+          null
+        } else {
+          cursor.readThumbnailFileInfo()
+        }
+      }
+  }
+
  fun getDataFilePath(attachmentId: AttachmentId): String? {
    return readableDatabase
      .select(DATA_FILE)
@@ -1320,8 +1389,10 @@ class AttachmentTable(
              uploadTimestamp = jsonObject.getLong(UPLOAD_TIMESTAMP),
              dataHash = jsonObject.getString(DATA_HASH_END),
              archiveCdn = jsonObject.getInt(ARCHIVE_CDN),
+              archiveThumbnailCdn = jsonObject.getInt(ARCHIVE_THUMBNAIL_CDN),
              archiveMediaName = jsonObject.getString(ARCHIVE_MEDIA_NAME),
-              archiveMediaId = jsonObject.getString(ARCHIVE_MEDIA_ID)
+              archiveMediaId = jsonObject.getString(ARCHIVE_MEDIA_ID),
+              hasArchiveThumbnail = !TextUtils.isEmpty(jsonObject.getString(THUMBNAIL_FILE))
            )
          }
        }
@@ -1361,13 +1432,14 @@ class AttachmentTable(
    return readableDatabase.rawQuery(query, null)
  }

-  fun setArchiveData(attachmentId: AttachmentId, archiveCdn: Int, archiveMediaName: String, archiveMediaId: String) {
+  fun setArchiveData(attachmentId: AttachmentId, archiveCdn: Int, archiveMediaName: String, archiveMediaId: String, archiveThumbnailMediaId: String) {
    writableDatabase
      .update(TABLE_NAME)
      .values(
        ARCHIVE_CDN to archiveCdn,
        ARCHIVE_MEDIA_ID to archiveMediaId,
        ARCHIVE_MEDIA_NAME to archiveMediaName,
+        ARCHIVE_THUMBNAIL_MEDIA_ID to archiveThumbnailMediaId,
        ARCHIVE_TRANSFER_STATE to ArchiveTransferState.FINISHED.value
      )
      .where("$ID = ?", attachmentId.id)
@@ -1375,13 +1447,14 @@ class AttachmentTable(
  }

  fun updateArchiveCdnByMediaId(archiveMediaId: String, archiveCdn: Int): Int {
-    return writableDatabase
-      .update(TABLE_NAME)
-      .values(
-        ARCHIVE_CDN to archiveCdn
-      )
-      .where("$ARCHIVE_MEDIA_ID = ?", archiveMediaId)
-      .run()
+    return writableDatabase.rawQuery(
+      "UPDATE $TABLE_NAME SET " +
+        "$ARCHIVE_THUMBNAIL_CDN = CASE WHEN $ARCHIVE_THUMBNAIL_MEDIA_ID = ? THEN ? ELSE $ARCHIVE_THUMBNAIL_CDN END," +
+        "$ARCHIVE_CDN = CASE WHEN $ARCHIVE_MEDIA_ID = ? THEN ? ELSE $ARCHIVE_CDN END " +
+        "WHERE $ARCHIVE_MEDIA_ID = ? OR $ARCHIVE_THUMBNAIL_MEDIA_ID = ? " +
+        "RETURNING $ARCHIVE_CDN, $ARCHIVE_THUMBNAIL_CDN",
+      SqlUtil.buildArgs(archiveMediaId, archiveCdn, archiveMediaId, archiveCdn, archiveMediaId, archiveMediaId)
+    ).count
  }

  fun clearArchiveData(attachmentIds: List<AttachmentId>) {
@@ -1485,6 +1558,21 @@ class AttachmentTable(
    }
  }

+  @Throws(FileNotFoundException::class)
+  private fun getThumbnailStream(attachmentId: AttachmentId, offset: Long): InputStream? {
+    val thumbnailInfo = getThumbnailFileInfo(attachmentId) ?: return null
+
+    return try {
+      ModernDecryptingPartInputStream.createFor(attachmentSecret, thumbnailInfo.random, thumbnailInfo.file, offset)
+    } catch (e: FileNotFoundException) {
+      Log.w(TAG, e)
+      throw e
+    } catch (e: IOException) {
+      Log.w(TAG, e)
+      null
+    }
+  }
+
  @Throws(IOException::class)
  private fun newTransferFile(): File {
    val partsDirectory = context.getDir(DIRECTORY, Context.MODE_PRIVATE)
@@ -1664,6 +1752,7 @@ class AttachmentTable(
        put(ARCHIVE_CDN, attachment.archiveCdn)
        put(ARCHIVE_MEDIA_NAME, attachment.archiveMediaName)
        put(ARCHIVE_MEDIA_ID, attachment.archiveMediaId)
+        put(ARCHIVE_THUMBNAIL_MEDIA_ID, attachment.archiveThumbnailMediaId)

        attachment.stickerLocator?.let { sticker ->
          put(STICKER_PACK_ID, sticker.packId)
@@ -1874,8 +1963,10 @@ class AttachmentTable(
      uploadTimestamp = cursor.requireLong(UPLOAD_TIMESTAMP),
      dataHash = cursor.requireString(DATA_HASH_END),
      archiveCdn = cursor.requireInt(ARCHIVE_CDN),
+      archiveThumbnailCdn = cursor.requireInt(ARCHIVE_THUMBNAIL_CDN),
      archiveMediaName = cursor.requireString(ARCHIVE_MEDIA_NAME),
-      archiveMediaId = cursor.requireString(ARCHIVE_MEDIA_ID)
+      archiveMediaId = cursor.requireString(ARCHIVE_MEDIA_ID),
+      hasArchiveThumbnail = !cursor.isNull(THUMBNAIL_FILE)
    )
  }

@@ -1900,6 +1991,14 @@ class AttachmentTable(
    )
  }

+  private fun Cursor.readThumbnailFileInfo(): ThumbnailFileInfo {
+    return ThumbnailFileInfo(
+      id = AttachmentId(this.requireLong(ID)),
+      file = File(this.requireNonNullString(THUMBNAIL_FILE)),
+      random = this.requireNonNullBlob(THUMBNAIL_RANDOM)
+    )
+  }
+
  private fun Cursor.readStickerLocator(): StickerLocator? {
    return if (this.requireInt(STICKER_ID) >= 0) {
      StickerLocator(
@@ -1954,6 +2053,13 @@ class AttachmentTable(
    val uploadTimestamp: Long
  )

+  @VisibleForTesting
+  class ThumbnailFileInfo(
+    val id: AttachmentId,
+    val file: File,
+    val random: ByteArray
+  )
+
  @Parcelize
  data class TransformProperties(
    @JsonProperty("skipTransform")
--- a/app/src/main/java/org/thoughtcrime/securesms/database/MessageTable.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/database/MessageTable.kt
@@ -357,7 +357,8 @@ open class MessageTable(context: Context?, databaseHelper: SignalDatabase) : Dat
            '${AttachmentTable.MESSAGE_ID}', ${AttachmentTable.TABLE_NAME}.${AttachmentTable.MESSAGE_ID},
            '${AttachmentTable.DATA_SIZE}', ${AttachmentTable.TABLE_NAME}.${AttachmentTable.DATA_SIZE}, 
            '${AttachmentTable.FILE_NAME}', ${AttachmentTable.TABLE_NAME}.${AttachmentTable.FILE_NAME}, 
-            '${AttachmentTable.DATA_FILE}', ${AttachmentTable.TABLE_NAME}.${AttachmentTable.DATA_FILE}, 
+            '${AttachmentTable.DATA_FILE}', ${AttachmentTable.TABLE_NAME}.${AttachmentTable.DATA_FILE},
+            '${AttachmentTable.THUMBNAIL_FILE}', ${AttachmentTable.TABLE_NAME}.${AttachmentTable.THUMBNAIL_FILE},
            '${AttachmentTable.CONTENT_TYPE}', ${AttachmentTable.TABLE_NAME}.${AttachmentTable.CONTENT_TYPE}, 
            '${AttachmentTable.CDN_NUMBER}', ${AttachmentTable.TABLE_NAME}.${AttachmentTable.CDN_NUMBER}, 
            '${AttachmentTable.REMOTE_LOCATION}', ${AttachmentTable.TABLE_NAME}.${AttachmentTable.REMOTE_LOCATION}, 
@@ -381,6 +382,7 @@ open class MessageTable(context: Context?, databaseHelper: SignalDatabase) : Dat
            '${AttachmentTable.UPLOAD_TIMESTAMP}', ${AttachmentTable.TABLE_NAME}.${AttachmentTable.UPLOAD_TIMESTAMP},
            '${AttachmentTable.DATA_HASH_END}', ${AttachmentTable.TABLE_NAME}.${AttachmentTable.DATA_HASH_END},
            '${AttachmentTable.ARCHIVE_CDN}', ${AttachmentTable.TABLE_NAME}.${AttachmentTable.ARCHIVE_CDN},
+            '${AttachmentTable.ARCHIVE_THUMBNAIL_CDN}', ${AttachmentTable.TABLE_NAME}.${AttachmentTable.ARCHIVE_THUMBNAIL_CDN},
            '${AttachmentTable.ARCHIVE_MEDIA_NAME}', ${AttachmentTable.TABLE_NAME}.${AttachmentTable.ARCHIVE_MEDIA_NAME},
            '${AttachmentTable.ARCHIVE_MEDIA_ID}', ${AttachmentTable.TABLE_NAME}.${AttachmentTable.ARCHIVE_MEDIA_ID}
          )
--- a/app/src/main/java/org/thoughtcrime/securesms/database/helpers/SignalDatabaseMigrations.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/database/helpers/SignalDatabaseMigrations.kt
@@ -88,6 +88,7 @@ import org.thoughtcrime.securesms.database.helpers.migration.V227_AddAttachmentA
 import org.thoughtcrime.securesms.database.helpers.migration.V228_AddNameCollisionTables
 import org.thoughtcrime.securesms.database.helpers.migration.V229_MarkMissedCallEventsNotified
 import org.thoughtcrime.securesms.database.helpers.migration.V230_UnreadCountIndices
+import org.thoughtcrime.securesms.database.helpers.migration.V231_ArchiveThumbnailColumns

 /**
 * Contains all of the database migrations for [SignalDatabase]. Broken into a separate file for cleanliness.
@@ -178,10 +179,11 @@ object SignalDatabaseMigrations {
    227 to V227_AddAttachmentArchiveTransferState,
    228 to V228_AddNameCollisionTables,
    229 to V229_MarkMissedCallEventsNotified,
-    230 to V230_UnreadCountIndices
+    230 to V230_UnreadCountIndices,
+    231 to V231_ArchiveThumbnailColumns
  )

-  const val DATABASE_VERSION = 230
+  const val DATABASE_VERSION = 231

  @JvmStatic
  fun migrate(context: Application, db: SQLiteDatabase, oldVersion: Int, newVersion: Int) {
--- a/app/src/main/java/org/thoughtcrime/securesms/database/helpers/migration/V231_ArchiveThumbnailColumns.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/database/helpers/migration/V231_ArchiveThumbnailColumns.kt
@@ -0,0 +1,18 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.database.helpers.migration
+
+import android.app.Application
+import net.zetetic.database.sqlcipher.SQLiteDatabase
+
+object V231_ArchiveThumbnailColumns : SignalDatabaseMigration {
+  override fun migrate(context: Application, db: SQLiteDatabase, oldVersion: Int, newVersion: Int) {
+    db.execSQL("ALTER TABLE attachment ADD COLUMN thumbnail_file TEXT DEFAULT NULL")
+    db.execSQL("ALTER TABLE attachment ADD COLUMN thumbnail_random BLOB DEFAULT NULL")
+    db.execSQL("ALTER TABLE attachment ADD COLUMN archive_thumbnail_cdn INTEGER DEFAULT 0")
+    db.execSQL("ALTER TABLE attachment ADD COLUMN archive_thumbnail_media_id TEXT DEFAULT NULL")
+  }
+}