From ac5e907d7242d97e7a939160b7ad73c03ddd30d0 Mon Sep 17 00:00:00 2001
From: tophf <tophf@gmx.com>
Date: Mon, 24 Feb 2020 18:20:11 +0300
Subject: [PATCH] fixup! track iframes via ports

---
 background/icon-manager.js | 24 ++++++++++++++++++++++--
 content/apply.js           | 32 +++++++++++++++++++++-----------
 2 files changed, 43 insertions(+), 13 deletions(-)

diff --git a/background/icon-manager.js b/background/icon-manager.js
index 59bf2cb5..53360fd5 100644
--- a/background/icon-manager.js
+++ b/background/icon-manager.js
@@ -4,6 +4,7 @@
 
 const iconManager = (() => {
   const ICON_SIZES = FIREFOX || CHROME >= 2883 && !VIVALDI ? [16, 32] : [19, 38];
+  const staleBadges = new Set();
 
   prefs.subscribe([
     'disableAll',
@@ -31,8 +32,10 @@ const iconManager = (() => {
     updateIconBadge(styleIds) {
       // FIXME: in some cases, we only have to redraw the badge. is it worth a optimization?
       const {frameId, tab: {id: tabId}} = this.sender;
-      tabManager.set(tabId, 'styleIds', frameId, styleIds.length ? styleIds.map(Number) : undefined);
-      refreshIconBadgeText(tabId);
+      const value = styleIds.length ? styleIds.map(Number) : undefined;
+      tabManager.set(tabId, 'styleIds', frameId, value);
+      debounce(refreshStaleBadges, frameId ? 250 : 0);
+      staleBadges.add(tabId);
       if (!frameId) refreshIcon(tabId, true);
     },
   });
@@ -41,6 +44,18 @@ const iconManager = (() => {
     if (!frameId) tabManager.set(tabId, 'styleIds', undefined);
   });
 
+  chrome.runtime.onConnect.addListener(port => {
+    if (port.name === 'iframe') {
+      port.onDisconnect.addListener(onPortDisconnected);
+    }
+  });
+
+  function onPortDisconnected({sender}) {
+    if (tabManager.get(sender.tab.id, 'styleIds')) {
+      API_METHODS.updateIconBadge.call({sender}, []);
+    }
+  }
+
   function refreshIconBadgeText(tabId) {
     const text = prefs.get('show-badge') ? `${getStyleCount(tabId)}` : '';
     iconUtil.setBadgeText({tabId, text});
@@ -110,4 +125,9 @@ const iconManager = (() => {
       refreshIconBadgeText(tabId);
     }
   }
+
+  function refreshStaleBadges() {
+    for (const tabId of staleBadges) refreshIconBadgeText(tabId);
+    staleBadges.clear();
+  }
 })();
diff --git a/content/apply.js b/content/apply.js
index 0e0d098a..895e9b9a 100644
--- a/content/apply.js
+++ b/content/apply.js
@@ -9,18 +9,21 @@
 self.INJECTED !== 1 && (() => {
   self.INJECTED = 1;
 
+  let IS_TAB = !chrome.tabs || location.pathname !== '/popup.html';
+  const IS_FRAME = window !== parent;
   const STYLE_VIA_API = !chrome.app && document instanceof XMLDocument;
   const styleInjector = createStyleInjector({
     compare: (a, b) => a.id - b.id,
     onUpdate: onInjectorUpdate,
   });
   const initializing = init();
+  /** @type chrome.runtime.Port */
+  let port;
 
-  // if chrome.tabs is absent it's a web page, otherwise we'll check for popup/options as those aren't tabs
-  let isTab = !chrome.tabs;
-  if (chrome.tabs) {
+  // the popup needs a check as it's not a tab but can be opened in a tab manually for whatever reason
+  if (!IS_TAB) {
     chrome.tabs.getCurrent(tab => {
-      isTab = Boolean(tab);
+      IS_TAB = Boolean(tab);
       if (tab && styleInjector.list.length) updateCount();
     });
   }
@@ -41,7 +44,7 @@ self.INJECTED !== 1 && (() => {
   let parentDomain;
 
   prefs.subscribe(['disableAll'], (key, value) => doDisableAll(value));
-  if (window !== parent) {
+  if (IS_FRAME) {
     prefs.subscribe(['exposeIframes'], updateExposeIframes);
   }
 
@@ -64,7 +67,7 @@ self.INJECTED !== 1 && (() => {
       // dynamic about: and javascript: iframes don't have an URL yet
       // so we'll try the parent frame which is guaranteed to have a real URL
       try {
-        if (window !== parent) {
+        if (IS_FRAME) {
           matchUrl = parent.location.href;
         }
       } catch (e) {}
@@ -162,11 +165,18 @@ self.INJECTED !== 1 && (() => {
   }
 
   function updateCount() {
-    if (isTab) {
-      (STYLE_VIA_API ?
-        API.styleViaAPI({method: 'updateCount'}) :
-        API.updateIconBadge(styleInjector.list.map(style => style.id))
-      ).catch(msg.ignoreError);
+    if (!IS_TAB) return;
+    if (STYLE_VIA_API) {
+      API.styleViaAPI({method: 'updateCount'}).catch(msg.ignoreError);
+    } else {
+      API.updateIconBadge(styleInjector.list.map(style => style.id)).catch(msg.ignoreError);
+    }
+    if (IS_FRAME) {
+      if (!port && styleInjector.list.length) {
+        port = chrome.runtime.connect({name: 'iframe'});
+      } else if (port && !styleInjector.list.length) {
+        port.disconnect();
+      }
     }
   }