Make boot-file architecture config distro-aware

toolkit/tools/internal/initrdutils/initrdread.go

@@ -0,0 +1,120 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package initrdutils
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"io"
+	"io/fs"
+	"os"
+
+	"github.com/cavaliergopher/cpio"
+	"github.com/klauspost/compress/zstd"
+	"github.com/klauspost/pgzip"
+)
+
+// Magic byte signatures for the compression formats we recognize at the start
+// of an initramfs stream.
+//
+//   - gzip: RFC 1952 section 2.3.1 ("ID1 ID2" = 1F 8B)
+//     https://www.rfc-editor.org/rfc/rfc1952#section-2.3.1
+//   - zstd: RFC 8478 section 3.1.1 ("Magic_Number" = 0xFD2FB528 little-endian)
+//     https://www.rfc-editor.org/rfc/rfc8478#section-3.1.1
+var (
+	magicGzip        = []byte{0x1f, 0x8b}
+	magicZstd        = []byte{0x28, 0xb5, 0x2f, 0xfd}
+	longestMagicSize = len(magicZstd)
+)
+
+// ReadFirstFileFromInitrd scans an initramfs cpio archive once and returns the content of the first candidate path in
+// the provided list that exists as a regular file. Returns an error wrapping fs.ErrNotExist if none of the candidates
+// is present.
+func ReadFirstFileFromInitrd(initrdPath string, candidates []string) (content []byte, foundPath string, err error) {
+	f, err := os.Open(initrdPath)
+	if err != nil {
+		return nil, "", fmt.Errorf("failed to open initrd (%s):\n%w", initrdPath, err)
+	}
+	defer f.Close()
+
+	decompressed, err := openInitrdDecompressor(f)
+	if err != nil {
+		return nil, "", fmt.Errorf("failed to decompress initrd (%s):\n%w", initrdPath, err)
+	}
+	defer decompressed.Close()
+
+	wanted := make(map[string]struct{}, len(candidates))
+	for _, c := range candidates {
+		wanted[c] = struct{}{}
+	}
+
+	found := make(map[string][]byte, len(candidates))
+	cpioReader := cpio.NewReader(decompressed)
+	for {
+		hdr, err := cpioReader.Next()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return nil, "", fmt.Errorf("failed to read cpio header from initrd (%s):\n%w", initrdPath, err)
+		}
+
+		if _, ok := wanted[hdr.Name]; !ok {
+			continue
+		}
+
+		// Only read regular files.
+		if hdr.Mode&cpio.ModeType != cpio.TypeReg {
+			continue
+		}
+
+		data, err := io.ReadAll(cpioReader)
+		if err != nil {
+			return nil, "", fmt.Errorf("failed to read (%s) from initrd (%s):\n%w", hdr.Name, initrdPath, err)
+		}
+
+		found[hdr.Name] = data
+	}
+
+	for _, candidate := range candidates {
+		if data, ok := found[candidate]; ok {
+			return data, candidate, nil
+		}
+	}
+
+	return nil, "", fmt.Errorf("failed to find any of %v in initrd (%s): %w", candidates, initrdPath, fs.ErrNotExist)
+}
+
+// openInitrdDecompressor auto-detects the compression format of an initramfs stream from its leading magic bytes and
+// returns a reader over the decompressed (cpio) content.
+func openInitrdDecompressor(r io.Reader) (io.ReadCloser, error) {
+	br := bufio.NewReader(r)
+	head, err := br.Peek(longestMagicSize)
+	if err != nil && err != io.EOF {
+		return nil, fmt.Errorf("failed to peek initrd magic bytes:\n%w", err)
+	}
+
+	switch {
+	case bytes.HasPrefix(head, magicGzip):
+		gz, err := pgzip.NewReader(br)
+		if err != nil {
+			return nil, fmt.Errorf("failed to open gzip reader for initrd:\n%w", err)
+		}
+		return gz, nil
+
+	case bytes.HasPrefix(head, magicZstd):
+		zr, err := zstd.NewReader(br)
+		if err != nil {
+			return nil, fmt.Errorf("failed to open zstd reader for initrd:\n%w", err)
+		}
+
+		// zstd.Decoder satisfies io.Reader but not io.Closer since its Close() returns no error, so wrap it to
+		// implement io.ReadCloser, like pgzip.Reader.
+		return zstdReadCloser{Decoder: zr}, nil
+
+	default:
+		return nil, fmt.Errorf("unrecognized initrd compression format (leading bytes: % x)", head)
+	}
+}

toolkit/tools/internal/initrdutils/zstdreadcloser.go

