nix/hosts/common/configs/system/impermanence/options.nix

{
  config,
  lib,
  pkgs,
  utils,
  ...
}:
let
  cfg = config.environment.persistence;

  # ["/home/user/" "/.screenrc"] -> ["home" "user" ".screenrc"]
  splitPath =
    paths:
    (builtins.filter (s: builtins.typeOf s == "string" && s != "") (
      builtins.concatMap (builtins.split "/") paths
    ));

  # ["/home/user/" "/.screenrc"] -> "/home/user/.screenrc"
  mergePaths =
    paths:
    let
      prefix = lib.strings.optionalString (lib.strings.hasPrefix "/" (builtins.head paths)) "/";
      path = lib.strings.concatStringsSep "/" (splitPath paths);
    in
    prefix + path;

  # "/home/user/.screenrc" -> ["/home", "/home/user"]
  parentsOf =
    path:
    let
      prefix = lib.strings.optionalString (lib.strings.hasPrefix "/" path) "/";
      split = splitPath [ path ];
      parents = lib.lists.take ((lib.lists.length split) - 1) split;
    in
    lib.lists.foldl' (
      state: item:
      state
      ++ [
        (mergePaths [
          (if state != [ ] then lib.lists.last state else prefix)
          item
        ])
      ]
    ) [ ] parents;
