Bug 1870848 - [remote] Introduce shared webdriver helper UserContextManager r=webdriver-reviewers,whimboo

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

remote/jar.mn

@@ -31,12 +31,14 @@ remote.jar:
   content/shared/Stack.sys.mjs (shared/Stack.sys.mjs)
   content/shared/Sync.sys.mjs (shared/Sync.sys.mjs)
   content/shared/TabManager.sys.mjs (shared/TabManager.sys.mjs)
+  content/shared/UserContextManager.sys.mjs (shared/UserContextManager.sys.mjs)
   content/shared/UUID.sys.mjs (shared/UUID.sys.mjs)
   content/shared/WebSocketConnection.sys.mjs (shared/WebSocketConnection.sys.mjs)
   content/shared/WindowManager.sys.mjs (shared/WindowManager.sys.mjs)
   content/shared/listeners/BrowsingContextListener.sys.mjs (shared/listeners/BrowsingContextListener.sys.mjs)
   content/shared/listeners/ConsoleAPIListener.sys.mjs (shared/listeners/ConsoleAPIListener.sys.mjs)
   content/shared/listeners/ConsoleListener.sys.mjs (shared/listeners/ConsoleListener.sys.mjs)
+  content/shared/listeners/ContextualIdentityListener.sys.mjs (shared/listeners/ContextualIdentityListener.sys.mjs)
   content/shared/listeners/LoadListener.sys.mjs (shared/listeners/LoadListener.sys.mjs)
   content/shared/listeners/NavigationListener.sys.mjs (shared/listeners/NavigationListener.sys.mjs)
   content/shared/listeners/NetworkEventRecord.sys.mjs (shared/listeners/NetworkEventRecord.sys.mjs)

remote/shared/UserContextManager.sys.mjs

