nix-config/modules/nixos/impermanence.nix

{
  config,
  pkgs,
  lib,
  ...
}:
let
  inherit (lib)
    attrNames
    flip
    isAttrs
    mapAttrs
    mkIf
    mkMerge
    mkOption
    optionals
    types
    ;
in
{
  boot.zfs.forceImportRoot = false;
  boot.initrd.systemd.enable = true;
  boot.initrd.systemd.services.impermanence-rollback =
    mkIf
      (config.environment.persistence."/persist".enable || config.environment.persistence."/state".enable)
      {
        description = "Rollback root filesystem";
        wantedBy = [ "initrd.target" ];
        after = [ "zfs-import-rpool.service" ];
        before = [ "sysroot.mount" ];
        unitConfig.DefaultDependencies = "no";
        serviceConfig = {
          Type = "oneshot";
          ExecStart = "${pkgs.zfs}/bin/zfs rollback -r rpool/local/root@blank";
        };
      };

  age.identityPaths = [ "/persist/etc/ssh/ssh_host_ed25519_key" ];

  fileSystems."/state" = mkIf config.environment.persistence."/state".enable {
    neededForBoot = true;
  };
  environment.persistence."/state" = {
    enable = false;
    hideMounts = true;
    directories = [
      "/var/lib/systemd"
      "/var/log"
      "/var/spool"
    ];
  };

  fileSystems."/persist" = mkIf config.environment.persistence."/persist".enable {
    neededForBoot = true;
  };
  environment.persistence."/persist" = {
    enable = false;
    hideMounts = true;
    files = [
      (mkIf (!config.boot.isContainer) "/etc/machine-id")
      "/etc/adjtime"
      "/etc/ssh/ssh_host_ed25519_key"
      "/etc/ssh/ssh_host_ed25519_key.pub"
    ];
    directories = [
      "/var/lib/nixos"
    ]
    ++ optionals config.security.acme.acceptTerms [
      {
        directory = "/var/lib/acme";
        user = "acme";
        group = "acme";
        mode = "0755";
      }
    ]
    ++ optionals config.services.printing.enable [
      {
        directory = "/var/lib/cups";
        mode = "0700";
      }
    ]
    ++ optionals config.hardware.bluetooth.enable [
      "/var/lib/bluetooth"
    ];
  };

  users.mutableUsers = !config.environment.persistence."/persist".enable;

  # For each user that has a home-manager config, merge the locally defined
  # persistence options that we defined above.
  imports =
    let
      mkUserFiles = map (
        x: { parentDirectory.mode = "700"; } // (if isAttrs x then x else { file = x; })
      );
      mkUserDirs = map (x: { mode = "700"; } // (if isAttrs x then x else { directory = x; }));
    in
    [
      {
        environment.persistence = mkMerge (
          flip map (attrNames config.home-manager.users) (
            user:
            let
              hmUserCfg = config.home-manager.users.${user};
            in
            flip mapAttrs hmUserCfg.home.persistence (
              _: sourceCfg: {
                users.${user} = {
                  files = mkUserFiles sourceCfg.files;
                  directories = mkUserDirs sourceCfg.directories;
                };
              }
            )
          )
        );
      }
    ];

  home-manager.sharedModules = [
    {
      options.home.persistence = mkOption {
        description = "Additional persistence config for the given source path";
        default = { };
        type = types.attrsOf (
          types.submodule {
            options = {
              files = mkOption {
                description = "Additional files to persist via NixOS impermanence.";
                type = types.listOf (types.either types.attrs types.str);
                default = [ ];
              };

              directories = mkOption {
                description = "Additional directories to persist via NixOS impermanence.";
                type = types.listOf (types.either types.attrs types.str);
                default = [ ];
              };
            };
          }
        );
      };
    }
  ];
}