From 9a98e8eb702eba27b376bd26afc121d38e93d81e Mon Sep 17 00:00:00 2001
From: Justin Baur <19896123+justindbaur@users.noreply.github.com>
Date: Fri, 20 Jan 2023 12:33:41 -0500
Subject: [PATCH] [PS-2303] MV3 Service Cache Enhancements (#4507)

* Add Should Listen Gate

* Fixed ServiceCache Being Sent to Wrong Arg

* Await Combined Listener
---
 apps/browser/src/background.ts                  |  7 ++++++-
 .../src/browser/cipher-context-menu-handler.ts  | 17 ++++++++++++++---
 .../src/browser/context-menu-clicked-handler.ts |  1 +
 apps/browser/src/listeners/combine.spec.ts      |  6 ++++--
 apps/browser/src/listeners/combine.ts           |  6 +++---
 apps/browser/src/listeners/index.ts             |  5 ++++-
 6 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/apps/browser/src/background.ts b/apps/browser/src/background.ts
index a6baca89c4..2c888cfd4b 100644
--- a/apps/browser/src/background.ts
+++ b/apps/browser/src/background.ts
@@ -21,7 +21,12 @@ if (BrowserApi.manifestVersion === 3) {
   chrome.tabs.onReplaced.addListener(tabsOnReplacedListener);
   chrome.tabs.onUpdated.addListener(tabsOnUpdatedListener);
   chrome.contextMenus.onClicked.addListener(contextMenusClickedListener);
-  BrowserApi.messageListener("runtime.background", runtimeMessageListener);
+  BrowserApi.messageListener(
+    "runtime.background",
+    (message: { command: string }, sender, sendResponse) => {
+      runtimeMessageListener(message, sender);
+    }
+  );
 } else {
   const bitwardenMain = ((window as any).bitwardenMain = new MainBackground());
   bitwardenMain.bootstrap().then(() => {
diff --git a/apps/browser/src/browser/cipher-context-menu-handler.ts b/apps/browser/src/browser/cipher-context-menu-handler.ts
index ea87f8da93..5942ce3b14 100644
--- a/apps/browser/src/browser/cipher-context-menu-handler.ts
+++ b/apps/browser/src/browser/cipher-context-menu-handler.ts
@@ -120,13 +120,24 @@ export class CipherContextMenuHandler {
     await cipherContextMenuHandler.update(tab.url);
   }
 
-  static async messageListener(message: { command: string }, cachedServices: CachedServices) {
+  static async messageListener(
+    message: { command: string },
+    sender: chrome.runtime.MessageSender,
+    cachedServices: CachedServices
+  ) {
+    if (!CipherContextMenuHandler.shouldListen(message)) {
+      return;
+    }
     const cipherContextMenuHandler = await CipherContextMenuHandler.create(cachedServices);
     await cipherContextMenuHandler.messageListener(message);
   }
 
-  async messageListener(message: { command: string }) {
-    if (!LISTENED_TO_COMMANDS.includes(message.command)) {
+  private static shouldListen(message: { command: string }) {
+    return LISTENED_TO_COMMANDS.includes(message.command);
+  }
+
+  async messageListener(message: { command: string }, sender?: chrome.runtime.MessageSender) {
+    if (!CipherContextMenuHandler.shouldListen(message)) {
       return;
     }
 
diff --git a/apps/browser/src/browser/context-menu-clicked-handler.ts b/apps/browser/src/browser/context-menu-clicked-handler.ts
index c11c6dbb94..f0043805da 100644
--- a/apps/browser/src/browser/context-menu-clicked-handler.ts
+++ b/apps/browser/src/browser/context-menu-clicked-handler.ts
@@ -126,6 +126,7 @@ export class ContextMenuClickedHandler {
 
   static async messageListener(
     message: { command: string; data: LockedVaultPendingNotificationsItem },
+    sender: chrome.runtime.MessageSender,
     cachedServices: CachedServices
   ) {
     if (
diff --git a/apps/browser/src/listeners/combine.spec.ts b/apps/browser/src/listeners/combine.spec.ts
index 3975d27017..10a0fd30d1 100644
--- a/apps/browser/src/listeners/combine.spec.ts
+++ b/apps/browser/src/listeners/combine.spec.ts
@@ -1,22 +1,24 @@
 import { combine } from "./combine";
 
 describe("combine", () => {
-  it("runs", () => {
+  it("runs", async () => {
     const combined = combine([
       (arg: Record<string, unknown>, serviceCache: Record<string, unknown>) => {
         arg["one"] = true;
         serviceCache["one"] = true;
+        return Promise.resolve();
       },
       (arg: Record<string, unknown>, serviceCache: Record<string, unknown>) => {
         if (serviceCache["one"] !== true) {
           throw new Error("One should have ran.");
         }
         arg["two"] = true;
+        return Promise.resolve();
       },
     ]);
 
     const arg: Record<string, unknown> = {};
-    combined(arg);
+    await combined(arg);
 
     expect(arg["one"]).toBeTruthy();
 
diff --git a/apps/browser/src/listeners/combine.ts b/apps/browser/src/listeners/combine.ts
index 278772ba21..b87631de97 100644
--- a/apps/browser/src/listeners/combine.ts
+++ b/apps/browser/src/listeners/combine.ts
@@ -1,15 +1,15 @@
 import { CachedServices } from "../background/service_factories/factory-options";
 
-type Listener<T extends unknown[]> = (...args: [...T, CachedServices]) => void;
+type Listener<T extends unknown[]> = (...args: [...T, CachedServices]) => Promise<void>;
 
 export const combine = <T extends unknown[]>(
   listeners: Listener<T>[],
   startingServices: CachedServices = {}
 ) => {
-  return (...args: T) => {
+  return async (...args: T) => {
     const cachedServices = { ...startingServices };
     for (const listener of listeners) {
-      listener(...[...args, cachedServices]);
+      await listener(...[...args, cachedServices]);
     }
   };
 };
diff --git a/apps/browser/src/listeners/index.ts b/apps/browser/src/listeners/index.ts
index b55ce884a1..0c00dea784 100644
--- a/apps/browser/src/listeners/index.ts
+++ b/apps/browser/src/listeners/index.ts
@@ -23,7 +23,10 @@ const tabsOnUpdatedListener = combine([
 
 const contextMenusClickedListener = ContextMenuClickedHandler.onClickedListener;
 
-const runtimeMessageListener = combine([
+// TODO: All message listeners should be RuntimeMessage in Notifications follow up then this type annotation can be inferred
+const runtimeMessageListener = combine<
+  [message: { command: string }, sender: chrome.runtime.MessageSender]
+>([
   UpdateBadge.messageListener,
   CipherContextMenuHandler.messageListener,
   ContextMenuClickedHandler.messageListener,