@@ -0,0 +1,180 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+const lazy = {};
+
+ChromeUtils.defineESModuleGetters(lazy, {
+  ContextualIdentityService:
+    "resource://gre/modules/ContextualIdentityService.sys.mjs",
+
+  ContextualIdentityListener:
+    "chrome://remote/content/shared/listeners/ContextualIdentityListener.sys.mjs",
+  generateUUID: "chrome://remote/content/shared/UUID.sys.mjs",
+});
+
+/**
+ * A UserContextManager instance keeps track of all public user contexts and
+ * maps their internal platform.
+ *
+ * This class is exported for test purposes. Otherwise the UserContextManager
+ * singleton should be used.
+ */
+export class UserContextManagerClass {
+  #contextualIdentityListener;
+  #userContextIds;
+
+  DEFAULT_CONTEXT_ID = "default";
+  DEFAULT_INTERNAL_ID = 0;
+
+  constructor() {
+    // Map from internal ids (numbers) from the ContextualIdentityService to
+    // opaque UUIDs (string).
+    this.#userContextIds = new Map();
+
+    // The default user context is always using 0 as internal user context id
+    // and should be exposed as "default" instead of a randomly generated id.
+    this.#userContextIds.set(this.DEFAULT_INTERNAL_ID, this.DEFAULT_CONTEXT_ID);
+
+    // Register other (non-default) public contexts.
+    lazy.ContextualIdentityService.getPublicIdentities().forEach(identity =>
+      this.#registerIdentity(identity)
+    );
+
+    this.#contextualIdentityListener = new lazy.ContextualIdentityListener();
+    this.#contextualIdentityListener.on("created", this.#onIdentityCreated);
+    this.#contextualIdentityListener.on("deleted", this.#onIdentityDeleted);
+    this.#contextualIdentityListener.startListening();
+  }
+
+  destroy() {
+    this.#contextualIdentityListener.off("created", this.#onIdentityCreated);
+    this.#contextualIdentityListener.off("deleted", this.#onIdentityDeleted);
+    this.#contextualIdentityListener.destroy();
+
+    this.#userContextIds = null;
+  }
+
+  /**
+   * Creates a new user context.
+   *
+   * @param {string} prefix
+   *     The prefix to use for the name of the user context.
+   *
+   * @returns {string}
+   *     The user context id of the new user context.
+   */
+  createContext(prefix = "remote") {
+    // Prepare the opaque id and name beforehand.
+    const userContextId = lazy.generateUUID();
+    const name = `${prefix}-${userContextId}`;
+
+    // Create the user context.
+    const identity = lazy.ContextualIdentityService.create(name);
+    const internalId = identity.userContextId;
+
+    // An id has been set already by the contextual-identity-created observer.
+    // Override it with `userContextId` to match the container name.
+    this.#userContextIds.set(internalId, userContextId);
+
+    return userContextId;
+  }
+
+  /**
+   * Retrieve the user context id corresponding to the provided internal id.
+   *
+   * @param {number} internalId
+   *     The internal user context id.
+   *
+   * @returns {string|null}
+   *     The corresponding user context id or null if the user context does not
+   *     exist.
+   */
+  getIdByInternalId(internalId) {
+    if (this.#userContextIds.has(internalId)) {
+      return this.#userContextIds.get(internalId);
+    }
+    return null;
+  }
+
+  /**
+   * Retrieve the internal id corresponding to the provided user
+   * context id.
+   *
+   * @param {string} userContextId
+   *     The user context id.
+   *
+   * @returns {number|null}
+   *     The internal user context id or null if the user context does not
+   *     exist.
+   */
+  getInternalIdById(userContextId) {
+    for (const [internalId, id] of this.#userContextIds) {
+      if (userContextId == id) {
+        return internalId;
+      }
+    }
+    return null;
+  }
+
+  /**
+   * Returns an array of all known user context ids.
+   *
+   * @returns {Array<string>}
+   *     The array of user context ids.
+   */
+  getUserContextIds() {
+    return Array.from(this.#userContextIds.values());
+  }
+
+  /**
+   * Checks if the provided user context id is known by this UserContextManager.
+   *
+   * @param {string} userContextId
+   *     The id of the user context to check.
+   */
+  hasUserContextId(userContextId) {
+    return this.getUserContextIds().includes(userContextId);
+  }
+
+  /**
+   * Removes a user context and closes all related container tabs.
+   *
+   * @param {string} userContextId
+   *     The id of the user context to remove.
+   * @param {object=} options
+   * @param {boolean=} options.closeContextTabs
+   *     Pass true if the tabs owned by the user context should also be closed.
+   *     Defaults to false.
+   */
+  removeUserContext(userContextId, options = {}) {
+    const { closeContextTabs = false } = options;
+
+    if (!this.hasUserContextId(userContextId)) {
+      return;
+    }
+
+    const internalId = this.getInternalIdById(userContextId);
+    if (closeContextTabs) {
+      lazy.ContextualIdentityService.closeContainerTabs(internalId);
+    }
+    lazy.ContextualIdentityService.remove(internalId);
+  }
+
+  #onIdentityCreated = (eventName, data) => {
+    this.#registerIdentity(data.identity);
+  };
+
+  #onIdentityDeleted = (eventName, data) => {
+    this.#userContextIds.delete(data.identity.userContextId);
+  };
+
+  #registerIdentity(identity) {
+    // Note: the id for identities created via UserContextManagerClass.createContext
+    // are overridden in createContext.
+    this.#userContextIds.set(identity.userContextId, lazy.generateUUID());
+  }
+}
+
+// Expose a shared singleton.
+export const UserContextManager = new UserContextManagerClass();

remote/shared/listeners/ContextualIdentityListener.sys.mjs

