Issue 4579 - Wrap runtime.onMessage

polyfill.js

Index: polyfill.js
===================================================================
--- a/polyfill.js
+++ b/polyfill.js
@@ -48,16 +48,22 @@
   // Since we add a callback for all messaging API calls in our wrappers,
   // Chrome assumes we're interested in the response; when there's no response,
   // it sets runtime.lastError
   const portClosedBeforeResponseError =
     // Older versions of Chrome have a typo:
     // https://crrev.com/c33f51726eacdcc1a487b21a13611f7eab580d6d
     /^The message port closed before a res?ponse was received\.$/;
 
+  // This is the error Firefox throws when a message listener is not a
+  // function.
+  const invalidMessageListenerError = "Invalid listener for runtime.onMessage.";
+
+  let messageListeners = new WeakMap();
+
   function wrapAPI(api)
   {
     let object = browser;
     let path = api.split(".");
     let name = path.pop();
 
     for (let node of path)
     {
@@ -65,19 +71,22 @@
 
       if (!object)
         return;
     }
 
     let func = object[name];
     if (!func)
       return;
+
     let descriptor = Object.getOwnPropertyDescriptor(object, name);
+
     delete descriptor["get"];
     delete descriptor["set"];
+
     descriptor.value = function(...args)
     {
       let callStack = new Error().stack;
 
       if (typeof args[args.length - 1] == "function")
         return func.apply(object, args);
 
       // If the last argument is undefined, we drop it from the list assuming
@@ -109,19 +118,84 @@
           }
           else
           {
             resolve(result);
           }
         });
       });
     };
+
     Object.defineProperty(object, name, descriptor);
   }
 
+  function wrapMessageAPIs()
+  {
+    let {onMessage} = browser.runtime;
+    let {addListener, removeListener, hasListener} = onMessage;
+
+    onMessage.addListener = function(listener)
+    {
+      if (typeof listener != "function")
+        throw new Error(invalidMessageListenerError);
+
+      // Don't add the same listener twice or we end up with multiple wrappers.
+      if (messageListeners.has(listener))
+        return;
+
+      let wrapper = (message, sender, sendResponse) =>
+      {
+        let wait = listener(message, sender, sendResponse);
+
+        if (wait instanceof Promise)
+        {
+          wait.then(sendResponse, reason =>
+          {
+            try
+            {
+              sendResponse();
+            }
+            finally
+            {
+              // sendResponse can throw if the internal port is closed; be sure
+              // to throw the original error.
+              throw reason;
+            }
+          });
+        }
+
+        return !!wait;
+      };
+
+      addListener.call(onMessage, wrapper);
+      messageListeners.set(listener, wrapper);
+    };
+
+    onMessage.removeListener = function(listener)
+    {
+      if (typeof listener != "function")
+        throw new Error(invalidMessageListenerError);
+
+      let wrapper = messageListeners.get(listener);
+      if (wrapper)
+      {
+        removeListener.call(onMessage, wrapper);
+        messageListeners.delete(listener);
+      }
+    };
+
+    onMessage.hasListener = function(listener)
+    {
+      if (typeof listener != "function")
+        throw new Error(invalidMessageListenerError);
+
+      return messageListeners.has(listener);
+    };
+  }
+
   function shouldWrapAPIs()
   {
     try
     {
       return !(browser.storage.local.get([]) instanceof Promise);
     }
     catch (error)
     {
@@ -135,16 +209,18 @@
     // Unlike Firefox and Microsoft Edge, Chrome doesn't have a "browser"
     // object, but provides the extension API through the "chrome" namespace
     // (non-standard).
     if (typeof browser == "undefined")
       window.browser = chrome;
 
     for (let api of asyncAPIs)
       wrapAPI(api);
+
+    wrapMessageAPIs();
   }
 
   // Workaround since HTMLCollection, NodeList, StyleSheetList, and CSSRuleList
   // didn't have iterator support before Chrome 51.
   // https://bugs.chromium.org/p/chromium/issues/detail?id=401699
   for (let object of [HTMLCollection, NodeList, StyleSheetList, CSSRuleList])
   {
     if (!(Symbol.iterator in object.prototype))