From 9907be9537a0cf1cecac00178ad369cfa6f6c55f Mon Sep 17 00:00:00 2001
From: Jonathan Kingston <jonathan@jooped.co.uk>
Date: Fri, 14 Jul 2017 14:46:41 +0100
Subject: [PATCH] Move async popup functions to message pass through the
 background script using a message queue because
 'extensions.webextensions.remote' pref breaks async messages. Fixes #670

---
 index.js                   | 14 +++++++--
 webextension/background.js |  7 +++++
 webextension/js/popup.js   | 58 ++++++++++++++++++++++++++++++++++++--
 3 files changed, 74 insertions(+), 5 deletions(-)

diff --git a/index.js b/index.js
index afdfa61..61bfc3e 100644
--- a/index.js
+++ b/index.js
@@ -274,9 +274,19 @@ const ContainerService = {
 
     try {
       const api = await webExtension.startup();
-      api.browser.runtime.onMessage.addListener((message, sender, sendReply) => {
+      api.browser.runtime.onMessage.addListener(async function (message, sender, sendReply) {
         if ("method" in message && methods.indexOf(message.method) !== -1) {
-          sendReply(this[message.method](message));
+          const response = ContainerService[message.method](message);
+          if (response instanceof Promise) {
+            const responseValue = await response;
+            ContainerService.triggerBackgroundCallback({
+              response: responseValue,
+              method: message.method,
+              uuid: message.uuid
+            }, "async-background-response");
+          } else {
+            sendReply(response);
+          }
         }
       });
 
diff --git a/webextension/background.js b/webextension/background.js
index 2884dfd..6916092 100644
--- a/webextension/background.js
+++ b/webextension/background.js
@@ -499,6 +499,13 @@ const messageHandler = {
     const port = browser.runtime.connect();
     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
+         });
+        break;
       case "lightweight-theme-changed":
         themeManager.update(m.message);
         break;
diff --git a/webextension/js/popup.js b/webextension/js/popup.js
index 1006dab..0ba99a5 100644
--- a/webextension/js/popup.js
+++ b/webextension/js/popup.js
@@ -83,6 +83,14 @@ const Logic = {
     const identitiesPromise = this.refreshIdentities();
     // 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;
+         }
+       }
+    });
 
     try {
       await Promise.all([identitiesPromise, variationPromise]);
@@ -183,12 +191,56 @@ const Logic = {
     return false;
   },
 
+  asyncTimeout: 2000,
+  asyncRefresh: 20,
+  asyncMessageQueue: {},
+
+  asyncId() {
+    // 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)
+    )
+  },
+
+  asyncCreateMessage() {
+    const guid = this.asyncId();
+    this.asyncMessageQueue[guid] = {
+      time: Date.now(),
+      response: null
+    };
+    return guid;
+  },
+
+  awaitMessage(methodName) {
+    return new Promise((resolve, reject) => {
+      const messageId = this.asyncCreateMessage();
+      browser.runtime.sendMessage({
+        method: methodName,
+        uuid: messageId
+      });
+      const checkState = () => {
+        const messageState = this.asyncMessageQueue[messageId];
+        if (messageState.response !== null) {
+          resolve(messageState.response);
+          delete this.asyncMessageQueue[messageId];
+          return;
+        }
+        if (Date.now() - this.asyncTimeout > messageState.time) {
+          reject(null);
+          delete this.asyncMessageQueue[messageId];
+          return;
+        } else {
+          setTimeout(checkState, this.asyncRefresh);
+        }
+      };
+      checkState();
+    });
+  },
+
   refreshIdentities() {
     return Promise.all([
       browser.contextualIdentities.query({}),
-      browser.runtime.sendMessage({
-        method: "queryIdentitiesState"
-      })
+      this.awaitMessage("queryIdentitiesState")
     ]).then(([identities, state]) => {
       this._identities = identities.map((identity) => {
         const stateObject = state[Logic.userContextId(identity.cookieStoreId)];