Follow up commit to add all messages to be promise safe including background.js messages. Fixe #670

index.js

@@ -276,7 +276,7 @@ const ContainerService = {
       const api = await webExtension.startup();
       api.browser.runtime.onMessage.addListener(async function (message, sender, sendReply) {
         if ("method" in message && methods.indexOf(message.method) !== -1) {
-          const response = ContainerService[message.method](message);
+          const response = ContainerService[message.method](message.message);
           if (response instanceof Promise) {
             const responseValue = await response;
             ContainerService.triggerBackgroundCallback({

webextension/background.js

@@ -419,7 +419,9 @@ const backgroundLogic = {
     // Unhide all hidden tabs
     browser.runtime.sendMessage({
       method: "showTabs",
-      userContextId: options.userContextId
+      message: {
+        userContextId: options.userContextId
+      }
     });
     return browser.tabs.create({
       url,
@@ -456,7 +458,7 @@ const messageHandler = {
 
   init() {
     // Handles messages from webextension code
-    browser.runtime.onMessage.addListener((m) => {
+    browser.runtime.onMessage.addListener(async function (m) {
       let response;
 
       switch (m.method) {
@@ -474,7 +476,7 @@ const messageHandler = {
         assignManager._neverAsk(m);
         break;
       case "getAssignment":
-        response = browser.tabs.get(m.tabId).then((tab) => {
+        response = browser.tabs.get(m.message.tabId).then((tab) => {
           return assignManager._getAssignment(tab);
         });
         break;
@@ -485,13 +487,22 @@ const messageHandler = {
         // m.tabId is used for where to place the in content message
         // m.url is the assignment to be removed/added
         response = browser.tabs.get(m.tabId).then((tab) => {
-          return assignManager._setOrRemoveAssignment(tab.id, m.url, m.userContextId, m.value);
+          return assignManager._setOrRemoveAssignment(tab.id, m.message.url, m.message.userContextId, m.message.value);
         });
         break;
       case "exemptContainerAssignment":
         response = assignManager._exemptTab(m);
         break;
       }
+      if (response instanceof Promise) {
+        const resolvedResponse = await response;
+        browser.runtime.sendMessage({
+          method: "async-popup-response",
+          message: resolvedResponse,
+          uuid: m.uuid
+        });
+        return true;
+      }
       return response;
     });
 
@@ -500,11 +511,11 @@ const messageHandler = {
     port.onMessage.addListener(m => {
       switch (m.type) {
       case "async-background-response":
-         browser.runtime.sendMessage({
-           method: "async-popup-response",
-           message: m.message.response,
-           uuid: m.message.uuid
-         });
+        browser.runtime.sendMessage({
+          method: "async-popup-response",
+          message: m.message.response,
+          uuid: m.message.uuid
+        });
         break;
       case "lightweight-theme-changed":
         themeManager.update(m.message);
@@ -707,7 +718,9 @@ messageHandler.init();
 
 browser.runtime.sendMessage({
   method: "getPreference",
-  pref: "browser.privatebrowsing.autostart"
+  message: {
+    pref: "browser.privatebrowsing.autostart"
+  }
 }).then(pbAutoStart => {
 
   // We don't want to disable the addon if we are in auto private-browsing.

webextension/js/popup.js

@@ -84,12 +84,12 @@ const Logic = {
     // Get the onboarding variation
     const variationPromise = this.getShieldStudyVariation();
     browser.runtime.onMessage.addListener((m) => {
-       if (m.method === "async-popup-response" && m.uuid) {
-         const uuid = m.uuid;
-         if (uuid in this.asyncMessageQueue) {
-           this.asyncMessageQueue[uuid].response = m.message;
-         }
-       }
+      if (m.method === "async-popup-response" && m.uuid) {
+        const uuid = m.uuid;
+        if (uuid in this.asyncMessageQueue) {
+          this.asyncMessageQueue[uuid].response = m.message;
+        }
+      }
     });
 
     try {
@@ -199,28 +199,30 @@ const Logic = {
     // UUID SO result, as you do
     return ([1e7]+-1e3+-4e3+-8e3+-1e11).replace(/[018]/g, c =>
       (c ^ crypto.getRandomValues(new Uint8Array(1))[0] & 15 >> c / 4).toString(16)
-    )
+    );
   },
+  asyncNotResponded: "@@_not_responded_@@",
 
   asyncCreateMessage() {
     const guid = this.asyncId();
     this.asyncMessageQueue[guid] = {
       time: Date.now(),
-      response: null
+      response: this.asyncNotResponded
     };
     return guid;
   },
 
-  awaitMessage(methodName) {
+  awaitMessage(methodName, message) {
     return new Promise((resolve, reject) => {
       const messageId = this.asyncCreateMessage();
       browser.runtime.sendMessage({
         method: methodName,
-        uuid: messageId
+        uuid: messageId,
+        message
       });
       const checkState = () => {
         const messageState = this.asyncMessageQueue[messageId];
-        if (messageState.response !== null) {
+        if (messageState.response !== this.asyncNotResponded) {
           resolve(messageState.response);
           delete this.asyncMessageQueue[messageId];
           return;
@@ -330,28 +332,25 @@ const Logic = {
       return Promise.reject("removeIdentity must be called with userContextId argument.");
     }
 
-    return browser.runtime.sendMessage({
-      method: "deleteContainer",
-      message: {userContextId}
+    return Logic.awaitMessage("deleteContainer", {
+      userContextId
     });
   },
 
   getAssignment(tab) {
-    return browser.runtime.sendMessage({
-      method: "getAssignment",
+    return Logic.awaitMessage("getAssignment", {
       tabId: tab.id
     });
   },
 
   getAssignmentObjectByContainer(userContextId) {
-    return browser.runtime.sendMessage({
-      method: "getAssignmentObjectByContainer",
-      message: {userContextId}
+    return Logic.awaitMessage("getAssignmentObjectByContainer", {
+      userContextId
     });
   },
 
   setOrRemoveAssignment(tabId, url, userContextId, value) {
-    return browser.runtime.sendMessage({
+    return Logic.awaitMessage("setOrRemoveAssignment", {
       method: "setOrRemoveAssignment",
       tabId,
       url,
@@ -525,9 +524,7 @@ Logic.registerPanel(P_CONTAINERS_LIST, {
     });
 
     Logic.addEnterHandler(document.querySelector("#sort-containers-link"), () => {
-      browser.runtime.sendMessage({
-        method: "sortTabs"
-      }).then(() => {
+      Logic.awaitMessage("sortTabs").then(() => {
         window.close();
       }).catch(() => {
         window.close();
@@ -659,12 +656,9 @@ Logic.registerPanel(P_CONTAINERS_LIST, {
         if (e.target.matches(".open-newtab")
             || e.target.parentNode.matches(".open-newtab")
             || e.type === "keydown") {
-          browser.runtime.sendMessage({
-            method: "openTab",
-            message: {
-              userContextId: Logic.userContextId(identity.cookieStoreId),
-              source: "pop-up"
-            }
+          Logic.awaitMessage("openTab", {
+            userContextId: Logic.userContextId(identity.cookieStoreId),
+            source: "pop-up"
           }).then(() => {
             window.close();
           }).catch(() => {
@@ -711,8 +705,8 @@ Logic.registerPanel(P_CONTAINER_INFO, {
 
     Logic.addEnterHandler(document.querySelector("#container-info-hideorshow"), () => {
       const identity = Logic.currentIdentity();
-      browser.runtime.sendMessage({
-        method: identity.hasHiddenTabs ? "showTabs" : "hideTabs",
+      const method = identity.hasHiddenTabs ? "showTabs" : "hideTabs";
+      Logic.awaitMessage(method, {
         userContextId: Logic.currentUserContextId()
       }).then(() => {
         window.close();
@@ -722,9 +716,7 @@ Logic.registerPanel(P_CONTAINER_INFO, {
     });
 
     // Check if the user has incompatible add-ons installed
-    browser.runtime.sendMessage({
-      method: "checkIncompatibleAddons"
-    }).then(incompatible => {
+    Logic.awaitMessage("checkIncompatibleAddons").then(incompatible => {
       const moveTabsEl = document.querySelector("#container-info-movetabs");
       if (incompatible) {
         const fragment = document.createDocumentFragment();
@@ -741,8 +733,7 @@ Logic.registerPanel(P_CONTAINER_INFO, {
         moveTabsEl.parentNode.insertBefore(fragment, moveTabsEl.nextSibling);
       } else {
         Logic.addEnterHandler(moveTabsEl, () => {
-          browser.runtime.sendMessage({
-            method: "moveTabsToWindow",
+          Logic.awaitMessage("moveTabsToWindow", {
             userContextId: Logic.userContextId(Logic.currentIdentity().cookieStoreId),
           }).then(() => {
             window.close();
@@ -783,9 +774,8 @@ Logic.registerPanel(P_CONTAINER_INFO, {
     }
 
     // Let's retrieve the list of tabs.
-    return browser.runtime.sendMessage({
-      method: "getTabs",
-      userContextId: Logic.currentUserContextId(),
+    return Logic.awaitMessage("getTabs", {
+      userContextId: Logic.currentUserContextId()
     }).then(this.buildInfoTable);
   },
 
@@ -805,9 +795,8 @@ Logic.registerPanel(P_CONTAINER_INFO, {
       if (tab.active) {
         tr.classList.add("clickable");
         Logic.addEnterHandler(tr, () => {
-          browser.runtime.sendMessage({
-            method: "showTab",
-            tabId: tab.id,
+          Logic.awaitMessage("showTab", {
+            tabId: tab.id
           }).then(() => {
             window.close();
           }).catch(() => {
@@ -925,15 +914,12 @@ Logic.registerPanel(P_CONTAINER_EDIT, {
 
   _submitForm() {
     const formValues = new FormData(this._editForm);
-    return browser.runtime.sendMessage({
-      method: "createOrUpdateContainer",
-      message: {
-        userContextId: formValues.get("container-id") || NEW_CONTAINER_ID,
-        params: {
-          name: document.getElementById("edit-container-panel-name-input").value || Logic.generateIdentityName(),
-          icon: formValues.get("container-icon") || DEFAULT_ICON,
-          color: formValues.get("container-color") || DEFAULT_COLOR,
-        }
+    return Logic.awaitMessage("createOrUpdateContainer", {
+      userContextId: formValues.get("container-id") || NEW_CONTAINER_ID,
+      params: {
+        name: document.getElementById("edit-container-panel-name-input").value || Logic.generateIdentityName(),
+        icon: formValues.get("container-icon") || DEFAULT_ICON,
+        color: formValues.get("container-color") || DEFAULT_COLOR,
       }
     }).then(() => {
       return Logic.refreshIdentities();