Bug 1301837 - Defer SwapDocShells event registration in MessageManagerProxy r=rpl

The "SwapDocShells" event should be deferred until "EndSwapDocShells".
Otherwise the event MessageManagerProxy may swap the event listeners
twice, and end up having the listeners on the incorrect message manager.

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

--HG--
extra : moz-landing-system : lando

toolkit/components/extensions/MessageManagerProxy.jsm

@@ -192,7 +192,12 @@ class MessageManagerProxy {
   handleEvent(event) {
     if (event.type == "SwapDocShells") {
       this.removeListeners(this.eventTarget);
-      this.addListeners(event.detail);
+      // The SwapDocShells event is dispatched for both browsers that are being
+      // swapped. To avoid double-swapping, register the event handler after
+      // both SwapDocShells events have fired.
+      this.eventTarget.addEventListener("EndSwapDocShells", () => {
+        this.addListeners(event.detail);
+      }, {once: true});
     }
   }
 }

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

@@ -0,0 +1,73 @@
+/* -*- Mode: indent-tabs-mode: nil; js-indent-level: 2 -*- */
+/* vim: set sts=2 sw=2 et tw=80: */
+"use strict";
+
+const {MessageManagerProxy} = ChromeUtils.import("resource://gre/modules/MessageManagerProxy.jsm");
+const {PromiseUtils} = ChromeUtils.import("resource://gre/modules/PromiseUtils.jsm");
+
+class TestMessageManagerProxy extends MessageManagerProxy {
+  constructor(contentPage, identifier) {
+    super(contentPage.browser);
+    this.identifier = identifier;
+    this.contentPage = contentPage;
+    this.deferred = null;
+  }
+
+  // Registers message listeners. Call dispose() once you've finished.
+  async setupPingPongListeners() {
+    await this.contentPage.loadFrameScript(`() => {
+      this.addMessageListener("test:MessageManagerProxy:Ping", ({data}) => {
+        this.sendAsyncMessage("test:MessageManagerProxy:Pong", "${this.identifier}:" + data);
+      });
+    }`);
+
+    // Register the listener here instead of during testPingPong, to make sure
+    // that the listener is correctly registered during the whole test.
+    this.addMessageListener("test:MessageManagerProxy:Pong", event => {
+      ok(this.deferred, `[${this.identifier}] expected to be waiting for ping-pong`);
+      this.deferred.resolve(event.data);
+      this.deferred = null;
+    });
+  }
+
+  async testPingPong(description) {
+    equal(this.deferred, null, "should not be waiting for a message");
+    this.deferred = PromiseUtils.defer();
+    this.sendAsyncMessage("test:MessageManagerProxy:Ping", description);
+    let result = await this.deferred.promise;
+    equal(result, `${this.identifier}:${description}`, "Expected ping-pong");
+  }
+}
+
+// Tests that MessageManagerProxy continues to proxy messages after docshells
+// have been swapped.
+add_task(async function test_message_after_swapdocshells() {
+  let page1 = await ExtensionTestUtils.loadContentPage("about:blank");
+  let page2 = await ExtensionTestUtils.loadContentPage("about:blank");
+
+  let testProxyOne = new TestMessageManagerProxy(page1, "page1");
+  let testProxyTwo = new TestMessageManagerProxy(page2, "page2");
+
+  await testProxyOne.setupPingPongListeners();
+  await testProxyTwo.setupPingPongListeners();
+
+  await testProxyOne.testPingPong("after setup (to 1)");
+  await testProxyTwo.testPingPong("after setup (to 2)");
+
+  page1.browser.swapDocShells(page2.browser);
+
+  await testProxyOne.testPingPong("after docshell swap (to 1)");
+  await testProxyTwo.testPingPong("after docshell swap (to 2)");
+
+  // Swap again to verify that listeners are repeatedly moved when needed.
+  page1.browser.swapDocShells(page2.browser);
+
+  await testProxyOne.testPingPong("after another docshell swap (to 1)");
+  await testProxyTwo.testPingPong("after another docshell swap (to 2)");
+
+  // Verify that dispose() works regardless of the browser's validity.
+  await testProxyOne.dispose();
+  await page1.close();
+  await page2.close();
+  await testProxyTwo.dispose();
+});

toolkit/components/extensions/test/xpcshell/xpcshell-common.ini

@@ -1,3 +1,5 @@
+[test_ext_MessageManagerProxy.js]
+skip-if = os == 'android' # Bug 1545439
 [test_ext_alarms.js]
 [test_ext_alarms_does_not_fire.js]
 [test_ext_alarms_periodic.js]