From b8cf955e1f3024cca355a47bc45361583f4477e2 Mon Sep 17 00:00:00 2001
From: Minigugus <43109623+Minigugus@users.noreply.github.com>
Date: Sun, 12 Jun 2022 11:37:21 +0200
Subject: [PATCH] Add "isIsolated" to JSON backups

---
 src/js/background/backgroundLogic.js | 47 +++++++++++++++++++++-------
 1 file changed, 35 insertions(+), 12 deletions(-)

diff --git a/src/js/background/backgroundLogic.js b/src/js/background/backgroundLogic.js
index 309cc1e..de59166 100644
--- a/src/js/background/backgroundLogic.js
+++ b/src/js/background/backgroundLogic.js
@@ -277,9 +277,21 @@ const backgroundLogic = {
     return Promise.all(
       identities.map(async ({ cookieStoreId, color, icon, name }) => {
         const userContextId = this.getUserContextIdFromCookieStoreId(cookieStoreId);
-        const sitesByContainer = await assignManager.storageArea.getAssignedSites(userContextId);
+        const [
+          { isIsolated },
+          sitesByContainer
+        ] = await Promise.all([
+          identityState.storageArea.get(cookieStoreId),
+          assignManager.storageArea.getAssignedSites(userContextId)
+        ]);
         const sites = Object.values(sitesByContainer).map(({ neverAsk, hostname }) => ({ neverAsk, hostname }));
-        return ({ color, icon, name, sites });
+        return ({
+          color,
+          icon,
+          name,
+          isolated: isIsolated && true, // either `true` or `undefined`
+          sites
+        });
       })
     );
   },
@@ -288,13 +300,19 @@ const backgroundLogic = {
     const backup = await browser.contextualIdentities.query({});
     const incomplete = [];
     let allSucceed = true;
-    const identitiesPromise = identities.map(async ({ color, icon, name, sites }) => {
+    const identitiesPromise = identities.map(async ({ color, icon, name, isolated, sites }) => {
       try {
-        if (typeof color !== "string" || typeof icon !== "string" || typeof name !== "string" || !Array.isArray((sites)))
+        if (
+          typeof color !== "string" || 
+          typeof icon !== "string" || 
+          typeof name !== "string" ||
+          (isolated !== true && isolated !== undefined) ||
+          !Array.isArray((sites))
+        )
           throw new Error("Corrupted container backup");
         const identity = await browser.contextualIdentities.create({ color, icon, name });
         try {
-          await identityState.storageArea.get(identity.cookieStoreId);
+          await identityState.storageArea.get(identity.cookieStoreId); // to create identity state
           const userContextId = this.getUserContextIdFromCookieStoreId(identity.cookieStoreId);
           for (const { neverAsk, hostname } of sites) {
             if (typeof neverAsk !== "boolean" || typeof hostname !== "string" || hostname === "")
@@ -305,6 +323,8 @@ const backgroundLogic = {
               userContextId
             });
           }
+          if (isolated)
+            await identityState.storageArea.set(identity.cookieStoreId, { isIsolated: "locked" });
         } catch (err) {
           incomplete.push(name); // site association damaged
         }
@@ -314,6 +334,7 @@ const backgroundLogic = {
         return null;
       }
     });
+
     const created = await Promise.all(identitiesPromise);
     if (!allSucceed) { // Importation failed, restore previous state
       await Promise.all(
@@ -325,14 +346,16 @@ const backgroundLogic = {
         })
       );
       throw new Error("Some containers couldn't be created");
-    } else { // Importation succeed, remove old identities
-      await Promise.all(
-        backup.map(async (identity) => {
-          await identityState.storageArea.remove(identity.cookieStoreId);
-          await browser.contextualIdentities.remove(identity.cookieStoreId);
-        })
-      );
     }
+
+    // Importation succeed, remove old identities
+    await Promise.all(
+      backup.map(async (identity) => {
+        await identityState.storageArea.remove(identity.cookieStoreId);
+        await browser.contextualIdentities.remove(identity.cookieStoreId);
+      })
+    );
+
     return { created: created.length, incomplete };
   },