From 1d330739ec25e76b42fa3106c9e06bfd48e917a3 Mon Sep 17 00:00:00 2001
From: hackademix <giorgio@maone.net>
Date: Thu, 15 Jun 2023 00:33:15 +0200
Subject: [PATCH] [TabGuard] Load with credentials when reloading from
 NoScript's UI.

---
 src/bg/TabGuard.js | 28 ++++++++++++++++++++++------
 src/bg/main.js     |  9 +++++++++
 src/ui/popup.js    |  4 ++--
 3 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/src/bg/TabGuard.js b/src/bg/TabGuard.js
index 40e7296..bd73aff 100644
--- a/src/bg/TabGuard.js
+++ b/src/bg/TabGuard.js
@@ -36,6 +36,18 @@ var TabGuard = (() => {
   };
   forget();
 
+  function mergeGroups(groups, {tabDomain, otherDomains} /* anonymizedTabInfo */) {
+    if (!(tabDomain in groups)) groups[tabDomain] = new Set();
+    let currentGroup = groups[tabDomain];
+    for (let d of otherDomains) {
+      if (!(d in groups)) groups[d] = new Set();
+      // add this domain to the allow/block group of the other tied ones...
+      groups[d].add(tabDomain);
+      // ... and vice-versa
+      currentGroup.add(d);
+    }
+  }
+
   const AUTH_HEADERS_RX = /^(?:authorization|cookie)/i;
 
   function getDomain(u) {
@@ -193,10 +205,9 @@ var TabGuard = (() => {
               if (ret.button !== 0) {
                 return {cancel: true};
               }
-              let list = ret.option === 0 ? filteredGroups : allowedGroups;
-              otherDomains.add(tabDomain);
-              for (let d of otherDomains) list[d] = otherDomains;
-              return list === filteredGroups ? filterAuth() : null;
+              const groups = ret.option === 0 ? filteredGroups : allowedGroups;
+              mergeGroups(groups, {tabDomain, otherDomains});
+              return groups === filteredGroups ? filterAuth() : null;
             })();
           }
         }
@@ -237,9 +248,14 @@ var TabGuard = (() => {
       // return a deep copy
       return JSON.parse(JSON.stringify(anonymizedTabs.get(tabId)));
     },
-    reloadNormally(tabId) {
+    async reloadNormally(tabId) {
       TabTies.cut(tabId);
-      browser.tabs.reload(tabId);
+      await browser.tabs.reload(tabId);
+    },
+    allow(tabId) {
+      if (!TabGuard.isAnonymizedTab(tabId)) return;
+      const info = this.getAnonymizedTabInfo(tabId);
+      mergeGroups(allowedGroups, info);
     }
   }
 })();
diff --git a/src/bg/main.js b/src/bg/main.js
index 7e3c883..75c6048 100644
--- a/src/bg/main.js
+++ b/src/bg/main.js
@@ -165,6 +165,7 @@
       let policy = ns.policy.dry(true);
       let seen = tabId !== -1 ? await ns.collectSeen(tabId) : null;
       let xssUserChoices = await XSS.getUserChoices();
+      let anonymyzedTabInfo =
       await Messages.send("settings", {
         policy,
         seen,
@@ -174,6 +175,7 @@
         unrestrictedTab: ns.unrestrictedTabs.has(tabId),
         tabId,
         xssBlockedInTab: XSS.getBlockedInTab(tabId),
+        anonymyzedTabInfo: TabGuard.isAnonymizedTab(tabId) && TabGuard.getAnonymizedTabInfo(tabId),
       });
     },
 
@@ -236,6 +238,13 @@
         matchAboutBlank: true,
         allFrames: true,
       });
+    },
+
+    async reloadWithCredentials({tabId, remember}) {
+      if (remember) {
+        TabGuard.allow(tabId);
+      }
+      await TabGuard.reloadNormally(tabId);
     }
   };
 
diff --git a/src/ui/popup.js b/src/ui/popup.js
index 27d913b..52d3362 100644
--- a/src/ui/popup.js
+++ b/src/ui/popup.js
@@ -389,9 +389,9 @@ addEventListener("unload", e => {
       window.scrollTo(0, 0);
     }
 
-    function reload() {
+    async function  reload() {
       if (sitesUI) sitesUI.clear();
-      browser.tabs.reload(tabId);
+      await Messages.send("reloadWithCredentials", {tabId});
       pendingReload(false);
     }