From d047723f04a8bc56787eb8adc753ca29ac68c3ad Mon Sep 17 00:00:00 2001
From: Jared Snider <116684653+JaredSnider-Bitwarden@users.noreply.github.com>
Date: Fri, 2 Feb 2024 16:23:15 -0500
Subject: [PATCH] Auth & Autofill / PM-5976 - Safari Browser SSO Initialization
 Race Condition Attempted Fix 2 (#7794)

* Implementing pinging system for SSO to address issue on Safari with race condition

* Implementing pinging system for SSO to address issue on Safari with race condition

* [PM-5976] Updating references within sso.ts

---------

Co-authored-by: Cesar Gonzalez <cgonzalez@bitwarden.com>
---
 .../content/content-message-handler.ts        |  4 +++
 apps/web/src/connectors/sso.ts                | 29 ++++++++++++++-----
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/apps/browser/src/autofill/content/content-message-handler.ts b/apps/browser/src/autofill/content/content-message-handler.ts
index caa2883ef8..a7a371af60 100644
--- a/apps/browser/src/autofill/content/content-message-handler.ts
+++ b/apps/browser/src/autofill/content/content-message-handler.ts
@@ -35,6 +35,10 @@ class ContentMessageHandler implements ContentMessageHandlerInterface {
     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.
diff --git a/apps/web/src/connectors/sso.ts b/apps/web/src/connectors/sso.ts
index e594a97d51..8218b0d9be 100644
--- a/apps/web/src/connectors/sso.ts
+++ b/apps/web/src/connectors/sso.ts
@@ -24,17 +24,30 @@ window.addEventListener("load", () => {
 });
 
 function initiateBrowserSsoIfDocumentReady(code: string, state: string, lastpass: boolean) {
-  if (document.readyState === "complete") {
-    initiateBrowserSso(code, state, lastpass);
-    return;
-  }
+  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);
 
-  const interval = setInterval(() => {
-    if (document.readyState === "complete") {
-      clearInterval(interval);
       initiateBrowserSso(code, state, lastpass);
     }
-  }, 50);
+  };
+
+  window.addEventListener("message", handleWindowMessage);
 }
 
 function initiateBrowserSso(code: string, state: string, lastpass: boolean) {