Bug 1248497 – Add promise support to the sendMessage APIs. r=billm

MozReview-Commit-ID: AZH9LUq8kGr

browser/components/extensions/test/browser/browser_ext_tabs_sendMessage.js

@@ -2,6 +2,95 @@
 /* vim: set sts=2 sw=2 et tw=80: */
 "use strict";
 
+add_task(function* tabsSendMessageReply() {
+  let extension = ExtensionTestUtils.loadExtension({
+    manifest: {
+      "permissions": ["tabs"],
+
+      "content_scripts": [{
+        "matches": ["http://example.com/"],
+        "js": ["content-script.js"],
+        "run_at": "document_start",
+      }],
+    },
+
+    background: function() {
+      let promiseResponse = new Promise(resolve => {
+        browser.runtime.onMessage.addListener((msg, sender, respond) => {
+          if (msg == "content-script-ready") {
+            let tabId = sender.tab.id;
+
+            browser.tabs.sendMessage(tabId, "respond-never", response => {
+              browser.test.fail("Got unexpected response callback");
+              browser.test.notifyFail("sendMessage");
+            });
+
+            Promise.all([
+              promiseResponse,
+              browser.tabs.sendMessage(tabId, "respond-now"),
+              new Promise(resolve => browser.tabs.sendMessage(tabId, "respond-soon", resolve)),
+              browser.tabs.sendMessage(tabId, "respond-promise"),
+              browser.tabs.sendMessage(tabId, "respond-never"),
+              browser.tabs.sendMessage(tabId, "respond-error").catch(error => Promise.resolve({error})),
+              browser.tabs.sendMessage(tabId, "throw-error").catch(error => Promise.resolve({error})),
+            ]).then(([response, respondNow, respondSoon, respondPromise, respondNever, respondError, throwError]) => {
+              browser.test.assertEq("expected-response", response, "Content script got the expected response");
+
+              browser.test.assertEq("respond-now", respondNow, "Got the expected immediate response");
+              browser.test.assertEq("respond-soon", respondSoon, "Got the expected delayed response");
+              browser.test.assertEq("respond-promise", respondPromise, "Got the expected promise response");
+              browser.test.assertEq(undefined, respondNever, "Got the expected no-response resolution");
+
+              browser.test.assertEq("respond-error", respondError.error.message, "Got the expected error response");
+              browser.test.assertEq("throw-error", throwError.error.message, "Got the expected thrown error response");
+
+              return browser.tabs.remove(tabId);
+            }).then(() => {
+              browser.test.notifyPass("sendMessage");
+            });
+
+            return Promise.resolve("expected-response");
+          } else if (msg[0] == "got-response") {
+            resolve(msg[1]);
+          }
+        });
+      });
+
+      browser.tabs.create({url: "http://example.com/"});
+    },
+
+    files: {
+      "content-script.js": function() {
+        browser.runtime.onMessage.addListener((msg, sender, respond) => {
+          if (msg == "respond-now") {
+            respond(msg);
+          } else if (msg == "respond-soon") {
+            setTimeout(() => { respond(msg); }, 0);
+            return true;
+          } else if (msg == "respond-promise") {
+            return Promise.resolve(msg);
+          } else if (msg == "respond-never") {
+            return;
+          } else if (msg == "respond-error") {
+            return Promise.reject(new Error(msg));
+          } else if (msg == "throw-error") {
+            throw new Error(msg);
+          }
+        });
+        browser.runtime.sendMessage("content-script-ready").then(response => {
+          browser.runtime.sendMessage(["got-response", response]);
+        });
+      },
+    },
+  });
+
+  yield extension.startup();
+
+  yield extension.awaitFinish("sendMessage");
+
+  yield extension.unload();
+});
+
 add_task(function* tabsSendMessageNoExceptionOnNonExistentTab() {
   let extension = ExtensionTestUtils.loadExtension({
     manifest: {

toolkit/components/extensions/ExtensionContent.jsm

@@ -99,7 +99,7 @@ var api = context => {
         }
 
         let recipient = extensionId ? {extensionId} : {extensionId: context.extensionId};
-        context.messenger.sendMessage(context.messageManager, message, recipient, responseCallback);
+        return context.messenger.sendMessage(context.messageManager, message, recipient, responseCallback);
       },
     },
 

toolkit/components/extensions/ExtensionUtils.jsm

@@ -862,25 +862,34 @@ Messenger.prototype = {
     recipient.messageId = id;
     this.broker.sendMessage(messageManager, "message", msg, this.sender, recipient);
 
-    let onClose;
-    let listener = ({data: response}) => {
-      messageManager.removeMessageListener(replyName, listener);
-      this.context.forgetOnClose(onClose);
-
-      if (response.gotData) {
-        // TODO: Handle failure to connect to the extension?
-        runSafe(this.context, responseCallback, response.data);
-      }
-    };
-    onClose = {
-      close() {
+    let promise = new Promise((resolve, reject) => {
+      let onClose;
+      let listener = ({data: response}) => {
         messageManager.removeMessageListener(replyName, listener);
-      },
-    };
-    if (responseCallback) {
+        this.context.forgetOnClose(onClose);
+
+        if (response.gotData) {
+          resolve(response.data);
+        } else if (response.error) {
+          reject(response.error);
+        } else if (!responseCallback) {
+          // As a special case, we don't call the callback variant if we
+          // receive no response, but the promise needs to resolve or
+          // reject in either case.
+          resolve();
+        }
+      };
+      onClose = {
+        close() {
+          messageManager.removeMessageListener(replyName, listener);
+        },
+      };
+
       messageManager.addMessageListener(replyName, listener);
       this.context.callOnClose(onClose);
-    }
+    });
+
+    return this.context.wrapPromise(promise, responseCallback);
   },
 
   onMessage(name) {
@@ -895,23 +904,33 @@ Messenger.prototype = {
         let mm = getMessageManager(target);
         let replyName = `Extension:Reply-${recipient.messageId}`;
 
-        let valid = true, sent = false;
-        let sendResponse = data => {
-          if (!valid) {
-            return;
-          }
-          sent = true;
-          mm.sendAsyncMessage(replyName, {data, gotData: true});
-        };
-        sendResponse = Cu.exportFunction(sendResponse, this.context.cloneScope);
+        new Promise((resolve, reject) => {
+          let sendResponse = Cu.exportFunction(resolve, this.context.cloneScope);
 
-        let result = runSafeSyncWithoutClone(callback, message, sender, sendResponse);
-        if (result !== true) {
-          valid = false;
-          if (!sent) {
-            mm.sendAsyncMessage(replyName, {gotData: false});
+          // Note: We intentionally do not use runSafe here so that any
+          // errors are propagated to the message sender.
+          let result = callback(message, sender, sendResponse);
+          if (result instanceof Promise) {
+            resolve(result);
+          } else if (result !== true) {
+            reject();
           }
-        }
+        }).then(
+          data => {
+            mm.sendAsyncMessage(replyName, {data, gotData: true});
+          },
+          error => {
+            if (error) {
+              // The result needs to be structured-clonable, which
+              // ordinary Error objects are not.
+              try {
+                error = {message: String(error.message), stack: String(error.stack)};
+              } catch (e) {
+                error = {message: String(error)};
+              }
+            }
+            mm.sendAsyncMessage(replyName, {error, gotData: false});
+          });
       };
 
       this.broker.addListener("message", listener, this.filter);