Moving bulk of removeIdentity code into popup.js

index.js

@@ -197,11 +197,11 @@ const ContainerService = {
       "queryIdentities",
       "getIdentity",
       "createIdentity",
-      "removeIdentity",
       "updateIdentity",
       "getPreference",
       "sendTelemetryPayload",
       "getTheme",
+      "forgetIdentityAndRefresh",
       "checkIncompatibleAddons"
     ];
 
@@ -309,7 +309,7 @@ const ContainerService = {
   },
 
   registerBackgroundConnection(api) {
-    // This is only used for theme and container deletion notifications
+    // This is only used for theme notifications
     api.browser.runtime.onConnect.addListener((port) => {
       this._onBackgroundConnectCallback = (message, topic) => {
         port.postMessage({
@@ -967,30 +967,6 @@ const ContainerService = {
     });
   },
 
-  removeIdentity(args) {
-    const eventName = "delete-container";
-    if (!("userContextId" in args)) {
-      return Promise.reject("removeIdentity must be called with userContextId argument.");
-    }
-
-    this.sendTelemetryPayload({
-      "event": eventName,
-      "userContextId": args.userContextId
-    });
-
-    const tabsToClose = [];
-    this._containerTabIterator(args.userContextId, tab => {
-      tabsToClose.push(tab);
-    });
-
-    return this._closeTabs(tabsToClose).then(() => {
-      const removed = ContextualIdentityProxy.remove(args.userContextId);
-      this.triggerBackgroundCallback({userContextId: args.userContextId}, eventName);
-      this._forgetIdentity(args.userContextId);
-      return this._refreshNeeded().then(() => removed );
-    });
-  },
-
   // Preferences
 
   getPreference(args) {
@@ -1169,6 +1145,11 @@ const ContainerService = {
     // End-Of-Hack
   },
 
+  forgetIdentityAndRefresh(args) {
+    this._forgetIdentity(args.userContextId);
+    return this._refreshNeeded();
+  },
+
   _forgetIdentity(userContextId = 0) {
     for (let window of windows.browserWindows) { // eslint-disable-line prefer-const
       window = viewFor(window);

webextension/background.js

@@ -59,18 +59,25 @@ const assignManager = {
   },
 
   init() {
-    browser.runtime.onMessage.addListener((neverAskMessage) => {
-      const pageUrl = neverAskMessage.pageUrl;
-      if (neverAskMessage.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) => {
-          if (siteSettings) {
-            siteSettings.neverAsk = true;
-            this.storageArea.set(pageUrl, siteSettings);
+    browser.runtime.onMessage.addListener((m) => {
+      switch (m.type) {
+        case "delete-container":
+          assignManager.deleteContainer(m.message.userContextId);
+          break;
+        case "never-ask":
+          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) => {
+              if (siteSettings) {
+                siteSettings.neverAsk = true;
+                this.storageArea.set(pageUrl, siteSettings);
+              }
+            }).catch((e) => {
+              throw e;
+            });
           }
-        }).catch((e) => {
-          throw e;
-        });
+          break;
       }
     });
 
@@ -254,9 +261,6 @@ const messageHandler = {
       case "lightweight-theme-changed":
         themeManager.update(m.message);
         break;
-      case "delete-container":
-        assignManager.deleteContainer(m.message.userContextId);
-        break;
       default:
         throw new Error(`Unhandled message type: ${m.message}`);
       }

webextension/js/confirm-page.js

@@ -9,6 +9,7 @@ document.getElementById("redirect-form").addEventListener("submit", (e) => {
   // Sending neverAsk message to background to store for next time we see this process
   if (neverAsk) {
     browser.runtime.sendMessage({
+      type: "never-ask",
       neverAsk: true,
       pageUrl: redirectUrl
     }).then(() => {

webextension/js/popup.js

@@ -141,6 +141,58 @@ const Logic = {
     return this._currentIdentity;
   },
 
+  cookieStoreId(userContextId) {
+    return `firefox-container-${userContextId}`;
+  },
+
+  _containerTabIterator(userContextId, cb) {
+    browser.tabs.query({
+      cookieStoreId: Logic.cookieStoreId(userContextId)
+    }).then((tabs) => {
+      tabs.forEach((tab) => {
+        cb(tab);
+      });
+    });
+  },
+
+  _containers(userContextId) {
+    return browser.tabs.query({
+      cookieStoreId: Logic.cookieStoreId(userContextId)
+    });
+  },
+
+  removeIdentity(userContextId) {
+    const eventName = "delete-container";
+    if (!userContextId) {
+      return Promise.reject("removeIdentity must be called with userContextId argument.");
+    }
+
+    browser.runtime.sendMessage({
+      method: "sendTelemetryPayload",
+      event: eventName,
+      userContextId
+    });
+
+    const removeTabsPromise = Logic._containers(userContextId).then((tabs) => {
+      const tabIds = tabs.map((tab) => tab.id);
+      return browser.tabs.remove(tabIds);
+    });
+
+    return removeTabsPromise.then(() => {
+      const removed = browser.contextualIdentities.remove(Logic.cookieStoreId(userContextId));
+      // Send delete event to webextension/background.js
+      browser.runtime.sendMessage({
+        type: eventName,
+        message: {userContextId}
+      });
+      browser.runtime.sendMessage({
+        method: "forgetIdentityAndRefresh"
+      }).then(() => {
+        return removed;
+      });
+    });
+  },
+
   generateIdentityName() {
     const defaultName = "Container #";
     const ids = [];
@@ -603,10 +655,7 @@ Logic.registerPanel(P_CONTAINER_DELETE, {
     });
 
     document.querySelector("#delete-container-ok-link").addEventListener("click", () => {
-      browser.runtime.sendMessage({
-        method: "removeIdentity",
-        userContextId: Logic.currentIdentity().userContextId,
-      }).then(() => {
+      Logic.removeIdentity(Logic.currentIdentity().userContextId).then(() => {
         return Logic.refreshIdentities();
       }).then(() => {
         Logic.showPreviousPanel();

webextension/manifest.json

@@ -23,6 +23,7 @@
     "activeTab",
     "cookies",
     "contextMenus",
+    "contextualIdentities",
     "history",
     "idle",
     "notifications",