From 54586daa59e4d46968fbec498b31bf992e557e43 Mon Sep 17 00:00:00 2001
From: Damien Cassou <damien@cassou.me>
Date: Fri, 13 Oct 2023 08:34:32 +0200
Subject: [PATCH] firefox: add container support

---
 modules/programs/firefox.nix                  | 104 +++++++++++++++++-
 tests/modules/programs/firefox/default.nix    |   1 +
 .../firefox/duplicate-container-ids.nix       |  35 ++++++
 .../profile-settings-expected-containers.json |   1 +
 .../programs/firefox/profile-settings.nix     |  17 ++-
 5 files changed, 155 insertions(+), 3 deletions(-)
 create mode 100644 tests/modules/programs/firefox/duplicate-container-ids.nix
 create mode 100644 tests/modules/programs/firefox/profile-settings-expected-containers.json

diff --git a/modules/programs/firefox.nix b/modules/programs/firefox.nix
index 568472c54..1bcb0307b 100644
--- a/modules/programs/firefox.nix
+++ b/modules/programs/firefox.nix
@@ -59,6 +59,24 @@ let
       ${extraPrefs}
     '';
 
+  mkContainersJson = containers:
+    let
+      containerToIdentity = _: container: {
+        userContextId = container.id;
+        name = container.name;
+        icon = container.icon;
+        color = container.color;
+        public = true;
+      };
+    in ''
+      ${builtins.toJSON {
+        version = 4;
+        lastUserContextId =
+          elemAt (mapAttrsToList (_: container: container.id) containers) 0;
+        identities = mapAttrsToList containerToIdentity containers;
+      }}
+    '';
+
   firefoxBookmarksFile = bookmarks:
     let
       indent = level:
@@ -457,6 +475,83 @@ in {
               };
             };
 
+            containers = mkOption {
+              type = types.attrsOf (types.submodule ({ name, ... }: {
+                options = {
+                  name = mkOption {
+                    type = types.str;
+                    default = name;
+                    description = "Container name, e.g., shopping.";
+                  };
+
+                  id = mkOption {
+                    type = types.ints.unsigned;
+                    default = 0;
+                    description = ''
+                      Container ID. This should be set to a unique number per container in this profile.
+                    '';
+                  };
+
+                  # List of colors at
+                  # https://searchfox.org/mozilla-central/rev/5ad226c7379b0564c76dc3b54b44985356f94c5a/toolkit/components/extensions/parent/ext-contextualIdentities.js#32
+                  color = mkOption {
+                    type = types.enum [
+                      "blue"
+                      "turquoise"
+                      "green"
+                      "yellow"
+                      "orange"
+                      "red"
+                      "pink"
+                      "purple"
+                      "toolbar"
+                    ];
+                    default = "pink";
+                    description = "Container color.";
+                  };
+
+                  icon = mkOption {
+                    type = types.enum [
+                      "briefcase"
+                      "cart"
+                      "circle"
+                      "dollar"
+                      "fence"
+                      "fingerprint"
+                      "gift"
+                      "vacation"
+                      "food"
+                      "fruit"
+                      "pet"
+                      "tree"
+                      "chill"
+                    ];
+                    default = "fruit";
+                    description = "Container icon.";
+                  };
+                };
+              }));
+              default = { };
+              example = {
+                "shopping" = {
+                  id = 1;
+                  color = "blue";
+                  icon = "cart";
+                };
+                "dangerous" = {
+                  id = 2;
+                  color = "red";
+                  icon = "fruit";
+                };
+              };
+              description = ''
+                Attribute set of container configurations. See
+                [Multi-Account
+                Containers](https://support.mozilla.org/en-US/kb/containers)
+                for more information.
+              '';
+            };
+
             extensions = mkOption {
               type = types.listOf types.package;
               default = [ ];
@@ -514,7 +609,9 @@ in {
       })
 
       (mkNoDuplicateAssertion cfg.profiles "profile")
-    ];
+    ] ++ (mapAttrsToList
+      (_: profile: mkNoDuplicateAssertion profile.containers "container")
+      cfg.profiles);
 
     warnings = optional (cfg.enableGnomeExtensions or false) ''
       Using 'programs.firefox.enableGnomeExtensions' has been deprecated and
