From de4a762bb2eb59967ce67dd420050c734c0bb0df Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dominik=20Luba=C5=84ski?= <dominik.lubanski@gmail.com>
Date: Fri, 10 Jan 2025 09:07:26 +0100
Subject: [PATCH] fix(session): Safari bug with delayed set of cookies

---
 src/background/session.js           | 51 ++++++++++++++++++++++-------
 src/pages/settings/views/account.js |  7 ++--
 src/utils/api.js                    | 24 ++++++++++----
 3 files changed, 60 insertions(+), 22 deletions(-)

diff --git a/src/background/session.js b/src/background/session.js
index 919e40d8c..f879852ef 100644
--- a/src/background/session.js
+++ b/src/background/session.js
@@ -8,18 +8,47 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at http://mozilla.org/MPL/2.0
  */
-import { UPDATE_SESSION_ACTION_NAME } from '/store/session.js';
-import { HOME_PAGE_URL, ACCOUNT_PAGE_URL } from '/utils/api.js';
+import { store } from 'hybrids';
+import Session, { UPDATE_SESSION_ACTION_NAME } from '/store/session.js';
+import { HOME_PAGE_URL, ACCOUNT_PAGE_URL, COOKIE_DOMAIN } from '/utils/api.js';
+
+function refreshSession() {
+  chrome.runtime
+    .sendMessage({ action: UPDATE_SESSION_ACTION_NAME })
+    .catch(() => null);
+}
 
 // Observe cookie changes (login/logout actions)
-chrome.webNavigation.onDOMContentLoaded.addListener(async ({ url = '' }) => {
-  if (url === HOME_PAGE_URL || url.includes(ACCOUNT_PAGE_URL)) {
-    // Send message to update session in other contexts
-    chrome.runtime
-      .sendMessage({ action: UPDATE_SESSION_ACTION_NAME })
-      // The function only throws if the other end does not exist. Mainly, it happens
-      // when the background process starts, but there is no other content script or
-      // extension page, which could receive a message.
-      .catch(() => null);
+chrome.cookies.onChanged.addListener(async ({ cookie }) => {
+  if (cookie.domain === COOKIE_DOMAIN && cookie.name === 'access_token') {
+    refreshSession();
   }
 });
+
+async function retry(fn, retries = 10, delay = 1000) {
+  if (!(await fn())) {
+    if (retries > 0) {
+      await new Promise((resolve) => setTimeout(resolve, delay));
+      return retry(fn, retries - 1, delay);
+    }
+  }
+}
+
+if (__PLATFORM__ === 'safari') {
+  chrome.webNavigation.onDOMContentLoaded.addListener(async ({ url = '' }) => {
+    // Safari on iOS 18.x has a bug where cookies are available to the extension with a significant delay
+    // However, there is no event to listen to when the cookies are available, we need to poll
+    // To avoid the same possible problem on macOS, we do it for all Safari platforms
+    // TODO: Check if this is still needed after the next Safari release
+
+    if (url === HOME_PAGE_URL) {
+      // Check for logged out state
+      await retry(() => store.resolve(Session).then(({ user }) => !user));
+      refreshSession();
+    } else if (url.includes(ACCOUNT_PAGE_URL)) {
+      // Check for logged in state
+      await retry(() => store.resolve(Session).then(({ user }) => user));
+      refreshSession();
+    }
+  });
+}
diff --git a/src/pages/settings/views/account.js b/src/pages/settings/views/account.js
index 15619fef4..396877e1d 100644
--- a/src/pages/settings/views/account.js
+++ b/src/pages/settings/views/account.js
@@ -26,19 +26,18 @@ function openGhosteryPage(url) {
         details.tabId === tab.id &&
         details.url.startsWith(ACCOUNT_PAGE_URL)
       ) {
-        chrome.webNavigation.onCommitted.removeListener(onSuccess);
+        chrome.webNavigation.onDOMContentLoaded.removeListener(onSuccess);
         chrome.tabs.remove(tab.id);
       }
     };
 
     const onRemove = (tabId) => {
       if (tabId === tab.id) {
-        chrome.webNavigation.onCommitted.removeListener(onSuccess);
+        chrome.webNavigation.onDOMContentLoaded.removeListener(onSuccess);
         chrome.tabs.onRemoved.removeListener(onRemove);
 
         // The tab is closed before the background listeners can catch the event
         // so we need to refresh session and trigger sync manually
-        store.clear(Session);
         chrome.runtime.sendMessage({ action: 'syncOptions' });
 
         // Restore the original tab
@@ -46,8 +45,8 @@ function openGhosteryPage(url) {
       }
     };
 
+    chrome.webNavigation.onDOMContentLoaded.addListener(onSuccess);
     chrome.tabs.onRemoved.addListener(onRemove);
-    chrome.webNavigation.onCommitted.addListener(onSuccess);
 
     tab = await chrome.tabs.create({ url });
   };
diff --git a/src/utils/api.js b/src/utils/api.js
index 4b71dd3dd..7638bfdbc 100644
--- a/src/utils/api.js
+++ b/src/utils/api.js
@@ -100,13 +100,23 @@ async function getCookie(name) {
       : {}),
   });
 
-  if (
-    cookie &&
-    (cookie.session ||
-      cookie.expirationDate * (__PLATFORM__ !== 'safari' ? 1000 : 1) >
-        Date.now())
-  ) {
-    return cookie;
+  if (cookie) {
+    if (cookie.session) return cookie;
+
+    let expirationDate = cookie.expirationDate;
+
+    // By the specs, the `expirationDate` should be in seconds since the epoch
+    // and we need to convert it to milliseconds, but Safari returns it in milliseconds
+    if (
+      __PLATFORM__ !== 'safari' ||
+      new Date(expirationDate).getFullYear() === 1970
+    ) {
+      expirationDate = expirationDate * 1000;
+    }
+
+    if (expirationDate > Date.now()) {
+      return cookie;
+    }
   }
 
   return null;