From 65bf5b2d4fa45eb27095eedf36595f0221d82c21 Mon Sep 17 00:00:00 2001
From: Andrew Kidd <contact@andrewiankidd.co.uk>
Date: Thu, 26 Dec 2024 13:38:13 +0000
Subject: [PATCH] Add wip netboot image gen

---
 flake.nix                  |   4 +-
 net-image/default.nix      |  71 ++++++++++++++++
 net-image/make-root-fs.nix |  54 ++++++++++++
 net-image/net-image.nix    | 163 +++++++++++++++++++++++++++++++++++++
 4 files changed, 291 insertions(+), 1 deletion(-)
 create mode 100644 net-image/default.nix
 create mode 100644 net-image/make-root-fs.nix
 create mode 100644 net-image/net-image.nix

diff --git a/flake.nix b/flake.nix
index c0acdc3..d0f1730 100644
--- a/flake.nix
+++ b/flake.nix
@@ -60,11 +60,12 @@
           libcamera-overlay = self.overlays.libcamera;
         };
         sd-image = import ./sd-image;
+        net-image = import ./net-image;
       };
       nixosConfigurations = {
         rpi-example = srcs.nixpkgs.lib.nixosSystem {
           system = "aarch64-linux";
-          modules = [ self.nixosModules.raspberry-pi self.nixosModules.sd-image ./example ];
+          modules = [ self.nixosModules.raspberry-pi self.nixosModules.sd-image self.nixosModules.net-image ./example ];
         };
       };
       checks.aarch64-linux = self.packages.aarch64-linux;
