1 files changed, 198 insertions, 0 deletions
diff --git a/toolkit/components/antitracking/test/browser/browser_staticPartition_websocket.js b/toolkit/components/antitracking/test/browser/browser_staticPartition_websocket.js
new file mode 100644
index 0000000000..11f53d69ed
--- /dev/null
+++ b/toolkit/components/antitracking/test/browser/browser_staticPartition_websocket.js
@@ -0,0 +1,198 @@
+/* Any copyright is dedicated to the Public Domain.
+ * http://creativecommons.org/publicdomain/zero/1.0/
+ */
+
+"use strict";
+
+const FIRST_PARTY_A = "http://example.com";
+const FIRST_PARTY_B = "http://example.org";
+const THIRD_PARTY = "http://example.net";
+const WS_ENDPOINT_HOST = "mochi.test:8888";
+
+function getWSTestUrlForHost(host) {
+  return (
+    getRootDirectory(gTestPath).replace(
+      "chrome://mochitests/content",
+      `ws://${host}`
+    ) + `file_ws_handshake_delay`
+  );
+}
+
+function connect(browsingContext, host, protocol) {
+  let url = getWSTestUrlForHost(host);
+  info("Creating websocket with endpoint " + url);
+
+  // Create websocket connection in third party iframe.
+  let createPromise = SpecialPowers.spawn(
+    browsingContext.children[0],
+    [url, protocol],
+    (url, protocol) => {
+      let ws = new content.WebSocket(url, [protocol]);
+      ws.addEventListener("error", () => {
+        ws._testError = true;
+      });
+      if (!content.ws) {
+        content.ws = {};
+      }
+      content.ws[protocol] = ws;
+    }
+  );
+
+  let openPromise = createPromise.then(() =>
+    SpecialPowers.spawn(
+      browsingContext.children[0],
+      [protocol],
+      async protocol => {
+        let ws = content.ws[protocol];
+        if (ws.readyState != 0) {
+          return !ws._testError;
+        }
+        // Still connecting.
+        let result = await Promise.race([
+          ContentTaskUtils.waitForEvent(ws, "open"),
+          ContentTaskUtils.waitForEvent(ws, "error"),
+        ]);
+        return result.type != "error";
+      }
+    )
+  );
+
+  let result = { createPromise, openPromise };
+  return result;
+}
+
+// Open 3 websockets which target the same ip/port combination, but have
+// different principals. We send a protocol identifier to the server to signal
+// how long the request should be delayed.
+//
+// When partitioning is disabled A blocks B and B blocks C. The timeline will
+// look like this:
+// A________
+//          B____
+//               C_
+//
+// When partitioning is enabled, connection handshakes for A and B will run
+// (semi-) parallel since they have different origin attributes. B and C share
+// origin attributes and therefore still run serially.
+// A________
+// B____
+//      C_
+//
+// By observing the order of the handshakes we can ensure that the queue
+// partitioning is working correctly.
+async function runTest(partitioned) {
+  await SpecialPowers.pushPrefEnv({
+    set: [["privacy.partition.network_state", partitioned]],
+  });
+
+  let tabA = BrowserTestUtils.addTab(gBrowser, FIRST_PARTY_A);
+  await BrowserTestUtils.browserLoaded(tabA.linkedBrowser);
+  let tabB = BrowserTestUtils.addTab(gBrowser, FIRST_PARTY_B);
+  await BrowserTestUtils.browserLoaded(tabB.linkedBrowser);
+
+  for (let tab of [tabA, tabB]) {
+    await SpecialPowers.spawn(tab.linkedBrowser, [THIRD_PARTY], async src => {
+      let frame = content.document.createElement("iframe");
+      frame.src = src;
+      let loadPromise = ContentTaskUtils.waitForEvent(frame, "load");
+      content.document.body.appendChild(frame);
+      await loadPromise;
+    });
+  }
+
+  // First ensure that we can open websocket connections to the test endpoint.
+  let { openPromise, createPromise } = await connect(
+    tabA.linkedBrowser.browsingContext,
+    WS_ENDPOINT_HOST,
+    false
+  );
+  await createPromise;
+  let openPromiseResult = await openPromise;
+  ok(openPromiseResult, "Websocket endpoint accepts connections.");
+
+  let openedA;
+  let openedB;
+  let openedC;
+
+  let { createPromise: createPromiseA, openPromise: openPromiseA } = connect(
+    tabA.linkedBrowser.browsingContext,
+    WS_ENDPOINT_HOST,
+    "test-6"
+  );
+
+  openPromiseA = openPromiseA.then(opened => {
+    openedA = opened;
+    info("Completed WS connection A");
+    if (partitioned) {
+      ok(openedA, "Should have opened A");
+      ok(openedB, "Should have opened B");
+    } else {
+      ok(openedA, "Should have opened A");
+      ok(openedB == null, "B should be pending");
+    }
+  });
+  await createPromiseA;
+
+  // The frame of connection B is embedded in a different first party as A.
+  let { createPromise: createPromiseB, openPromise: openPromiseB } = connect(
+    tabB.linkedBrowser.browsingContext,
+    WS_ENDPOINT_HOST,
+    "test-3"
+  );
+  openPromiseB = openPromiseB.then(opened => {
+    openedB = opened;
+    info("Completed WS connection B");
+    if (partitioned) {
+      ok(openedA == null, "A should be pending");
+      ok(openedB, "Should have opened B");
+      ok(openedC == null, "C should be pending");
+    } else {
+      ok(openedA, "Should have opened A");
+      ok(openedB, "Should have opened B");
+      ok(openedC == null, "C should be pending");
+    }
+  });
+  await createPromiseB;
+
+  // The frame of connection C is embedded in the same first party as B.
+  let { createPromise: createPromiseC, openPromise: openPromiseC } = connect(
+    tabB.linkedBrowser.browsingContext,
+    WS_ENDPOINT_HOST,
+    "test-0"
+  );
+  openPromiseC = openPromiseC.then(opened => {
+    openedC = opened;
+    info("Completed WS connection C");
+    if (partitioned) {
+      ok(openedB, "Should have opened B");
+      ok(openedC, "Should have opened C");
+    } else {
+      ok(opened, "Should have opened B");
+      ok(opened, "Should have opened C");
+    }
+  });
+  await createPromiseC;
+
+  // Wait for all connections to complete before closing the tabs.
+  await Promise.all([openPromiseA, openPromiseB, openPromiseC]);
+
+  BrowserTestUtils.removeTab(tabA);
+  BrowserTestUtils.removeTab(tabB);
+
+  await SpecialPowers.popPrefEnv();
+}
+
+add_setup(async function () {
+  // This test relies on a WS connection timeout > 6 seconds.
+  await SpecialPowers.pushPrefEnv({
+    set: [["network.websocket.timeout.open", 20]],
+  });
+});
+
+add_task(async function test_non_partitioned() {
+  await runTest(false);
+});
+
+add_task(async function test_partitioned() {
+  await runTest(true);
+});