[PM-5976] Safari Browser SSO Initialization Race Condition Attempted Fix 3 (#7800)

* [PM-5976] Safari Browser SSO Initialization Race Condition Attempted Fix 3

* [PM-5976] Safari Browser SSO Initialization Race Condition Attempted Fix 3

* [PM-5976] Removing usage of pinging system and keeping reworked top-level registration of window message listener events

* [PM-5976] Pulling the implementation of the static content script delcaration for the content-message-handler file to the top of the list of content_scripts

* [PM-5976] Pulling the implementation of the static content script delcaration for the content-message-handler file to the top of the list of content_scripts

* [PM-5976] Removing the useCapture value within the window message event listener

apps/browser/src/autofill/content/abstractions/content-message-handler.ts

@@ -1,6 +1,24 @@
-interface ContentMessageHandler {
-  init(): void;
-  destroy(): void;
-}
+type ContentMessageWindowData = {
+  command: string;
+  lastpass?: boolean;
+  code?: string;
+  state?: string;
+  data?: string;
+  remember?: boolean;
+};
+type ContentMessageWindowEventParams = {
+  data: ContentMessageWindowData;
+  referrer: string;
+};
 
-export { ContentMessageHandler };
+type ContentMessageWindowEventHandlers = {
+  [key: string]: ({ data, referrer }: ContentMessageWindowEventParams) => void;
+  authResult: ({ data, referrer }: ContentMessageWindowEventParams) => void;
+  webAuthnResult: ({ data, referrer }: ContentMessageWindowEventParams) => void;
+};
+
+export {
+  ContentMessageWindowData,
+  ContentMessageWindowEventParams,
+  ContentMessageWindowEventHandlers,
+};

apps/browser/src/autofill/content/bootstrap-content-message-handler.ts

@@ -1,12 +0,0 @@
-import { setupExtensionDisconnectAction } from "../utils";
-
-import ContentMessageHandler from "./content-message-handler";
-
-(function (windowContext) {
-  if (!windowContext.bitwardenContentMessageHandler) {
-    windowContext.bitwardenContentMessageHandler = new ContentMessageHandler();
-    setupExtensionDisconnectAction(() => windowContext.bitwardenContentMessageHandler.destroy());
-
-    windowContext.bitwardenContentMessageHandler.init();
-  }
-})(window);

apps/browser/src/autofill/content/content-message-handler.spec.ts

@@ -1,37 +1,31 @@
+import { mock } from "jest-mock-extended";
+
 import { postWindowMessage, sendExtensionRuntimeMessage } from "../jest/testing-utils";
 
-import ContentMessageHandler from "./content-message-handler";
-
 describe("ContentMessageHandler", () => {
-  let contentMessageHandler: ContentMessageHandler;
   const sendMessageSpy = jest.spyOn(chrome.runtime, "sendMessage");
+  let portOnDisconnectAddListenerCallback: CallableFunction;
+  chrome.runtime.connect = jest.fn(() =>
+    mock<chrome.runtime.Port>({
+      onDisconnect: {
+        addListener: jest.fn((callback) => {
+          portOnDisconnectAddListenerCallback = callback;
+        }),
+        removeListener: jest.fn(),
+      },
+    }),
+  );
 
   beforeEach(() => {
-    contentMessageHandler = new ContentMessageHandler();
+    require("./content-message-handler");
   });
 
   afterEach(() => {
+    jest.resetModules();
     jest.clearAllMocks();
-    contentMessageHandler.destroy();
   });
 
-  describe("init", () => {
-    it("should add event listeners", () => {
-      const addEventListenerSpy = jest.spyOn(window, "addEventListener");
-      const addListenerSpy = jest.spyOn(chrome.runtime.onMessage, "addListener");
-
-      contentMessageHandler.init();
-
-      expect(addEventListenerSpy).toHaveBeenCalledTimes(1);
-      expect(addListenerSpy).toHaveBeenCalledTimes(1);
-    });
-  });
-
-  describe("handleWindowMessage", () => {
-    beforeEach(() => {
-      contentMessageHandler.init();
-    });
-
+  describe("handled window messages", () => {
     it("ignores messages from other sources", () => {
       postWindowMessage({ command: "authResult" }, "https://localhost/", null);
 
@@ -47,7 +41,6 @@ describe("ContentMessageHandler", () => {
     it("sends an authResult message", () => {
       postWindowMessage({ command: "authResult", lastpass: true, code: "code", state: "state" });
 
-      expect(sendMessageSpy).toHaveBeenCalledTimes(1);
       expect(sendMessageSpy).toHaveBeenCalledWith({
         command: "authResult",
         code: "code",
@@ -60,7 +53,6 @@ describe("ContentMessageHandler", () => {
     it("sends a webAuthnResult message", () => {
       postWindowMessage({ command: "webAuthnResult", data: "data", remember: true });
 
-      expect(sendMessageSpy).toHaveBeenCalledTimes(1);
       expect(sendMessageSpy).toHaveBeenCalledWith({
         command: "webAuthnResult",
         data: "data",
@@ -70,11 +62,7 @@ describe("ContentMessageHandler", () => {
     });
   });
 
-  describe("handleExtensionMessage", () => {
-    beforeEach(() => {
-      contentMessageHandler.init();
-    });
-
+  describe("handled extension messages", () => {
     it("ignores the message to the extension background if it is not present in the forwardCommands list", () => {
       sendExtensionRuntimeMessage({ command: "someOtherCommand" });
 
@@ -88,4 +76,17 @@ describe("ContentMessageHandler", () => {
       expect(sendMessageSpy).toHaveBeenCalledWith({ command: "bgUnlockPopoutOpened" });
     });
   });
+
+  describe("extension disconnect action", () => {
+    it("removes the window message listener and the extension message listener", () => {
+      const removeEventListenerSpy = jest.spyOn(window, "removeEventListener");
+
+      portOnDisconnectAddListenerCallback(mock<chrome.runtime.Port>());
+
+      expect(removeEventListenerSpy).toHaveBeenCalledTimes(1);
+      expect(removeEventListenerSpy).toHaveBeenCalledWith("message", expect.any(Function));
+      expect(chrome.runtime.onMessage.removeListener).toHaveBeenCalledTimes(1);
+      expect(chrome.runtime.onMessage.removeListener).toHaveBeenCalledWith(expect.any(Function));
+    });
+  });
 });

apps/browser/src/autofill/content/content-message-handler.ts

@@ -1,84 +1,110 @@
-import { ContentMessageHandler as ContentMessageHandlerInterface } from "./abstractions/content-message-handler";
+import {
+  ContentMessageWindowData,
+  ContentMessageWindowEventHandlers,
+} from "./abstractions/content-message-handler";
 
-class ContentMessageHandler implements ContentMessageHandlerInterface {
-  private forwardCommands = [
-    "bgUnlockPopoutOpened",
-    "addToLockedVaultPendingNotifications",
-    "unlockCompleted",
-    "addedCipher",
-  ];
+/**
+ * IMPORTANT: Safari seems to have a bug where it doesn't properly handle
+ * window message events from content scripts when the listener these events
+ * is registered within a class. This is why these listeners are registered
+ * at the top level of this file.
+ */
+window.addEventListener("message", handleWindowMessageEvent, false);
+chrome.runtime.onMessage.addListener(handleExtensionMessage);
+setupExtensionDisconnectAction(() => {
+  window.removeEventListener("message", handleWindowMessageEvent);
+  chrome.runtime.onMessage.removeListener(handleExtensionMessage);
+});
 
-  /**
-   * Initialize the content message handler. Sets up
-   * a window message listener and a chrome runtime
-   * message listener.
-   */
-  init() {
-    window.addEventListener("message", this.handleWindowMessage, false);
-    chrome.runtime.onMessage.addListener(this.handleExtensionMessage);
-  }
+/**
+ * Handlers for window messages from the content script.
+ */
+const windowMessageHandlers: ContentMessageWindowEventHandlers = {
+  authResult: ({ data, referrer }: { data: any; referrer: string }) =>
+    handleAuthResultMessage(data, referrer),
+  webAuthnResult: ({ data, referrer }: { data: any; referrer: string }) =>
+    handleWebAuthnResultMessage(data, referrer),
+};
 
-  /**
-   * Handle a message from the window. This implementation
-   * specifically handles the authResult and webAuthnResult
-   * commands. This facilitates single sign-on.
-   *
-   * @param event - The message event.
-   */
-  private handleWindowMessage = (event: MessageEvent) => {
-    const { source, data } = event;
-
-    if (source !== window || !data?.command) {
-      return;
-    }
-
-    const { command } = data;
-    const referrer = source.location.hostname;
-
-    if (command === "checkIfReadyForAuthResult") {
-      window.postMessage({ command: "readyToReceiveAuthResult" }, "*");
-    }
-
-    if (command === "authResult") {
-      const { lastpass, code, state } = data;
-      // FIXME: Verify that this floating promise is intentional. If it is, add an explanatory comment and ensure there is proper error handling.
-      // eslint-disable-next-line @typescript-eslint/no-floating-promises
-      chrome.runtime.sendMessage({ command, code, state, lastpass, referrer });
-    }
-
-    if (command === "webAuthnResult") {
-      const { remember } = data;
-      // FIXME: Verify that this floating promise is intentional. If it is, add an explanatory comment and ensure there is proper error handling.
-      // eslint-disable-next-line @typescript-eslint/no-floating-promises
-      chrome.runtime.sendMessage({ command, data: data.data, remember, referrer });
-    }
-  };
-
-  /**
-   * Handle a message from the extension. This
-   * implementation forwards the message to the
-   * extension background so that it can  be received
-   * in other contexts of the background script.
-   *
-   * @param message - The message from the extension.
-   */
-  private handleExtensionMessage = (message: any) => {
-    if (this.forwardCommands.includes(message.command)) {
-      // FIXME: Verify that this floating promise is intentional. If it is, add an explanatory comment and ensure there is proper error handling.
-      // eslint-disable-next-line @typescript-eslint/no-floating-promises
-      chrome.runtime.sendMessage(message);
-    }
-  };
-
-  /**
-   * Destroy the content message handler. Removes
-   * the window message listener and the chrome
-   * runtime message listener.
-   */
-  destroy = () => {
-    window.removeEventListener("message", this.handleWindowMessage);
-    chrome.runtime.onMessage.removeListener(this.handleExtensionMessage);
-  };
+/**
+ * Handles the auth result message from the window.
+ *
+ * @param data - Data from the window message
+ * @param referrer - The referrer of the window
+ */
+async function handleAuthResultMessage(data: ContentMessageWindowData, referrer: string) {
+  const { command, lastpass, code, state } = data;
+  await chrome.runtime.sendMessage({ command, code, state, lastpass, referrer });
 }
 
-export default ContentMessageHandler;
+/**
+ * Handles the webauthn result message from the window.
+ *
+ * @param data - Data from the window message
+ * @param referrer - The referrer of the window
+ */
+async function handleWebAuthnResultMessage(data: ContentMessageWindowData, referrer: string) {
+  const { command, remember } = data;
+  await chrome.runtime.sendMessage({ command, data: data.data, remember, referrer });
+}
+
+/**
+ * Handles the window message event.
+ *
+ * @param event - The window message event
+ */
+function handleWindowMessageEvent(event: MessageEvent) {
+  const { source, data } = event;
+  if (source !== window || !data?.command) {
+    return;
+  }
+
+  const referrer = source.location.hostname;
+  const handler = windowMessageHandlers[data.command];
+  if (handler) {
+    handler({ data, referrer });
+  }
+}
+
+/**
+ * Commands to forward from this script to the extension background.
+ */
+const forwardCommands = new Set([
+  "bgUnlockPopoutOpened",
+  "addToLockedVaultPendingNotifications",
+  "unlockCompleted",
+  "addedCipher",
+]);
+
+/**
+ * Handles messages from the extension. Currently, this is
+ * used to forward messages from the background context to
+ * other scripts within the extension.
+ *
+ * @param message - The message from the extension
+ */
+async function handleExtensionMessage(message: any) {
+  if (forwardCommands.has(message.command)) {
+    await chrome.runtime.sendMessage(message);
+  }
+}
+
+/**
+ * Duplicate implementation of the same named method within `apps/browser/src/autofill/utils/index.ts`.
+ * This is done due to some strange observed compilation behavior present when importing the method from
+ * the utils file.
+ *
+ * TODO: Investigate why webpack tree shaking is not removing other methods when importing from the utils file.
+ * Possible cause can be seen below:
+ * @see https://stackoverflow.com/questions/71679366/webpack5-does-not-seem-to-tree-shake-unused-exports
+ *
+ * @param callback - Callback function to run when the extension disconnects
+ */
+function setupExtensionDisconnectAction(callback: (port: chrome.runtime.Port) => void) {
+  const port = chrome.runtime.connect({ name: "autofill-injected-script-port" });
+  const onDisconnectCallback = (disconnectedPort: chrome.runtime.Port) => {
+    callback(disconnectedPort);
+    port.onDisconnect.removeListener(onDisconnectCallback);
+  };
+  port.onDisconnect.addListener(onDisconnectCallback);
+}

apps/browser/src/autofill/globals.d.ts

@@ -1,9 +1,7 @@
 import { AutofillInit } from "./content/abstractions/autofill-init";
-import ContentMessageHandler from "./content/content-message-handler";
 
 declare global {
   interface Window {
     bitwardenAutofillInit?: AutofillInit;
-    bitwardenContentMessageHandler?: ContentMessageHandler;
   }
 }

apps/browser/src/autofill/services/autofill.service.spec.ts

@@ -207,11 +207,11 @@ describe("AutofillService", () => {
       });
     });
 
-    it("injects the bootstrap-content-message-handler script if not injecting on page load", async () => {
+    it("injects the content-message-handler script if not injecting on page load", async () => {
       await autofillService.injectAutofillScripts(sender.tab, sender.frameId, false);
 
       expect(BrowserApi.executeScriptInTab).toHaveBeenCalledWith(tabMock.id, {
-        file: "content/bootstrap-content-message-handler.js",
+        file: "content/content-message-handler.js",
         ...defaultExecuteScriptOptions,
       });
     });

apps/browser/src/autofill/services/autofill.service.ts

@@ -105,7 +105,7 @@ export default class AutofillService implements AutofillServiceInterface {
       injectedScripts.push("autofiller.js");
     } else {
       await BrowserApi.executeScriptInTab(tab.id, {
-        file: "content/bootstrap-content-message-handler.js",
+        file: "content/content-message-handler.js",
         runAt: "document_start",
       });
     }

apps/browser/src/manifest.json

@@ -15,6 +15,12 @@
     "128": "images/icon128.png"
   },
   "content_scripts": [
+    {
+      "all_frames": false,
+      "js": ["content/content-message-handler.js"],
+      "matches": ["http://*/*", "https://*/*", "file:///*"],
+      "run_at": "document_start"
+    },
     {
       "all_frames": true,
       "js": [
@@ -24,12 +30,6 @@
       "matches": ["http://*/*", "https://*/*", "file:///*"],
       "run_at": "document_start"
     },
-    {
-      "all_frames": false,
-      "js": ["content/bootstrap-content-message-handler.js"],
-      "matches": ["http://*/*", "https://*/*", "file:///*"],
-      "run_at": "document_start"
-    },
     {
       "all_frames": true,
       "css": ["content/autofill.css"],

apps/browser/src/manifest.v3.json

@@ -17,14 +17,14 @@
   },
   "content_scripts": [
     {
-      "all_frames": true,
-      "js": ["content/trigger-autofill-script-injection.js"],
+      "all_frames": false,
+      "js": ["content/content-message-handler.js"],
       "matches": ["http://*/*", "https://*/*", "file:///*"],
       "run_at": "document_start"
     },
     {
-      "all_frames": false,
-      "js": ["content/bootstrap-content-message-handler.js"],
+      "all_frames": true,
+      "js": ["content/trigger-autofill-script-injection.js"],
       "matches": ["http://*/*", "https://*/*", "file:///*"],
       "run_at": "document_start"
     },

apps/browser/webpack.config.js

@@ -169,8 +169,7 @@ const mainConfig = {
     "content/autofiller": "./src/autofill/content/autofiller.ts",
     "content/notificationBar": "./src/autofill/content/notification-bar.ts",
     "content/contextMenuHandler": "./src/autofill/content/context-menu-handler.ts",
-    "content/bootstrap-content-message-handler":
-      "./src/autofill/content/bootstrap-content-message-handler.ts",
+    "content/content-message-handler": "./src/autofill/content/content-message-handler.ts",
     "content/fido2/trigger-fido2-content-script-injection":
       "./src/vault/fido2/content/trigger-fido2-content-script-injection.ts",
     "content/fido2/content-script": "./src/vault/fido2/content/content-script.ts",

apps/web/src/connectors/sso.ts

@@ -8,9 +8,9 @@ window.addEventListener("load", () => {
   const lastpass = getQsParam("lp");
 
   if (lastpass === "1") {
-    initiateBrowserSsoIfDocumentReady(code, state, true);
+    initiateBrowserSso(code, state, true);
   } else if (state != null && state.includes(":clientId=browser")) {
-    initiateBrowserSsoIfDocumentReady(code, state, false);
+    initiateBrowserSso(code, state, false);
   } else {
     window.location.href = window.location.origin + "/#/sso?code=" + code + "&state=" + state;
     // Match any characters between "_returnUri='" and the next "'"
@@ -23,33 +23,6 @@ window.addEventListener("load", () => {
   }
 });
 
-function initiateBrowserSsoIfDocumentReady(code: string, state: string, lastpass: boolean) {
-  const MAX_ATTEMPTS = 200;
-  const TIMEOUT_MS = 50;
-  let attempts = 0;
-
-  const pingInterval = setInterval(() => {
-    if (attempts >= MAX_ATTEMPTS) {
-      clearInterval(pingInterval);
-      throw new Error("Failed to initiate browser SSO");
-    }
-
-    attempts++;
-    window.postMessage({ command: "checkIfReadyForAuthResult" }, "*");
-  }, TIMEOUT_MS);
-
-  const handleWindowMessage = (event: MessageEvent) => {
-    if (event.source === window && event.data?.command === "readyToReceiveAuthResult") {
-      clearInterval(pingInterval);
-      window.removeEventListener("message", handleWindowMessage);
-
-      initiateBrowserSso(code, state, lastpass);
-    }
-  };
-
-  window.addEventListener("message", handleWindowMessage);
-}
-
 function initiateBrowserSso(code: string, state: string, lastpass: boolean) {
   window.postMessage({ command: "authResult", code: code, state: state, lastpass: lastpass }, "*");
   const handOffMessage = ("; " + document.cookie)