From 05e6ac6671c9a224f922cb2c5621fcab6134af18 Mon Sep 17 00:00:00 2001
From: Marija Gerasimovska <marija@toddle.dev>
Date: Fri, 7 Mar 2025 10:50:06 +0100
Subject: [PATCH] Update session rules when adding or removing cookie

---
 chrome/manifest.json     | 15 ++++++++----
 chrome/src/background.ts | 52 +++++++++++++++++++++-------------------
 chrome/src/helpers.ts    | 32 +++++++++++++++++++++++++
 shared/setCookies.ts     | 38 ++++++++++++++++-------------
 4 files changed, 93 insertions(+), 44 deletions(-)
 create mode 100644 chrome/src/helpers.ts

diff --git a/chrome/manifest.json b/chrome/manifest.json
index 7ec3d98..1b1ba67 100644
--- a/chrome/manifest.json
+++ b/chrome/manifest.json
@@ -2,7 +2,7 @@
   "manifest_version": 3,
   "name": "toddle",
   "description": "Browser extension for the toddle editor.",
-  "version": "2.3.2",
+  "version": "2.3.3",
   "icons": {
     "32": "icons/toddle_32x32.png",
     "64": "icons/toddle_64x64.png",
@@ -19,11 +19,18 @@
     "declarativeNetRequest",
     "webRequest"
   ],
-  "host_permissions": ["https://*.toddle.site/*"],
+  "host_permissions": [
+    "https://*.toddle.site/*"
+  ],
   "content_scripts": [
     {
-      "matches": ["https://*.toddle.dev/*", "https://*.toddle.site/*"],
-      "js": ["chrome/src/install_notifier.js"],
+      "matches": [
+        "https://*.toddle.dev/*",
+        "https://*.toddle.site/*"
+      ],
+      "js": [
+        "chrome/src/install_notifier.js"
+      ],
       "run_at": "document_start"
     }
   ]
diff --git a/chrome/src/background.ts b/chrome/src/background.ts
index 9cb06ef..9f1fbb5 100644
--- a/chrome/src/background.ts
+++ b/chrome/src/background.ts
@@ -1,6 +1,7 @@
 // The .js extension is necessary for Chrome to pickup the import correctly
 import { setCookies } from '../../shared/setCookies.js'
 import type { RequireFields } from '../../shared/setCookies.js'
+import { updateSessionRules } from './helpers.js'
 
 console.log('toddle extension loaded')
 
@@ -58,28 +59,7 @@ chrome.webNavigation.onBeforeNavigate.addListener(
     })
 
     if (domainCookies.length > 0) {
-      const cookieValue =
-        domainCookies.map((c) => `${c.name}=${c.value}`).join('; ') + ';'
-
-      await chrome.declarativeNetRequest.updateSessionRules({
-        removeRuleIds: [RULE_ID],
-        addRules: [
-          {
-            id: RULE_ID,
-            condition: {},
-            action: {
-              type: 'modifyHeaders',
-              requestHeaders: [
-                {
-                  header: 'Cookie',
-                  operation: 'set',
-                  value: cookieValue,
-                },
-              ],
-            },
-          },
-        ],
-      })
+      await updateSessionRules({ domainCookies, RULE_ID })
     }
   },
   {
@@ -103,8 +83,32 @@ chrome.webRequest.onHeadersReceived.addListener(
       setCookies({
         setCookieHeaders,
         requestUrl: info.url,
-        setCookie: (cookie) => chrome.cookies.set(cookie),
-        removeCookie: (cookie) => chrome.cookies.remove(cookie),
+        setCookie: async (cookie, domain) => {
+          await chrome.cookies.set(cookie)
+
+          const domainCookies = await chrome.cookies.getAll({
+            domain,
+          })
+
+          if (domainCookies.length > 0) {
+            await updateSessionRules({ domainCookies, RULE_ID })
+          }
+        },
+        removeCookie: async (cookie, domain) => {
+          await chrome.cookies.remove(cookie)
+
+          const domainCookies = await chrome.cookies.getAll({
+            domain,
+          })
+
+          if (domainCookies.length > 0) {
+            await updateSessionRules({ domainCookies, RULE_ID })
+          } else {
+            await chrome.declarativeNetRequest.updateSessionRules({
+              removeRuleIds: [RULE_ID],
+            })
+          }
+        },
         notifyUser: async (requestedUrl) => {
           const url = new URL(info.url)
           const domainCookies = await chrome.cookies.getAll({
diff --git a/chrome/src/helpers.ts b/chrome/src/helpers.ts
new file mode 100644
index 0000000..ddf97d9
--- /dev/null
+++ b/chrome/src/helpers.ts
@@ -0,0 +1,32 @@
+export interface updateSessionRulesArguments {
+  domainCookies: chrome.cookies.Cookie[]
+  RULE_ID: number
+}
+
+export async function updateSessionRules({
+  domainCookies,
+  RULE_ID,
+}: updateSessionRulesArguments) {
+  const cookieValue =
+    domainCookies.map((c) => `${c.name}=${c.value}`).join('; ') + ';'
+
+  await chrome.declarativeNetRequest.updateSessionRules({
+    removeRuleIds: [RULE_ID],
+    addRules: [
+      {
+        id: RULE_ID,
+        condition: {},
+        action: {
+          type: 'modifyHeaders',
+          requestHeaders: [
+            {
+              header: 'Cookie',
+              operation: 'set',
+              value: cookieValue,
+            },
+          ],
+        },
+      },
+    ],
+  })
+}
diff --git a/shared/setCookies.ts b/shared/setCookies.ts
index faa3a78..6171f5d 100644
--- a/shared/setCookies.ts
+++ b/shared/setCookies.ts
@@ -1,8 +1,8 @@
 export interface SetCookiesArguments {
   requestUrl: string
   setCookieHeaders: string[]
-  setCookie: (cookie: ParsedCookie & { url: string }) => void
-  removeCookie: (cookie: { name: string; url: string }) => void
+  setCookie: (cookie: ParsedCookie & { url: string }, domain?: string) => void
+  removeCookie: (cookie: { name: string; url: string }, domain?: string) => void
   notifyUser: (requestedUrl: string) => void
 }
 
@@ -41,22 +41,28 @@ export function setCookies({
           // eslint-disable-next-line no-empty
         } catch {}
         if (value === '') {
-          removeCookie({
-            name,
-            url,
-          })
+          removeCookie(
+            {
+              name,
+              url,
+            },
+            domain,
+          )
         } else {
-          setCookie({
-            name,
-            value,
-            url,
-            secure,
-            httpOnly,
-            path,
-            sameSite,
-            expirationDate,
+          setCookie(
+            {
+              name,
+              value,
+              url,
+              secure,
+              httpOnly,
+              path,
+              sameSite,
+              expirationDate,
+              domain,
+            },
             domain,
-          })
+          )
         }
         try {
           notifyUser(url)