Bug 1643214 - Wait for explict responses from runtime.onMessage r=mixedpuppy

When there are `runtime.onMessage` listeners in multiple extension contexts, and one promises a response, make sure we wait for that instead of resolving with the first `undefined` from the other context which didn't explicitly respond.

Differential Revision: https://phabricator.services.mozilla.com/D78652

toolkit/components/extensions/ConduitsParent.jsm

@@ -275,16 +275,24 @@ class BroadcastConduit extends BaseConduit {
    */
   _raceResponses(promises) {
     return new Promise((resolve, reject) => {
+      let result;
       promises.map(p =>
         p
-          // Ignore responses from child Messengers without any listeners.
-          .then(value => value && resolve(value))
-
+          .then(value => {
+            if (value.response) {
+              // We have an explicit response, resolve immediately.
+              resolve(value);
+            } else if (value.received) {
+              // Message was received, but no response.
+              // Resolve with this only if there is no other explicit response.
+              result = value;
+            }
+          })
           // Ignore errors trying to query child Messengers being destroyed.
           .catch(err => err.result !== Cr.NS_ERROR_NOT_AVAILABLE && reject(err))
       );
-      // Ensure resolveing when there are no (actual) responses.
-      Promise.allSettled(promises).then(() => resolve());
+      // Ensure resolving when there are no responses.
+      Promise.allSettled(promises).then(() => resolve(result));
     });
   }
 

toolkit/components/extensions/ExtensionChild.jsm

@@ -171,17 +171,19 @@ class SimpleEventAPI extends EventManager {
 class MessageEvent extends SimpleEventAPI {
   emit(holder, sender) {
     if (!this.fires.size || !this.context.active) {
-      return;
+      return { received: false };
     }
 
     sender = Cu.cloneInto(sender, this.context.cloneScope);
     let message = holder.deserialize(this.context.cloneScope);
 
-    let all = [...this.fires]
+    let responses = [...this.fires]
       .map(fire => this.wrapResponse(fire, message, sender))
       .filter(x => x !== undefined);
 
-    return all.length ? Promise.race(all).then(v => [v]) : [];
+    return !responses.length
+      ? { received: true, response: false }
+      : Promise.race(responses).then(value => ({ response: true, value }));
   }
 
   wrapResponse(fire, message, sender) {

toolkit/components/extensions/ExtensionParent.jsm

@@ -329,7 +329,9 @@ const ProxyMessenger = {
     arg.firstResponse = true;
     let kind = await this.normalizeArgs(arg, sender);
     let result = await this.conduit.castRuntimeMessage(kind, arg);
-    return result ? result[0] : Promise.reject({ message: ERROR_NO_RECEIVERS });
+    return result
+      ? result.value
+      : Promise.reject({ message: ERROR_NO_RECEIVERS });
   },
 
   async recvPortConnect(arg, { sender }) {

toolkit/components/extensions/test/xpcshell/test_ext_runtime_sendMessage.js

@@ -2,6 +2,10 @@
 /* vim: set sts=2 sw=2 et tw=80: */
 "use strict";
 
+const server = createHttpServer();
+server.registerDirectory("/data/", do_get_file("data"));
+const BASE_URL = `http://localhost:${server.identity.primaryPort}/data`;
+
 add_task(async function runtimeSendMessageReply() {
   function background() {
     browser.runtime.onMessage.addListener((msg, sender, respond) => {
@@ -309,3 +313,84 @@ add_task(async function sendMessageResponseGC() {
   ok("Long running tasks responded");
   await extension.unload();
 });
+
+add_task(async function sendMessage_async_response_multiple_contexts() {
+  let extension = ExtensionTestUtils.loadExtension({
+    background() {
+      browser.runtime.onMessage.addListener((msg, _, respond) => {
+        browser.test.log(`Background got request: ${msg}`);
+
+        switch (msg) {
+          case "ask-bg-fast":
+            respond("bg-respond");
+            return true;
+
+          case "ask-bg-slow":
+            return new Promise(r => setTimeout(() => r("bg-promise")), 1000);
+        }
+      });
+      browser.test.sendMessage("bg-ready");
+    },
+
+    manifest: {
+      content_scripts: [
+        {
+          matches: ["http://localhost/*/file_sample.html"],
+          js: ["cs.js"],
+        },
+      ],
+    },
+
+    files: {
+      "page.html":
+        "<!DOCTYPE html><meta charset=utf-8><script src=page.js></script>",
+      "page.js"() {
+        browser.runtime.onMessage.addListener((msg, _, respond) => {
+          browser.test.log(`Page got request: ${msg}`);
+
+          switch (msg) {
+            case "ask-page-fast":
+              respond("page-respond");
+              return true;
+
+            case "ask-page-slow":
+              return new Promise(r => setTimeout(() => r("page-promise")), 500);
+          }
+        });
+        browser.test.sendMessage("page-ready");
+      },
+
+      "cs.js"() {
+        Promise.all([
+          browser.runtime.sendMessage("ask-bg-fast"),
+          browser.runtime.sendMessage("ask-bg-slow"),
+          browser.runtime.sendMessage("ask-page-fast"),
+          browser.runtime.sendMessage("ask-page-slow"),
+        ]).then(responses => {
+          browser.test.assertEq(
+            responses.join(),
+            ["bg-respond", "bg-promise", "page-respond", "page-promise"].join(),
+            "Got all expected responses from correct contexts"
+          );
+          browser.test.notifyPass("cs-done");
+        });
+      },
+    },
+  });
+
+  await extension.startup();
+  await extension.awaitMessage("bg-ready");
+
+  let url = `moz-extension://${extension.uuid}/page.html`;
+  let page = await ExtensionTestUtils.loadContentPage(url, { extension });
+  await extension.awaitMessage("page-ready");
+
+  let content = await ExtensionTestUtils.loadContentPage(
+    BASE_URL + "/file_sample.html"
+  );
+  await extension.awaitFinish("cs-done");
+  await content.close();
+
+  await page.close();
+  await extension.unload();
+});