From 606b952304ae01b643331bf80402fce31f5c3fa0 Mon Sep 17 00:00:00 2001
From: Tom Alexander <tom@fizz.buzz>
Date: Mon, 24 Nov 2025 20:14:45 -0500
Subject: [PATCH] Make rollback datasets configurable.

---
 nix/configuration/hosts/odo/default.nix     |  5 +++
 nix/configuration/hosts/odowork/default.nix |  5 +++
 nix/configuration/hosts/quark/default.nix   |  5 +++
 nix/configuration/roles/boot/default.nix    | 34 +++++++++++++++++----
 4 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/nix/configuration/hosts/odo/default.nix b/nix/configuration/hosts/odo/default.nix
index b2f55e67..222858ad 100644
--- a/nix/configuration/hosts/odo/default.nix
+++ b/nix/configuration/hosts/odo/default.nix
@@ -32,6 +32,11 @@
     boot.loader.generic-extlinux-compatible.enable = true;
     boot.loader.systemd-boot.enable = lib.mkForce false;
 
+    me.rollback.dataset = [
+      "zroot/linux/nix/root@blank"
+      "zroot/linux/nix/home@blank"
+    ];
+
     me.optimizations = {
       enable = true;
       arch = "znver4";
diff --git a/nix/configuration/hosts/odowork/default.nix b/nix/configuration/hosts/odowork/default.nix
index a606add8..4c1fbb63 100644
--- a/nix/configuration/hosts/odowork/default.nix
+++ b/nix/configuration/hosts/odowork/default.nix
@@ -33,6 +33,11 @@
     boot.loader.generic-extlinux-compatible.enable = true;
     boot.loader.systemd-boot.enable = lib.mkForce false;
 
+    me.rollback.dataset = [
+      "zroot/linux/nixwork/root@blank"
+      "zroot/linux/nixwork/home@blank"
+    ];
+
     me.optimizations = {
       enable = true;
       arch = "znver4";
diff --git a/nix/configuration/hosts/quark/default.nix b/nix/configuration/hosts/quark/default.nix
index 6440cda4..b3be6ffa 100644
--- a/nix/configuration/hosts/quark/default.nix
+++ b/nix/configuration/hosts/quark/default.nix
@@ -30,6 +30,11 @@
     boot.loader.generic-extlinux-compatible.enable = true;
     boot.loader.systemd-boot.enable = lib.mkForce false;
 
+    me.rollback.dataset = [
+      "zroot/linux/nix/root@blank"
+      "zroot/linux/nix/home@blank"
+    ];
+
     me.optimizations = {
       enable = true;
       arch = "znver4";
diff --git a/nix/configuration/roles/boot/default.nix b/nix/configuration/roles/boot/default.nix
index 2b3c2123..035f747e 100644
--- a/nix/configuration/roles/boot/default.nix
+++ b/nix/configuration/roles/boot/default.nix
@@ -23,6 +23,27 @@
       example = true;
       description = "Enable to use secure boot.";
     };
+
+    rollback.enable = lib.mkOption {
+      type = lib.types.bool;
+      default = true;
+      example = true;
+      description = "Whether we want to enable rolling back during boot.";
+    };
+
+    rollback.dataset = lib.mkOption {
+      default = { };
+      example = lib.literalExpression ''
+        {
+          "zroot/linux/nix/root@blank" = true;
+          "zroot/linux/nix/home@blank" = lib.mkForce false;
+        }
+      '';
+      type = lib.types.coercedTo (lib.types.listOf lib.types.str) (
+        enabled: lib.listToAttrs (map (fs: lib.nameValuePair fs true) enabled)
+      ) (lib.types.attrsOf lib.types.bool);
+      description = "List of ZFS datasets to rollback to during boot.";
+    };
   };
 
   config = lib.mkIf config.me.boot.enable (
@@ -51,7 +72,7 @@
 
         # Check what will be lost with `zfs diff zroot/linux/root@blank`
         boot.initrd.systemd.enable = lib.mkDefault true;
-        boot.initrd.systemd.services.zfs-rollback = {
+        boot.initrd.systemd.services.zfs-rollback = lib.mkIf config.me.rollback.enable {
           description = "Rollback ZFS root dataset to blank snapshot";
           wantedBy = [
             "initrd.target"
@@ -64,11 +85,12 @@
           ];
           unitConfig.DefaultDependencies = "no";
           serviceConfig.Type = "oneshot";
-          script = ''
-            ${config.boot.zfs.package}/sbin/zfs rollback -r zroot/linux/nix/root@blank
-            ${config.boot.zfs.package}/sbin/zfs rollback -r zroot/linux/nix/home@blank
-            echo "rollback complete"
-          '';
+          script = lib.concatStringsSep "\n" (
+            (builtins.map (ds: "${config.boot.zfs.package}/sbin/zfs rollback -r '${ds}'") (
+              builtins.attrNames config.me.rollback.dataset
+            ))
+            ++ [ ''echo "rollback complete"'' ]
+          );
         };
 
         # boot.loader.systemd-boot.extraEntries = {