From f48a13afc0e6a9e44ec21f13f1189a0bd53cbd9d Mon Sep 17 00:00:00 2001
From: Greyson Parrelli <greyson@signal.org>
Date: Fri, 28 Mar 2025 10:40:08 -0400
Subject: [PATCH] Fix handling of missing files during archive upload.

---
 .../securesms/database/AttachmentTable.kt     | 26 +++++++++++++++----
 .../jobs/UploadAttachmentToArchiveJob.kt      | 12 ++++++---
 2 files changed, 29 insertions(+), 9 deletions(-)

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 718e2b7a73..70de72dadf 100644
--- a/app/src/main/java/org/thoughtcrime/securesms/database/AttachmentTable.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/database/AttachmentTable.kt
@@ -308,11 +308,7 @@ class AttachmentTable(
 
   @Throws(IOException::class)
   fun getAttachmentStream(attachmentId: AttachmentId, offset: Long): InputStream {
-    return try {
-      getDataStream(attachmentId, offset)
-    } catch (e: FileNotFoundException) {
-      throw IOException("No stream for: $attachmentId", e)
-    } ?: throw IOException("No stream for: $attachmentId")
+    return getDataStream(attachmentId, offset) ?: throw FileNotFoundException("No stream for: $attachmentId")
   }
 
   @Throws(IOException::class)
@@ -671,6 +667,26 @@ class AttachmentTable(
     }
   }
 
+  /**
+   * Sets the archive transfer state for the given attachment and all other attachments that share the same data file.
+   */
+  fun setArchiveTransferStateUnlessPermanentFailure(id: AttachmentId, state: ArchiveTransferState) {
+    writableDatabase.withinTransaction {
+      val dataFile: String = readableDatabase
+        .select(DATA_FILE)
+        .from(TABLE_NAME)
+        .where("$ID = ?", id.id)
+        .run()
+        .readToSingleObject { it.requireString(DATA_FILE) } ?: return@withinTransaction
+
+      writableDatabase
+        .update(TABLE_NAME)
+        .values(ARCHIVE_TRANSFER_STATE to state.value)
+        .where("$DATA_FILE = ? AND $ARCHIVE_TRANSFER_STATE != ${ArchiveTransferState.PERMANENT_FAILURE.value}", dataFile)
+        .run()
+    }
+  }
+
   /**
    * Marks eligible attachments as offloaded based on their received at timestamp, their last restore time,
    * presence of thumbnail if media, and the full file being available in the archive.
diff --git a/app/src/main/java/org/thoughtcrime/securesms/jobs/UploadAttachmentToArchiveJob.kt b/app/src/main/java/org/thoughtcrime/securesms/jobs/UploadAttachmentToArchiveJob.kt
index 9ba6dda0d1..524e122f7e 100644
--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/UploadAttachmentToArchiveJob.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/UploadAttachmentToArchiveJob.kt
@@ -26,6 +26,7 @@ import org.whispersystems.signalservice.api.NetworkResult
 import org.whispersystems.signalservice.api.archive.ArchiveMediaUploadFormStatusCodes
 import org.whispersystems.signalservice.api.attachment.AttachmentUploadResult
 import org.whispersystems.signalservice.api.messages.SignalServiceAttachment
+import java.io.FileNotFoundException
 import java.io.IOException
 import java.net.ProtocolException
 import kotlin.random.Random
@@ -74,7 +75,7 @@ class UploadAttachmentToArchiveJob private constructor(
 
     if (transferStatus == AttachmentTable.ArchiveTransferState.NONE) {
       Log.d(TAG, "[$attachmentId] Updating archive transfer state to ${AttachmentTable.ArchiveTransferState.UPLOAD_IN_PROGRESS}")
-      SignalDatabase.attachments.setArchiveTransferState(attachmentId, AttachmentTable.ArchiveTransferState.UPLOAD_IN_PROGRESS)
+      SignalDatabase.attachments.setArchiveTransferStateUnlessPermanentFailure(attachmentId, AttachmentTable.ArchiveTransferState.UPLOAD_IN_PROGRESS)
     }
   }
 
@@ -110,7 +111,6 @@ class UploadAttachmentToArchiveJob private constructor(
 
     if (attachment.remoteKey == null || attachment.remoteIv == null) {
       Log.w(TAG, "[$attachmentId] Attachment is missing remote key or IV! Cannot upload.")
-      SignalDatabase.attachments.setArchiveTransferState(attachmentId, AttachmentTable.ArchiveTransferState.NONE)
       return Result.failure()
     }
 
@@ -143,8 +143,12 @@ class UploadAttachmentToArchiveJob private constructor(
           override fun shouldCancel() = this@UploadAttachmentToArchiveJob.isCanceled
         }
       )
+    } catch (e: FileNotFoundException) {
+      Log.w(TAG, "[$attachmentId] No file exists for this attachment! Marking as a permanent failure.", e)
+      SignalDatabase.attachments.setArchiveTransferState(attachmentId, AttachmentTable.ArchiveTransferState.PERMANENT_FAILURE)
+      return Result.failure()
     } catch (e: IOException) {
-      Log.e(TAG, "[$attachmentId] Failed to get attachment stream.", e)
+      Log.w(TAG, "[$attachmentId] Failed while reading the stream.", e)
       return Result.retry(defaultBackoff())
     }
 
@@ -188,7 +192,7 @@ class UploadAttachmentToArchiveJob private constructor(
       Log.w(TAG, "[$attachmentId] Job was canceled, updating archive transfer state to ${AttachmentTable.ArchiveTransferState.NONE}.")
       SignalDatabase.attachments.setArchiveTransferState(attachmentId, AttachmentTable.ArchiveTransferState.NONE)
     } else {
-      Log.w(TAG, "[$attachmentId] Job failed, updating archive transfer state to ${AttachmentTable.ArchiveTransferState.TEMPORARY_FAILURE}.")
+      Log.w(TAG, "[$attachmentId] Job failed, updating archive transfer state to ${AttachmentTable.ArchiveTransferState.TEMPORARY_FAILURE} (if not already a permanent failure).")
       SignalDatabase.attachments.setArchiveTransferState(attachmentId, AttachmentTable.ArchiveTransferState.TEMPORARY_FAILURE)
     }
   }