Issue 4579 - Wrap runtime.onMessage

polyfill.js

Index: polyfill.js
===================================================================
--- a/polyfill.js
+++ b/polyfill.js
@@ -47,16 +47,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)

[kzar, 2018/01/08 14:35:16]

Since we're now wrapping the message APIs as well perhaps we should rename this
to `wrapAsyncAPI` which would be consistent with the `asyncAPIs` variable name
as well?

[Manish Jethani, 2018/01/09 09:37:35]

Done.

   {
     let object = browser;
     let path = api.split(".");
     let name = path.pop();
 
     for (let node of path)
     {
@@ -105,39 +111,105 @@
           {
             resolve(result);
           }
         });
       });
     };
   }
 
+  function wrapMessageAPIs()

[kzar, 2018/01/08 14:35:16]

Nit: `wrapRuntimeOnMessage`?

[Manish Jethani, 2018/01/09 09:37:35]

Done.

+  {
+    let {onMessage} = browser.runtime;
+    let {addListener, removeListener, hasListener} = onMessage;
+
+    onMessage.addListener = function(listener)
+    {
+      if (typeof listener != "function")

[Manish Jethani, 2017/10/23 12:21:04]

We must check the type here to be consistent with Firefox.

Note that typeof also works for Proxy instances.

[kzar, 2018/01/08 14:35:16]

Acknowledged.

+        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

[Manish Jethani, 2017/10/23 12:21:04]

I haven't verified this, but in any case adding a finally here seems like a good
idea.

[kzar, 2018/01/08 14:35:16]

Acknowledged.

+              // 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)
     {
     }
 
     return true;
   }
 
   if (shouldWrapAPIs())
   {
-    // Unlike Firefox and Microsoft Edge, Chrome doesn't have a "browser" object,
-    // but provides the extension API through the "chrome" namespace
+    // 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))