Issue 4579 - Wrap runtime.onMessage

polyfill.js

 /*
  * This file is part of Adblock Plus <https://adblockplus.org/>,
  * Copyright (C) 2006-present eyeo GmbH
  *
  * Adblock Plus is free software: you can redistribute it and/or modify
  * it under the terms of the GNU General Public License version 3 as
  * published by the Free Software Foundation.
  *
  * Adblock Plus is distributed in the hope that it will be useful,
  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  * GNU General Public License for more details.
  *
  * You should have received a copy of the GNU General Public License
  * along with Adblock Plus.  If not, see <http://www.gnu.org/licenses/>.
  */
 
 "use strict";
 
 {
   const asyncAPIs = [
     "contextMenus.removeAll",
     "devtools.panels.create",
     "notifications.clear",
     "notifications.create",
+    "runtime.getBrowserInfo",
     "runtime.openOptionsPage",
     "runtime.sendMessage",
     "runtime.setUninstallURL",
     "storage.local.get",
     "storage.local.remove",
     "storage.local.set",
     "storage.managed.get",
     "tabs.create",
     "tabs.get",
     "tabs.getCurrent",
     "tabs.insertCSS",
     "tabs.query",
     "tabs.reload",
     "tabs.sendMessage",
     "tabs.update",
     "webNavigation.getAllFrames",
     "webRequest.handlerBehaviorChanged",
     "windows.create",
     "windows.update"
   ];
 
   // 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\.$/;
 
-  function wrapAPI(api)
+  // 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 wrapAsyncAPI(api)
   {
     let object = browser;
     let path = api.split(".");
     let name = path.pop();
 
     for (let node of path)
     {
       object = object[node];
 
       if (!object)
         return;
     }
 
     let func = object[name];
-    object[name] = function(...args)
+    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
       // it stands for the optional callback. We must do this, because we have
       // to replace it with our own callback. If we simply append our own
       // callback to the list, it won't match the signature of the function and
@@ skipping 20 matching lines @@
 
             reject(error);
           }
           else
           {
             resolve(result);
           }
         });
       });
     };
+
+    Object.defineProperty(object, name, descriptor);
+  }
+
+  function wrapRuntimeOnMessage()
+  {
+    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)
     {
     }
 
     return true;
   }
 
   if (shouldWrapAPIs())
   {
     // Unlike Firefox and Microsoft Edge, Chrome doesn't have a "browser"

[Manish Jethani, 2017/10/23 02:32:45]

Unrelated, but there's an ESLint error here because of line length.

     // 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);
-
-    let {addListener} = browser.runtime.onMessage;
-    browser.runtime.onMessage.addListener = function(listener)
-    {
-      addListener.call(browser.runtime.onMessage, (
-        message, sender, sendResponse
-      ) =>
-      {
-        let wait = listener(message, sender, sendResponse);
-
-        if (wait instanceof Promise)

[Manish Jethani, 2017/10/23 02:32:45]

There's no need to add a catch handler, Chrome (and probably Edge) handles this
automagically. The error will be logged to the console, the sender will get
undefined as a response.

Firefox actually doesn't send an undefined response to the sender, instead it
sends an "unexpected error" in the sender's catch handler. We cannot mimic this
part, unfortunately. We could probably override the Firefox version of the API
to behave like Chrome's (send an undefined response), but I don't know if it
matters so much.

[Wladimir Palant, 2017/10/23 08:53:26]

Not really, the message clearly says that not handling promise rejects is
deprecated. So you should handle rejections explicitly. In particular, you need
to make sure that a response is always sent and be it undefined.
We *could*, by adjusting the behavior of the sendMessage wrapper so that it
recognizes some special values as rejections. Probably not worth doing however.

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

Did you mean that Chrome will no longer do this automatically in future
versions?

Done.
Ack.

-          wait.then(sendResponse);
-
-        return !!wait;
-      });
-    };

[Wladimir Palant, 2017/10/23 08:53:26]

You have a problem here, removeListener will no longer work. The best solution
should be a WeakMap mapping wrapers to original listeners. You will have to wrap
removeListener then and remove the wrapper if you have it in the map. Ideally,
the same should be done with hasListener.

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

Done.

+      wrapAsyncAPI(api);
+
+    wrapRuntimeOnMessage();
   }
 
   // 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))
       object.prototype[Symbol.iterator] = Array.prototype[Symbol.iterator];
   }
 }