@@ -560,6 +657,11 @@ in {
             mkUserJs profile.settings profile.extraConfig profile.bookmarks;
         };
 
+      "${profilesPath}/${profile.path}/containers.json" =
+        mkIf (profile.containers != { }) {
+          text = mkContainersJson profile.containers;
+        };
+
       "${profilesPath}/${profile.path}/search.json.mozlz4" = mkIf
         (profile.search.default != null || profile.search.order != [ ]
           || profile.search.engines != { }) {
diff --git a/tests/modules/programs/firefox/default.nix b/tests/modules/programs/firefox/default.nix
index 1df80c5bc..5492e383c 100644
--- a/tests/modules/programs/firefox/default.nix
+++ b/tests/modules/programs/firefox/default.nix
@@ -3,4 +3,5 @@
   firefox-state-version-19_09 = ./state-version-19_09.nix;
   firefox-deprecated-native-messenger = ./deprecated-native-messenger.nix;
   firefox-duplicate-profile-ids = ./duplicate-profile-ids.nix;
+  firefox-duplicate-container-ids = ./duplicate-container-ids.nix;
 }
diff --git a/tests/modules/programs/firefox/duplicate-container-ids.nix b/tests/modules/programs/firefox/duplicate-container-ids.nix
new file mode 100644
index 000000000..fce91fa02
--- /dev/null
+++ b/tests/modules/programs/firefox/duplicate-container-ids.nix
@@ -0,0 +1,35 @@
+{ config, lib, ... }:
+
+{
+  imports = [ ./setup-firefox-mock-overlay.nix ];
+
+  config = lib.mkIf config.test.enableBig {
+    test.asserts.assertions.expected = [''
+      Must not have a Firefox container with an existing ID but
+        - ID 9 is used by dangerous, shopping''];
+
+    programs.firefox = {
+      enable = true;
+
+      profiles = {
+        my-profile = {
+          isDefault = true;
+          id = 1;
+
+          containers = {
+            "shopping" = {
+              id = 9;
+              color = "blue";
+              icon = "circle";
+            };
+            "dangerous" = {
+              id = 9;
+              color = "red";
+              icon = "circle";
+            };
+          };
+        };
+      };
+    };
+  };
+}
diff --git a/tests/modules/programs/firefox/profile-settings-expected-containers.json b/tests/modules/programs/firefox/profile-settings-expected-containers.json
new file mode 100644
index 000000000..d1727cc2a
--- /dev/null
+++ b/tests/modules/programs/firefox/profile-settings-expected-containers.json
@@ -0,0 +1 @@
+{"identities":[{"color":"yellow","icon":"circle","name":"shopping","public":true,"userContextId":6}],"lastUserContextId":6,"version":4}
diff --git a/tests/modules/programs/firefox/profile-settings.nix b/tests/modules/programs/firefox/profile-settings.nix
index 72a44e23e..6529a1bb2 100644
--- a/tests/modules/programs/firefox/profile-settings.nix
+++ b/tests/modules/programs/firefox/profile-settings.nix
@@ -1,7 +1,5 @@
 { config, lib, pkgs, ... }:
 
-with lib;
-
 {
   imports = [ ./setup-firefox-mock-overlay.nix ];
 
@@ -142,6 +140,17 @@ with lib;
           };
         };
       };
+
+      profiles.containers = {
+        id = 5;
+        containers = {
+          "shopping" = {
+            id = 6;
+            icon = "circle";
+            color = "yellow";
+          };
+        };
+      };
     };
 
     nmt.script = ''
@@ -155,6 +164,10 @@ with lib;
         home-files/.mozilla/firefox/test/user.js \
         ${./profile-settings-expected-user.js}
 
+      assertFileContent \
+        home-files/.mozilla/firefox/containers/containers.json \
+        ${./profile-settings-expected-containers.json}
+
       bookmarksUserJs=$(normalizeStorePaths \
         home-files/.mozilla/firefox/bookmarks/user.js)