From 29dc255fd3aed8fc192e373133a6a9ef8c8b26f2 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/to/device-tree-file.dtb
2) /absolute/path/%{kernel}/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
values of this property are handled in the same order/priority as
EFI.DeviceTree.
---
 bin/generate-zbm                    | 43 ++++++++++++++++++++++++++---
 zfsbootmenu/lib/zfsbootmenu-core.sh | 32 +++++++++++++++++++++
 2 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/bin/generate-zbm b/bin/generate-zbm
index 0c33633a9..08d716314 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,38 @@ sub createUEFIBundle {
     $uki_offset = addBundleSection( \@cmd, ".splash", $config{EFI}{SplashImage}, $uki_offset, $uki_alignment );
   }
 
+  my $dtb;
+
+  if ( has_value $config{EFI}{DeviceTree} ) {
+    my @dtb_tests;
+
+    # Possibly an absolute path to a file
+    if ( $config{EFI}{DeviceTree} =~ m,^/, ) {
+      push( @dtb_tests, $config{EFI}{DeviceTree} );
+    }
+
+    # Possibly a templated file path
+    if ( $config{EFI}{DeviceTree} =~ m/\Q%{kernel}/ ) {
+      my $dtb_template = $config{EFI}{DeviceTree};
+      $dtb_template =~ s/\Q%{kernel}/$runConf{kernel_version}/;
+      push( @dtb_tests, $dtb_template );
+    }
+
+    # Possibly a partial path
+    push( @dtb_tests, sprintf( "/boot/dtbs/dtbs-%s/%s", $runConf{kernel_version}, $config{EFI}{DeviceTree} ) );
+
+    foreach my $dtb_test (@dtb_tests) {
+      if ( -f $dtb_test ) {
+        $dtb = $dtb_test;
+        last;
+      }
+    }
+  }
+
+  if ( has_value $dtb and -f $dtb ) {
+    $uki_offset = addBundleSection( \@cmd, ".dtb", $dtb, $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..8b363dcb8 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_tries dtb_try dtb_file kver
 
   selected="${1}"
   if [ -z "${selected}" ]; then
@@ -390,8 +391,38 @@ 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
+      # possibly an absolute path to a file
+      dtb_tries+=( "${mnt}/${dtb_prop}" )
+      if [[ "${dtb_prop}" =~ %{kernel} ]] ; then
+        # possibly a templated absolute path to a file
+        dtb_tries+=( "${mnt}/${dtb_prop//%\{kernel\}/$kver}" )
+      fi
+    else
+      # possibly a suffix to a well-known path
+      dtb_tries+=( "${mnt}/boot/dtbs/dtbs${kver}/${dtb_prop}" )
+    fi
+
+    for dtb_try in "${dtb_tries[@]}" ; do
+      zdebug "Checking if '${dtb_try}' exists"
+      if [ -f "${dtb_try}" ] ; then
+        zdebug "exists: '${dtb_try}'"
+        dtb_file="${dtb_try}"
+        break
+      fi
+    done
+  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 +2050,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