From 9502718e5c39acf8f7f84c5170e789c68b88d8b4 Mon Sep 17 00:00:00 2001
From: Zach Dykstra <dykstra.zachary@gmail.com>
Date: Mon, 23 Oct 2023 20:31:17 -0500
Subject: [PATCH] bin/generate-zbm: use mkinitcpio method to calc offsets

Newer versions of the EFI stub provided by systemd-boot are incompatible
with the method used by generate-zbm to add sections to the output
bundle. mkinitcpio uses a process that evaluates the sizes of each
component being combined into the EFI bundle and adjusts their section
VMA accordingly.

This has been tested with the latest systemd-boot EFI stub and the most
recent Gummiboot EFI stub.
---
 bin/generate-zbm | 93 ++++++++++++++++++++++++++++++++++++------------
 1 file changed, 71 insertions(+), 22 deletions(-)

diff --git a/bin/generate-zbm b/bin/generate-zbm
index cb2ce5fb3..539436298 100755
--- a/bin/generate-zbm
+++ b/bin/generate-zbm
@@ -7,7 +7,7 @@ use warnings;
 our $VERSION = '2.2.1';
 
 use Getopt::Long qw(:config no_ignore_case auto_version);
-use Pod::Usage qw(pod2usage);
+use Pod::Usage   qw(pod2usage);
 use File::Basename;
 use File::Temp qw(tempfile tempdir);
 use File::Copy;
@@ -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 bigint qw(hex);
 
 use Pod::Usage qw(pod2usage);
 
@@ -53,7 +54,7 @@ BEGIN {
 
 my ( %runConf, %config );
 
-$runConf{config} = "/etc/zfsbootmenu/config.yaml";
+$runConf{config}  = "/etc/zfsbootmenu/config.yaml";
 $runConf{bootdir} = "/boot";
 
 GetOptions(
@@ -713,8 +714,27 @@ EOF
   return $namever;
 }
 
+sub increaseBundleOffset {
+  my ( $step, $offset, $alignment ) = @_;
+  $offset += ( ( $step + $alignment - 1 ) / $alignment * $alignment );
+  return $offset;
+}
+
+sub addBundleSection {
+  my ( $cmds, $secname, $filename, $offset, $alignment ) = @_;
+
+  my $hex_offset = sprintf( "0x%X", $offset );
+  push( @$cmds, ( "--add-section", "$secname=\"$filename\"" ), qw(--change-section-vma), ("$secname=\"$hex_offset\""),
+  );
+
+  my $sb = stat($filename);
+  return increaseBundleOffset( $sb->size, $offset, $alignment );
+
+}
+
 # Creates a UEFI bundle from an initramfs and kernel
 # Returns the path to the bundle or dies with an error
+
 sub createUEFIBundle {
   my ( $imagedir, $kernel, $initramfs ) = @_;
 
@@ -734,6 +754,7 @@ sub createUEFIBundle {
       exit 1;
     }
   } else {
+
     # For now, default stub locations are x86_64 only
     my @uefi_stub_defaults = qw(
       /usr/lib/gummiboot/linuxx64.efi.stub
@@ -756,17 +777,53 @@ sub createUEFIBundle {
     }
   }
 
-  my @cmd = qw(objcopy);
+  my ( $uki_alignment, $uki_offset );
+
+  # Determine stub alignment
+  my @cmd = qw(objdump -p);
+  push( @cmd, $uefi_stub );
+
+  my @output = execute(@cmd);
+  my $status = pop(@output);
+  if ( $status eq 0 ) {
+    foreach my $line (@output) {
+      if ( $line =~ m/SectionAlignment\s+(\d+)/ ) {
+        Log( "Alignment is: " . hex($1) );
+        $uki_alignment = hex($1);
+      }
+    }
+  } else {
+    print "Unable to determine stub alignment!\n";
+    exit 1;
+  }
+
+  # Determine initial UKI offset value
+  @cmd = qw(objdump -w -h);
+  push( @cmd, $uefi_stub );
+
+  @output = execute(@cmd);
+  $status = pop(@output);
+  if ( $status eq 0 ) {
+    my @sizes = split( /\s+/, @output[ scalar @output - 1 ] );
+
+    my $size  = "0x" . $sizes[3];
+    my $vma   = "0x" . $sizes[4];
+    my $sum = hex($size) + hex($vma);
+
+    $uki_offset = increaseBundleOffset( $sum, 0, $uki_alignment );
+  } else {
+    print "Unable to determine initial stub offset!\n";
+    exit 1;
+  }
+
+  @cmd = qw(objcopy);
+
+  my ( $hex_offset, $sb );
 
-  # Add os-release, if it exists
   if ( -f "/etc/os-release" ) {
-    push( @cmd,
-      qw(--add-section .osrel=/etc/os-release),
-      qw(--change-section-vma .osrel=0x20000)
-    );
+    $uki_offset = addBundleSection( \@cmd, ".osrel", "/etc/os-release", $uki_offset, $uki_alignment );
   }
 
-  # Add cmdline, if it exists
   if ( nonempty $runConf{cmdline} ) {
     my $cmdline = join( '/', $imagedir, "cmdline.txt" );
 
@@ -774,27 +831,19 @@ sub createUEFIBundle {
     print $fh $runConf{cmdline};
     close($fh);
 
-    push( @cmd,
-      ( "--add-section", ".cmdline=\"$cmdline\"" ),
-      qw(--change-section-vma .cmdline=0x30000)
-    );
+    $uki_offset = addBundleSection( \@cmd, ".cmdline", $cmdline, $uki_offset, $uki_alignment );
   }
 
-  # Mandatory kernel and initramfs images
-  push( @cmd, (
-    ( "--add-section", ".linux=\"$kernel\"" ),
-    qw(--change-section-vma .linux=0x2000000),
-    ( "--add-section", ".initrd=\"$initramfs\"" ),
-    qw(--change-section-vma .initrd=0x3000000)
-  ));
+  $uki_offset = addBundleSection( \@cmd, ".linux",  $kernel,    $uki_offset, $uki_alignment );
+  $uki_offset = addBundleSection( \@cmd, ".initrd", $initramfs, $uki_offset, $uki_alignment );
 
   push( @cmd, ( $uefi_stub, $output_file ) );
 
   my $command = join( ' ', @cmd );
   Log("Executing: $command");
 
-  my @output = execute(@cmd);
-  my $status = pop(@output);
+  @output = execute(@cmd);
+  $status = pop(@output);
   if ( $status eq 0 ) {
     foreach my $line (@output) {
       Log($line);