From e72a5c3320aa5b6b655a4029150fe75210340fa5 Mon Sep 17 00:00:00 2001
From: automated-signal <37887102+automated-signal@users.noreply.github.com>
Date: Thu, 13 Jun 2024 12:16:26 -0500
Subject: [PATCH] Avoid race condition in AttachmentDownloadQueue

---
 ts/util/attachmentDownloadQueue.ts | 86 ++++++++++++++++++------------
 1 file changed, 53 insertions(+), 33 deletions(-)

diff --git a/ts/util/attachmentDownloadQueue.ts b/ts/util/attachmentDownloadQueue.ts
index 09bca78fa..b15715c2d 100644
--- a/ts/util/attachmentDownloadQueue.ts
+++ b/ts/util/attachmentDownloadQueue.ts
@@ -1,16 +1,16 @@
 // Copyright 2023 Signal Messenger, LLC
 // SPDX-License-Identifier: AGPL-3.0-only
 
-import type { MessageAttributesType } from '../model-types.d';
 import type { MessageModel } from '../models/messages';
 import * as log from '../logging/log';
 import { isMoreRecentThan } from './timestamp';
+import { isNotNil } from './isNotNil';
 
 const MAX_ATTACHMENT_DOWNLOAD_AGE = 3600 * 72 * 1000;
 const MAX_ATTACHMENT_MSGS_TO_DOWNLOAD = 250;
 
 let isEnabled = true;
-let attachmentDownloadQueue: Array<MessageModel> | undefined = [];
+let attachmentDownloadQueue: Array<string> | undefined = [];
 const queueEmptyCallbacks: Set<() => void> = new Set();
 
 export function shouldUseAttachmentDownloadQueue(): boolean {
@@ -25,6 +25,11 @@ export function registerQueueEmptyCallback(callback: () => void): void {
   queueEmptyCallbacks.add(callback);
 }
 
+function onQueueEmpty() {
+  queueEmptyCallbacks.forEach(callback => callback());
+  queueEmptyCallbacks.clear();
+}
+
 export function addToAttachmentDownloadQueue(
   idLog: string,
   message: MessageModel
@@ -33,7 +38,7 @@ export function addToAttachmentDownloadQueue(
     return;
   }
 
-  attachmentDownloadQueue.unshift(message);
+  attachmentDownloadQueue.unshift(message.id);
 
   log.info(
     `${idLog}: Adding to attachmentDownloadQueue`,
@@ -42,48 +47,63 @@ export function addToAttachmentDownloadQueue(
 }
 
 export async function flushAttachmentDownloadQueue(): Promise<void> {
-  if (!attachmentDownloadQueue) {
-    return;
-  }
-
   // NOTE: ts/models/messages.ts expects this global to become undefined
   // once we stop processing the queue.
   isEnabled = false;
 
-  const attachmentsToDownload = attachmentDownloadQueue.filter(
-    (message, index) =>
-      index <= MAX_ATTACHMENT_MSGS_TO_DOWNLOAD ||
-      isMoreRecentThan(message.getReceivedAt(), MAX_ATTACHMENT_DOWNLOAD_AGE) ||
-      // Stickers and long text attachments has to be downloaded for UI
-      // to display the message properly.
-      message.hasRequiredAttachmentDownloads()
+  if (!attachmentDownloadQueue?.length) {
+    onQueueEmpty();
+    return;
+  }
+
+  const messageIdsToDownload = attachmentDownloadQueue.slice(
+    0,
+    MAX_ATTACHMENT_MSGS_TO_DOWNLOAD
+  );
+
+  const messageIdsToSave: Array<string> = [];
+  let numMessagesQueued = 0;
+  await Promise.all(
+    messageIdsToDownload.map(async messageId => {
+      const message = window.MessageCache.__DEPRECATED$getById(messageId);
+      if (!message) {
+        log.warn(
+          'attachmentDownloadQueue: message not found in messageCache, maybe it was deleted?'
+        );
+        return;
+      }
+
+      if (
+        isMoreRecentThan(
+          message.getReceivedAt(),
+          MAX_ATTACHMENT_DOWNLOAD_AGE
+        ) ||
+        // Stickers and long text attachments has to be downloaded for UI
+        // to display the message properly.
+        message.hasRequiredAttachmentDownloads()
+      ) {
+        const shouldSave = await message.queueAttachmentDownloads();
+        if (shouldSave) {
+          messageIdsToSave.push(messageId);
+        }
+        numMessagesQueued += 1;
+      }
+    })
   );
 
   log.info(
-    'Downloading recent attachments of total attachments',
-    attachmentsToDownload.length,
-    attachmentDownloadQueue.length
+    `Downloading recent attachments for ${numMessagesQueued} ` +
+      `of ${attachmentDownloadQueue.length} total messages`
   );
 
-  const messagesWithDownloads = await Promise.all(
-    attachmentsToDownload.map(message => {
-      const updatedMessage =
-        window.MessageCache.__DEPRECATED$getById(message.id) ?? message;
-      return updatedMessage.queueAttachmentDownloads();
-    })
-  );
-  const messagesToSave: Array<MessageAttributesType> = [];
-  messagesWithDownloads.forEach((shouldSave, messageKey) => {
-    if (shouldSave) {
-      const message = attachmentsToDownload[messageKey];
-      messagesToSave.push(message.attributes);
-    }
-  });
+  const messagesToSave = messageIdsToSave
+    .map(messageId => window.MessageCache.accessAttributes(messageId))
+    .filter(isNotNil);
+
   await window.Signal.Data.saveMessages(messagesToSave, {
     ourAci: window.storage.user.getCheckedAci(),
   });
 
   attachmentDownloadQueue = undefined;
-  queueEmptyCallbacks.forEach(callback => callback());
-  queueEmptyCallbacks.clear();
+  onQueueEmpty();
 }