From 584b39baa8341a640d41f5fbbdc9df52200b113c Mon Sep 17 00:00:00 2001
From: Josh Perez <60019601+josh-signal@users.noreply.github.com>
Date: Thu, 11 Aug 2022 18:26:26 -0400
Subject: [PATCH] Avoid race conditions when queueing a story for download

---
 ts/state/ducks/stories.ts                    | 42 +++++++++++++++++++-
 ts/test-electron/state/ducks/stories_test.ts |  4 +-
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/ts/state/ducks/stories.ts b/ts/state/ducks/stories.ts
index 315a9a5b7..1e863dcc0 100644
--- a/ts/state/ducks/stories.ts
+++ b/ts/state/ducks/stories.ts
@@ -52,6 +52,7 @@ import { viewedReceiptsJobQueue } from '../../jobs/viewedReceiptsJobQueue';
 export type StoryDataType = {
   attachment?: AttachmentType;
   messageId: string;
+  startedDownload?: boolean;
 } & Pick<
   MessageAttributesType,
   | 'canReplyToStory'
@@ -92,6 +93,7 @@ export type StoriesStateType = {
 const DOE_STORY = 'stories/DOE';
 const LOAD_STORY_REPLIES = 'stories/LOAD_STORY_REPLIES';
 const MARK_STORY_READ = 'stories/MARK_STORY_READ';
+const QUEUE_STORY_DOWNLOAD = 'stories/QUEUE_STORY_DOWNLOAD';
 const REPLY_TO_STORY = 'stories/REPLY_TO_STORY';
 export const RESOLVE_ATTACHMENT_URL = 'stories/RESOLVE_ATTACHMENT_URL';
 const STORY_CHANGED = 'stories/STORY_CHANGED';
@@ -116,6 +118,11 @@ type MarkStoryReadActionType = {
   payload: string;
 };
 
+type QueueStoryDownloadActionType = {
+  type: typeof QUEUE_STORY_DOWNLOAD;
+  payload: string;
+};
+
 type ReplyToStoryActionType = {
   type: typeof REPLY_TO_STORY;
   payload: MessageAttributesType;
@@ -155,6 +162,7 @@ export type StoriesActionType =
   | MessageChangedActionType
   | MessageDeletedActionType
   | MessagesAddedActionType
+  | QueueStoryDownloadActionType
   | ReplyToStoryActionType
   | ResolveAttachmentUrlActionType
   | StoryChangedActionType
@@ -431,7 +439,7 @@ function queueStoryDownload(
   void,
   RootStateType,
   unknown,
-  NoopActionType | ResolveAttachmentUrlActionType
+  NoopActionType | QueueStoryDownloadActionType | ResolveAttachmentUrlActionType
 > {
   return async (dispatch, getState) => {
     const { stories } = getState().stories;
@@ -477,7 +485,11 @@ function queueStoryDownload(
       return;
     }
 
-    if (isDownloading(attachment)) {
+    // isDownloading checks for the downloadJobId which is set by
+    // queueAttachmentDownloads but we optimistically set story.startedDownload
+    // in redux to prevent race conditions from queuing up multiple attachment
+    // downloads before the attachment save takes place.
+    if (isDownloading(attachment) || story.startedDownload) {
       return;
     }
 
@@ -488,7 +500,13 @@ function queueStoryDownload(
       // completed attachment download.
       message.set({ storyReplyContext: undefined });
 
+      dispatch({
+        type: QUEUE_STORY_DOWNLOAD,
+        payload: storyId,
+      });
+
       await queueAttachmentDownloads(message.attributes);
+      return;
     }
 
     dispatch({
@@ -1205,5 +1223,25 @@ export function reducer(
     };
   }
 
+  if (action.type === QUEUE_STORY_DOWNLOAD) {
+    const storyIndex = state.stories.findIndex(
+      story => story.messageId === action.payload
+    );
+
+    if (storyIndex < 0) {
+      return state;
+    }
+
+    const existingStory = state.stories[storyIndex];
+
+    return {
+      ...state,
+      stories: replaceIndex(state.stories, storyIndex, {
+        ...existingStory,
+        startedDownload: true,
+      }),
+    };
+  }
+
   return state;
 }
diff --git a/ts/test-electron/state/ducks/stories_test.ts b/ts/test-electron/state/ducks/stories_test.ts
index 2161b9a99..d3a15852a 100644
--- a/ts/test-electron/state/ducks/stories_test.ts
+++ b/ts/test-electron/state/ducks/stories_test.ts
@@ -202,8 +202,8 @@ describe('both/state/ducks/stories', () => {
       await queueStoryDownload(storyId)(dispatch, getState, null);
 
       sinon.assert.calledWith(dispatch, {
-        type: 'NOOP',
-        payload: null,
+        type: 'stories/QUEUE_STORY_DOWNLOAD',
+        payload: storyId,
       });
     });
   });