Fixed possible surprises in background script message handling.

2 files changed, 311 insertions, 333 deletions

diff --git a/src/bg/RequestGuard.js b/src/bg/RequestGuard.js
index 69c6348..ccf3c70 100644
--- a/src/bg/RequestGuard.js
+++ b/src/bg/RequestGuard.js
@@ -213,71 +213,71 @@ var RequestGuard = (() => {
   if (!("setIcon" in browser.browserAction)) { // unsupported on Android
     TabStatus._updateTabNow = TabStatus.updateTab = () => {};
   }
-
-  const Content = {
-
-
-    async hearFrom(message, sender) {
-      debug("Received message from content", message, sender);
-      switch (message.type) {
-        case "pageshow":
-          TabStatus.recordAll(sender.tab.id, message.seen);
-          return true;
-        case "enable": {
-          let {url, documentUrl, policyType} = message;
-          let TAG = `<${policyType.toUpperCase()}>`;
-          let origin = Sites.origin(url);
-          let {siteKey} = Sites.parse(url);
-          let options;
-          if (siteKey === origin) {
-            TAG += `@${siteKey}`;
-          } else {
-            options = [
-              {label: _("allowLocal", siteKey), checked: true},
-              {label: _("allowLocal", origin)}
-            ];
-          }
-          // let parsedDoc = Sites.parse(documentUrl);
-          let t = u => `${TAG}@${u}`;
-          let ret = await Prompts.prompt({
-            title: _("BlockedObjects"),
-            message: _("allowLocal", TAG),
-            options});
-          debug(`Prompt returned %o`);
-          if (ret.button !== 0) return;
-          let key = [siteKey, origin][ret.option || 0];
-          if (!key) return;
-          let {siteMatch, contextMatch, perms} = ns.policy.get(key, documentUrl);
-          let {capabilities} = perms;
-          if (!capabilities.has(policyType)) {
-            perms = new Permissions(new Set(capabilities), false);
-            perms.capabilities.add(policyType);
-
-            /* TODO: handle contextual permissions
-            if (documentUrl) {
-              let context = new URL(documentUrl).origin;
-              let contextualSites = new Sites([context, perms]);
-              perms = new Permissions(new Set(capabilities), false, contextualSites);
-            }
-            */
-            ns.policy.set(key, perms);
-            ns.savePolicy();
-          }
-          return true;
-        }
-        case "docStatus": {
-          let {frameId, tab} = sender;
-          let {url} = message;
-          let tabId = tab.id;
-          let records = TabStatus.map.get(tabId);
-          let noscriptFrames = records && records.noscriptFrames;
-          let canScript = !(noscriptFrames && noscriptFrames[sender.frameId]);
-          let shouldScript = !ns.isEnforced(tabId) || !url.startsWith("http") || ns.policy.can(url, "script");
-          debug("Frame %s %s of  %o, canScript: %s, shouldScript: %s", frameId, url, noscriptFrames, canScript, shouldScript);
-          return {canScript, shouldScript};
+  
+  
+  let messageHandler = {
+    async pageshow(message, sender) {
+      TabStatus.recordAll(sender.tab.id, message.seen);
+      return true;
+    },
+    async enable(message, sender) {
+      let {url, documentUrl, policyType} = message;
+      let TAG = `<${policyType.toUpperCase()}>`;
+      let origin = Sites.origin(url);
+      let {siteKey} = Sites.parse(url);
+      let options;
+      if (siteKey === origin) {
+        TAG += `@${siteKey}`;
+      } else {
+        options = [
+          {label: _("allowLocal", siteKey), checked: true},
+          {label: _("allowLocal", origin)}
+        ];
+      }
+      // let parsedDoc = Sites.parse(documentUrl);
+      let t = u => `${TAG}@${u}`;
+      let ret = await Prompts.prompt({
+        title: _("BlockedObjects"),
+        message: _("allowLocal", TAG),
+        options});
+      debug(`Prompt returned %o`);
+      if (ret.button !== 0) return;
+      let key = [siteKey, origin][ret.option || 0];
+      if (!key) return;
+      let {siteMatch, contextMatch, perms} = ns.policy.get(key, documentUrl);
+      let {capabilities} = perms;
+      if (!capabilities.has(policyType)) {
+        perms = new Permissions(new Set(capabilities), false);
+        perms.capabilities.add(policyType);
+
+        /* TODO: handle contextual permissions
+        if (documentUrl) {
+          let context = new URL(documentUrl).origin;
+          let contextualSites = new Sites([context, perms]);
+          perms = new Permissions(new Set(capabilities), false, contextualSites);
         }
+        */
+        ns.policy.set(key, perms);
+        await ns.savePolicy();
       }
+      return true;
     },
+    
+    async docStatus(message, sender) {
+      let {frameId, tab} = sender;
+      let {url} = message;
+      let tabId = tab.id;
+      let records = TabStatus.map.get(tabId);
+      let noscriptFrames = records && records.noscriptFrames;
+      let canScript = !(noscriptFrames && noscriptFrames[sender.frameId]);
+      let shouldScript = !ns.isEnforced(tabId) || !url.startsWith("http") || ns.policy.can(url, "script");
+      debug("Frame %s %s of  %o, canScript: %s, shouldScript: %s", frameId, url, noscriptFrames, canScript, shouldScript);
+      return {canScript, shouldScript};
+    }
+    
+  }
+  
+  const Content = {
 
     async reportTo(request, allowed, policyType) {
       let {requestId, tabId, frameId, type, url, documentUrl, originUrl} = request;
@@ -312,7 +312,6 @@ var RequestGuard = (() => {
       }
     }
   };
-  browser.runtime.onMessage.addListener(Content.hearFrom);
 
   const pendingRequests = new Map();
   function initPendingRequest(request) {
@@ -551,6 +550,8 @@ var RequestGuard = (() => {
   
   const RequestGuard = {
     async start() {
+      Messages.addHandler(messageHandler);
+      
       let wr = browser.webRequest;
       let listen = (what, ...args) => wr[what].addListener(listeners[what], ...args);
       
@@ -608,6 +609,7 @@ var RequestGuard = (() => {
         }
       }
       wr.onBeforeRequest.removeListener(onViolationReport);
+      Messages.removeHandler(messageHandler);
     }
   };
 
diff --git a/src/bg/main.js b/src/bg/main.js
index 75ea868..3882191 100644
--- a/src/bg/main.js
+++ b/src/bg/main.js
@@ -1,272 +1,248 @@
- var ns = (() => {
-   'use strict';
-
-   const popupURL = browser.extension.getURL("/ui/popup.html");
-   let popupFor = tabId => `${popupURL}#tab${tabId}`;
-
-   let ctxMenuId = "noscript-ctx-menu";
-
-   async function toggleCtxMenuItem(show = ns.local.showCtxMenuItem) {
-     if (!"contextMenus" in browser) return;
-     let id = ctxMenuId;
-     try {
-         await browser.contextMenus.remove(id);
-     } catch (e) {}
-
-     if (show) {
-       browser.contextMenus.create({
-         id,
-         title: "NoScript",
-         contexts: ["all"]
-       });
-     }
-   }
-
-   async function init() {
-     let policyData = (await Storage.get("sync", "policy")).policy;
-     if (policyData && policyData.DEFAULT) {
-       ns.policy = new Policy(policyData);
-       await ChildPolicies.update(policyData);
-     } else {
-       await include("/legacy/Legacy.js");
-       ns.policy = await Legacy.createOrMigratePolicy();
-       ns.savePolicy();
-     }
-     
-     
-     await include("/bg/defaults.js");
-     await ns.defaults;
-     await include("/bg/RequestGuard.js");
-     await RequestGuard.start();
-     await XSS.start(); // we must start it anyway to initialize sub-objects
-     if (!ns.sync.xss) {
-       XSS.stop();
-     }
-     Commands.install();
-   };
-
-   var Commands = {
-     openPageUI() {
-       try {
-         browser.browserAction.openPopup();
-         return;
-       } catch (e) {
-         debug(e);
-       }
-       browser.windows.create({
-         url: popupURL,
-         width: 800,
-         height: 600,
-         type: "panel"
-       });
-     },
-
-     togglePermissions() {},
-     install() {
-
-
-       if ("command" in browser) {
-         // keyboard shortcuts
-         browser.commands.onCommand.addListener(cmd => {
-           if (cmd in Commands) {
-             Commands[cmd]();
-           }
-         });
-       }
-
-       if ("contextMenus" in browser) {
-         toggleCtxMenuItem();
-         browser.contextMenus.onClicked.addListener((info, tab) => {
-           if (info.menuItemId == ctxMenuId) {
-             this.openPageUI();
-           }
-         });
-       }
-
-       // wiring main UI
-       let ba = browser.browserAction;
-       if ("setIcon" in ba) {
-         //desktop
-         ba.setPopup({
-           popup: popupURL
-         });
-       } else {
-         // mobile
-         ba.onClicked.addListener(async tab => {
-           try {
-             await browser.tabs.remove(await browser.tabs.query({
-               url: popupURL
-             }));
-           } catch (e) {}
-           await browser.tabs.create({
-             url: popupFor(tab.id)
-           });
-         });
-       }
-     }
-   }
-
-   var MessageHandler = {
-     responders: {
-
-       async updateSettings(settings, sender) {
-         await Settings.update(settings);
-         toggleCtxMenuItem();
-       },
-       async broadcastSettings({
-         tabId = -1
-       }) {
-         let policy = ns.policy.dry(true);
-         let seen = tabId !== -1 ? await ns.collectSeen(tabId) : null;
-         let xssUserChoices = await XSS.getUserChoices();
-         browser.runtime.sendMessage({
-           type: "settings",
-           policy,
-           seen,
-           xssUserChoices,
-           local: ns.local,
-           sync: ns.sync,
-           unrestrictedTab: ns.unrestrictedTabs.has(tabId),
-         });
-       },
-
-       exportSettings(m, sender, sendResponse) {
-         sendResponse(Settings.export());
-         return false;
-       },
-
-       async importSettings({
-         data
-       }) {
-         return await Settings.import(data);
-       },
-
-       async openStandalonePopup() {
-         let win = await browser.windows.getLastFocused();
-         let [tab] = (await browser.tabs.query({
-           lastFocusedWindow: true,
-           active: true
-         }));
-
-         if (!tab || tab.id === -1) {
-           log("No tab found to open the UI for");
-           return;
-         }
-         browser.windows.create({
-           url: popupFor(tab.id),
-           width: 800,
-           height: 600,
-           top: win.top + 48,
-           left: win.left + 48,
-           type: "panel"
-         });
-       }
-     },
-     onMessage(m, sender, sendResponse) {
-       let {
-         type
-       } = m;
-       let {
-         responders
-       } = MessageHandler;
-
-
-       if (type && (type = type.replace(/^NoScript\./, '')) in responders) {
-         return responders[type](m, sender, sendResponse);
-       } else {
-         debug("Received unkown message", m, sender);
-       }
-       return false;
-     },
-
-     listen() {
-       browser.runtime.onMessage.addListener(this.onMessage);
-     },
-   }
-
-
-
-   return {
-     running: false,
-     policy: null,
-     local: null,
-     sync: null,
-     unrestrictedTabs: new Set(),
-     isEnforced(tabId = -1) {
-       return this.policy.enforced && (tabId === -1 || !this.unrestrictedTabs.has(tabId));
-     },
-
-     start() {
-       if (this.running) return;
-       this.running = true;
-
-       deferWebTraffic(init(), 
-          async () => {
-
-           await include("/bg/Settings.js");
-           MessageHandler.listen();
-
-           log("STARTED");
-
-           this.devMode = (await browser.management.getSelf()).installType === "development";
-           if (this.local.debug) {
-             if (this.devMode) {
-               include("/test/run.js");
-             }
-           } else {
-             debug = () => {}; // suppress verbosity
-           }
-         });
-     },
-
-     stop() {
-       if (!this.running) return;
-       this.running = false;
-       RequestGuard.stop();
-       log("STOPPED");
-     },
-
-     async savePolicy() {
-       if (this.policy) {
-         await ChildPolicies.update(this.policy);
-         await Storage.set("sync", {
-           policy: this.policy.dry()
-         });
-         await browser.webRequest.handlerBehaviorChanged()
-       }
-       return this.policy;
-     },
-
-
-
-     async save(obj) {
-       if (obj && obj.storage) {
-         let toBeSaved = {
-           [obj.storage]: obj
-         };
-         Storage.set(obj.storage, toBeSaved);
-       }
-       return obj;
-     },
-
-     async collectSeen(tabId) {
-
-       try {
-         let seen = Array.from(await browser.tabs.sendMessage(tabId, {
-           type: "collect"
-         }, {
-           frameId: 0
-         }));
-         debug("Collected seen", seen);
-         return seen;
-       } catch (e) {
-         // probably a page where content scripts cannot run, let's open the options instead
-         error(e, "Cannot collect noscript activity data");
-       }
-
-       return null;
-     },
-   };
- })();
+ {
+  'use strict';
+
+  let popupURL = browser.extension.getURL("/ui/popup.html");
+  let popupFor = tabId => `${popupURL}#tab${tabId}`;
+
+  let ctxMenuId = "noscript-ctx-menu";
+
+  async function toggleCtxMenuItem(show = ns.local.showCtxMenuItem) {
+    if (!"contextMenus" in browser) return;
+    let id = ctxMenuId;
+    try {
+      await browser.contextMenus.remove(id);
+    } catch (e) {}
+
+    if (show) {
+      browser.contextMenus.create({
+        id,
+        title: "NoScript",
+        contexts: ["all"]
+      });
+    }
+  }
+
+  async function init() {
+    let policyData = (await Storage.get("sync", "policy")).policy;
+    if (policyData && policyData.DEFAULT) {
+      ns.policy = new Policy(policyData);
+      await ChildPolicies.update(policyData);
+    } else {
+      await include("/legacy/Legacy.js");
+      ns.policy = await Legacy.createOrMigratePolicy();
+      ns.savePolicy();
+    }
+    
+    
+    await include("/bg/defaults.js");
+    await ns.defaults;
+    await include("/bg/RequestGuard.js");
+    await RequestGuard.start();
+    await XSS.start(); // we must start it anyway to initialize sub-objects
+    if (!ns.sync.xss) {
+      XSS.stop();
+    }
+    Commands.install();
+  };
+
+  let Commands = {
+    openPageUI() {
+      try {
+        browser.browserAction.openPopup();
+        return;
+      } catch (e) {
+        debug(e);
+      }
+      browser.windows.create({
+        url: popupURL,
+        width: 800,
+        height: 600,
+        type: "panel"
+      });
+    },
+
+    togglePermissions() {},
+    install() {
+
+
+      if ("command" in browser) {
+        // keyboard shortcuts
+        browser.commands.onCommand.addListener(cmd => {
+          if (cmd in Commands) {
+            Commands[cmd]();
+          }
+        });
+      }
+
+      if ("contextMenus" in browser) {
+        toggleCtxMenuItem();
+        browser.contextMenus.onClicked.addListener((info, tab) => {
+          if (info.menuItemId == ctxMenuId) {
+            this.openPageUI();
+          }
+        });
+      }
+
+      // wiring main UI
+      let ba = browser.browserAction;
+      if ("setIcon" in ba) {
+        //desktop
+        ba.setPopup({
+          popup: popupURL
+        });
+      } else {
+        // mobile
+        ba.onClicked.addListener(async tab => {
+          try {
+            await browser.tabs.remove(await browser.tabs.query({
+              url: popupURL
+            }));
+          } catch (e) {}
+          await browser.tabs.create({
+            url: popupFor(tab.id)
+          });
+        });
+      }
+    }
+  }
+
+  let messageHandler = {
+    async updateSettings(settings, sender) {
+      await Settings.update(settings);
+      toggleCtxMenuItem();
+    },
+      
+    async broadcastSettings({
+      tabId = -1
+    }) {
+      let policy = ns.policy.dry(true);
+      let seen = tabId !== -1 ? await ns.collectSeen(tabId) : null;
+      let xssUserChoices = await XSS.getUserChoices();
+      browser.runtime.sendMessage({
+        type: "settings",
+        policy,
+        seen,
+        xssUserChoices,
+        local: ns.local,
+        sync: ns.sync,
+        unrestrictedTab: ns.unrestrictedTabs.has(tabId),
+      });
+    },
+
+    async exportSettings() {
+      return Settings.export();
+    },
+
+    async importSettings({data}) {
+      return await Settings.import(data);
+    },
+
+    async openStandalonePopup() {
+      let win = await browser.windows.getLastFocused();
+      let [tab] = (await browser.tabs.query({
+        lastFocusedWindow: true,
+        active: true
+      }));
+
+      if (!tab || tab.id === -1) {
+        log("No tab found to open the UI for");
+        return;
+      }
+      browser.windows.create({
+        url: popupFor(tab.id),
+        width: 800,
+        height: 600,
+        top: win.top + 48,
+        left: win.left + 48,
+        type: "panel"
+      });
+    },
+  };
+
+
+
+  var ns = {
+    running: false,
+    policy: null,
+    local: null,
+    sync: null,
+    unrestrictedTabs: new Set(),
+    isEnforced(tabId = -1) {
+      return this.policy.enforced && (tabId === -1 || !this.unrestrictedTabs.has(tabId));
+    },
+
+    start() {
+      if (this.running) return;
+      this.running = true;
+
+      deferWebTraffic(init(), 
+        async () => {
+
+          await include("/bg/Settings.js");
+          Messages.addHandler(messageHandler);
+
+          log("STARTED");
+
+          this.devMode = (await browser.management.getSelf()).installType === "development";
+          if (this.local.debug) {
+            if (this.devMode) {
+              include("/test/run.js");
+            }
+          } else {
+            debug = () => {}; // suppress verbosity
+          }
+        });
+    },
+
+    stop() {
+      if (!this.running) return;
+      this.running = false;
+      Messages.removeHandler(messageHandler);
+      RequestGuard.stop();
+      log("STOPPED");
+    },
+
+    async savePolicy() {
+      if (this.policy) {
+        await ChildPolicies.update(this.policy);
+        await Storage.set("sync", {
+          policy: this.policy.dry()
+        });
+        await browser.webRequest.handlerBehaviorChanged()
+      }
+      return this.policy;
+    },
+
+
+
+    async save(obj) {
+      if (obj && obj.storage) {
+        let toBeSaved = {
+          [obj.storage]: obj
+        };
+        Storage.set(obj.storage, toBeSaved);
+      }
+      return obj;
+    },
+
+    async collectSeen(tabId) {
+
+      try {
+        let seen = Array.from(await browser.tabs.sendMessage(tabId, {
+          type: "collect"
+        }, {
+          frameId: 0
+        }));
+        debug("Collected seen", seen);
+        return seen;
+      } catch (e) {
+        // probably a page where content scripts cannot run, let's open the options instead
+        error(e, "Cannot collect noscript activity data");
+      }
+
+      return null;
+    },
+  };
+ }
 
  ns.start();