in
{
  options.environment =
    with lib;
    with types;
    {
      impermanence = {
        enable = mkEnableOption "Impermanence";

        device = mkOption {
          type = str;
          default = config.disko.devices.disk.main.content.partitions.root.content.content.device;
          description = ''
            LUKS BTRFS partition to wipe on boot.
          '';
        };
      };

      persistence =
        let
          isPathLike = strings.hasPrefix "/";
        in
        mkOption {
          type = (
            addCheck (attrsOf (
              attrsOf (
                submodule (
                  { name, config, ... }:
                  {
                    options = {
                      enable = mkOption {
                        type = bool;
                        default = true;
                        description = "Whether to enable the item.";
                      };

                      service = mkOption {
                        type = str;
                        readOnly = true;
                        description = ''
                          Systemd service that prepares and syncs the item.
                          Can be used as a dependency in other units.
                        '';
                      };

                      mount = mkOption {
                        type = str;
                        readOnly = true;
                        description = ''
                          Systemd mount that binds the item.
                          Can be used as a dependency in other units.
                        '';
                      };

                      path = mkOption {
                        type = str;
                        readOnly = true;
                        default = name;
                      };

                      _sourceRoot = mkOption {
                        type = str;
                        internal = true;
                      };

                      source = mkOption {
                        type = str;
                        readOnly = true;
                      };

                      _targetRoot = mkOption {
                        type = str;
                        internal = true;
                      };

                      target = mkOption {
                        type = str;
                        readOnly = true;
                      };

                      create = mkOption {
                        type = enum [
                          "none"
                          "file"
                          "directory"
                        ];
                        default = "none";
                        description = ''
                          Whether to create the file or directory
                          in persistence if it does not exist.
                        '';
                      };
                    };
                  }
                )
              )
            )) (attrs: lists.all isPathLike (builtins.attrNames attrs))
          );
          apply =
            ps:
            builtins.mapAttrs (
              persistence: items:
              builtins.mapAttrs (
                _: config:
                let
                  path = config.path;

                  _sourceRoot = persistence;

                  source = mergePaths [
                    _sourceRoot
                    path
                  ];

                  _targetRoot =
                    let
                      parents = lists.reverseList (parentsOf path);
                    in
                    lists.foldl' (
                      acc: parent:
                      if acc == "/" then
                        lists.findFirst (
                          otherPersistence: lists.any (other: parent == other) (builtins.attrNames ps.${otherPersistence})
                        ) "/" (builtins.attrNames ps)
                      else
                        acc
                    ) "/" parents;

                  target = mergePaths [
                    _targetRoot
                    path
                  ];
                in
                config
                // {
                  inherit
                    _sourceRoot
                    source
                    _targetRoot
                    target
                    ;
                  service = "${utils.escapeSystemdPath target}.service";
                  mount = "${utils.escapeSystemdPath target}.mount";
                }
              ) items
            ) ps;
          default = { };
          description = "Persistence config.";
        };
    };

  config =
    let
      all = lib.lists.flatten (builtins.concatMap builtins.attrValues (builtins.attrValues cfg));
    in
    lib.mkIf config.environment.impermanence.enable {
      boot.initrd.systemd = {
        enable = true;

        initrdBin = with pkgs; [
          coreutils
          util-linux
          findutils
          btrfs-progs
        ];

        services.impermanence = {
          description = "Rollback BTRFS subvolumes to a pristine state";
          wantedBy = [ "initrd.target" ];
          before = [ "sysroot.mount" ];
          after = [
            "cryptsetup.target"
            "local-fs-pre.target"
          ];
          unitConfig.DefaultDependencies = false;
          serviceConfig.Type = "oneshot";
          environment.DEVICE = config.environment.impermanence.device;
          script = builtins.readFile ./scripts/wipe.sh;
        };
      };

      systemd = {
        mounts = builtins.map (c: {
          description = c.path;
          requiredBy = [ "local-fs.target" ];
          requires = [ c.service ];
          bindsTo = [ c.service ];
          after = [ c.service ];
          unitConfig.ConditionPathExists = [ (lib.strings.escape [ " " ] c.source) ];
          what = c.source;
          where = c.target;
          options = lib.strings.concatStringsSep "," [
            "bind"
            "X-fstrim.notrim"
            "x-gvfs-hide"
          ];
        }) all;

        services = builtins.listToAttrs (
          builtins.map (c: {
            name = utils.escapeSystemdPath c.target;
            value = {
              description = c.path;
              after = [ "local-fs-pre.target" ];
              requiredBy = [
                "local-fs.target"
                c.mount
              ];
              before = [
                "local-fs.target"
                c.mount
                "umount.target"
              ];
              conflicts = [ "umount.target" ];
              unitConfig = {
                DefaultDependencies = false;
                RefuseManualStart = true;
                RefuseManualStop = true;
              };
              serviceConfig = {
                Type = "oneshot";
                RemainAfterExit = true;
                TimeoutStopSec = "1800";
              };
              script = ''
                source=${lib.strings.escapeShellArg c._sourceRoot}
                target=${lib.strings.escapeShellArg c._targetRoot}
                path=${lib.strings.escapeShellArg c.path}
                create=${lib.strings.escapeShellArg c.create}

                ${builtins.readFile ./scripts/start.sh}
              '';
              preStop = ''
                source=${lib.strings.escapeShellArg c._sourceRoot}
                target=${lib.strings.escapeShellArg c._targetRoot}
                path=${lib.strings.escapeShellArg c.path}
                create=${lib.strings.escapeShellArg c.create}

                ${builtins.readFile ./scripts/stop.sh}
              '';
            };
          }) all
        );
      };

      fileSystems = builtins.mapAttrs (_: _: { neededForBoot = true; }) cfg // {
        "/persist".neededForBoot = true;
      };

      environment.persistence = {
        "/persist/user"."/etc/nixos" = { };
        "/persist/state" = {
          "/var/lib/nixos" = { };
          "/var/lib/systemd" = { };
          "/var/log" = { };
        };
      };

      assertions =
        let
          paths = builtins.map (c: c.path) all;
          duplicates = lib.lists.filter (t: lib.lists.count (o: o == t) paths > 1) (lib.lists.unique paths);
        in
        [
          {
            assertion = lib.lists.length duplicates == 0;
            message = "Each target must be defined under a single persistence. Duplicate targets found: ${lib.concatStringsSep ", " duplicates}";
          }
        ];
    };
}