Rework (and fix) impermanence

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
-        '';
-      };
-    };
   };
 }