@@ -0,0 +1,85 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+const lazy = {};
+
+ChromeUtils.defineESModuleGetters(lazy, {
+  EventEmitter: "resource://gre/modules/EventEmitter.sys.mjs",
+});
+
+const OBSERVER_TOPIC_CREATED = "contextual-identity-created";
+const OBSERVER_TOPIC_DELETED = "contextual-identity-deleted";
+
+/**
+ * The ContextualIdentityListener can be used to listen for notifications about
+ * contextual identities (containers) being created or deleted.
+ *
+ * Example:
+ * ```
+ * const listener = new ContextualIdentityListener();
+ * listener.on("created", onCreated);
+ * listener.startListening();
+ *
+ * const onCreated = (eventName, data = {}) => {
+ *   const { identity } = data;
+ *   ...
+ * };
+ * ```
+ *
+ * @fires message
+ *    The ContextualIdentityListener emits "created" and "deleted" events,
+ *    with the following object as payload:
+ *      - {object} identity
+ *            The contextual identity which was created or deleted.
+ */
+export class ContextualIdentityListener {
+  #listening;
+
+  /**
+   * Create a new BrowsingContextListener instance.
+   */
+  constructor() {
+    lazy.EventEmitter.decorate(this);
+
+    this.#listening = false;
+  }
+
+  destroy() {
+    this.stopListening();
+  }
+
+  observe(subject, topic, data) {
+    switch (topic) {
+      case OBSERVER_TOPIC_CREATED:
+        this.emit("created", { identity: subject.wrappedJSObject });
+        break;
+
+      case OBSERVER_TOPIC_DELETED:
+        this.emit("deleted", { identity: subject.wrappedJSObject });
+        break;
+    }
+  }
+
+  startListening() {
+    if (this.#listening) {
+      return;
+    }
+
+    Services.obs.addObserver(this, OBSERVER_TOPIC_CREATED);
+    Services.obs.addObserver(this, OBSERVER_TOPIC_DELETED);
+
+    this.#listening = true;
+  }
+
+  stopListening() {
+    if (!this.#listening) {
+      return;
+    }
+
+    Services.obs.removeObserver(this, OBSERVER_TOPIC_CREATED);
+    Services.obs.removeObserver(this, OBSERVER_TOPIC_DELETED);
+
+    this.#listening = false;
+  }
+}

remote/shared/listeners/test/browser/browser.toml

@@ -14,6 +14,8 @@ prefs = ["remote.messagehandler.modulecache.useBrowserTestRoot=true"]
 
 ["browser_ConsoleListener_cached_messages.js"]
 
+["browser_ContextualIdentityListener.js"]
+
 ["browser_NetworkListener.js"]
 
 ["browser_PromptListener.js"]

remote/shared/listeners/test/browser/browser_ContextualIdentityListener.js

@@ -0,0 +1,38 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this file,
+ * You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+const { ContextualIdentityListener } = ChromeUtils.importESModule(
+  "chrome://remote/content/shared/listeners/ContextualIdentityListener.sys.mjs"
+);
+
+add_task(async function test_createdOnNewContextualIdentity() {
+  const listener = new ContextualIdentityListener();
+  const created = listener.once("created");
+
+  listener.startListening();
+
+  ContextualIdentityService.create("test_name");
+
+  const { identity } = await created;
+  is(identity.name, "test_name", "Received expected identity");
+
+  listener.stopListening();
+
+  ContextualIdentityService.remove(identity.userContextId);
+});
+
+add_task(async function test_deletedOnRemovedContextualIdentity() {
+  const listener = new ContextualIdentityListener();
+  const deleted = listener.once("deleted");
+
+  listener.startListening();
+
+  const testIdentity = ContextualIdentityService.create("test_name");
+  ContextualIdentityService.remove(testIdentity.userContextId);
+
+  const { identity } = await deleted;
+  is(identity.name, "test_name", "Received expected identity");
+
+  listener.stopListening();
+});

remote/shared/test/browser/browser.toml

@@ -12,3 +12,5 @@ support-files = ["head.js"]
 ["browser_NavigationManager_notify.js"]
 
 ["browser_TabManager.js"]
+
+["browser_UserContextManager.js"]

remote/shared/test/browser/browser_UserContextManager.js

