Implemented site isolation

Added feature to isolate (lock) assigned sites: When you are in
a container with site isolation enabled, navigating to a site
outside of the assignments will open that site in a new default
container tab.

Co-authored-by: Francis McKenzie <francis.mckenzie@gmail.com>

src/js/background/assignManager.js

@@ -205,10 +205,33 @@ window.assignManager = {
       return {};
     }
     const userContextId = this.getUserContextIdFromCookieStore(tab);
-    if (!siteSettings
+
-        || userContextId === siteSettings.userContextId
+    // https://github.com/mozilla/multi-account-containers/issues/847
-        || this.storageArea.isExempted(options.url, tab.id)) {
+    //    
-      return {};
+    // Handle the case where this request's URL is not assigned to any particular
+    // container. We must do the following check:
+    //
+    // If the current tab's container is "unlocked", we can just go ahead
+    // and open the URL in the current tab, since an "unlocked" container accepts
+    // any-and-all sites.
+    //
+    // But if the current tab's container has been "locked" by the user, then we must
+    // re-open the page in the default container, because the user doesn't want random
+    // sites polluting their locked container.
+    //
+    // For example:
+    //   - the current tab's container is locked and only allows "www.google.com"
+    //   - the incoming request is for "www.amazon.com", which has no specific container assignment
+    //   - in this case, we must re-open "www.amazon.com" in a new tab in the default container
+    const siteIsolatedReloadInDefault = 
+      await this._maybeSiteIsolatedReloadInDefault(siteSettings, tab);
+
+    if (!siteIsolatedReloadInDefault) {
+      if (!siteSettings
+          || userContextId === siteSettings.userContextId
+          || this.storageArea.isExempted(options.url, tab.id)) {
+        return {};
+      }
     }
     const removeTab = backgroundLogic.NEW_TAB_PAGES.has(tab.url)
       || (messageHandler.lastCreatedTab
@@ -257,15 +280,24 @@ window.assignManager = {
       }
     }
 
-    this.reloadPageInContainer(
+    if (siteIsolatedReloadInDefault) {
-      options.url,
+      this.reloadPageInDefaultContainer(
-      userContextId,
+        options.url,
-      siteSettings.userContextId,
+        tab.index + 1,
-      tab.index + 1,
+        tab.active,
-      tab.active,
+        openTabId
-      siteSettings.neverAsk,
+      );
-      openTabId
+    } else {
-    );
+      this.reloadPageInContainer(
+        options.url,
+        userContextId,
+        siteSettings.userContextId,
+        tab.index + 1,
+        tab.active,
+        siteSettings.neverAsk,
+        openTabId
+      );
+    }
     this.calculateContextMenu(tab);
 
     /* Removal of existing tabs:
@@ -299,6 +331,29 @@ window.assignManager = {
     };
   },
 
+  async _maybeSiteIsolatedReloadInDefault(siteSettings, tab) {
+    // Tab doesn't support cookies, so containers not supported either.
+    if (!("cookieStoreId" in tab)) {
+      return false;
+    }
+
+    // Requested page has been assigned to a specific container.
+    // I.e. it will be opened in that container anyway, so we don't need to check if the
+    // current tab's container is locked or not.
+    if (siteSettings) {
+      return false;  
+    }
+
+    //tab is alredy reopening in the default container
+    if (tab.cookieStoreId === "firefox-default") {
+      return false;
+    }
+    // Requested page is not assigned to a specific container. If the current tab's container
+    // is locked, then the page must be reloaded in the default container.
+    const currentContainerState = await identityState.storageArea.get(tab.cookieStoreId);
+    return currentContainerState && currentContainerState.isIsolated;
+  },
+
   init() {
     browser.contextMenus.onClicked.addListener((info, tab) => {
       info.bookmarkId ? 
@@ -502,8 +557,13 @@ window.assignManager = {
       }, exemptedTabIds);
       actionName = "assigned site to always open in this container";
     } else {
+      // Remove assignment
       await this.storageArea.remove(pageUrl);
+
       actionName = "removed from assigned sites list";
+
+      // remove site isolation if now empty
+      await this._maybeRemoveSiteIsolation(userContextId);   
     }
 
     if (tabId) {
@@ -519,6 +579,18 @@ window.assignManager = {
     }
   },
 
+  async _maybeRemoveSiteIsolation(userContextId) {
+    const assignments = await this.storageArea.getByContainer(userContextId);
+    const hasAssignments = assignments && Object.keys(assignments).length > 0;
+    if (hasAssignments) {
+      return;
+    }
+    await backgroundLogic.addRemoveSiteIsolation(
+      backgroundLogic.cookieStoreId(userContextId),
+      true
+    );
+  },
+
   async _getAssignment(tab) {
     const cookieStore = this.getUserContextIdFromCookieStore(tab);
     // Ensure we have a cookieStore to assign to
@@ -596,6 +668,28 @@ window.assignManager = {
     });
   },
 
+  reloadPageInDefaultContainer(url, index, active, openerTabId) {
+    // To create a new tab in the default container, it is easiest just to omit the
+    // cookieStoreId entirely.
+    // 
+    // Unfortunately, if you create a new tab WITHOUT a cookieStoreId but WITH an openerTabId,
+    // then the new tab automatically inherits the opener tab's cookieStoreId.
+    // I.e. it opens in the wrong container!
+    // 
+    // So we have to explicitly pass in a cookieStoreId when creating the tab, since we
+    // are specifying the openerTabId. There doesn't seem to be any way
+    // to look up the default container's cookieStoreId programatically, so sadly
+    // we have to hardcode it here as "firefox-default". This is potentially
+    // not cross-browser compatible.
+    // 
+    // Note that we could have just omitted BOTH cookieStoreId and openerTabId. But the
+    // drawback then is that if the user later closes the newly-created tab, the browser
+    // does not automatically return to the original opener tab. To get this desired behaviour,
+    // we MUST specify the openerTabId when creating the new tab.
+    const cookieStoreId = "firefox-default";
+    browser.tabs.create({url, cookieStoreId, index, active, openerTabId});
+  },
+
   reloadPageInContainer(url, currentUserContextId, userContextId, index, active, neverAsk = false, openerTabId = null) {
     const cookieStoreId = backgroundLogic.cookieStoreId(userContextId);
     const loadPage = browser.extension.getURL("confirm-page.html");

src/js/background/backgroundLogic.js

@@ -136,6 +136,20 @@ const backgroundLogic = {
     }
   },
 
+  // https://github.com/mozilla/multi-account-containers/issues/847
+  async addRemoveSiteIsolation(cookieStoreId, remove = false) {
+    const containerState = await identityState.storageArea.get(cookieStoreId);
+    try {
+      if ("isIsolated" in containerState || remove) {
+        delete containerState.isIsolated;
+      } else {
+        containerState.isIsolated = "locked";        
+      }
+      return await identityState.storageArea.set(cookieStoreId, containerState);
+    } catch (error) {
+      console.error(`No container: ${cookieStoreId}`);
+    }
+  },
 
   async moveTabsToWindow(options) {
     const requiredArguments = ["cookieStoreId", "windowId"];
@@ -242,7 +256,8 @@ const backgroundLogic = {
         hasHiddenTabs: !!containerState.hiddenTabs.length,
         hasOpenTabs: !!openTabs.length,
         numberOfHiddenTabs: containerState.hiddenTabs.length,
-        numberOfOpenTabs: openTabs.length
+        numberOfOpenTabs: openTabs.length,
+        isIsolated: !!containerState.isIsolated
       };
       return;
     });

src/js/background/identityState.js

@@ -28,7 +28,7 @@ window.identityState = {
         await this.set(cookieStoreId, defaultContainerState);
         return defaultContainerState;
       }
-      throw new Error (`${cookieStoreId} not found`);
+      return false;
     },
 
     set(cookieStoreId, data) {

src/js/background/messageHandler.js

@@ -32,6 +32,9 @@ const messageHandler = {
       case "neverAsk":
         assignManager._neverAsk(m);
         break;
+      case "addRemoveSiteIsolation":
+        response = backgroundLogic.addRemoveSiteIsolation(m.cookieStoreId);
+        break;
       case "getAssignment":
         response = browser.tabs.get(m.tabId).then((tab) => {
           return assignManager._getAssignment(tab);

src/js/popup.js

@@ -202,6 +202,7 @@ const Logic = {
         identity.hasHiddenTabs = stateObject.hasHiddenTabs;
         identity.numberOfHiddenTabs = stateObject.numberOfHiddenTabs;
         identity.numberOfOpenTabs = stateObject.numberOfOpenTabs;
+        identity.isIsolated = stateObject.isIsolated;
       }
       return identity;
     });
@@ -1298,6 +1299,14 @@ Logic.registerPanel(P_CONTAINER_EDIT, {
       containerName.select();
       containerName.focus();
     });
+    const siteIsolation = document.querySelector("#site-isolation");
+    siteIsolation.checked = !!identity.isIsolated;
+    siteIsolation.addEventListener( "change", function() {
+      browser.runtime.sendMessage({ 
+        method: "addRemoveSiteIsolation",
+        cookieStoreId: identity.cookieStoreId
+      });
+    });
     [...document.querySelectorAll("[name='container-color']")].forEach(colorInput => {
       colorInput.checked = colorInput.value === identity.color;
     });

test/features/assignment.test.js

@@ -14,7 +14,7 @@ describe("Assignment Reopen Feature", function () {
     this.webExt.destroy();
   });
 
-  describe("click the 'Always open in' checkbox in the popup", function () {
+  describe("set to 'Always open in' firefox-container-4", function () {
     beforeEach(async function () {
       // popup click to set assignment for activeTab.url
       await this.webExt.popup.helper.clickElementById("always-open-in");