Wildcard subdomains - e.g. *.google.com

src/css/popup.css

@@ -818,6 +818,11 @@ span ~ .panel-header-text {
   flex: 1;
 }
 
+/* Wildcard subdomains: https://github.com/mozilla/multi-account-containers/issues/473 */
+.assigned-sites-list .hostname .subdomain:hover {
+  text-decoration: underline;
+}
+
 .radio-choice > .radio-container {
   align-items: center;
   block-size: 29px;

src/js/background/assignManager.js

@@ -1,3 +1,183 @@
+/**
+  Utils for dealing with hosts.
+  
+  E.g. www.google.com:443
+  */
+const HostUtils = {
+  getHost(pageUrl) {
+    const url = new window.URL(pageUrl);
+    if (url.port === "80" || url.port === "443") {
+      return `${url.hostname}`;
+    } else {
+      return `${url.hostname}${url.port}`;
+    }
+  },
+  
+  // Wildcard subdomains: https://github.com/mozilla/multi-account-containers/issues/473
+  hasSubdomain(host) {
+    return host.indexOf(".") >= 0;
+  },
+  
+  removeSubdomain(host) {
+    const indexOfDot = host.indexOf(".");
+    if (indexOfDot < 0) {
+      return null;
+    } else {
+      return host.substring(indexOfDot + 1);
+    }
+  }    
+};
+
+/**
+  Store data in 'named stores'.
+  
+  (In actual fact, all data for all stores is stored in the same storage area,
+  but this class provides accessor methods to get/set only the data that applies
+  to one specific named store, as identified in the constructor.)
+ */
+class AssignStore {
+  constructor(name) {
+    this.prefix = `${name}@@_`;
+  }
+
+  _storeKeyForKey(key) {
+    if (Array.isArray(key)) {
+      return key.map(oneKey => oneKey.startsWith(this.prefix) ? oneKey : `${this.prefix}${oneKey}`);
+    } else if (key) {
+      return key.startsWith(this.prefix) ? key : `${this.prefix}${key}`;
+    } else {
+      return null;
+    }
+  }
+  
+  _keyForStoreKey(storeKey) {
+    if (Array.isArray(storeKey)) {
+      return storeKey.map(oneStoreKey => oneStoreKey.startsWith(this.prefix) ? oneStoreKey.substring(this.prefix.length) : null);
+    } else if (storeKey) {
+      return storeKey.startsWith(this.prefix) ? storeKey.substring(this.prefix.length) : null;
+    } else {
+      return null;
+    }
+  }
+  
+  get(key) {
+    if (typeof key !== "string") { return Promise.reject(new Error(`[AssignStore.get] Invalid key: ${key}`)); }  
+    const storeKey = this._storeKeyForKey(key);
+    return new Promise((resolve, reject) => {
+      browser.storage.local.get([storeKey]).then((storageResponse) => {
+        if (storeKey in storageResponse) {
+          resolve(storageResponse[storeKey]);
+        } else {
+          resolve(null);
+        }
+      }).catch((e) => {
+        reject(e);
+      });
+    });
+  }
+
+  getAll(keys) {
+    if (keys && !Array.isArray(keys)) { return Promise.reject(new Error(`[AssignStore.getAll] Invalid keys: ${keys}`)); }
+    const storeKeys = this._storeKeyForKey(keys);
+    return new Promise((resolve, reject) => {
+      browser.storage.local.get(storeKeys).then((storageResponse) => {
+        if (storageResponse) {
+          resolve(Object.assign({}, ...Object.entries(storageResponse).map(([oneStoreKey, data]) => {
+            const key = this._keyForStoreKey(oneStoreKey);
+            return key ? { [key]: data } : null;
+          })));
+        } else {
+          resolve(null);
+        }
+      }).catch((e) => {
+        reject(e);
+      });
+    });
+  }
+  
+  set(key, data) {
+    if (typeof key !== "string") { return Promise.reject(new Error(`[AssignStore.set] Expected String, but received ${key}`)); }
+    const storeKey = this._storeKeyForKey(key);
+    return browser.storage.local.set({
+      [storeKey]: data
+    });
+  }
+
+  remove(key) {
+    if (typeof key !== "string") { return Promise.reject(new Error(`[AssignStore.remove] Expected String, but received ${key}`)); }
+    const storeKey = this._storeKeyForKey(key);
+    return browser.storage.local.remove(storeKey);
+  }
+  
+  removeAll(keys) {
+    if (keys && !Array.isArray(keys)) { return Promise.reject(new Error(`[AssignStore.removeAll] Invalid keys: ${keys}`)); }
+    const storeKeys = this._storeKeyForKey(keys);
+    return browser.storage.local.remove(storeKeys);
+  }
+}
+
+/**
+  Manages mappings of Site Host <-> Wildcard Host.
+  
+  E.g. drive.google.com <-> google.com
+  
+  Wildcard subdomains: https://github.com/mozilla/multi-account-containers/issues/473
+ */
+const WildcardManager = {
+  bySite:     new AssignStore("siteToWildcardMap"),
+  byWildcard: new AssignStore("wildcardToSiteMap"),
+
+  // Site -> Wildcard
+  get(site) {
+    return this.bySite.get(site);
+  },
+  
+  async getAll(sites) {
+    return this.bySite.getAll(sites);
+  },
+  
+  async set(site, wildcard) {
+    // Remove existing site -> wildcard
+    const oldSite = await this.byWildcard.get(wildcard);
+    if (oldSite) { await this.bySite.remove(oldSite); }
+    
+    // Set new mappings site <-> wildcard
+    await this.bySite.set(site, wildcard);
+    await this.byWildcard.set(wildcard, site);
+  },
+
+  async remove(site) {
+    const wildcard = await this.bySite.get(site);
+    if (!wildcard) { return; }
+    
+    await this.bySite.remove(site);
+    await this.byWildcard.remove(wildcard);
+  },
+  
+  async removeAll(sites) {
+    const data = await this.bySite.getAll(sites);
+    const existingSites = Object.keys(data);
+    const existingWildcards = Object.values(data);
+    
+    await this.bySite.removeAll(existingSites);
+    await this.byWildcard.removeAll(existingWildcards);
+  },
+    
+  // Site -> Site that owns Wildcard
+  async match(site) {
+    // Keep stripping subdomains off site domain until match a wildcard domain
+    do {
+      // Use the ever-shortening site hostname as if it is a wildcard
+      const siteHavingWildcard = await this.byWildcard.get(site);
+      if (siteHavingWildcard) { return siteHavingWildcard; }
+    } while ((site = HostUtils.removeSubdomain(site)));
+    return null;
+  }        
+};
+
+/**
+  Main interface for managing assignments.
+ */
 const assignManager = {
   MENU_ASSIGN_ID: "open-in-this-container",
   MENU_REMOVE_ID: "remove-open-in-this-container",
@@ -6,91 +186,123 @@ const assignManager = {
   MENU_MOVE_ID: "move-to-new-window-container",
 
   storageArea: {
-    area: browser.storage.local,
+    store: new AssignStore("siteContainerMap"),
     exemptedTabs: {},
 
-    getSiteStoreKey(pageUrl) {
+    setExempted(host, tabId) {
-      const url = new window.URL(pageUrl);
+      if (!(host in this.exemptedTabs)) {
-      const storagePrefix = "siteContainerMap@@_";
+        this.exemptedTabs[host] = [];
-      if (url.port === "80" || url.port === "443") {
-        return `${storagePrefix}${url.hostname}`;
-      } else {
-        return `${storagePrefix}${url.hostname}${url.port}`;
       }
+      this.exemptedTabs[host].push(tabId);
     },
 
-    setExempted(pageUrl, tabId) {
+    removeExempted(host) {
-      const siteStoreKey = this.getSiteStoreKey(pageUrl);
+      this.exemptedTabs[host] = [];
-      if (!(siteStoreKey in this.exemptedTabs)) {
-        this.exemptedTabs[siteStoreKey] = [];
-      }
-      this.exemptedTabs[siteStoreKey].push(tabId);
     },
 
-    removeExempted(pageUrl) {
+    isExemptedUrl(pageUrl, tabId) {
-      const siteStoreKey = this.getSiteStoreKey(pageUrl);
+      const host = HostUtils.getHost(pageUrl);
-      this.exemptedTabs[siteStoreKey] = [];
+      if (!(host in this.exemptedTabs)) {
-    },
-
-    isExempted(pageUrl, tabId) {
-      const siteStoreKey = this.getSiteStoreKey(pageUrl);
-      if (!(siteStoreKey in this.exemptedTabs)) {
         return false;
       }
-      return this.exemptedTabs[siteStoreKey].includes(tabId);
+      return this.exemptedTabs[host].includes(tabId);
     },
 
-    get(pageUrl) {
+    async matchUrl(pageUrl) {
-      const siteStoreKey = this.getSiteStoreKey(pageUrl);
+      const host = HostUtils.getHost(pageUrl);
-      return new Promise((resolve, reject) => {
+      
-        this.area.get([siteStoreKey]).then((storageResponse) => {
+      // Try exact match
-          if (storageResponse && siteStoreKey in storageResponse) {
+      const result = await this.get(host);
-            resolve(storageResponse[siteStoreKey]);
+      if (result) { return result; }
-          }
+      
-          resolve(null);
+      // Try wildcard match
-        }).catch((e) => {
+      const wildcard = await WildcardManager.match(host);
-          reject(e);
+      if (wildcard) { return await this.get(wildcard); }
-        });
+      
-      });
+      return null;
     },
 
-    set(pageUrl, data, exemptedTabIds) {
+    async get(host) {
-      const siteStoreKey = this.getSiteStoreKey(pageUrl);
+      const result = await this.store.get(host);
+      if (result) {
+        if (result.host !== host) { result.host = host; }
+        result.wildcard = await WildcardManager.get(host);
+      }
+      return result;
+    },
+
+    async set(host, data, exemptedTabIds, wildcard) {
+      // Store exempted tabs
       if (exemptedTabIds) {
         exemptedTabIds.forEach((tabId) => {
-          this.setExempted(pageUrl, tabId);
+          this.setExempted(host, tabId);
         });
       }
-      return this.area.set({
+      // Store wildcard mapping
-        [siteStoreKey]: data
+      if (wildcard) {
-      });
+        if (wildcard === host) {
+          await WildcardManager.remove(host);
+        } else {
+          await WildcardManager.set(host, wildcard);        
+        }
+      }
+      // Do not store wildcard property
+      if (data.wildcard) {
+        data = Object.assign(data);
+        delete data.wildcard;
+      }
+      // Store assignment
+      return this.store.set(host, data);
     },
 
-    remove(pageUrl) {
+    async remove(host) {
-      const siteStoreKey = this.getSiteStoreKey(pageUrl);
       // When we remove an assignment we should clear all the exemptions
-      this.removeExempted(pageUrl);
+      this.removeExempted(host);
-      return this.area.remove([siteStoreKey]);
+      // ...and also clear the wildcard mapping
+      await WildcardManager.remove(host);
+      
+      return this.store.remove(host);
     },
 
     async deleteContainer(userContextId) {
       const sitesByContainer = await this.getByContainer(userContextId);
-      this.area.remove(Object.keys(sitesByContainer));
+      const sites = Object.keys(sitesByContainer);
+      
+      sites.forEach((site) => {
+        // When we remove an assignment we should clear all the exemptions
+        this.removeExempted(site);
+      });
+      
+      // ...and also clear the wildcard mappings
+      await WildcardManager.removeAll(sites);
+
+      return this.store.removeAll(sites);
     },
 
     async getByContainer(userContextId) {
-      const sites = {};
+      // Get sites
-      const siteConfigs = await this.area.get();
+      const sitesConfig = await this.store.getAll();
-      Object.keys(siteConfigs).forEach((key) => {
+      const sites = Object.assign({}, ...Object.entries(sitesConfig).map(([host, data]) => {
         // For some reason this is stored as string... lets check them both as that
-        if (String(siteConfigs[key].userContextId) === String(userContextId)) {
+        if (String(data.userContextId) === String(userContextId)) {
-          const site = siteConfigs[key];
           // In hindsight we should have stored this
           // TODO file a follow up to clean the storage onLoad
-          site.hostname = key.replace(/^siteContainerMap@@_/, "");
+          data.host = host;
-          sites[key] = site;
+          return { [host]: data };
-        }
+        } else {
-      });
+          return null;
+        }      
+      }));
+      
+      // Add wildcards
+      const hosts = Object.keys(sites);
+      if (hosts.length > 0) {
+        const sitesToWildcards = await WildcardManager.getAll(hosts);
+        Object.entries(sitesToWildcards).forEach(([site, wildcard]) => {
+          sites[site].wildcard = wildcard;
+        });
+      }
+      
       return sites;
     }
   },
@@ -99,10 +311,10 @@ const assignManager = {
     const pageUrl = m.pageUrl;
     if (m.neverAsk === true) {
       // If we have existing data and for some reason it hasn't been deleted etc lets update it
-      this.storageArea.get(pageUrl).then((siteSettings) => {
+      this.storageArea.matchUrl(pageUrl).then((siteSettings) => {
         if (siteSettings) {
           siteSettings.neverAsk = true;
-          this.storageArea.set(pageUrl, siteSettings);
+          return this.storageArea.set(siteSettings.host, siteSettings);
         }
       }).catch((e) => {
         throw e;
@@ -113,7 +325,8 @@ const assignManager = {
   // We return here so the confirm page can load the tab when exempted
   async _exemptTab(m) {
     const pageUrl = m.pageUrl;
-    this.storageArea.setExempted(pageUrl, m.tabId);
+    const host = HostUtils.getHost(pageUrl);
+    this.storageArea.setExempted(host, m.tabId);
     return true;
   },
 
@@ -125,7 +338,7 @@ const assignManager = {
     this.removeContextMenu();
     const [tab, siteSettings] = await Promise.all([
       browser.tabs.get(options.tabId),
-      this.storageArea.get(options.url)
+      this.storageArea.matchUrl(options.url)
     ]);
     let container;
     try {
@@ -144,7 +357,7 @@ const assignManager = {
     if (!siteSettings
         || userContextId === siteSettings.userContextId
         || tab.incognito
-        || this.storageArea.isExempted(options.url, tab.id)) {
+        || this.storageArea.isExemptedUrl(options.url, tab.id)) {
       return {};
     }
     const removeTab = backgroundLogic.NEW_TAB_PAGES.has(tab.url)
@@ -299,7 +512,7 @@ const assignManager = {
     return true;
   },
 
-  async _setOrRemoveAssignment(tabId, pageUrl, userContextId, remove) {
+  async _setOrRemoveAssignment(tabId, pageUrl, userContextId, remove, options) {
     let actionName;
 
     // https://github.com/mozilla/testpilot-containers/issues/626
@@ -307,13 +520,15 @@ const assignManager = {
     // the value to a string for accurate checking
     userContextId = String(userContextId);
 
+    const assignmentHost = HostUtils.getHost(pageUrl);
     if (!remove) {
       const tabs = await browser.tabs.query({});
-      const assignmentStoreKey = this.storageArea.getSiteStoreKey(pageUrl);
+      const wildcardHost = options && options.wildcard ? options.wildcard : null;
       const exemptedTabIds = tabs.filter((tab) => {
-        const tabStoreKey = this.storageArea.getSiteStoreKey(tab.url);
+        const tabHost = HostUtils.getHost(tab.url);
         /* Auto exempt all tabs that exist for this hostname that are not in the same container */
-        if (tabStoreKey === assignmentStoreKey &&
+        if ( (tabHost === assignmentHost ||
+              (wildcardHost && tabHost.endsWith(wildcardHost))) &&
             this.getUserContextIdFromCookieStore(tab) !== userContextId) {
           return true;
         }
@@ -321,29 +536,38 @@ const assignManager = {
       }).map((tab) => {
         return tab.id;
       });
-
+      
-      await this.storageArea.set(pageUrl, {
+      await this.storageArea.set(assignmentHost, {
         userContextId,
         neverAsk: false
-      }, exemptedTabIds);
+      }, exemptedTabIds, (wildcardHost || assignmentHost));
       actionName = "added";
     } else {
-      await this.storageArea.remove(pageUrl);
+      await this.storageArea.remove(assignmentHost);
       actionName = "removed";
     }
-    browser.tabs.sendMessage(tabId, {
+    if (!options || !options.silent) {
-      text: `Successfully ${actionName} site to always open in this container`
+      browser.tabs.sendMessage(tabId, {
-    });
+        text: `Successfully ${actionName} site to always open in this container`
+      });
+    }
     const tab = await browser.tabs.get(tabId);
     this.calculateContextMenu(tab);
   },
+  
+  async _setOrRemoveWildcard(tabId, pageUrl, userContextId, wildcard) {
+    // Remove assignment
+    await this._setOrRemoveAssignment(tabId, pageUrl, userContextId, true, {silent:true});
+    // Add assignment
+    await this._setOrRemoveAssignment(tabId, pageUrl, userContextId, false, {wildcard:wildcard, silent:true});  
+  },
 
   async _getAssignment(tab) {
     const cookieStore = this.getUserContextIdFromCookieStore(tab);
     // Ensure we have a cookieStore to assign to
     if (cookieStore
         && this.isTabPermittedAssign(tab)) {
-      return await this.storageArea.get(tab.url);
+      return await this.storageArea.matchUrl(tab.url);
     }
     return false;
   },

src/js/background/messageHandler.js

@@ -34,6 +34,12 @@ const messageHandler = {
           return assignManager._setOrRemoveAssignment(tab.id, m.url, m.userContextId, m.value);
         });
         break;
+      case "setOrRemoveWildcard":
+        // Wildcard subdomains: https://github.com/mozilla/multi-account-containers/issues/473
+        response = browser.tabs.get(m.tabId).then((tab) => {
+          return assignManager._setOrRemoveWildcard(tab.id, m.url, m.userContextId, m.wildcard);
+        });      
+        break;
       case "sortTabs":
         backgroundLogic.sortTabs();
         break;

src/js/popup.js

@@ -361,6 +361,17 @@ const Logic = {
     });
   },
 
+  // Wildcard subdomains: https://github.com/mozilla/multi-account-containers/issues/473
+  setOrRemoveWildcard(tabId, url, userContextId, wildcard) {
+    return browser.runtime.sendMessage({
+      method: "setOrRemoveWildcard",
+      tabId,
+      url,
+      userContextId,
+      wildcard
+    });
+  },
+  
   generateIdentityName() {
     const defaultName = "Container #";
     const ids = [];
@@ -381,7 +392,7 @@ const Logic = {
         return defaultName + (id < 10 ? "0" : "") + id;
       }
     }
-  },
+  }
 };
 
 // P_ONBOARDING_1: First page for Onboarding.
@@ -985,18 +996,40 @@ Logic.registerPanel(P_CONTAINER_EDIT, {
         const trElement = document.createElement("div");
         /* As we don't have the full or correct path the best we can assume is the path is HTTPS and then replace with a broken icon later if it doesn't load.
            This is pending a better solution for favicons from web extensions */
-        const assumedUrl = `https://${site.hostname}`;
+        const assumedUrl = `https://${site.host}`;
         trElement.innerHTML = escaped`
         <img class="icon" src="${assumedUrl}/favicon.ico">
-        <div title="${site.hostname}" class="truncate-text hostname">
+        <div title="${site.host}" class="truncate-text hostname"></div>
-          ${site.hostname}
-        </div>
         <img
           class="pop-button-image delete-assignment"
           src="/img/container-delete.svg"
-        />`;
+        />`;        
-        const deleteButton = trElement.querySelector(".delete-assignment");
+        trElement.querySelector(".hostname").appendChild(this.assignmentElement(site));
+        
         const that = this;
+        
+        // Wildcard subdomains: https://github.com/mozilla/multi-account-containers/issues/473
+        trElement.querySelectorAll(".subdomain").forEach(function(subdomainLink) {
+          subdomainLink.addEventListener("click", async (e) => {
+            const userContextId = Logic.currentUserContextId();
+            // Wildcard hostname is stored in id attribute
+            const wildcard = e.target.id;
+            if (wildcard) {
+              // Remove wildcard from other site that has same wildcard
+              Object.values(assignments).forEach((site) => {
+                if (site.wildcard === wildcard) { delete site.wildcard; }
+              });
+              site.wildcard = wildcard;
+            } else {
+              delete site.wildcard;
+            }
+            const currentTab = await Logic.currentTab();
+            Logic.setOrRemoveWildcard(currentTab.id, assumedUrl, userContextId, wildcard);
+            that.showAssignedContainers(assignments);
+          });
+        });
+        
+        const deleteButton = trElement.querySelector(".delete-assignment");
         Logic.addEnterHandler(deleteButton, async () => {
           const userContextId = Logic.currentUserContextId();
           // Lets show the message to the current tab
@@ -1006,11 +1039,46 @@ Logic.registerPanel(P_CONTAINER_EDIT, {
           delete assignments[siteKey];
           that.showAssignedContainers(assignments);
         });
+        
         trElement.classList.add("container-info-tab-row", "clickable");
         tableElement.appendChild(trElement);
       });
     }
   },
+  
+  // Wildcard subdomains: https://github.com/mozilla/multi-account-containers/issues/473
+  assignmentElement(site) {
+    const result = document.createElement("span");
+  
+    // Remove wildcard subdomain
+    if (site.wildcard && site.wildcard !== site.host) {
+      result.appendChild(this.assignmentSubdomainLink(null, "___"));
+      result.appendChild(document.createTextNode("."));
+    }
+  
+    // Add wildcard subdomain
+    let host = site.wildcard ? site.wildcard : site.host;
+    let indexOfDot;
+    while ((indexOfDot = host.indexOf(".")) >= 0) {
+      const subdomain = host.substring(0, indexOfDot);
+      host = host.substring(indexOfDot + 1);
+      result.appendChild(this.assignmentSubdomainLink(host, subdomain));
+      result.appendChild(document.createTextNode("."));
+    }
+
+    // Root domain
+    result.appendChild(document.createTextNode(host));
+  
+    return result;
+  },
+  
+  assignmentSubdomainLink(wildcard, text) {
+    const result = document.createElement("a");
+    if (wildcard) { result.id = wildcard; }
+    result.className = "subdomain";
+    result.appendChild(document.createTextNode(text));
+    return result;
+  },
 
   initializeRadioButtons() {
     const colorRadioTemplate = (containerColor) => {

test/features/wildcard.test.js

@@ -0,0 +1,55 @@
+// Wildcard subdomains: https://github.com/mozilla/multi-account-containers/issues/473
+describe("Wildcard Subdomains Feature", () => {
+  const activeTab = {
+    id: 1,
+    cookieStoreId: "firefox-container-1",
+    url: "http://www.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("click the assigned URL's subdomain to convert it to a wildcard", () => {
+      beforeEach(async () => {
+        await helper.popup.setWildcard(activeTab, "example.com");
+      });
+
+      describe("open new Tab with a different subdomain in the default container", () => {
+        const newTab = {
+          id: 2,
+          cookieStoreId: "firefox-default",
+          url: "http://mail.example.com",
+          index: 1,
+          active: true
+        };
+        beforeEach(async () => {
+          await helper.browser.openNewTab(newTab);
+        });
+
+        it("should open the confirm page", async () => {
+          // should have created a new tab with the confirm page
+          background.browser.tabs.create.should.have.been.calledWith({
+            url: "moz-extension://multi-account-containers/confirm-page.html?" +
+                 `url=${encodeURIComponent(newTab.url)}` +
+                 `&cookieStoreId=${activeTab.cookieStoreId}`,
+            cookieStoreId: undefined,
+            openerTabId: null,
+            index: 2,
+            active: true
+          });
+        });
+
+        it("should remove the new Tab that got opened in the default container", () => {
+          background.browser.tabs.remove.should.have.been.calledWith(newTab.id);
+        });
+      });
+    });
+  });
+});

test/helper.js

@@ -42,6 +42,15 @@ module.exports = {
       clickEvent.initEvent("click");
       popup.document.getElementById(id).dispatchEvent(clickEvent);
       await nextTick();
+    },
+    
+    // Wildcard subdomains: https://github.com/mozilla/multi-account-containers/issues/473
+    async setWildcard(tab, wildcard) {
+      const site = new URL(tab.url).hostname;
+      const siteToWildcardKey = `siteToWildcardMap@@_${site}`;
+      const wildcardToSiteKey = `wildcardToSiteMap@@_${wildcard}`;
+      await background.browser.storage.local.set({[siteToWildcardKey]: wildcard});
+      await background.browser.storage.local.set({[wildcardToSiteKey]: site});
     }
   },
 };