From ff0ef825a6dbcd74c23e8def32752c46a6842426 Mon Sep 17 00:00:00 2001
From: Daniel Barlow <dan@telent.net>
Date: Fri, 22 Dec 2023 22:15:54 +0000
Subject: [PATCH] tftpboot: add option for kernel image format

---
 modules/outputs/tftpboot.nix | 28 +++++++++++++++++++++-------
 tests/tftpboot/test.nix      | 16 +++++++++++-----
 2 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/modules/outputs/tftpboot.nix b/modules/outputs/tftpboot.nix
index de464cf..87ae34f 100644
--- a/modules/outputs/tftpboot.nix
+++ b/modules/outputs/tftpboot.nix
@@ -9,9 +9,15 @@ let
   cfg = config.boot.tftp;
 in {
   imports = [ ../ramdisk.nix ];
-  options.boot.tftp.freeSpaceBytes = mkOption {
-    type = types.int;
-    default = 0;
+  options.boot.tftp = {
+    freeSpaceBytes = mkOption {
+      type = types.int;
+      default = 0;
+    };
+    kernelFormat = mkOption {
+      type = types.enum [ "zimage" "uimage" ];
+      default = "uimage";
+    };
   };
   options.system.outputs = {
     tftpboot = mkOption {
@@ -44,6 +50,14 @@ in {
         let
           inherit (pkgs.lib.trivial) toHexString;
           o = config.system.outputs;
+          image = let choices = {
+            uimage = o.uimage;
+            zimage = o.zimage;
+          }; in choices.${cfg.kernelFormat};
+          bootCommand = let choices = {
+            uimage = "bootm";
+            zimage = "bootz";
+          }; in choices.${cfg.kernelFormat};
           cmdline = concatStringsSep " " config.boot.commandLine;
         in
           pkgs.runCommand "tftpboot" { nativeBuildInputs = [ pkgs.pkgsBuildBuild.dtc ];  } ''
@@ -53,8 +67,8 @@ in {
             ln -s ${o.kernel} vmlinux
             ln -s ${o.manifest} manifest
             ln -s ${o.kernel.headers} build
-            ln -s ${o.uimage} uimage
-            uimageSize=$(($(stat -L -c %s ${o.uimage}) + 0x1000 &(~0xfff)))
+            ln -s ${image} image
+            uimageSize=$(($(stat -L -c %s ${image}) + 0x1000 &(~0xfff)))
             rootfsStart=$(printf %x $((${toString cfg.loadAddress} + 0x100000 + $uimageSize   &(~0xfffff) )))
             rootfsBytes=$(($(stat -L -c %s ${o.rootfs}) + 0x100000 &(~0xfffff)))
             rootfsBytes=$(($rootfsBytes + ${toString cfg.freeSpaceBytes} ))
@@ -82,8 +96,8 @@ in {
             setenv serverip ${cfg.serverip}
             setenv ipaddr ${cfg.ipaddr}
             setenv bootargs 'liminix ${cmdline} $cmd'
-            tftpboot 0x${lib.toHexString cfg.loadAddress} result/uimage ; tftpboot 0x$rootfsStart result/rootfs ; tftpboot 0x$dtbStart result/dtb
-            bootm 0x${lib.toHexString cfg.loadAddress} - 0x$dtbStart
+            tftpboot 0x${lib.toHexString cfg.loadAddress} result/image ; tftpboot 0x$rootfsStart result/rootfs ; tftpboot 0x$dtbStart result/dtb
+            ${bootCommand} 0x${lib.toHexString cfg.loadAddress} - 0x$dtbStart
             EOF
          '';
 
diff --git a/tests/tftpboot/test.nix b/tests/tftpboot/test.nix
index bb9c3af..d4682f6 100644
--- a/tests/tftpboot/test.nix
+++ b/tests/tftpboot/test.nix
@@ -1,10 +1,13 @@
 {
   liminix
 }:
-let check = deviceName : ubootName :
+let check = deviceName : ubootName : config :
 let derivation = (import liminix {
       device = import "${liminix}/devices/${deviceName}/";
-      liminix-config = ./configuration.nix;
+      liminix-config = { pkgs, ... } : {
+        imports = [./configuration.nix];
+        inherit config;
+      };
     });
     img = derivation.outputs.tftpboot;
     uboot = derivation.pkgs.${ubootName};
@@ -30,7 +33,10 @@ run-liminix-vm \
 expect ${./script.expect} 2>&1 |tee $out
 '';
 in {
-  aarch64 = check "qemu-aarch64" "ubootQemuAarch64";
-  arm = check  "qemu-armv7l" "ubootQemuArm";
-  mips = check  "qemu" "ubootQemuMips";
+  aarch64 = check "qemu-aarch64" "ubootQemuAarch64" {};
+  arm = check  "qemu-armv7l" "ubootQemuArm" {};
+  armZimage = check  "qemu-armv7l" "ubootQemuArm" {
+    boot.tftp.kernelFormat = "zimage";
+  };
+  mips = check  "qemu" "ubootQemuMips" {};
 }