ZFSBootMenu on NixOS

Since the end of 2023 or so, I’ve migrated all my workstations from Void Linux to NixOS. Since I was using ZFS with a dedicated boot partition with Void Linux, migration to NixOS was easy: boot the NixOS distro image, load the ZFS pool, create a new ZFS partition, and install ZFS onto that new root. Then comes the problem: I really want to avoid using GRUB2 with ZFS (due to loads of incompatibility issues over the years). The answer to that is always ZFSBootMenu and this is how I get ZFSBootMenu working with NixOS.

Understanding ZFSBootMenu

The way ZFSBootMenu works is actually pretty simple. It creates a small initramfs image on a host machine using a standard initramfs creation methods such as dracut or mkinitcpio. During the boot process, this image is loaded by an intermediate bootloader (such as rEFInd/gummiboot or syslinux).

In case ZFSBootMenu is being used with an encrypted ZFS pool, it then performs the following:

• ZFSBootMenu prompts for an encryption passphrase, decrypts the pool, locates the kernel and initramfs on the ZFS dataset, then displays the boot menu.
• ZFSBootMenu appends the corresponding root=zfs:... spl.spl_hostid=... parameter to the kernel boot parameter, then kexec into the final kernel using the chosen kernel and initramfs.
• The final kernel decrypts the ZFS pool again using the encryption key stored on the final initramfs¹ then boots into the final system.

This allows ZFSBootMenu to have the same feature level as any ZFS system out there, instead of relying on a custom implementation like with GRUB. This also means all we need to make ZFSBootMenu work is:

• The ZFSBootMenu initramfs image itself (thankfully, ZFSBootMenu has a binary release)
• A way to locate the kernel image and initramfs during nixos-rebuild (which means boot.loader.external.installHook).
• A way to install an intermediate bootloader (easy in case of rEFInd).

ZFSBootMenu module

That resulted in this ZFSBootMenu module for NixOS:

{ config, pkgs, lib, ... }:

let
  inherit (lib) concatStringsSep mkDefault mkIf mkOption types pipe;

  currentZfs = config.boot.zfs.package;
in
{
  options = {
    boot = {
      loader = {
        zfsbootmenu = {
          enable = mkOption {
            type = types.bool;
            default = false;
            description = ''
              Enable support for ZFSBootMenu. Note that you need to install
              ZFSBootMenu and EFI bootloader (such as rEFInd) manually.
            '';
          };

          bootfs = mkOption {
            type = types.str;
            default = "zroot/ROOT/nixos";
            description = ''
              The bootfs dataset where ZFSBootMenu will boot from.
            '';
          };

          keyfile = mkOption {
            type = types.nullOr types.path;
            default = null;
            description = ''
              A keyfile to include in the initramfs image as /etc/zfs/zroot.key.
            '';
          };
        };
      };
    };
  };

  config = mkIf config.boot.loader.zfsbootmenu.enable {
    boot.supportedFilesystems = [ "zfs" ];
    boot.zfs.devNodes = mkDefault "/dev/disk/by-id";
    boot.loader.external = {
      enable = true;
      installHook = pkgs.writeShellScript "zfsbootinstall" ''
        #!{pkgs.bash}/bin/bash
        echo "updating the boot generations directory..."
        ${pkgs.coreutils}/bin/mkdir -p /boot
        bootfs=${config.boot.loader.zfsbootmenu.bootfs}

        declare -A kernels

        kernver() {
          local path=$1
          echo $path | ${pkgs.gnused}/bin/sed -E 's|.*linux-([0-9]+\.[0-9]+\.[0-9]+).*|\1|'
        }

        copyToBoot() {
          local src=$1
          local dstname=$2
          local generation=$3
          local dst=/boot/$dstname-$(kernver "$src")-g$generation

          if ! test -e "$dst"; then
            local dstTmp=$dst.tmp.''$''$
            ${pkgs.coreutils}/bin/cp $src $dstTmp
            ${pkgs.coreutils}/bin/mv $dstTmp $dst
          fi
        }

        addEntry() {
          local path=$1
          local generation=$2

          if ! test -e $path/kernel -a -e $path/initrd; then
            return
          fi

          local kernel=$(${pkgs.coreutils}/bin/readlink -f $path/kernel)
          local initrd=$(${pkgs.coreutils}/bin/readlink -f $path/initrd)

          copyToBoot $kernel vmlinuz $generation
          copyToBoot $initrd initramfs $generation
        }

        updateInit() {
          local dataset=$1
          local init=$2/init

          if ! test -f $init; then
            echo "invalid init path given: $init"
            exit 1
          fi

          local oldval=$(${currentZfs}/bin/zfs get -H org.zfsbootmenu:commandline -o value $dataset)
          local newval

          newval="init=$init ${concatStringsSep " " config.boot.kernelParams}"
          ${currentZfs}/bin/zfs set org.zfsbootmenu:commandline="$newval" $dataset
        }

        # NixOS reads init from cmdline, but ZBM only allows setting
        # cmdline via either ZBM EFI executable or ZFS dataset. This
        # means we can only have a single environment as we need to
        # resolve the path to init in init=... cmdline.
        ${pkgs.coreutils}/bin/rm /boot/vmlinuz-* || true
        ${pkgs.coreutils}/bin/rm /boot/initramfs-* || true

        for generation in $(
          (cd /nix/var/nix/profiles && ${pkgs.coreutils}/bin/ls -d system-*-link) \
          | ${pkgs.gnused}/bin/sed 's/system-\([0-9]\+\)-link/\1/' \
          | ${pkgs.coreutils}/bin/sort -n -r); do
          link=/nix/var/nix/profiles/system-''${generation}-link
          addEntry $link $generation
          updateInit $bootfs $link
          break
        done
      '';
    };

    boot.initrd.extraFiles =
      if !builtins.isNull config.boot.loader.zfsbootmenu.keyfile then {
        "etc/zfs/zroot.key".source = pkgs.runCommandLocal "zroot.key" { } ''
          ${pkgs.coreutils}/bin/cp ${config.boot.loader.zfsbootmenu.keyfile} $out
          ${pkgs.coreutils}/bin/chmod 0000 $out
        '';
      } else { };

    environment.etc =
      if !builtins.isNull config.boot.loader.zfsbootmenu.keyfile then {
        "zfs/zroot.key" = {
          source = config.boot.loader.zfsbootmenu.keyfile;
          mode = "0000";
        };
      } else { };
  };
}

