Use asyncrhonous messages to deliver SyncMessage payloads on Firefox.

1 files changed, 115 insertions, 86 deletions

diff --git a/src/lib/SyncMessage.js b/src/lib/SyncMessage.js
index 1df9411..de1d784 100644
--- a/src/lib/SyncMessage.js
+++ b/src/lib/SyncMessage.js
@@ -7,14 +7,33 @@
     if (typeof browser.runtime.onSyncMessage !== "object") {
       // Background Script side
 
-      // cache of senders from early async messages to track tab ids in Firefox
       let pending = new Map();
       if (MOZILLA) {
         // we don't care this is async, as long as it get called before the
         // sync XHR (we are not interested in the response on the content side)
         browser.runtime.onMessage.addListener((m, sender) => {
-          if (!m.___syncMessageId) return;
-          pending.set(m.___syncMessageId, sender);
+          let wrapper = m.__syncMessage__;
+          if (!wrapper) return;
+          let {id} = wrapper;
+          pending.set(id, wrapper);
+          let result;
+          let unsuspend = result => {
+            pending.delete(id);
+            if (wrapper.unsuspend) {
+              setTimeout(wrapper.unsuspend, 0);
+            }
+            return result;
+          }
+          try {
+            result = notifyListeners(JSON.stringify(wrapper.payload), sender);
+          } catch(e) {
+            unsuspend();
+            throw e;
+          }
+          console.debug("sendSyncMessage: returning", result);
+          return (result instanceof Promise ? result
+            : new Promise(resolve => resolve(result))
+          ).then(result => unsuspend(result));
         });
       }
 
@@ -26,6 +45,7 @@
 
 
       let obrListener = request => {
+        try {
         let {url, tabId} = request;
         let params = new URLSearchParams(url.split("?")[1]);
         let msgId = params.get("id");
@@ -33,91 +53,73 @@
           return asyncRet(msgId);
         }
         let msg = params.get("msg");
+
+        if (MOZILLA || tabId === TAB_ID_NONE) {
+          // this shoud be a mozilla suspension request
+          return params.get("suspend") ? new Promise(resolve => {
+            if (pending.has(msgId)) {
+              let wrapper = pending.get(msgId);
+              if (!wrapper.unsuspend) {
+                wrapper.unsuspend = resolve;
+                return;
+              }
+            }
+            resolve();
+          }).then(() => ret("go on"))
+          : CANCEL; // otherwise, bail
+        }
+        // CHROME from now on
         let documentUrl = params.get("url");
-        let suspension = !!params.get("suspend");
-        let sender;
-        if (tabId === TAB_ID_NONE || suspension) {
-          // Firefox sends privileged content script XHR without valid tab ids
-          // so we cache sender info from unprivileged XHR correlated by msgId
-          if (pending.has(msgId)) {
-            sender = pending.get(msgId);
-            if (suspension) { // we hold any script execution / DOM modification on this promise
-              return new Promise(resolve => {
-                sender.unsuspend = resolve;
+        let {frameAncestors, frameId} = request;
+        let isTop = frameId === 0 || !!params.get("top");
+        let tabUrl = frameAncestors && frameAncestors.length
+          && frameAncestors[frameAncestors.length - 1].url;
+
+        if (!tabUrl) {
+          if (isTop) {
+            tabUrlCache.set(tabId, tabUrl = documentUrl);
+            if (!tabRemovalListener) {
+              browser.tabs.onRemoved.addListener(tabRemovalListener = tab => {
+                tabUrlCache.delete(tab.id);
               });
             }
-            if (sender.unsuspend) {
-              let {unsuspend} = sender;
-              delete sender.unsuspend;
-              setTimeout(unsuspend(ret("unsuspend")), 0);
-            }
-            pending.delete(msgId);
           } else {
-            throw new Error(`sendSyncMessage: cannot correlate sender info for ${msgId}.`);
-          }
-        } else {
-          let {frameAncestors, frameId} = request;
-          let isTop = frameId === 0 || !!params.get("top");
-          let tabUrl = frameAncestors && frameAncestors.length
-            && frameAncestors[frameAncestors.length - 1].url;
-
-          if (!tabUrl) {
-            if (isTop) {
-              tabUrlCache.set(tabId, tabUrl = documentUrl);
-              if (!tabRemovalListener) {
-                browser.tabs.onRemoved.addListener(tabRemovalListener = tab => {
-                  tabUrlCache.delete(tab.id);
-                });
-              }
-            } else {
-              tabUrl = tabUrlCache.get(tabId);
-            }
+            tabUrl = tabUrlCache.get(tabId);
           }
-          sender = {
-            tab: {
-              id: tabId,
-              url: tabUrl
-            },
-            frameId,
-            url: documentUrl,
-            timeStamp: Date.now()
-          };
         }
+        let sender = {
+          tab: {
+            id: tabId,
+            url: tabUrl
+          },
+          frameId,
+          url: documentUrl,
+          timeStamp: Date.now()
+        };
+
         if (!(msg !== null && sender)) {
           return CANCEL;
         }
-        // Just like in the async runtime.sendMessage() API,
-        // we process the listeners in order until we find a not undefined
-        // result, then we return it (or undefined if none returns anything).
-        let result;
-        for (let l of listeners) {
-          try {
-            if ((result = l(JSON.parse(msg), sender)) !== undefined) break;
-          } catch (e) {
-            console.error("%o processing message %o from %o", e, msg, sender);
-          }
-        }
+        let result = notifyListeners(msg, sender);
         if (result instanceof Promise) {
-          if (MOZILLA) {
-            // Firefox supports asynchronous webRequest listeners, so we can
-            // just defer the return
-            return (async () => ret(await result))();
-          } else {
-            // On Chromium, if the promise is not resolved yet,
-            // we redirect the XHR to the same URL (hence same msgId)
-            // while the result get cached for asynchronous retrieval
-            result.then(r => {
-              asyncResults.set(msgId, result = r);
-            });
-            return asyncResults.has(msgId)
-            ? asyncRet(msgId) // promise was already resolved
-            : {redirectUrl: url.replace(
-                /&redirects=(\d+)|$/, // redirects count to avoid loop detection
-                (all, count) => `&redirects=${parseInt(count) + 1 || 1}`)};
-          }
+
+          // On Chromium, if the promise is not resolved yet,
+          // we redirect the XHR to the same URL (hence same msgId)
+          // while the result get cached for asynchronous retrieval
+          result.then(r => {
+            asyncResults.set(msgId, result = r);
+          });
+          return asyncResults.has(msgId)
+          ? asyncRet(msgId) // promise was already resolved
+          : {redirectUrl: url.replace(
+              /&redirects=(\d+)|$/, // redirects count to avoid loop detection
+              (all, count) => `&redirects=${parseInt(count) + 1 || 1}`)};
         }
         return ret(result);
-      };
+      } catch(e) {
+        console.error(e);
+        return CANCEL;
+      } };
 
       let ret = r => ({redirectUrl:  `data:application/json,${JSON.stringify(r)}`})
       let asyncRet = msgId => {
@@ -127,6 +129,19 @@
       }
 
       let listeners = new Set();
+      function notifyListeners(msg, sender) {
+        // Just like in the async runtime.sendMessage() API,
+        // we process the listeners in order until we find a not undefined
+        // result, then we return it (or undefined if none returns anything).
+        for (let l of listeners) {
+          try {
+            let result = l(JSON.parse(msg), sender);
+            if (result !== undefined) return result;
+          } catch (e) {
+            console.error("%o processing message %o from %o", e, msg, sender);
+          }
+        }
+      }
       browser.runtime.onSyncMessage = {
         ENDPOINT_PREFIX,
         addListener(l) {
@@ -163,30 +178,46 @@
         // about frameAncestors
         url += "&top=true";
       }
-      let finalizers = [];
       if (MOZILLA) {
         // on Firefox we first need to send an async message telling the
         // background script about the tab ID, which does not get sent
         // with "privileged" XHR
-        browser.runtime.sendMessage({___syncMessageId: msgId});
+        let result;
+        browser.runtime.sendMessage(
+          {__syncMessage__: {id: msgId, payload: msg}}
+        ).then(r => {
+          result = r;
+        }).catch(e => {
+          throw e;
+        });
 
         // In order to cope with inconsistencies in XHR synchronicity,
         // allowing DOM element to be inserted and script to be executed
         // (seen with file:// and ftp:// loads) we additionally suspend on
         // Mutation notifications and beforescriptexecute events
-        let suspendURL = url + "&suspend";
+        let suspendURL = url + "&suspend=true";
+        let suspended = false;
         let suspend = () => {
-          let r = new XMLHttpRequest();
-          r.open("GET", url, false);
-          r.send(null);
+          if (result || suspended) return;
+          suspended = true;
+          try {
+            let r = new XMLHttpRequest();
+            r.open("GET", suspendURL, false);
+            r.send(null);
+          } finally {
+            suspended = false;
+          }
         };
         let domSuspender = new MutationObserver(suspend);
         domSuspender.observe(document.documentElement, {childList: true});
         addEventListener("beforescriptexecute", suspend, true);
-        finalizers.push(() => {
+        try {
+          suspend();
+        } finally {
           removeEventListener("beforescriptexecute", suspend, true);
           domSuspender.disconnect();
-        });
+        }
+        return result;
       }
       // then we send the payload using a privileged XHR, which is not subject
       // to CORS but unfortunately doesn't carry any tab id except on Chromium
@@ -199,8 +230,6 @@
         return JSON.parse(r.responseText);
       } catch(e) {
         console.error(`syncMessage error in ${document.URL}: ${e.message} (response ${r.responseText})`);
-      } finally {
-        for (let f of finalizers) try { f(); } catch(e) { console.error(e); }
       }
       return null;
     };