@@ -82,6 +83,7 @@
         in
         {
           example-sd-image = self.nixosConfigurations.rpi-example.config.system.build.sdImage;
+          example-net-image = self.nixosConfigurations.rpi-example.config.system.build.netImage;
           firmware = pinned.raspberrypifw;
           libcamera = pinned.libcamera;
           wireless-firmware = pinned.raspberrypiWirelessFirmware;
diff --git a/net-image/default.nix b/net-image/default.nix
new file mode 100644
index 0000000..ef196d9
--- /dev/null
+++ b/net-image/default.nix
@@ -0,0 +1,71 @@
+{ config, lib, pkgs, ... }:
+
+{
+  imports = [ ./net-image.nix ];
+
+  config = {
+    boot.loader.grub.enable = false;
+
+    boot.consoleLogLevel = lib.mkDefault 7;
+
+    boot.kernelParams = [
+      "rw"
+      "nfsroot=${config.netImage.nfsRoot}"
+      "ip=dhcp"
+      "root=/dev/nfs"
+      "rootwait"
+      "elevator=deadline"
+      "console=tty1"
+      "console=serial0,115200n8"
+      "init=/sbin/init"
+      "loglevel=7"
+    ];
+
+    netImage =
+      let
+        kernel-params = pkgs.writeTextFile {
+          name = "cmdline.txt";
+          text = ''
+            ${lib.strings.concatStringsSep " " config.boot.kernelParams}
+          '';
+        };
+        cfg = config.raspberry-pi-nix;
+        version = cfg.kernel-version;
+        board = cfg.board;
+        kernel = "${config.system.build.kernel}/${config.system.boot.loader.kernelFile}";
+        initrd = "${config.system.build.initialRamdisk}/${config.system.boot.loader.initrdFile}";
+        populate-kernel =
+          if cfg.uboot.enable
+          then ''
+            cp ${cfg.uboot.package}/u-boot.bin firmware/u-boot-rpi-arm64.bin
+          ''
+          else ''
+            cp "${kernel}" firmware/kernel.img
+            cp "${initrd}" firmware/initrd
+            cp "${kernel-params}" firmware/cmdline.txt
+          '';
+      in
+      {
+        populateFirmwareCommands = ''
+          ${populate-kernel}
+          cp -r ${pkgs.raspberrypifw}/share/raspberrypi/boot/{start*.elf,*.dtb,bootcode.bin,fixup*.dat,overlays} firmware
+          cp ${config.hardware.raspberry-pi.config-output} firmware/config.txt
+        '';
+        populateRootCommands =
+          if cfg.uboot.enable
+          then ''
+            mkdir -p ./files/boot
+            ${config.boot.loader.generic-extlinux-compatible.populateCmd} -c ${config.system.build.toplevel} -d ./files/boot
+          ''
+          else ''
+            mkdir -p ./files/sbin
+            content="$(
+              echo "#!${pkgs.bash}/bin/bash"
+              echo "exec ${config.system.build.toplevel}/init"
+            )"
+            echo "$content" > ./files/sbin/init
+            chmod 744 ./files/sbin/init
+          '';
+      };
+  };
+}
diff --git a/net-image/make-root-fs.nix b/net-image/make-root-fs.nix
new file mode 100644
index 0000000..34d30c8
--- /dev/null
+++ b/net-image/make-root-fs.nix
@@ -0,0 +1,54 @@
+# Builds an ext4 image containing a populated /nix/store with the closure
+# of store paths passed in the storePaths parameter, in addition to the
+# contents of a directory that can be populated with commands. The
+# generated image is sized to only fit its contents, with the expectation
+# that a script resizes the filesystem at boot time.
+{
+  pkgs,
+  lib,
+  # List of derivations to be included
+  storePaths,
+  # Shell commands to populate the ./files directory.
+  # All files in that directory are copied to the root of the FS.
+  populateImageCommands ? "",
+  perl
+}:
+
+let
+  netbootClosureInfo = pkgs.buildPackages.closureInfo { rootPaths = storePaths; };
+in
+pkgs.stdenv.mkDerivation {
+  name = "root-fs";
+
+  nativeBuildInputs = [
+    perl
+  ];
+
+  buildCommand = ''
+    echo "Populating image with command: ${populateImageCommands}"
+    mkdir -p ./files
+    ${populateImageCommands}
+
+    echo "Preparing store paths for image..."
+    # Create nix/store before copying path
+    mkdir -p ./rootImage/nix/store
+
+    xargs -I % cp -a --reflink=auto % -t ./rootImage/nix/store/ < ${netbootClosureInfo}/store-paths
+    (
+      GLOBIGNORE=".:.."
+      shopt -u dotglob
+
+      for f in ./files/*; do
+          cp -a --reflink=auto -t ./rootImage/ "$f"
+      done
+    )
+
+    # Also include a manifest of the closures in a format suitable for nix-store --load-db
+    cp ${netbootClosureInfo}/registration ./rootImage/nix-path-registration
+
+    # done
+    echo "Image populated."
+    ls -aR ./rootImage
+    cp -r ./rootImage/ $out
+  '';
+}
diff --git a/net-image/net-image.nix b/net-image/net-image.nix
new file mode 100644
index 0000000..8ce8b5b
--- /dev/null
+++ b/net-image/net-image.nix
@@ -0,0 +1,163 @@
+# This module creates a bootable netboot image containing the given NixOS
+# configuration. The generated image consists of a `/boot` partition
+# (for TFTP boot) and a `/root` partition (for NFS root). The goal is to
+# allow the system to boot over the network, using TFTP to retrieve the
+# boot files and NFS to mount the root filesystem, enabling fully
+# headless deployment of the NixOS system.
+
+# The generated netboot image consists of two directories:
+# - `/boot`: Contains the necessary bootloader, kernel, and initrd files
+#   required for booting the system. These files will be served via TFTP
+#   to the target machine.
+# - `/root`: Contains the root filesystem that will be mounted by the
+#   target machine over NFS. This is typically an ext4 root partition
+#   populated with the necessary NixOS configuration.
+
+# The image is generated in such a way that it can be used to netboot a
+# Raspberry Pi (or any other compatible hardware) directly, as long as
+# the appropriate network boot infrastructure (TFTP server for `/boot`
+# and NFS server for `/root`) is configured.
+
+# The image does not include a bootable SD card but instead prepares the
+# filesystem and boot files for network-based booting. The NixOS
+# configuration will be automatically applied when the system boots.
+
+# The generated image will be placed in
+# config.system.build.netImage. This image is intended to be deployed
+# to a TFTP server (for the boot files) and an NFS server (for the root
+# filesystem) for a fully headless, network-booted NixOS system.
+
+# Note: This module assumes that you have already set up the TFTP and
+# NFS servers on your network, and the target machine is configured
+# for network booting.
+
+{ modulesPath, config, lib, pkgs, ... }:
+
+with lib;
+
+let
+  rootfsImage = pkgs.callPackage (builtins.path { path = ./make-root-fs.nix; }) ({
+    inherit (config.netImage) storePaths;
+    populateImageCommands = config.netImage.populateRootCommands;
+  });
+in
+{
+  imports = [ ];
+
+  options.netImage = {
+    rootDirectoryName = mkOption {
+      default =
+        "${config.netImage.imageBaseName}-${config.system.nixos.label}-${pkgs.stdenv.hostPlatform.system}";
+      description = ''
+        Name of the generated root directory.
+      '';
+    };
+
+    imageBaseName = mkOption {
+      default = "nixos-net-image";
+      description = ''
+        Prefix of the name of the generated image file.
+      '';
+    };
+
+    storePaths = mkOption {
+      type = with types; listOf package;
+      example = literalExpression "[ pkgs.stdenv ]";
+      description = ''
+        Derivations to be included in the Nix store in the generated Netboot image.
+      '';
+    };
+
+    nfsRoot = mkOption {
+      type = types.str;
+      default = "192.168.0.108:/mnt/nfsshare/${config.netImage.imageBaseName}-${config.system.nixos.label}-${pkgs.stdenv.hostPlatform.system},v3";
+      description = ''
+        cmdline.txt nfs parameter for the root filesystem.
+      '';
+    };
+
+    populateFirmwareCommands = mkOption {
+      example =
+        literalExpression "'' cp \${pkgs.myBootLoader}/u-boot.bin firmware/ ''";
+      description = ''
+        Shell commands to populate the ./firmware directory.
+        All files in that directory are copied to the
+        /boot/firmware partition on the Netboot image.
+      '';
+    };
+
+    populateRootCommands = mkOption {
+      example = literalExpression
+        "''\${config.boot.loader.generic-extlinux-compatible.populateCmd} -c \${config.system.build.toplevel} -d ./files/boot''";
+      description = ''
+        Shell commands to populate the ./files directory.
+        All files in that directory are copied to the
+        root (/) partition on the Netboot image. Use this to
+        populate the ./files/boot (/boot) directory.
+      '';
+    };
+
+    postBuildCommands = mkOption {
+      example = literalExpression
+        "'' dd if=\${pkgs.myBootLoader}/SPL of=$img bs=1024 seek=1 conv=notrunc ''";
+      default = "";
+      description = ''
+        Shell commands to run after the image is built.
+        Can be used for boards requiring to dd u-boot SPL before actual partitions.
+      '';
+    };
+  };
+
+  config = {
+    netImage.storePaths = [ config.system.build.toplevel ];
+
+    system.build.netImage = pkgs.callPackage
+      ({ stdenv, util-linux }:
+        stdenv.mkDerivation {
+          name = config.netImage.rootDirectoryName;
+
+          nativeBuildInputs = [ util-linux ];
+
+          buildCommand = ''
+            set -e
+            set -x
+            mkdir -p $out/nix-support $out/net-image
+            export rootfs=$out/net-image/${config.netImage.rootDirectoryName}
+            export bootfs=$out/net-image/boot
+
+            echo "${pkgs.stdenv.buildPlatform.system}" > $out/nix-support/system
+
+            echo "Exporting rootfs image"
+            mkdir -p $rootfs
+            mv ${rootfsImage} $rootfs
+
+            # Populate the files intended for /boot/firmware
+            mkdir -p firmware
+            ${config.netImage.populateFirmwareCommands}
+            mkdir -p $bootfs
+            mv firmware $bootfs
+
+            ${config.netImage.postBuildCommands}
+          '';
+        })
+      { };
+
+    boot.postBootCommands = ''
+      # On the first boot do some maintenance tasks
+      if [ -f /nix-path-registration ]; then
+        set -euo pipefail
+        set -x
+
+        # Register the contents of the initial Nix store
+        ${config.nix.package.out}/bin/nix-store --load-db < /nix-path-registration
+
+        # nixos-rebuild also requires a "system" profile and an /etc/NIXOS tag.
+        touch /etc/NIXOS
+        ${config.nix.package.out}/bin/nix-env -p /nix/var/nix/profiles/system --set /run/current-system
+
+        # Prevents this from running on later boots.
+        rm -f /nix-path-registration
+      fi
+    '';
+  };
+}