Rework (and fix) impermanence

hosts/herugrim/default.nix

@@ -23,16 +23,18 @@
       hostname = "herugrim";
       username = "itsdrike";
 
-      impermanence.root = {
+      impermanence = {
+        root = {
           enable = true;
           # Some people use /nix/persist/system for this, leaving persistent files in /nix subvolume
           # I much prefer using a standalone subvolume for this though.
           persistentMountPoint = "/persist";
+        };
+
         # Configure automatic root subvolume wiping on boot from initrd
-        autoBtrfsWipe = {
+        autoWipeBtrfs = {
-          devicePath = "/dev/disk/by-label/NIXROOT";
+          enable = true;
-          subvolumePath = "root";
+          devices."/dev/disk/by-label/NIXROOT".subvolumes = [ "root" ];
-          cleanSnapshotPath = "root-blank";
         };
       };
     };

options/system/impermanence.nix

@@ -1,5 +1,5 @@
 { lib, config, ... }: with lib; let
-  inherit (lib) mkEnableOption mkOption literalExpression;
+  inherit (lib) mkEnableOption mkOption literalExpression types;
 
   cfg = config.myOptions.system.impermanence;
 in
@@ -12,6 +12,7 @@ in
 
       extraFiles = mkOption {
         default = [];
+        type = types.listOf types.path;
         example = literalExpression ''["/etc/nix/id_rsa"]'';
         description = ''
           Additional files in root to link to persistent storage.
@@ -20,6 +21,7 @@ in
 
       extraDirectories = mkOption {
         default = [];
+        type = types.listOf types.path;
         example = literalExpression ''["/etc/nix/id_rsa"]'';
         description = ''
           Additional directories in root to link to persistent storage.
@@ -34,48 +36,53 @@ in
           system state files.
         '';
       };
-
-      autoBtrfsWipe = {
-        enable = mkOption {
-          default = true;
-          description = ''
-            Enable automatic wiping of the root BTRFS subvolume from initrd.
-
-            Generally, you will want to keep this enabled, as otherwise setting up
-            impermanence is pointless. However in case you're using a non-BTRFS
-            system, or you wish to set up a custom handling for this auto-wiping,
-            which the current handling doesn't support, disable this.
-          '';
     };
 
-        devicePath = mkOption {
+    autoWipeBtrfs = let
-          default = "/dev/mapper/cryptfs";
+      btrfsDeviceOptionType = types.submodule {
+        options = {
+          subvolumes = mkOption {
+            type = types.listOf types.str;
+            default = [];
             description = ''
-            Path to the BTRFS block device containing the subvolume to be wiped.
+              List of BTRFS subvolumes to be wiped from the device.
 
-            This device will be mounted from initrd.
+              These subvolumes will be wiped from initrd, before the subvolumes are mounted.
             '';
+            example = literalExpression ''[ "root" "home" ]'';
           };
+        };
+      };
+    in {
+      enable = mkEnableOption ''
+          automatic wiping of specified BTRFS subvolumes from initrd.
 
-        subvolumePath = mkOption {
+          If you're using BTRFS, you will generally want to enable this, however
-          default = "root";
+          with a non-BTRFS system, or in case you wish to set up some custom handling
+          which this module doesn't support, you will need to write your own logic
+          for automatic root wiping.
+
+          One option is is to simply have your root get mounted from tmpfs, making it 
+          live in RAM. This does however require dedicating a concrete chunk of RAM.
+        '';
+
+      devices = mkOption {
+        default = {};
+        type = types.attrsOf btrfsDeviceOptionType;
         description = ''
-            Path to the BTRFS subvolume to be wiped.
+          BTRFS devices and their subvolumes to be wiped.
-
-            This is a relative path, starting from the BTRFS root.
         '';
+        example = literalExpression ''
+          {
+            "/dev/sda1" = {
+              subvolumes = [ "root" ];
             };
-
+            "/dev/mapper/cryptfs" = {
-        cleanSnapshotPath = mkOption {
+              subvolumes = [ "homeJohn" "homeBob" ];
-          default = "root-blank";
+            };
-          description = ''
+          }
-            Path to the BTRFS snapshot (subvolume) to be restore 
-            `myOptions.system.impermanence.root.autoWipe.btrfsSubvolume` to.
-
-            This should be a blank snapshot to achieve a complete wipe.
         '';
       };
     };
   };
-  };
 }

system/impermanence/autowipe.nix

@@ -0,0 +1,54 @@
+{ config, lib, ... }: let
+  inherit (lib) mkIf concatStringsSep flatten mapAttrsToList;
+  cfg = config.myOptions.system.impermanence.autoWipeBtrfs;
+in
+{
+  config = mkIf cfg.enable {
+    boot.initrd.systemd = {
+      enable = true; # This enables systemd support in stage 1 - required for below setup
+
+      services.rollback = {
+        description = "Rollback BTRFS subvolumes to a pristine state";
+        enable = true;
+        wantedBy = [ "initrd.target" ];
+        # Make sure it's done after decryption (i.e. LUKS/TPM process)
+        after = [ "systemd-cryptsetup@cryptfs.service" ];
+        # mount the root fs before clearing
+        before = [ "sysroot.mount" ];
+        unitConfig.DefaultDependencies = "no";
+        serviceConfig.Type = "oneshot";
+        script = let
+          wipeScript = devicePath: subvolumes: ''
+            # Mount the BTRFS device root to a temporary mount point
+            echo "Mounting BTRFS root from ${devicePath} to /mnt"
+            mount --mkdir "${devicePath}" /mnt
+
+            # Recreate each specified subvolume
+            ${concatStringsSep "\n" (map (subvolume: ''
+              delete_subvolume_recursively "/mnt/${subvolume}"
+              btrfs subvolume create "/mnt/${subvolume}"
+            '') subvolumes)}
+
+            # Cleanup: unmount the device
+            echo "Unmounting BTRFS root from ${devicePath}"
+            umount /mnt
+          '';
+        in ''
+          delete_subvolume_recursively() {
+            IFS=$'\n'
+            for i in $(btrfs subvolume list -o "$1" | cut -f 9- -d ' '); do
+              delete_subvolume_recursively "/mnt/$i"
+            done
+
+            echo "Deleting subvolume $1"
+            btrfs subvolume delete "$1"
+          }
+
+          ${concatStringsSep "\n" (mapAttrsToList (devicePath: deviceOpts:
+            wipeScript devicePath deviceOpts.subvolumes
+          ) cfg.devices)}
+        '';
+      };
+    };
+  };
+}

system/impermanence/default.nix

@@ -2,5 +2,6 @@
 {
   imports = [
     ./root.nix
+    ./autowipe.nix
   ];
 }

system/impermanence/root.nix

@@ -1,4 +1,6 @@
-{ config }: let
+{ config, lib, ... }: let
+  inherit (lib) mkIf mkForce;
+
   cfgSystem = config.myOptions.system;
   cfg = config.myOptions.system.impermanence.root;
 in
@@ -60,51 +62,5 @@ in
         type = "ed25519";
       }
     ];
-
-    boot.initrd.systemd = let
-      cfgWipe = cfg.autoBtrfsWipe;
-    in {
-      enable = true; # This enables systemd support in stage 1 - required for below setup
-      services.rollback-root = {
-        description = "Rollback BTRFS root subvolume to a pristine state";
-        enable = cfgWipe.enable;
-        wantedby = [ "initrd.target" ];
-        # Make sure it's done after decryption (i.e. LUKS/TPM process)
-        after = [ "systemd-cryptsetup@cryptfs.service" ];
-        # mount the root fs before clearing
-        before = [ "sysroot.mount" ];
-        unitConfig.DefaultDependencies = "no";
-        serviceConfig.Type = "oneshot";
-        script = ''
-          # Mount the BTRFS root to /mnt, so we can manipulate the subvolumes
-          mount --mkdir ${cfgWipe.devicePath} /mnt
-          
-          # To restore the root subvolume, we will first delete it, and then create
-          # a new snapshot from the blank snapshot, which will become our new root subvolume
-
-          # However, at this point, root subvol is already populated and contains a number
-          # of subvolumes, which would make `btrfs subvolume delete` fail.
-          #
-          # These existing subvolumes get created automatically, and we can safely remove
-          # them. They are: /srv, /var/lib/portables, /var/lib/machines, /var/tmp
-          sudo btrfs subvolume list -o "/mnt/${cfgWipe.subvolumePath}" | cut -f9 -d' ' |
-            while read subvolme; do
-              echo "deleting $subvolume subvolume..." &&
-              btrfs subvolume delete "/mnt/$subvolume"
-            done
-
-          # Now we can remove the root subvolume, and restore it from a snapshot
-          echo "deleting ${cfgWipe.subvolumePath} (root) subvolume..."
-          btrfs subvolume delete "/mnt/${cfg.subvolumePath}"
-          echo "restoring ${cfgWipe.subvolumePath} (root) subvolume..."
-          btrfs subvolume snapshot "/mnt/${cfgWipe.cleanSnapshotPath}"
-          "/mnt/${cfgWipe.subvolumePath}"
-
-          # Once we're done rolling back to a blank snapshot,
-          # we can unmount /mnt and continue on the boot process
-          umount /mnt
-        '';
-      };
-    };
   };
 }