@@ -0,0 +1,236 @@
+/* Any copyright is dedicated to the Public Domain.
+ * http://creativecommons.org/publicdomain/zero/1.0/ */
+
+"use strict";
+
+const { UserContextManagerClass } = ChromeUtils.importESModule(
+  "chrome://remote/content/shared/UserContextManager.sys.mjs"
+);
+
+add_task(async function test_invalid() {
+  const userContextManager = new UserContextManagerClass();
+
+  // Check invalid types for hasUserContextId/getInternalIdById which expects
+  // a string.
+  for (const value of [null, undefined, 1, [], {}]) {
+    is(userContextManager.hasUserContextId(value), false);
+    is(userContextManager.getInternalIdById(value), null);
+  }
+
+  // Check an invalid value for hasUserContextId/getInternalIdById which expects
+  // either "default" or a UUID from Services.uuid.generateUUID.
+  is(userContextManager.hasUserContextId("foo"), false);
+  is(userContextManager.getInternalIdById("foo"), null);
+
+  // Check invalid types for getIdByInternalId which expects a number.
+  for (const value of [null, undefined, "foo", [], {}]) {
+    is(userContextManager.getIdByInternalId(value), null);
+  }
+
+  userContextManager.destroy();
+});
+
+add_task(async function test_default_context() {
+  const userContextManager = new UserContextManagerClass();
+  ok(
+    userContextManager.hasUserContextId("default"),
+    `Context id default is known by the manager`
+  );
+  ok(
+    userContextManager.getUserContextIds().includes("default"),
+    `Context id default is listed by the manager`
+  );
+  is(
+    userContextManager.getInternalIdById("default"),
+    0,
+    "Default user context has the expected internal id"
+  );
+
+  userContextManager.destroy();
+});
+
+add_task(async function test_new_internal_contexts() {
+  info("Create a new user context with ContextualIdentityService");
+  const beforeInternalId =
+    ContextualIdentityService.create("before").userContextId;
+
+  info("Create the UserContextManager");
+  const userContextManager = new UserContextManagerClass();
+
+  const beforeContextId =
+    userContextManager.getIdByInternalId(beforeInternalId);
+  assertContextAvailable(userContextManager, beforeContextId, beforeInternalId);
+
+  info("Create another user context with ContextualIdentityService");
+  const afterInternalId =
+    ContextualIdentityService.create("after").userContextId;
+  const afterContextId = userContextManager.getIdByInternalId(afterInternalId);
+  assertContextAvailable(userContextManager, afterContextId, afterInternalId);
+
+  info("Delete both user contexts");
+  ContextualIdentityService.remove(beforeInternalId);
+  ContextualIdentityService.remove(afterInternalId);
+  assertContextRemoved(userContextManager, afterContextId, afterInternalId);
+  assertContextRemoved(userContextManager, beforeContextId, beforeInternalId);
+
+  userContextManager.destroy();
+});
+
+add_task(async function test_create_remove_context() {
+  const userContextManager = new UserContextManagerClass();
+
+  for (const closeContextTabs of [true, false]) {
+    info("Create two contexts via createContext");
+    const userContextId1 = userContextManager.createContext();
+    const internalId1 = userContextManager.getInternalIdById(userContextId1);
+    assertContextAvailable(userContextManager, userContextId1);
+
+    const userContextId2 = userContextManager.createContext();
+    const internalId2 = userContextManager.getInternalIdById(userContextId2);
+    assertContextAvailable(userContextManager, userContextId2);
+
+    info("Create tabs in various user contexts");
+    const url = "https://example.com/document-builder.sjs?html=tab";
+    const tabDefault = await addTab(gBrowser, url);
+    const tabContext1 = await addTab(gBrowser, url, {
+      userContextId: internalId1,
+    });
+    const tabContext2 = await addTab(gBrowser, url, {
+      userContextId: internalId2,
+    });
+
+    info("Remove the user context 1 via removeUserContext");
+    userContextManager.removeUserContext(userContextId1, { closeContextTabs });
+
+    assertContextRemoved(userContextManager, userContextId1, internalId1);
+    if (closeContextTabs) {
+      ok(!gBrowser.tabs.includes(tabContext1), "Tab context 1 is closed");
+    } else {
+      ok(gBrowser.tabs.includes(tabContext1), "Tab context 1 is not closed");
+    }
+    ok(gBrowser.tabs.includes(tabDefault), "Tab default is not closed");
+    ok(gBrowser.tabs.includes(tabContext2), "Tab context 2 is not closed");
+
+    info("Remove the user context 2 via removeUserContext");
+    userContextManager.removeUserContext(userContextId2, { closeContextTabs });
+    assertContextRemoved(userContextManager, userContextId2, internalId2);
+    if (closeContextTabs) {
+      ok(!gBrowser.tabs.includes(tabContext2), "Tab context 2 is closed");
+    } else {
+      ok(gBrowser.tabs.includes(tabContext2), "Tab context 2 is not closed");
+    }
+    ok(gBrowser.tabs.includes(tabDefault), "Tab default is not closed");
+  }
+
+  userContextManager.destroy();
+});
+
+add_task(async function test_create_context_prefix() {
+  const userContextManager = new UserContextManagerClass();
+
+  info("Create a context with a custom prefix via createContext");
+  const userContextId = userContextManager.createContext("test_prefix");
+  const internalId = userContextManager.getInternalIdById(userContextId);
+  const identity =
+    ContextualIdentityService.getPublicIdentityFromId(internalId);
+  ok(
+    identity.name.startsWith("test_prefix"),
+    "The new identity used the provided prefix"
+  );
+
+  userContextManager.removeUserContext(userContextId);
+  userContextManager.destroy();
+});
+
+add_task(async function test_several_managers() {
+  const manager1 = new UserContextManagerClass();
+  const manager2 = new UserContextManagerClass();
+
+  info("Create a context via manager1");
+  const contextId1 = manager1.createContext();
+  const internalId = manager1.getInternalIdById(contextId1);
+  assertContextUnknown(manager2, contextId1);
+
+  info("Retrieve the corresponding user context id in manager2");
+  const contextId2 = manager2.getIdByInternalId(internalId);
+  is(
+    typeof contextId2,
+    "string",
+    "manager2 has a valid id for the user context created by manager 1"
+  );
+
+  ok(
+    contextId1 != contextId2,
+    "manager1 and manager2 have different ids for the same internal context id"
+  );
+
+  info("Remove the user context via manager2");
+  manager2.removeUserContext(contextId2);
+
+  info("Check that the user context is removed from both managers");
+  assertContextRemoved(manager1, contextId1, internalId);
+  assertContextRemoved(manager2, contextId2, internalId);
+
+  manager1.destroy();
+  manager2.destroy();
+});
+
+function assertContextAvailable(manager, contextId, expectedInternalId = null) {
+  ok(
+    manager.getUserContextIds().includes(contextId),
+    `Context id ${contextId} is listed by the manager`
+  );
+  ok(
+    manager.hasUserContextId(contextId),
+    `Context id ${contextId} is known by the manager`
+  );
+
+  const internalId = manager.getInternalIdById(contextId);
+  if (expectedInternalId != null) {
+    is(internalId, expectedInternalId, "Internal id has the expected value");
+  }
+
+  is(
+    typeof internalId,
+    "number",
+    `Context id ${contextId} corresponds to a valid internal id (${internalId})`
+  );
+  is(
+    manager.getIdByInternalId(internalId),
+    contextId,
+    `Context id ${contextId} is returned for internal id ${internalId}`
+  );
+  ok(
+    ContextualIdentityService.getPublicUserContextIds().includes(internalId),
+    `User context for context id ${contextId} is found by ContextualIdentityService`
+  );
+}
+
+function assertContextUnknown(manager, contextId) {
+  ok(
+    !manager.getUserContextIds().includes(contextId),
+    `Context id ${contextId} is not listed by the manager`
+  );
+  ok(
+    !manager.hasUserContextId(contextId),
+    `Context id ${contextId} is not known by the manager`
+  );
+  is(
+    manager.getInternalIdById(contextId),
+    null,
+    `Context id ${contextId} does not match any internal id`
+  );
+}
+
+function assertContextRemoved(manager, contextId, internalId) {
+  assertContextUnknown(manager, contextId);
+  is(
+    manager.getIdByInternalId(internalId),
+    null,
+    `Internal id ${internalId} cannot be converted to user context id`
+  );
+  ok(
+    !ContextualIdentityService.getPublicUserContextIds().includes(internalId),
+    `Internal id ${internalId} is not found in ContextualIdentityService`
+  );
+}

remote/shared/test/browser/head.js

@@ -12,11 +12,13 @@
  *     The browser element where the tab should be added.
  * @param {string} url
  *     The URL for the tab.
+ * @param {object=} options
+ *     Options object to forward to BrowserTestUtils.addTab.
  * @returns {Tab}
  *     The created tab.
  */
-function addTab(browser, url) {
-  const tab = BrowserTestUtils.addTab(browser, url);
+function addTab(browser, url, options) {
+  const tab = BrowserTestUtils.addTab(browser, url, options);
   registerCleanupFunction(() => browser.removeTab(tab));
   return tab;
 }