Updates suggested by groovecoder on https://github.com/mozilla/multi-account-containers/pull/1474

- Renamed "isLockedOut" constant so it's easier to tell what it means.
- Moved inline code blocks into separate functions, for improved readability.
- Fixed unnecessary use of "const that = this;"
- Improved code comments

src/js/background/assignManager.js

@@ -142,22 +142,31 @@ const assignManager = {
     }
     const userContextId = this.getUserContextIdFromCookieStore(tab);
     
-    // 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) {
+    // Handle the case where this request's URL is not assigned to any particular
-      const currentContainerState = await identityState.storageArea.get(tab.cookieStoreId);
+    // container. We must do the following check:
-      isLockedOut = !!currentContainerState.isLocked;
+    //
-    } else {
+    // If the current tab's container is "unlocked", we can just go ahead
-      isLockedOut = false;
+    // 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 mustReloadPageInDefaultContainerDueToLocking = await this._determineMustReloadPageInDefaultContainerDueToLocking(siteSettings, tab);
     
-    if ((!siteSettings && !isLockedOut)
+    if (!mustReloadPageInDefaultContainerDueToLocking) {
-        || (siteSettings && userContextId === siteSettings.userContextId)
+      if (!siteSettings
-        || this.storageArea.isExempted(options.url, tab.id)) {
+          || userContextId === siteSettings.userContextId
-      return {};
+          || this.storageArea.isExempted(options.url, tab.id)) {
+        return {};
+      }
     }
     const removeTab = backgroundLogic.NEW_TAB_PAGES.has(tab.url)
       || (messageHandler.lastCreatedTab
@@ -201,12 +210,11 @@ const assignManager = {
       }
     }
 
-    if (isLockedOut) {
+    if (mustReloadPageInDefaultContainerDueToLocking) {
-      // Open new tab in default context
+      // Open new tab in default container
-      // https://github.com/mozilla/multi-account-containers/issues/847
       browser.tabs.create({url: options.url});
     } else {
-      // Open new tab in specific context
+      // Open new tab in assigned container
       this.reloadPageInContainer(
         options.url,
         userContextId,
@@ -238,6 +246,25 @@ const assignManager = {
     };
   },
   
+  async _determineMustReloadPageInDefaultContainerDueToLocking(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;
+    }
+    
+    // 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.isLocked;
+  },
+
   init() {
     browser.contextMenus.onClicked.addListener((info, tab) => {
       info.bookmarkId ? this._onClickedBookmark(info) : this._onClickedHandler(info, tab);
@@ -417,20 +444,12 @@ const assignManager = {
       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";
+      
+      // Unlock container if now empty
+      await this._unlockContainerIfHasNoAssignments(userContextId);    
     }
     browser.tabs.sendMessage(tabId, {
       text: `Successfully ${actionName} site to always open in this container`
@@ -439,6 +458,17 @@ const assignManager = {
     this.calculateContextMenu(tab);
   },
   
+  async _unlockContainerIfHasNoAssignments(userContextId) {
+    const assignments = await this.storageArea.getByContainer(userContextId);
+    const hasAssignments = assignments && Object.keys(assignments).length > 0;
+    if (!hasAssignments) {
+      await backgroundLogic.lockOrUnlockContainer({
+        userContextId: userContextId,
+        isLocked: false
+      });
+    }
+  },
+
   async _getAssignment(tab) {
     const cookieStore = this.getUserContextIdFromCookieStore(tab);
     // Ensure we have a cookieStore to assign to

src/js/background/backgroundLogic.js

@@ -246,7 +246,6 @@ const backgroundLogic = {
         hasOpenTabs: !!openTabs.length,
         numberOfHiddenTabs: containerState.hiddenTabs.length,
         numberOfOpenTabs: openTabs.length,
-        // https://github.com/mozilla/multi-account-containers/issues/847
         isLocked: !!containerState.isLocked
       };
       return;

src/js/background/messageHandler.js

@@ -20,7 +20,6 @@ const messageHandler = {
         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":

src/js/popup.js

@@ -4,7 +4,6 @@
 
 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";
 
@@ -255,7 +254,6 @@ 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;
@@ -1041,7 +1039,6 @@ Logic.registerPanel(P_CONTAINER_EDIT, {
       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({
@@ -1051,12 +1048,12 @@ Logic.registerPanel(P_CONTAINER_EDIT, {
               isLocked: !isLocked
             }
           });
-          that.showAssignedContainers(assignments, !isLocked);
+          this.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 */
+      /* Container locking end */
       
       /* Assignment list */
       assignmentKeys.forEach((siteKey) => {