2 files changed, 116 insertions, 1 deletions
diff --git a/infra/modules/base/persist.nix b/infra/modules/base/persist.nix
new file mode 100644
index 0000000..634230e
--- /dev/null
+++ b/infra/modules/base/persist.nix
@@ -0,0 +1,115 @@
+{ inputs, ... }:
+{
+  flake.modules.nixos.base =
+    {
+      config,
+      lib,
+      pkgs,
+      ...
+    }:
+    {
+      imports = [ inputs.impermanence.nixosModules.impermanence ];
+
+      options.infra.persist = {
+        dir = lib.mkOption {
+          type = lib.types.path;
+          default = "/persist";
+          description = "The root directory to use for persistence.";
+        };
+
+        directories = lib.mkOption {
+          type = lib.types.listOf lib.types.str;
+          default = [ ];
+          description = "The list of root owned directories to persist across reboots.";
+        };
+
+        files = lib.mkOption {
+          type = lib.types.listOf lib.types.str;
+          default = [ ];
+          description = "The list of root owned files to persist across reboots.";
+        };
+
+        users = lib.mkOption {
+          type = lib.types.attrsOf (
+            lib.types.submodule {
+              options = {
+                directories = lib.mkOption {
+                  type = lib.types.listOf lib.types.str;
+                  default = [ ];
+                  description = "The list of directories to persist within the users home directory.";
+                };
+
+                files = lib.mkOption {
+                  type = lib.types.listOf lib.types.str;
+                  default = [ ];
+                  description = "The list of files to persist within the users home directory.";
+                };
+              };
+            }
+          );
+          default = { };
+          description = "The directories and files to persist for a user.";
+        };
+      };
+
+      config =
+        let
+          cfg = config.infra.persist;
+          blank = "${config.infra.zfs.pool}/local/root@blank";
+        in
+        {
+          boot = {
+            initrd = {
+              supportedFilesystems.zfs = true;
+              systemd = {
+                storePaths = [ pkgs.zfs ];
+                services.infra-rollback-root = {
+                  description = "Rollback root ZFS dataset to blank snapshot";
+                  wantedBy = [ "initrd.target" ];
+                  after = [ "zfs-import-${config.infra.zfs.pool}.service" ];
+                  before = [ "sysroot.mount" ];
+                  enableStrictShellChecks = true;
+
+                  serviceConfig = {
+                    Type = "oneshot";
+                  };
+
+                  script = ''
+                    if "${pkgs.zfs}/bin/zfs" list -H -o name "${blank}" >/dev/null 2>&1; then
+                      "${pkgs.zfs}/bin/zfs" rollback -r "${blank}"
+                    fi
+                  '';
+                };
+              };
+            };
+          };
+
+          fileSystems."${cfg.dir}".neededForBoot = true;
+          environment.persistence."${cfg.dir}" = {
+            directories = [
+              # keep-sorted start
+              "/var/lib/nixos"
+              "/var/lib/systemd"
+              "/var/log/journal"
+              # keep-sorted end
+            ]
+            ++ cfg.directories;
+            files = [
+              "/etc/machine-id"
+
+              # Needed for sops-nix to be able to decrypt secrets.
+              "/etc/ssh/ssh_host_ed25519_key"
+              "/etc/ssh/ssh_host_ed25519_key.pub"
+            ]
+            ++ cfg.files;
+          };
+
+          assertions = [
+            {
+              assertion = config.boot.initrd.systemd.enable;
+              message = "initrd systemd must be enabled for persistence (impermanence)";
+            }
+          ];
+        };
+    };
+}
diff --git a/infra/modules/base/zfs.nix b/infra/modules/base/zfs.nix
index e2cb007..4a818e7 100644
--- a/infra/modules/base/zfs.nix
+++ b/infra/modules/base/zfs.nix
@@ -99,7 +99,7 @@
 
             "persist" = {
               type = "zfs_fs";
-              mountpoint = "/persist";
+              mountpoint = config.infra.persist.dir;
               options = {
                 canmount = "on";
                 mountpoint = "legacy";