From 2c03698a2f85ea12916c8f4de08a7bca73b9e652 Mon Sep 17 00:00:00 2001
From: Dmitriy Kholkin <ataraxiadev@ataraxiadev.com>
Date: Mon, 10 Mar 2025 19:06:21 +0300
Subject: [PATCH] feat: add rollback service for btrfs

---
 modules/nixos/filesystems/btrfs.nix | 110 +++++++++++++++++++++++++++-
 1 file changed, 108 insertions(+), 2 deletions(-)

diff --git a/modules/nixos/filesystems/btrfs.nix b/modules/nixos/filesystems/btrfs.nix
index 0e5b327..91d44ac 100644
--- a/modules/nixos/filesystems/btrfs.nix
+++ b/modules/nixos/filesystems/btrfs.nix
@@ -1,16 +1,122 @@
 {
   config,
   lib,
+  pkgs,
   ...
 }:
 let
-  inherit (lib) mkIf mkEnableOption;
+  inherit (builtins) map;
+  inherit (lib)
+    concatStringsSep
+    mkIf
+    mkEnableOption
+    mkOption
+    mkBefore
+    ;
+  inherit (lib.types)
+    bool
+    str
+    listOf
+    submodule
+    ;
   cfg = config.ataraxia.filesystems.btrfs;
+
+  eraseVolumesOpts =
+    { ... }:
+    {
+      options = {
+        vol = mkOption {
+          type = str;
+          example = "rootfs";
+          description = "Name of submodule to erase";
+        };
+        blank = mkOption {
+          type = str;
+          example = "rootfs-blank";
+          description = "Name of submodule to clone into `vol`";
+        };
+      };
+    };
 in
 {
   options.ataraxia.filesystems.btrfs = {
     enable = mkEnableOption "Root on btrfs";
+    # Btrfs clean root
+    eraseOnBoot = {
+      enable = mkOption {
+        type = bool;
+        default = config.persist.enable;
+        description = "Clean btrfs subvolumes on boot";
+      };
+      device = mkOption {
+        type = str;
+        description = "Device on which is btrfs partititon";
+      };
+      systemdDevice = mkOption {
+        type = str;
+        description = "Escaped string with name of .device service";
+        example = "dev-disk-by\\x2did-ata\\x2dPhison_SATA_SSD_2165.device";
+      };
+      eraseVolumes = mkOption {
+        type = listOf (submodule eraseVolumesOpts);
+        default = [ ];
+        example = [
+          {
+            vol = "rootfs";
+            blank = "rootfs-blank";
+          }
+        ];
+        description = ''
+          A list of subvolumes to erase on boot.
+        '';
+      };
+    };
   };
 
-  config = mkIf cfg.enable { };
+  config =
+    let
+      script = ''
+        mkdir -p /mnt
+        mount -t btrfs -o subvol=/ ${cfg.eraseOnBoot.device} /mnt
+
+        ${concatStringsSep "\n" (
+          map (x: ''
+            btrfs subvolume list -o /mnt/${x.vol} |
+            cut -f9 -d' ' |
+            while read subvolume; do
+              echo "deleting /$subvolume subvolume..."
+              btrfs subvolume delete "/mnt/$subvolume"
+            done &&
+
+            echo "deleting /${x.vol} subvolume..."
+            btrfs subvolume delete /mnt/${x.vol}
+            echo "restoring blank ${x.blank} subvolume..."
+            btrfs subvolume snapshot /mnt/snapshots/${x.blank} /mnt/${x.vol}
+          '') cfg.eraseOnBoot.eraseVolumes
+        )}
+
+        umount /mnt
+      '';
+    in
+    mkIf cfg.enable {
+      boot.initrd = mkIf cfg.eraseOnBoot.enable {
+        postDeviceCommands = mkIf (!config.boot.initrd.systemd.enable) (mkBefore script);
+
+        systemd.services.rollback = mkIf config.boot.initrd.systemd.enable {
+          description = "Rollback btrfs root subvolume to a pristine state on boot";
+          wantedBy = [ "initrd.target" ];
+          requires = [ cfg.eraseOnBoot.systemdDevice ];
+          after = [ cfg.eraseOnBoot.systemdDevice ];
+          before = [ "sysroot.mount" ];
+          path = [
+            pkgs.btrfs-progs
+            pkgs.coreutils
+            pkgs.util-linuxMinimal.mount
+          ];
+          unitConfig.DefaultDependencies = "no";
+          serviceConfig.Type = "oneshot";
+          script = script;
+        };
+      };
+    };
 }