From ec869e2b2164a25ed20a4862afbf547e14d73b1b Mon Sep 17 00:00:00 2001
From: Zach Dykstra <dykstra.zachary@gmail.com>
Date: Sat, 1 Feb 2025 19:02:04 -0600
Subject: [PATCH] aarch64 support: efi stub and devicetree selection

generate-zbm now understands the EFI.DeviceTree key. This key can
contain one of three possible values:

1) /absolute/path/%{kernel}/to/device-tree-file.dtb
2) /absolute/path/to/device-tree-file.dtb
3) partial/path/device-tree-file.dtb

If found, %{kernel} is replaced with the kernel version being built in
to the EFI asset. Partial paths are appended to /boot/dtbs/dtb-$kernel/,
a well-known path for device tree files used by Void Linux.

Inside ZFSBootMenu, the org.zfsbootmenu:devicetree property is read. The
value of this property are handled in the same order/priority as
EFI.DeviceTree. %{kernel} is replaced with the kernel version that kexec
is loading / executing.
---
 bin/generate-zbm                    | 36 +++++++++++++++++++++++++----
 zfsbootmenu/lib/zfsbootmenu-core.sh | 31 +++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/bin/generate-zbm b/bin/generate-zbm
index 0c33633a9..3ef87f7cb 100755
--- a/bin/generate-zbm
+++ b/bin/generate-zbm
@@ -15,6 +15,7 @@ use File::stat;
 use File::Path qw(make_path remove_tree);
 use File::Glob qw(:globally :nocase);
 use Sort::Versions;
+use Config;
 use bigint qw(hex);
 
 use Pod::Usage qw(pod2usage);
@@ -684,10 +685,12 @@ sub createUEFIBundle {
     }
   } else {
 
-    # For now, default stub locations are x86_64 only
-    my @uefi_stub_defaults = qw(
-      /usr/lib/systemd/boot/efi/linuxx64.efi.stub
-    );
+    my @uefi_stub_defaults;
+    if ( $Config{archname} =~ m/x86_64/ ) {
+      push( @uefi_stub_defaults, '/usr/lib/systemd/boot/efi/linuxx64.efi.stub' );
+    } elsif ( $Config{archname} =~ m/aarch64/ ) {
+      push( @uefi_stub_defaults, '/usr/lib/systemd/boot/efi/linuxaa64.efi.stub' );
+    }
 
     foreach my $stubloc (@uefi_stub_defaults) {
       if ( -f $stubloc ) {
@@ -771,6 +774,31 @@ sub createUEFIBundle {
     $uki_offset = addBundleSection( \@cmd, ".splash", $config{EFI}{SplashImage}, $uki_offset, $uki_alignment );
   }
 
+  if ( has_value $config{EFI}{DeviceTree} ) {
+    my $dtb_try;
+
+    if ( $config{EFI}{DeviceTree} =~ m,^/, ) {
+      if ( $config{EFI}{DeviceTree} =~ m/\Q%{kernel}/ ) {
+
+        # Possibly a templated file path
+        $dtb_try = $config{EFI}{DeviceTree};
+        $dtb_try =~ s/\Q%{kernel}/$runConf{kernel_version}/;
+      } else {
+
+        # Possibly an absolute path to a file
+        $dtb_try = $config{EFI}{DeviceTree};
+      }
+    } else {
+
+      # Possibly a partial path
+      $dtb_try = sprintf( "/boot/dtbs/dtbs-%s/%s", $runConf{kernel_version}, $config{EFI}{DeviceTree} );
+    }
+
+    if ( has_value $dtb_try && -f $dtb_try ) {
+      $uki_offset = addBundleSection( \@cmd, ".dtb", $dtb_try, $uki_offset, $uki_alignment );
+    }
+  }
+
   $uki_offset = addBundleSection( \@cmd, ".initrd", $initramfs, $uki_offset, $uki_alignment );
 
   # Add the kernel last, so that it can decompress without overflowing other sections
diff --git a/zfsbootmenu/lib/zfsbootmenu-core.sh b/zfsbootmenu/lib/zfsbootmenu-core.sh
index 4321e4afb..40353e853 100644
--- a/zfsbootmenu/lib/zfsbootmenu-core.sh
+++ b/zfsbootmenu/lib/zfsbootmenu-core.sh
@@ -351,6 +351,7 @@ mount_zfs() {
 
 kexec_kernel() {
   local selected fs kernel initramfs output hook_envs
+  local dtb_prop dtb_try dtb_file kver
 
   selected="${1}"
   if [ -z "${selected}" ]; then
@@ -390,8 +391,37 @@ kexec_kernel() {
   cli_args="$( load_be_cmdline "${fs}" )"
   root_prefix="$( find_root_prefix "${fs}" "${mnt}" )"
 
+  dtb_prop="$( zfs get -H -o value org.zfsbootmenu:devicetree "${fs}" 2>/dev/null )"
+  if [ -n "${dtb_prop}" ] && [ "${dtb_prop}" != '-' ] ; then
+    zdebug "devicetree property set to '${dtb_prop}'"
+
+    kver="${kernel#"${kernel%%-*}"}"
+    if [ "${dtb_prop:0:1}" = "/" ] ; then
+      if [[ "${dtb_prop}" =~ %{kernel} ]] ; then
+        # possibly a templated absolute path to a file
+        dtb_try="${mnt}/${dtb_prop//%\{kernel\}/$kver}"
+      else
+        # possibly an absolute path to a file
+        dtb_try="${mnt}/${dtb_prop}"
+      fi
+    else
+      # possibly a suffix to a well-known path
+      dtb_try="${mnt}/boot/dtbs/dtbs${kver}/${dtb_prop}"
+    fi
+
+    if [ -n "${dtb_try}" ] && [ -f "${dtb_try}" ] ; then
+      zdebug "devicetree file exists: '${dtb_try}'"
+      dtb_file="${dtb_try}"
+    else
+      zdebug "devicetree file missing: '${dtb_try}'"
+    fi
+  else
+    zdebug "no devicetree property set"
+  fi
+
   if ! output="$( kexec -a -l "${mnt}${kernel}" \
     --initrd="${mnt}${initramfs}" \
+    ${dtb_file:+--dtb="${dtb_file}"} \
     --command-line="${root_prefix}${fs} ${cli_args}" 2>&1 )"
   then
     zerror "unable to load ${mnt}${kernel} and ${mnt}${initramfs} into memory"
@@ -2019,6 +2049,7 @@ zreport() {
       read -r INITRD_VERSION < /VERSION
       INITRD_VERSION="mkinitcpio ${INITRD_VERSION}"
     elif [ -f /etc/initrd-release ] ; then
+      # shellcheck disable=SC1091
       source /etc/initrd-release
       [ -n "${DRACUT_VERSION}" ] && INITRD_VERSION="Dracut ${DRACUT_VERSION}"
     fi