Container locking - restrict container to currently assigned domains

src/css/popup.css

@@ -925,6 +925,15 @@ span ~ .panel-header-text {
   padding-block-end: 6px;
 }
 
+/* https://github.com/mozilla/multi-account-containers/issues/847 */
+.container-lockorunlock.container-locked * {
+  filter: invert(0.5) sepia(1) saturate(127) hue-rotate(360deg);
+}
+
+.container-lockorunlock.container-unlocked * {
+  filter: invert(0.5);
+}
+
 /* Achievement panel elements */
 .share-ctas {
   padding-block-end: 0.5em;

src/img/container-lock.svg

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="utf-8"?>
+
+<!-- Adapted from https://www.materialui.co/icon/lock -->
+
+<svg version="1.1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px"
+	 viewBox="0 0 24 24" style="enable-background:new 0 0 24 24;" xml:space="preserve">
+  <path d="M18 8h-1V6c0-2.76-2.24-5-5-5S7 3.24 7 6v2H6c-1.1 0-2 .9-2 2v10c0 1.1.9 2 2 2h12c1.1 0 2-.9 2-2V10c0-1.1-.9-2-2-2zm-6 9c-1.1 0-2-.9-2-2s.9-2 2-2 2 .9 2 2-.9 2-2 2zm3.1-9H8.9V6c0-1.71 1.39-3.1 3.1-3.1 1.71 0 3.1 1.39 3.1 3.1v2z"/>
+</svg> 

src/img/container-unlock.svg

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="utf-8"?>
+
+<!-- Adapted from https://www.materialui.co/icon/lock-open -->
+
+<svg version="1.1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px"
+	 viewBox="0 0 24 24" style="enable-background:new 0 0 24 24;" xml:space="preserve">
+  <path d="M12 17c1.1 0 2-.9 2-2s-.9-2-2-2-2 .9-2 2 .9 2 2 2zm6-9h-1V6c0-2.76-2.24-5-5-5S7 3.24 7 6h1.9c0-1.71 1.39-3.1 3.1-3.1 1.71 0 3.1 1.39 3.1 3.1v2H6c-1.1 0-2 .9-2 2v10c0 1.1.9 2 2 2h12c1.1 0 2-.9 2-2V10c0-1.1-.9-2-2-2zm0 12H6V10h12v10z"/>
+</svg> 

src/js/background/assignManager.js

@@ -141,8 +141,21 @@ const assignManager = {
       return {};
     }
     const userContextId = this.getUserContextIdFromCookieStore(tab);
-    if (!siteSettings
-        || userContextId === siteSettings.userContextId
+    
+    // Determine if "locked out", i.e.:
+    // This request's URL is not associated with any particular contextualIdentity.
+    // But the current tab's contextualIdentity is locked. So must open request in new tab.
+    // https://github.com/mozilla/multi-account-containers/issues/847
+    let isLockedOut;
+    if (!siteSettings && "cookieStoreId" in tab) {
+      const currentContainerState = await identityState.storageArea.get(tab.cookieStoreId);
+      isLockedOut = !!currentContainerState.isLocked;
+    } else {
+      isLockedOut = false;
+    }
+    
+    if ((!siteSettings && !isLockedOut)
+        || (siteSettings && userContextId === siteSettings.userContextId)
         || this.storageArea.isExempted(options.url, tab.id)) {
       return {};
     }
@@ -188,15 +201,22 @@ const assignManager = {
       }
     }
 
-    this.reloadPageInContainer(
-      options.url,
-      userContextId,
-      siteSettings.userContextId,
-      tab.index + 1,
-      tab.active,
-      siteSettings.neverAsk,
-      openTabId
-    );
+    if (isLockedOut) {
+      // Open new tab in default context
+      // https://github.com/mozilla/multi-account-containers/issues/847
+      browser.tabs.create({url: options.url});
+    } else {
+      // Open new tab in specific context
+      this.reloadPageInContainer(
+        options.url,
+        userContextId,
+        siteSettings.userContextId,
+        tab.index + 1,
+        tab.active,
+        siteSettings.neverAsk,
+        openTabId
+      );
+    }
     this.calculateContextMenu(tab);
 
     /* Removal of existing tabs:
@@ -395,7 +415,20 @@ const assignManager = {
         neverAsk: false
       }, exemptedTabIds);
       actionName = "added";
+    
     } else {
+      // Unlock container if no more assignments after this one is removed.
+      // https://github.com/mozilla/multi-account-containers/issues/847
+      const assignments = await this.storageArea.getByContainer(userContextId);
+      const assignmentKeys = Object.keys(assignments);
+      if (!(assignmentKeys.length > 1)) {
+        await backgroundLogic.lockOrUnlockContainer({
+          userContextId: userContextId,
+          isLocked: false
+        });
+      }
+    
+      // Remove assignment
       await this.storageArea.remove(pageUrl);
       actionName = "removed";
     }

src/js/background/backgroundLogic.js

@@ -123,6 +123,22 @@ const backgroundLogic = {
     }
   },
 
+  // https://github.com/mozilla/multi-account-containers/issues/847
+  async lockOrUnlockContainer(options) {
+    if (!("userContextId" in options)) {
+      return Promise.reject("lockOrUnlockContainer must be called with userContextId argument.");
+    }
+  
+    const cookieStoreId = this.cookieStoreId(options.userContextId);
+    const containerState = await identityState.storageArea.get(cookieStoreId);
+    if (options.isLocked) {
+      containerState.isLocked = "locked";
+    } else {
+      delete containerState.isLocked;
+    }
+    return await identityState.storageArea.set(cookieStoreId, containerState);
+  },
+  
 
   async moveTabsToWindow(options) {
     const requiredArguments = ["cookieStoreId", "windowId"];
@@ -229,7 +245,9 @@ const backgroundLogic = {
         hasHiddenTabs: !!containerState.hiddenTabs.length,
         hasOpenTabs: !!openTabs.length,
         numberOfHiddenTabs: containerState.hiddenTabs.length,
-        numberOfOpenTabs: openTabs.length
+        numberOfOpenTabs: openTabs.length,
+        // https://github.com/mozilla/multi-account-containers/issues/847
+        isLocked: !!containerState.isLocked
       };
       return;
     });

src/js/background/messageHandler.js

@@ -19,6 +19,10 @@ const messageHandler = {
       case "createOrUpdateContainer":
         response = backgroundLogic.createOrUpdateContainer(m.message);
         break;
+      case "lockOrUnlockContainer":
+        // https://github.com/mozilla/multi-account-containers/issues/847
+        response = backgroundLogic.lockOrUnlockContainer(m.message);
+        break;
       case "neverAsk":
         assignManager._neverAsk(m);
         break;

src/js/popup.js

@@ -4,6 +4,9 @@
 
 const CONTAINER_HIDE_SRC = "/img/container-hide.svg";
 const CONTAINER_UNHIDE_SRC = "/img/container-unhide.svg";
+// https://github.com/mozilla/multi-account-containers/issues/847
+const CONTAINER_LOCKED_SRC = "/img/container-lock.svg";
+const CONTAINER_UNLOCKED_SRC = "/img/container-unlock.svg";
 
 const DEFAULT_COLOR = "blue";
 const DEFAULT_ICON = "circle";
@@ -252,6 +255,8 @@ const Logic = {
         identity.hasHiddenTabs = stateObject.hasHiddenTabs;
         identity.numberOfHiddenTabs = stateObject.numberOfHiddenTabs;
         identity.numberOfOpenTabs = stateObject.numberOfOpenTabs;
+        // https://github.com/mozilla/multi-account-containers/issues/847
+        identity.isLocked = stateObject.isLocked;
       }
       return identity;
     });
@@ -1011,7 +1016,7 @@ Logic.registerPanel(P_CONTAINER_EDIT, {
     }
   },
 
-  showAssignedContainers(assignments) {
+  showAssignedContainers(assignments, isLocked) {
     const assignmentPanel = document.getElementById("edit-sites-assigned");
     const assignmentKeys = Object.keys(assignments);
     assignmentPanel.hidden = !(assignmentKeys.length > 0);
@@ -1022,6 +1027,38 @@ Logic.registerPanel(P_CONTAINER_EDIT, {
       while (tableElement.firstChild) {
         tableElement.firstChild.remove();
       }
+      
+      /* Container locking: https://github.com/mozilla/multi-account-containers/issues/847 */
+      const lockOrUnlockIcon = isLocked ? CONTAINER_LOCKED_SRC : CONTAINER_UNLOCKED_SRC;
+      const lockOrUnlockLabel = isLocked ? "Locked" : "Unlocked";
+      const lockOrUnlockClass = isLocked ? "container-locked" : "container-unlocked";
+      const lockElement = document.createElement("div");
+      lockElement.innerHTML = escaped`
+        <img class="icon" src="${lockOrUnlockIcon}">
+        <div title="${lockOrUnlockLabel}">
+          ${lockOrUnlockLabel}
+        </div>`;
+      lockElement.classList.add("container-info-tab-row", "clickable", "container-lockorunlock", lockOrUnlockClass);
+      tableElement.appendChild(lockElement);
+      
+      const that = this;
+      Logic.addEnterHandler(lockElement, async () => {
+        try {
+          await browser.runtime.sendMessage({
+            method: "lockOrUnlockContainer",
+            message: {
+              userContextId: Logic.currentUserContextId(),
+              isLocked: !isLocked
+            }
+          });
+          that.showAssignedContainers(assignments, !isLocked);
+        } catch (e) {
+          throw new Error("Failed to lock/unlock. ", e.message);
+        }
+      });
+      /* Container locking: https://github.com/mozilla/multi-account-containers/issues/847 */
+      
+      /* Assignment list */
       assignmentKeys.forEach((siteKey) => {
         const site = assignments[siteKey];
         const trElement = document.createElement("div");
@@ -1047,7 +1084,7 @@ Logic.registerPanel(P_CONTAINER_EDIT, {
           const currentTab = await Logic.currentTab();
           Logic.setOrRemoveAssignment(currentTab.id, assumedUrl, userContextId, true);
           delete assignments[siteKey];
-          that.showAssignedContainers(assignments);
+          that.showAssignedContainers(assignments, isLocked);
         });
         trElement.classList.add("container-info-tab-row", "clickable");
         tableElement.appendChild(trElement);
@@ -1091,7 +1128,7 @@ Logic.registerPanel(P_CONTAINER_EDIT, {
 
     const userContextId = Logic.currentUserContextId();
     const assignments = await Logic.getAssignmentObjectByContainer(userContextId);
-    this.showAssignedContainers(assignments);
+    this.showAssignedContainers(assignments, identity.isLocked);
     document.querySelector("#edit-container-panel .panel-footer").hidden = !!userContextId;
 
     document.querySelector("#edit-container-panel-name-input").value = identity.name || "";

test/features/lock.test.js

@@ -0,0 +1,54 @@
+// https://github.com/mozilla/multi-account-containers/issues/847
+describe("Lock Feature", () => {
+  const activeTab = {
+    id: 1,
+    cookieStoreId: "firefox-container-1",
+    url: "http://example.com",
+    index: 0
+  };
+  beforeEach(async () => {
+    await helper.browser.initializeWithTab(activeTab);
+  });
+
+  describe("click the 'Always open in' checkbox in the popup", () => {
+    beforeEach(async () => {
+      // popup click to set assignment for activeTab.url
+      await helper.popup.clickElementById("container-page-assigned");
+    });
+    
+    describe("open different URL in same tab", () => {
+      const differentURL = "http://example2.com";
+      beforeEach(async () => {
+        await helper.browser.updateTab(activeTab, {
+          url: differentURL,
+          resetHistory: true
+        });
+      });
+
+      it("should not open a new tab", () => {
+        background.browser.tabs.create.should.not.have.been.called;
+      });
+    
+      describe("lock the container", () => {
+        beforeEach(async () => {
+          await helper.popup.setContainerIsLocked(activeTab.cookieStoreId, true);
+        });
+    
+        describe("open different URL in same tab", () => {
+          beforeEach(async () => {
+            await helper.browser.updateTab(activeTab, {
+              url: differentURL,
+              resetHistory: true
+            });
+          });
+
+          it("should open a new tab in the default container", () => {
+            background.browser.tabs.create.should.have.been.calledWith({
+              url: differentURL
+            });
+          });
+        });
+      });
+    });
+  });
+});

test/helper.js

@@ -29,7 +29,19 @@ module.exports = {
 
     async openNewTab(tab, options = {}) {
       return background.browser.tabs._create(tab, options);
-    }
+    },
+
+    // https://github.com/mozilla/multi-account-containers/issues/847
+    async updateTab(tab, options = {}) {
+      const updatedTab = {};
+      for (const key in tab) {
+        updatedTab[key] = tab[key];
+      }
+      for (const key in options) {
+        updatedTab[key] = options[key];
+      }
+      return this.openNewTab(updatedTab);
+    },
   },
 
   popup: {
@@ -39,6 +51,25 @@ module.exports = {
 
     async clickLastMatchingElementByQuerySelector(querySelector) {
       await popup.helper.clickElementByQuerySelectorAll(querySelector, "last");
+    },
+    
+    // https://github.com/mozilla/multi-account-containers/issues/847
+    async setContainerIsLocked(cookieStoreId, isLocked) {
+      const identityStateKey = this.getIdentityStateContainerStoreKey(cookieStoreId);
+      const identityState = await background.browser.storage.local.get([identityStateKey]) || {};
+      if (isLocked) {
+        identityState.isLocked = "locked";
+      } else {
+        delete identityState.isLocked;
+      }
+      // Must have valid 'hiddenTabs', otherwise backgroundLogic.showTabs() throws error
+      if (!identityState.hiddenTabs) { identityState.hiddenTabs = []; }
+      await background.browser.storage.local.set({[identityStateKey]: identityState});
+    },
+    
+    getIdentityStateContainerStoreKey(cookieStoreId) {
+      const storagePrefix = "identitiesState@@_";
+      return `${storagePrefix}${cookieStoreId}`;      
     }
   }
 };