refactor(nix): extract orion's btrfs rollback to NixOS module

flake.nix

@@ -226,6 +226,7 @@
               };
               modules = [
                 defaults
+                ./modules/btrfs-rollback.nix
                 inputs.impermanence.nixosModules.impermanence
                 inputs.agenix.nixosModules.default
                 inputs.disko.nixosModules.disko

hosts/orion/os/boot.nix

@@ -6,6 +6,13 @@
   security.tpm2.enable = true;
   environment.systemPackages = with pkgs; [ tpm2-tss ];
 
+  services.btrfs-rollback = {
+    enable = true;
+    diskLabel = "NixOS-Primary";
+    subvolume = "root";
+    snapshot = "root-base";
+  };
+
   boot = {
     loader = {
       systemd-boot.enable = true;
@@ -27,55 +34,6 @@
       systemd = {
         enable = true;
         enableTpm2 = true;
-        initrdBin = [
-          pkgs.libuuid
-          pkgs.gawk
-        ];
-        services.rollback = {
-          description = "Rollback btrfs root subvolume";
-          wantedBy = [ "initrd.target" ];
-          before = [ "sysroot.mount" ];
-          after = [ "initrd-root-device.target" ];
-          unitConfig.DefaultDependencies = "no";
-          serviceConfig.Type = "oneshot";
-          script = ''
-            mkdir -p /mnt
-            DISK_LABEL="NixOS-Primary"
-            FOUND_DISK=0
-            ATTEMPTS=50
-            printf "Attempting to find disk with label '%s'\n" "$DISK_LABEL"
-            while ((ATTEMPTS > 0)); do
-              if findfs LABEL="$DISK_LABEL"; then
-                FOUND_DISK=1
-                printf "Found disk!\n"
-                break;
-              fi
-              ((ATTEMPTS--))
-              sleep .1
-              printf "Remaining disk discovery attempts: %s\n" "$ATTEMPTS"
-            done
-            if (( FOUND_DISK == 0 )); then
-              printf "Discovery of disk with label '%s' failed! Cannot rollback!\n" "$DISK_LABEL"
-              exit 1
-            fi
-
-            mount -t btrfs -o subvol=/ $(findfs LABEL="$DISK_LABEL") /mnt
-            btrfs subvolume list -to /mnt/root \
-              | awk 'NR>2 { printf $4"\n" }' \
-              | while read subvol; do
-              printf "Removing Subvolume: %s\n" "$subvol";
-              btrfs subvolume delete "/mnt/$subvol"
-            done
-
-            printf "Removing /root subvolume\n"
-            btrfs subvolume delete /mnt/root
-
-            printf "Restoring base /root subvolume\n"
-            btrfs subvolume snapshot /mnt/root-base /mnt/root
-
-            umount /mnt
-          '';
-        };
       };
     };
   };

modules/btrfs-rollback.nix

@@ -0,0 +1,93 @@
+# TODO: Allow the rollback of multiple different subvolumes. Probably utilizing service templates (https://www.freedesktop.org/software/systemd/man/latest/systemd.service.html#Service%20Templates)
+{
+  lib,
+  config,
+  pkgs,
+  ...
+}:
+let
+  cfg = config.services.btrfs-rollback;
+in
+{
+  options.services.btrfs-rollback = {
+    enable = lib.mkEnableOption "BTRFS Rollback on boot";
+    diskLabel = lib.mkOption {
+      description = "The disk with the given label to rollback";
+      type = lib.types.str;
+    };
+    subvolume = lib.mkOption {
+      description = "The subvolume to rollback";
+      type = lib.types.str;
+    };
+    snapshot = lib.mkOption {
+      description = "The base snapshot to restore on rollback";
+      type = lib.types.str;
+    };
+    searchAttempts = lib.mkOption {
+      description = ''
+        The number of attempts that should be made to locate the labeled disk.
+
+        This may need to be increased if the given disk does in fact exist, but the rollback is
+        failing.
+      '';
+      type = lib.types.ints.unsigned;
+      default = 50;
+    };
+  };
+
+  config = lib.mkIf (cfg.enable && config.boot.initrd.systemd.enable) {
+    boot.initrd.systemd = {
+      initrdBin = with pkgs; [
+        libuuid
+        gawk
+      ];
+      services.rollback = {
+        description = "Rollback btrfs root subvolume";
+        wantedBy = [ "initrd.target" ];
+        before = [ "sysroot.mount" ];
+        after = [ "initrd-root-device.target" ];
+        unitConfig.DefaultDependencies = "no";
+        serviceConfig.Type = "oneshot";
+        script = # bash
+          ''
+            mkdir -p /mnt
+            DISK_LABEL="${cfg.diskLabel}"
+            FOUND_DISK=0
+            SEARCH_ATTEMPTS="${builtins.toString cfg.searchAttempts}"
+            printf "Attempting to find disk with label '%s'\n" "$DISK_LABEL"
+            while ((SEARCH_ATTEMPTS > 0)); do
+              if findfs LABEL="$DISK_LABEL"; then
+                FOUND_DISK=1
+                printf "Found disk!\n"
+                break;
+              fi
+              ((SEARCH_ATTEMPTS--))
+              sleep .1
+              printf "Remaining disk discovery attempts: %s\n" "$SEARCH_ATTEMPTS"
+            done
+
+            if (( FOUND_DISK == 0 )); then
+              printf "Discovery of disk with label '%s' failed! Cannot rollback!\n" "$DISK_LABEL" >&2
+              exit 1
+            fi
+
+            mount -t btrfs -o subvol=/ $(findfs LABEL="$DISK_LABEL") /mnt
+            btrfs subvolume list -to "/mnt/${cfg.subvolume}" \
+              | awk 'NR>2 { printf $4"\n" }' \
+              | while read subvol; do
+              printf "Removing Subvolume: %s\n" "$subvol";
+              btrfs subvolume delete "/mnt/$subvol"
+            done
+
+            printf "Removing ${cfg.subvolume} subvolume\n"
+            btrfs subvolume delete "/mnt/${cfg.subvolume}"
+
+            printf "Restoring base ${cfg.subvolume} subvolume from snapshot ${cfg.snapshot}\n"
+            btrfs subvolume snapshot "/mnt/${cfg.snapshot}" "/mnt/${cfg.subvolume}"
+
+            umount /mnt
+          '';
+      };
+    };
+  };
+}