config/modules/base/boot/default.nix

{
  config,
  lib,
  pkgs,
  ...
}:
let
  cfg = config.zfs-root.boot;
  inherit (lib)
    mkIf
    types
    mkDefault
    mkOption
    mkMerge
    strings
    ;
  inherit (builtins)
    head
    toString
    map
    tail
    ;
in
{
  options.zfs-root.boot = {
    enable = mkOption {
      description = "Enable root on ZFS support";
      type = types.bool;
      default = false; # TODO: change by @motiejus
    };
    devNodes = mkOption {
      description = "Specify where to discover ZFS pools";
      type = types.str;
      apply =
        x:
        assert (strings.hasSuffix "/" x || abort "devNodes '${x}' must have trailing slash!");
        x;
      default = "/dev/disk/by-id/";
    };
    bootDevices = mkOption {
      description = "Specify boot devices";
      type = types.nonEmptyListOf types.str;
    };
    availableKernelModules = mkOption {
      type = types.nonEmptyListOf types.str;
      default = [
        "uas"
        "nvme"
        "ahci"
      ];
    };
    kernelParams = mkOption {
      type = types.listOf types.str;
      default = [ ];
    };
    immutable = mkOption {
      description = "Enable root on ZFS immutable root support";
      type = types.bool;
      default = false;
    };
    removableEfi = mkOption {
      description = "install bootloader to fallback location";
      type = types.bool;
      default = true;
    };
    partitionScheme = mkOption {
      default = {
        biosBoot = "-part5";
        efiBoot = "-part1";
        swap = "-part4";
        bootPool = "-part2";
        rootPool = "-part3";
      };
      description = "Describe on disk partitions";
      type = types.attrsOf types.str;
    };
    sshUnlock = {
      enable = mkOption {
        type = types.bool;
        default = false;
      };
      authorizedKeys = mkOption {
        type = types.listOf types.str;
        default = [ ];
      };
    };
  };
  config = mkIf (cfg.enable) (mkMerge [
    {
      zfs-root.fileSystems.datasets = {
        "rpool/nixos/home" = mkDefault "/home";
        "rpool/nixos/var/lib" = mkDefault "/var/lib";
        "rpool/nixos/var/log" = mkDefault "/var/log";
        "rpool/nixos/nix" = mkDefault "/nix";
        "bpool/nixos/root" = "/boot";
      };
    }
    (mkIf (!cfg.immutable) {
      zfs-root.fileSystems.datasets = {
        "rpool/nixos/root" = "/";
      };
    })
    (mkIf cfg.immutable {
      zfs-root.fileSystems = {
        datasets = {
          "rpool/nixos/empty" = "/";
          "rpool/nixos/root" = "/oldroot";
        };
        bindmounts = {
          "/oldroot/nix" = "/nix";
          "/oldroot/etc/nixos" = "/etc/nixos";
        };
      };
      boot.initrd.postDeviceCommands = ''
        if ! grep -q zfs_no_rollback /proc/cmdline; then
          zpool import -N rpool
          zfs rollback -r rpool/nixos/empty@start
          zpool export -a
        fi
      '';
    })
    {
      zfs-root.fileSystems = {
        efiSystemPartitions = map (diskName: diskName + cfg.partitionScheme.efiBoot) cfg.bootDevices;
        swapPartitions =
          if cfg.partitionScheme ? swap then
            map (diskName: diskName + cfg.partitionScheme.swap) cfg.bootDevices
          else
            [ ];
      };
      boot = {
        initrd.availableKernelModules = cfg.availableKernelModules;
        kernelParams = cfg.kernelParams;
        supportedFilesystems = [ "zfs" ];
        zfs = {
          devNodes = cfg.devNodes;
          forceImportRoot = mkDefault false;
        };
        loader = {
          efi = {
            canTouchEfiVariables = if cfg.removableEfi then false else true;
            efiSysMountPoint = "/boot/efis/" + (head cfg.bootDevices) + cfg.partitionScheme.efiBoot;
          };
          generationsDir.copyKernels = true;
          grub = {
            enable = true;
            devices = map (diskName: cfg.devNodes + diskName) cfg.bootDevices;
            efiInstallAsRemovable = cfg.removableEfi;
            copyKernels = true;
            efiSupport = true;
            zfsSupport = true;
            extraInstallCommands = toString (
              map (diskName: ''
                set -x
                ${pkgs.coreutils-full}/bin/cp -r ${config.boot.loader.efi.efiSysMountPoint}/EFI /boot/efis/${diskName}${cfg.partitionScheme.efiBoot}
                set +x
              '') (tail cfg.bootDevices)
            );
          };
        };
      };
    }
    (mkIf cfg.sshUnlock.enable {
      boot.initrd = {
        network = {
          enable = true;
          ssh = {
            enable = true;
            hostKeys = [
              "/var/lib/ssh_unlock_zfs_ed25519_key"
              "/var/lib/ssh_unlock_zfs_rsa_key"
            ];
            authorizedKeys = cfg.sshUnlock.authorizedKeys;
          };
          postCommands = ''
            tee -a /root/.profile >/dev/null <<EOF
            if zfs load-key rpool/nixos; then
               pkill zfs
            fi
            exit
            EOF'';
        };
      };
    })
  ]);
}