@@ -0,0 +1,16 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package initrdutils
+
+import (
+	"github.com/klauspost/compress/zstd"
+)
+
+// zstdReadCloser adapts *zstd.Decoder (whose Close() returns no error) to io.ReadCloser.
+type zstdReadCloser struct{ *zstd.Decoder }
+
+func (z zstdReadCloser) Close() error {
+	z.Decoder.Close()
+	return nil
+}

toolkit/tools/internal/targetos/targetos.go

@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 
 	"github.com/microsoft/azure-linux-image-tools/toolkit/tools/internal/envfile"
+	"github.com/microsoft/azure-linux-image-tools/toolkit/tools/internal/initrdutils"
 	"github.com/microsoft/azure-linux-image-tools/toolkit/tools/internal/version"
 )
 
@@ -74,10 +75,27 @@ var (
 		Version:   []int{24, 4},
 	}
 
+	// OsReleaseCandidates lists the on-rootfs paths to probe for the os-release(5) file, in preference order.
 	OsReleaseFileCandidates = []string{
 		"/etc/os-release",
 		"/usr/lib/os-release",
 	}
+
+	// initrdReleaseCandidates lists the in-initrd paths to probe for an os-release(5)-type file, in preference order.
+	// These are matched against cpio member names, which are stored relative to the archive root, so they must not have
+	// a leading slash.
+	//
+	// os-release is the canonical filename and is the only candidate present in full-OS initrds built by Image
+	// Customizer (which pack the source rootfs directly, in which initrd-release does not exist at all).
+	//
+	// initrd-release is the dracut-runtime variant emitted by dracut's 99base module. os-release symlinks to this file
+	// in such cases, so it is a necessary fallback.
+	initrdReleaseCandidates = []string{
+		"etc/os-release",
+		"usr/lib/os-release",
+		"etc/initrd-release",
+		"usr/lib/initrd-release",
+	}
 )
 
 func New(distroId Distro, versionId string) TargetOs {
@@ -90,6 +108,9 @@ func New(distroId Distro, versionId string) TargetOs {
 	}
 }
 
+// GetInstalledTargetOs reads the os-release(5) file from the given rootfs and resolves it to a TargetOs.
+//
+// Returns an error wrapping fs.ErrNotExist when none of the candidates exist on the rootfs.
 func GetInstalledTargetOs(rootfs string) (TargetOs, error) {
 	var err error
 	var fields map[string]string
@@ -113,11 +134,34 @@ func GetInstalledTargetOs(rootfs string) (TargetOs, error) {
 	}
 
 	targetOs, err := GetInstalledTargetOsFromEnvFields(fields)
+	if err != nil {
+		return TargetOs{}, fmt.Errorf("failed to determine target OS from os-release file:\n%w", err)
+	}
+
+	return targetOs, nil
+}
+
+// GetInitrdTargetOs reads an os-release(5)-type file (os-release or dracut's initrd-release) from the given initrd and
+// resolves it to a TargetOs.
+//
+// Returns an error wrapping fs.ErrNotExist when none of the candidates exist in the initrd.
+func GetInitrdTargetOs(initrdPath string) (TargetOs, error) {
+	content, foundPath, err := initrdutils.ReadFirstFileFromInitrd(initrdPath, initrdReleaseCandidates)
 	if err != nil {
 		return TargetOs{}, err
 	}
 
-	return targetOs, err
+	fields, err := envfile.ParseEnv(string(content))
+	if err != nil {
+		return TargetOs{}, fmt.Errorf("failed to read (%s) file from initrd (%s):\n%w", foundPath, initrdPath, err)
+	}
+
+	targetOs, err := GetInstalledTargetOsFromEnvFields(fields)
+	if err != nil {
+		return TargetOs{}, fmt.Errorf("failed to determine target OS from initrd-release file:\n%w", err)
+	}
+
+	return targetOs, nil
 }
 
 func GetInstalledTargetOsFromEnvFields(fields map[string]string) (TargetOs, error) {
@@ -171,7 +215,7 @@ func GetInstalledTargetOsFromEnvFields(fields map[string]string) (TargetOs, erro
 		}, nil
 
 	default:
-		return TargetOs{}, fmt.Errorf("unknown ID (%s) in os-release", distroId)
+		return TargetOs{}, fmt.Errorf("unknown ID (%s)", distroId)
 	}
 }
 

toolkit/tools/pkg/imagecustomizerlib/artifactsinputoutput.go

