From bd20f3e419deae894f7a4804263adf78cc428a4a Mon Sep 17 00:00:00 2001
From: Daniel Barlow <dan@telent.net>
Date: Sun, 8 Oct 2023 22:35:30 +0100
Subject: [PATCH] uimage: make fit optional

---
 devices/belkin-rt3200/default.nix |  8 ++--
 modules/base.nix                  |  4 ++
 modules/outputs.nix               |  1 +
 pkgs/kernel/uimage.nix            | 68 ++++++++++++++++++-------------
 4 files changed, 49 insertions(+), 32 deletions(-)

diff --git a/devices/belkin-rt3200/default.nix b/devices/belkin-rt3200/default.nix
index 554e20f..6c07084 100644
--- a/devices/belkin-rt3200/default.nix
+++ b/devices/belkin-rt3200/default.nix
@@ -131,9 +131,11 @@
         SERIAL_OF_PLATFORM="y";
       };
     };
-    boot.commandLine = [
-      "console=ttyS0,115200"
-    ];
+    boot = {
+      commandLine = [ "console=ttyS0,115200" ];
+      tftp.loadAddress =  "0x4007ff28";
+      imageFormat = "fit";
+    };
     filesystem =
       let inherit (pkgs.pseudofile) dir symlink;
            in
diff --git a/modules/base.nix b/modules/base.nix
index a39c689..fb059d2 100644
--- a/modules/base.nix
+++ b/modules/base.nix
@@ -47,6 +47,10 @@ in {
         default = [];
         description = "Kernel command line";
       };
+      imageFormat = mkOption {
+        type = types.enum ["fit" "uimage"];
+        default = "uimage";
+      };
       tftp = {
         loadAddress = mkOption {
           type = types.str;
diff --git a/modules/outputs.nix b/modules/outputs.nix
index 3232d80..ae92bbf 100644
--- a/modules/outputs.nix
+++ b/modules/outputs.nix
@@ -78,6 +78,7 @@ in
       uimage = liminix.builders.uimage {
         commandLine = concatStringsSep " " config.boot.commandLine;
         inherit (config.hardware) loadAddress entryPoint;
+        inherit (config.boot) imageFormat;
         inherit kernel;
         inherit dtb;
       };
diff --git a/pkgs/kernel/uimage.nix b/pkgs/kernel/uimage.nix
index a5c56fb..e6ac685 100644
--- a/pkgs/kernel/uimage.nix
+++ b/pkgs/kernel/uimage.nix
@@ -7,21 +7,27 @@
 } :
 let
   objcopy = "${stdenv.cc.bintools.targetPrefix}objcopy";
+  arch = stdenv.hostPlatform.linuxArch;
+  stripAndZip = ''
+    ${objcopy} -O binary -R .reginfo -R .notes -R .note -R .comment -R .mdebug -R .note.gnu.build-id -S vmlinux.elf vmlinux.bin
+    rm -f vmlinux.bin.lzma ; lzma -k -z  vmlinux.bin
+  '';
 in {
   kernel
 , commandLine
 , entryPoint
 , extraName ? ""                # e.g. socFamily
 , loadAddress
+, imageFormat
 , dtb ? null
-} :
-stdenv.mkDerivation {
+} : stdenv.mkDerivation {
   name = "kernel.image";
   phases = [
     "preparePhase"
-    (if dtb != null then "dtbPhase" else ":")
-    "buildPhase"
-    "installPhase" ];
+    (if commandLine != null then assert dtb != null; "mungeDtbPhase" else ":")
+    (if imageFormat == "fit" then "buildPhaseFIT" else "buildPhaseUImage")
+    "installPhase"
+  ];
   nativeBuildInputs = [
     lzma
     dtc
@@ -31,35 +37,39 @@ stdenv.mkDerivation {
   preparePhase = ''
     cp ${kernel} vmlinux.elf; chmod +w vmlinux.elf
   '';
-  dtbPhase = ''
+  mungeDtbPhase = ''
     dtc -I dtb -O dts -o tmp.dts ${dtb}
     echo '/{ chosen { bootargs = ${builtins.toJSON commandLine}; }; };'  >> tmp.dts
     dtc -I dts -O dtb -o tmp.dtb tmp.dts
   '';
 
-  buildPhase =
-    let arch =
-          # per output of "mkimage -A list". I *think* these
-          # are the same as the kernel arch convention, but
-          # maybe that's coincidence
-          if stdenv.isMips
-          then "mips"
-          else if stdenv.isAarch64
-          then "arm64"
-          else throw "unknown arch";
-    in ''
-      ${objcopy} -O binary -R .reginfo -R .notes -R .note -R .comment -R .mdebug -R .note.gnu.build-id -S vmlinux.elf vmlinux.bin
-      rm -f vmlinux.bin.lzma ; lzma -k -z  vmlinux.bin
-      cat ${./kernel_fdt.its} > mkimage.its
-      echo '/ { images { kernel { data = /incbin/("./vmlinux.bin.lzma"); }; }; };' >> mkimage.its
-      echo '/ { images { kernel { load = <${loadAddress}>; }; }; };' >> mkimage.its
-      echo '/ { images { kernel { entry = <${entryPoint}>; }; }; };' >> mkimage.its
-      echo '/ { images { kernel { compression = "lzma"; }; }; };' >> mkimage.its
-      echo '/ { images { fdt-1 { data = /incbin/("./tmp.dtb"); }; }; }; ' >> mkimage.its
-      mkimage -f mkimage.its mkimage.itb
-      mkimage -l mkimage.itb
-    '';
+  buildPhaseUImage = ''
+    test -f tmp.dtb && ${objcopy} --update-section .appended_dtb=tmp.dtb vmlinux.elf || ${objcopy} --add-section .appended_dtb=tmp.dtb vmlinux.elf
+    ${stripAndZip}
+    mkimage -A ${arch} -O linux -T kernel -C lzma -a ${loadAddress} -e ${entryPoint} -n '${lib.toUpper arch} Liminix Linux ${extraName}' -d vmlinux.bin.lzma kernel.uimage
+  '';
+
+  buildPhaseFIT = ''
+    ${stripAndZip}
+    cat ${./kernel_fdt.its} > mkimage.its
+    cat << _VARS  >> mkimage.its
+    / {
+        images {
+            kernel {
+                data = /incbin/("./vmlinux.bin.lzma");
+                load = <${loadAddress}>;
+                entry = <${entryPoint}>;
+                compression = "lzma";
+            };
+            fdt-1 { data = /incbin/("./tmp.dtb"); };
+        };
+    };
+    _VARS
+    mkimage -f mkimage.its kernel.uimage
+    mkimage -l kernel.uimage
+  '';
+
   installPhase = ''
-    cp mkimage.itb $out
+    cp kernel.uimage $out
   '';
 }