nixfiles/modules/tmproot.nix

{ lib, pkgs, inputs, config, ... }:
  let
    inherit (builtins) elem;
    inherit (lib) concatStringsSep concatMap concatMapStringsSep mkIf mkDefault mkMerge mkForce mkVMOverride;
    inherit (lib.my) mkOpt mkBoolOpt mkVMOverride' dummyOption;

    cfg = config.my.tmproot;

    showUnsaved =
      ''
        #!${pkgs.python310}/bin/python
        import stat
        import sys
        import os

        ignored = [
          ${concatStringsSep ",\n  " (map (p: "'${p}'") cfg.unsaved.ignore)}
        ]

        base = '/'
        base_dev = os.stat(base).st_dev

        def recurse(p, link=None):
          try:
            for ignore in ignored:
              if p.startswith(ignore):
                return

            st = os.lstat(p)
            if st.st_dev != base_dev:
              return

            if stat.S_ISLNK(st.st_mode):
              target = os.path.realpath(p, strict=False)
              if os.access(target, os.F_OK):
                recurse(target, link=p)
                return
            elif stat.S_ISDIR(st.st_mode):
              for e in os.listdir(p):
                recurse(os.path.join(p, e))
              return

            print(link or p)
          except PermissionError as ex:
            print(f'{p}: {ex.strerror}', file=sys.stderr)

        recurse(base)
      '';

    rootDef = {
      device = "yeet";
      fsType = "tmpfs";
      options = [ "size=${cfg.size}" ];
    };
  in {
    imports = [ inputs.impermanence.nixosModule ];

    options = with lib.types; {
      my.tmproot = {
        enable = mkBoolOpt true;
        persistDir = mkOpt str "/persist";
        size = mkOpt str "2G";
        unsaved = {
          showMotd = mkBoolOpt true;
          ignore = mkOpt (listOf str) [];
        };
      };

      # Forward declare options that won't exist until the VM module is actually imported
      virtualisation = {
        diskImage = dummyOption;
      };
    };

    config = mkIf cfg.enable (mkMerge [
      {
        assertions = [
          {
            assertion = config.fileSystems ? "${cfg.persistDir}";
            message = "The 'fileSystems' option does not specify your persistence file system (${cfg.persistDir}).";
          }
          {
            # I mean you probably _could_, but if you're doing tmproot... come on
            assertion = !config.users.mutableUsers;
            message = "users.mutableUsers is incompatible with tmproot";
          }
        ];

        my.tmproot.unsaved.ignore = [
          "/tmp"

          # setup-etc.pl will create this for us
          "/etc/NIXOS"

          # Once mutableUsers is disabled, we should be all clear here
          "/etc/passwd"
          "/etc/group"
          "/etc/shadow"
          "/etc/subuid"
          "/etc/subgid"

          # Lock file for /etc/{passwd,shadow}
          "/etc/.pwd.lock"

          # systemd last updated? I presume they'll get updated on boot...
          "/etc/.updated"
          "/var/.updated"

          # Specifies obsolete files that should be deleted on activation - we'll never have those!
          "/etc/.clean"
        ];

        environment.systemPackages = [
          (pkgs.writeScriptBin "tmproot-unsaved" showUnsaved)
        ];

        # Catch non-existent source directories that are needed for boot (see `pathsNeededForBoot` in
        # nixos/lib/util.nix). We do this by monkey-patching the `waitDevice` function that would otherwise hang.
        boot.initrd.postDeviceCommands =
          ''
            ensurePersistSource() {
              [ -e "/mnt-root$1" ] && return
              echo "Persistent source directory $1 does not exist, creating..."
              install -dm "$2" "/mnt-root$1" || fail
            }

            _waitDevice() {
              local device="$1"

              ${concatMapStringsSep " || \\\n  " (d:
                let
                  sourceDir = "${d.persistentStoragePath}${d.directory}";
                in
                  ''([ "$device" = "/mnt-root${sourceDir}" ] && ensurePersistSource "${sourceDir}" "${d.mode}")'')
                config.environment.persistence."${cfg.persistDir}".directories}

              waitDevice "$@"
            }

            type waitDevice > /dev/null || (echo "waitDevice is missing!"; fail)
            alias waitDevice=_waitDevice
          '';

        environment.persistence."${cfg.persistDir}" = {
          hideMounts = mkDefault true;
          directories = [
            "/var/log"
            # In theory we'd include only the files needed individually (i.e. the {U,G}ID map files that track deleted
            # users and groups), but `update-users-groups.pl` actually deletes the original files for "atomic update".
            # Also the script runs before impermanence does.
            "/var/lib/nixos"
            "/var/lib/systemd"
          ];
          files = [
            "/etc/machine-id"
          ];
        };

        my.dynamic-motd.script = mkIf cfg.unsaved.showMotd
          ''
            tmprootUnsaved() {
              local count="$(tmproot-unsaved | wc -l)"
              [ $count -eq 0 ] && return

              echo
              echo -e "\t\e[31;1;4mWarning:\e[0m $count file(s) on / will be lost on shutdown!"
              echo -e '\tTo see them, run `tmproot-unsaved` as root.'
              echo -e '\tAdd these files to `environment.persistence."${cfg.persistDir}"` to keep them!'
              echo -e '\tOtherwise, they can be ignored by adding to `my.tmproot.unsaved.ignore`.'
              echo
            }

            tmprootUnsaved
          '';

        fileSystems."/" = rootDef;

        virtualisation = {
          diskImage = "./.vms/${config.system.name}-persist.qcow2";
        };
      }
      (mkIf config.services.openssh.enable {
        environment.persistence."${cfg.persistDir}".files =
          concatMap (k: [ k.path "${k.path}.pub" ]) config.services.openssh.hostKeys;
      })
      (mkIf config.networking.resolvconf.enable {
        my.tmproot.unsaved.ignore = [ "/etc/resolv.conf" ];
      })
      (mkIf config.security.doas.enable {
        my.tmproot.unsaved.ignore = [ "/etc/doas.conf" ];
      })
      (mkIf config.my.boot.isDevVM {
        my.tmproot.unsaved.ignore = [ "/nix" ];

        fileSystems = mkVMOverride {
          "/" = mkVMOverride' rootDef;
          # Hijack the "root" device for persistence in the VM
          "${cfg.persistDir}" = {
            device = config.virtualisation.bootDevice;
            neededForBoot = true;
          };
        };
      })
    ]);
  }