@@ -96,7 +96,7 @@ func outputArtifacts(ctx context.Context, items []imagecustomizerapi.OutputArtif
 	defer systemBootPartitionMount.Close()
 
 	// Detect system architecture
-	_, bootConfig, err := getBootArchConfig()
+	bootConfig, err := distroHandler.GetBootArchConfig()
 	if err != nil {
 		return err
 	}

toolkit/tools/pkg/imagecustomizerlib/artifactsinputoutput_test.go

@@ -55,7 +55,7 @@ func TestOutputAndInjectArtifacts(t *testing.T) {
 		return
 	}
 
-	espFiles := verifyAndSignOutputtedArtifacts(t, outputArtifactsDir, false)
+	espFiles := verifyAndSignOutputtedArtifacts(t, baseImageInfo, outputArtifactsDir, false)
 
 	// Use new buildDir to ensure the buildDir is created if it doesn't exist.
 	buildDirInject := filepath.Join(buildDir, "inject")
@@ -163,7 +163,7 @@ func TestOutputAndInjectArtifactsCosi(t *testing.T) {
 		return
 	}
 
-	espFiles := verifyAndSignOutputtedArtifacts(t, outputArtifactsDir, true)
+	espFiles := verifyAndSignOutputtedArtifacts(t, baseImageInfo, outputArtifactsDir, true)
 
 	// Inject artifacts into image.
 	options := InjectFilesOptions{
@@ -305,7 +305,19 @@ func TestOutputAndInjectArtifactsCosi(t *testing.T) {
 		"root", buildDir, "", "restart-on-corruption", false /*inlineVerity*/)
 }
 
-func verifyAndSignOutputtedArtifacts(t *testing.T, outputArtifactsDir string, expectVerityHash bool) []string {
+func verifyAndSignOutputtedArtifacts(t *testing.T, baseImageInfo testBaseImageInfo, outputArtifactsDir string, expectVerityHash bool) []string {
+	// Resolve the distro-version's boot file layout so artifact destinations can be checked against the distro's
+	// actual ESP paths. The boot file names differ per distro-version (for example, Azure Linux 4 uses the Fedora-style
+	// uppercase BOOTX64.EFI shim, while Azure Linux 3 uses lowercase bootx64.efi).
+	distroHandler, err := NewDistroHandler(baseImageInfo.TargetOs())
+	if !assert.NoError(t, err) {
+		return nil
+	}
+	bootConfig, err := distroHandler.GetBootArchConfig()
+	if !assert.NoError(t, err) {
+		return nil
+	}
+
 	// Confirm inject-files.yaml was generated
 	injectConfigPath := filepath.Join(outputArtifactsDir, "inject-files.yaml")
 	exists, err := file.PathExists(injectConfigPath)
@@ -333,15 +345,15 @@ func verifyAndSignOutputtedArtifacts(t *testing.T, outputArtifactsDir string, ex
 
 		switch entry.Type {
 		case imagecustomizerapi.OutputArtifactsItemShim:
-			assert.True(t, strings.HasPrefix(entry.Destination, "/EFI/BOOT/boot"), "Expected shim destination to start with /EFI/BOOT/boot")
-			assert.True(t, strings.HasSuffix(entry.Destination, ".efi"), "Expected shim destination to end with .efi")
+			expectedShimDestination := filepath.Join("/", bootConfig.espBootBinaryPath)
+			assert.Equal(t, expectedShimDestination, entry.Destination, "Expected shim destination to match the distro ESP boot binary path")
 			assert.True(t, strings.HasPrefix(entry.Source, "./shim/"), "Expected shim source to be in shim/ subdirectory")
 			hasShim = true
 			espFiles = append(espFiles, entry.Destination)
 
 		case imagecustomizerapi.OutputArtifactsItemBootloader:
-			assert.True(t, strings.HasPrefix(entry.Destination, "/EFI/BOOT/grub"), "Expected bootloader destination to start with /EFI/BOOT/grub")
-			assert.True(t, strings.HasSuffix(entry.Destination, ".efi"), "Expected bootloader destination to end with .efi")
+			expectedBootloaderDestination := filepath.Join("/", bootConfig.espGrubBinaryPath)
+			assert.Equal(t, expectedBootloaderDestination, entry.Destination, "Expected bootloader destination to match the distro ESP grub binary path")
 			assert.True(t, strings.HasPrefix(entry.Source, "./bootloader/"), "Expected bootloader source to be in bootloader/ subdirectory")
 			hasBootloader = true
 			espFiles = append(espFiles, entry.Destination)

toolkit/tools/pkg/imagecustomizerlib/customizeuki.go

@@ -258,7 +258,7 @@ func prepareUkiHelper(ctx context.Context, buildDir string, uki *imagecustomizer
 	if uki.Mode == imagecustomizerapi.UkiModeModify {
 		logger.Log.Infof("UKI mode is 'modify', skipping UKI preparation (will modify addon only)")
 
-		_, bootConfig, err := getBootArchConfig()
+		bootConfig, err := distroHandler.GetBootArchConfig()
 		if err != nil {
 			return err
 		}
@@ -295,7 +295,7 @@ func prepareUkiHelper(ctx context.Context, buildDir string, uki *imagecustomizer
 	}
 
 	// Detect system architecture.
-	_, bootConfig, err := getBootArchConfig()
+	bootConfig, err := distroHandler.GetBootArchConfig()
 	if err != nil {
 		return err
 	}
@@ -503,13 +503,13 @@ func findKernelsAndInitramfs(bootDir string) (map[string]string, error) {
 }
 
 // createUkiInModifyMode modifies UKI addons without touching the main UKI files.
-func createUkiInModifyMode(ctx context.Context, rc *ResolvedConfig) error {
+func createUkiInModifyMode(ctx context.Context, rc *ResolvedConfig, distroHandler DistroHandler) error {
 	_, span := otel.GetTracerProvider().Tracer(OtelTracerName).Start(ctx, "customize_uki_modify_mode")
 	defer span.End()
 
 	var err error
 
-	_, bootConfig, err := getBootArchConfig()
+	bootConfig, err := distroHandler.GetBootArchConfig()
 	if err != nil {
 		return err
 	}
@@ -618,7 +618,7 @@ func modifyUkiAddon(ukiFilePath string, stubPath string, rc *ResolvedConfig) err
 	return nil
 }
 
-func createUki(ctx context.Context, rc *ResolvedConfig) error {
+func createUki(ctx context.Context, rc *ResolvedConfig, distroHandler DistroHandler) error {
 	logger.Log.Infof("Creating UKIs")
 
 	// If mode is 'passthrough', skip UKI creation to preserve existing UKIs
@@ -630,7 +630,7 @@ func createUki(ctx context.Context, rc *ResolvedConfig) error {
 	// If mode is 'modify', only modify UKI addons (preserve main UKI)
 	if rc.Uki != nil && rc.Uki.Mode == imagecustomizerapi.UkiModeModify {
 		logger.Log.Infof("UKI mode is 'modify', modifying UKI addons only")
-		err := createUkiInModifyMode(ctx, rc)
+		err := createUkiInModifyMode(ctx, rc, distroHandler)
 		if err != nil {
 			return fmt.Errorf("failed to modify UKI addons in modify mode:\n%w", err)
 		}
@@ -642,7 +642,7 @@ func createUki(ctx context.Context, rc *ResolvedConfig) error {
 
 	var err error
 
-	_, bootConfig, err := getBootArchConfig()
+	bootConfig, err := distroHandler.GetBootArchConfig()
 	if err != nil {
 		return err
 	}

toolkit/tools/pkg/imagecustomizerlib/distrohandler.go

@@ -133,6 +133,10 @@ type DistroHandler interface {
 
 	// Distro has a root partition that is missing placeholder directories for special mounts like /dev.
 	RootMissingMountDirectories() bool
+
+	// GetBootArchConfig returns the boot-files configuration appropriate for this distro on the current runtime
+	// architecture.
+	GetBootArchConfig() (BootFilesArchConfig, error)
 }
 
 // NewDistroHandler creates a distro handler directly from TargetOs
@@ -185,3 +189,20 @@ func handleUnsupportedDistroVersion(rc *ResolvedConfig, targetOs targetos.Target
 
 	return nil
 }
+
+// NewDistroHandlerFromInitrd creates a distro handler by detecting the OS from the dracut-emitted initrd-release file
+// inside the initramfs at initrdPath. Used in ISO-to-ISO pipelines where no rootfs is mounted yet but a canonical
+// distro identity is needed to pick the correct boot-file layout.
+func NewDistroHandlerFromInitrd(initrdPath string) (DistroHandler, error) {
+	targetOs, err := targetos.GetInitrdTargetOs(initrdPath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to determine the target OS from initrd (%s):\n%w", initrdPath, err)
+	}
+
+	distroHandler, err := NewDistroHandler(targetOs)
+	if err != nil {
+		return nil, err
+	}
+
+	return distroHandler, nil
+}

toolkit/tools/pkg/imagecustomizerlib/distrohandler_acl.go

@@ -287,3 +287,7 @@ func (d *aclDistroHandler) GrubEfiPackage() string {
 func (d *aclDistroHandler) RootMissingMountDirectories() bool {
 	return true
 }
+
+func (d *aclDistroHandler) GetBootArchConfig() (BootFilesArchConfig, error) {
+	return bootArchConfigFromMap(bootloaderFilesConfigAzureLinux)
+}