This ZFSBootMenu module hooks into the boot.loader.external and:

• Clear out older kernel versions in /boot (files matching /boot/vmlinuz-* and /boot/initramfs-*)
• Look for the latest generation in /nix/var/nix/profiles
• Use that generation to copy the kernel image and initramfs to /boot (as /boot/<filename>-<kernver>-g<generation>)
• Configure the init kernel parameter via zfs set org.zfsbootmenu:commandline for NixOS.

The biggest caveat right now is that we can only have a single environment due to how NixOS reads init from a kernel parameter, and ZFSBootMenu only allows setting a kernel parameter via a ZFS property (zfs set). This means we cannot easily roll back NixOS generation during boot, and booting an older snapshot is not possible without manually editing init=... cmdline in the boot menu.

Putting it all together

Even though the module was created during migration from Void Linux to NixOS, an instruction during clean installation is probably more useful in the case of a blog post. With the above script, here’s how to set up a new machine with the above module.

Before everything: follow the instructions in Installing ZFSBootMenu for Void Linux with UEFI in ZFSBootMenu documentation until before the “Install Void” section.

Once ZFS pool is mounted as /mnt, continue to the next steps.

Install NixOS

Do the usual steps to install NixOS:

$ sudo nixos-generate-config --root /mnt

Normally, I would convert the configuration to Flake (see also my NixOS configuration), but for simplicity’s sake, put the ZFSBootMenu module in, e.g., /etc/nixos/lib/zfsbootmenu.nix:

$ mkdir -p /etc/nixos/lib
$ curl -sSL -o /etc/nixos/lib/zfsbootmenu.nix https://git.sr.ht/~sirn/nixos/blob/main/lib/zfsbootmenu.nix

Import the module in configuration.nix and configure NixOS:

{ config, pkgs, ... }:

{
  imports =
    [
      ./hardware-configuration.nix
      ./lib/zfsbootmenu.nix
    ];

  # Set this to your hostname.
  networking.hostName = "machine";

  # Set this to the value of `head -c 8 /etc/machine-id`
  networking.hostId = "deadbeef";

  boot.kernelPackages = pkgs.linuxPackages_6_6;
  boot.loader.zfsbootmenu = {
    enable = true;

    # Change this to your root dataset
    bootfs = "zroot/ROOT/nixos";

    # In case you're using ZFS encryption, put zroot.key in /etc/nixos and update this path.
    # (remember to gitignore it if you're version controlling your configuration)
    #keyfile = secrets/zroot.key;
  };
}

Then perform the usual nixos-install:

$ sudo nixos-install

Installing ZFSBootMenu and rEFInd

Make sure /boot is mounted to /mnt/boot (e.g., mount /dev/sda1 /mnt/boot) and enter chroot:

$ sudo nixos-enter

Install pre-built ZFSBootMenu:

$ sudo mkdir -p /boot/efi/EFI/ZBM
$ sudo curl -o /boot/efi/EFI/ZBM/VMLINUZ.EFI -L https://get.zfsbootmenu.org/efi
$ sudo cp /boot/efi/EFI/ZBM/VMLINUZ.EFI /boot/efi/EFI/ZBM/VMLINUZ-BACKUP.EFI

Install rEFInd:

$ sudo nix-shell -p refind efibootmgr
$ sudo refind-install

Configure rEFInd to boot ZFSBootMenu:

$ rm /boot/refind_linux.conf
$ cat << EOF > /boot/efi/EFI/ZBM/refind_linux.conf
"Boot default"  "quiet loglevel=0 zbm.skip"
"Boot to menu"  "quiet loglevel=0 zbm.show"
EOF

Exit chroot and reboot.

NixOS rebuild

Every time you run nixos-generate-config, NixOS will append our ZFS filesystems to its list of systemd-managed filesystem. This can cause mounting errors during boot (though it can generally be ignored). There are few ways to fix this:

• Set mount=legacy to all NixOS managed ZFS filesystems then rebuild²; or
• Set canmount=noauto to / mount (and /nix mount if you have one) and delete the rest from hardware-configuration.nix then rebuild.

Is it stable?

I don’t know, I’ve been using this module across multiple systems since 23.05 or so, and it worked great (even across NixOS upgrades). Though, the installation steps were written from a very old note of mine, it may require a bit of debugging to get it working. Feel free to contact me if you need help!

Changes

• Rev 1.1 (Dec 22, 2024)
  • Clarify more ways to handle NixOS rebuild