Image Customizer - Support modifying Verity protected base image. (#83)

<!-- Description: Please provide a summary of the changes and the
motivation behind them. -->

---

### **Checklist**
- [x] Tests added/updated
- [ ] Documentation updated (if needed)
- [x] Code conforms to style guidelines

Author: liulanze

toolkit/tools/pkg/imagecustomizerlib/customizepartitionsuuids.go

@@ -12,7 +12,6 @@ import (
 	"strings"
 
 	"github.com/google/uuid"
-	"github.com/microsoft/azurelinux/toolkit/tools/imagecustomizerapi"
 	"github.com/microsoft/azurelinux/toolkit/tools/imagegen/diskutils"
 	"github.com/microsoft/azurelinux/toolkit/tools/internal/logger"
 	"github.com/microsoft/azurelinux/toolkit/tools/internal/safeloopback"
@@ -167,20 +166,20 @@ func fixPartitionUuidsInFstabFile(partitions []diskutils.PartitionInfo, newUuids
 		}
 
 		// Find the partition.
-		// Note: The 'partitions' list was collected before all the changes were made. So, the fstab entires will still
+		// Note: The 'partitions' list was collected before all the changes were made. So, the fstab entries will still
 		// match the values in the `partitions` list.
-		mountIdType, _, partitionIndex, err := findSourcePartitionHelper(fstabEntry.Source, partitions)
+		mountIdType, _, partitionIndex, err := findSourcePartitionHelper(fstabEntry.Source, partitions, buildDir)
 		if err != nil {
 			return err
 		}
 
 		// Create a new value for the source.
 		newSource := fstabEntry.Source
 		switch mountIdType {
-		case imagecustomizerapi.MountIdentifierTypeUuid:
+		case ExtendedMountIdentifierTypeUuid:
 			newSource = fmt.Sprintf("UUID=%s", newUuids[partitionIndex])
 
-		case imagecustomizerapi.MountIdentifierTypePartUuid:
+		case ExtendedMountIdentifierTypePartUuid:
 			newSource = fmt.Sprintf("PARTUUID=%s", newPartUuids[partitionIndex])
 		}
 

toolkit/tools/pkg/imagecustomizerlib/customizeuki.go

@@ -290,7 +290,7 @@ func createUki(uki *imagecustomizerapi.Uki, buildDir string, buildImageFile stri
 		return err
 	}
 
-	systemBootPartitionTmpDir := filepath.Join(buildDir, tmpEspParitionDirName)
+	systemBootPartitionTmpDir := filepath.Join(buildDir, tmpEspPartitionDirName)
 	systemBootPartitionMount, err := safemount.NewMount(systemBootPartition.Path, systemBootPartitionTmpDir, systemBootPartition.FileSystemType, 0, "", true)
 	if err != nil {
 		return fmt.Errorf("failed to mount esp partition (%s):\n%w", bootPartition.Path, err)

toolkit/tools/pkg/imagecustomizerlib/customizeverity.go

@@ -106,11 +106,11 @@ func prepareGrubConfigForVerity(imageChroot *safechroot.Chroot) error {
 }
 
 func updateGrubConfigForVerity(rootfsVerity imagecustomizerapi.Verity, rootHash string, grubCfgFullPath string,
-	partIdToPartUuid map[string]string, partitions []diskutils.PartitionInfo,
+	partIdToPartUuid map[string]string, partitions []diskutils.PartitionInfo, buildDir string,
 ) error {
 	var err error
 
-	newArgs, err := constructVerityKernelCmdlineArgs(rootfsVerity, rootHash, partIdToPartUuid, partitions)
+	newArgs, err := constructVerityKernelCmdlineArgs(rootfsVerity, rootHash, partIdToPartUuid, partitions, buildDir)
 	if err != nil {
 		return fmt.Errorf("failed to generate verity kernel arguments:\n%w", err)
 	}
@@ -155,16 +155,17 @@ func updateGrubConfigForVerity(rootfsVerity imagecustomizerapi.Verity, rootHash
 }
 
 func constructVerityKernelCmdlineArgs(rootfsVerity imagecustomizerapi.Verity, rootHash string,
-	partIdToPartUuid map[string]string, partitions []diskutils.PartitionInfo) ([]string, error) {
+	partIdToPartUuid map[string]string, partitions []diskutils.PartitionInfo, buildDir string,
+) ([]string, error) {
 	// Format the dataPartitionId and hashPartitionId using the helper function.
 	formattedDataPartition, err := systemdFormatPartitionId(rootfsVerity.DataDeviceId,
-		rootfsVerity.DataDeviceMountIdType, partIdToPartUuid, partitions)
+		rootfsVerity.DataDeviceMountIdType, partIdToPartUuid, partitions, buildDir)
 	if err != nil {
 		return nil, err
 	}
 
 	formattedHashPartition, err := systemdFormatPartitionId(rootfsVerity.HashDeviceId,
-		rootfsVerity.HashDeviceMountIdType, partIdToPartUuid, partitions)
+		rootfsVerity.HashDeviceMountIdType, partIdToPartUuid, partitions, buildDir)
 	if err != nil {
 		return nil, err
 	}
@@ -218,11 +219,11 @@ func partitionMatchesDeviceId(configDeviceId string, partition diskutils.Partiti
 
 // systemdFormatPartitionId formats the partition ID based on the ID type following systemd dm-verity style.
 func systemdFormatPartitionId(configDeviceId string, mountIdType imagecustomizerapi.MountIdentifierType,
-	partIdToPartUuid map[string]string, partitions []diskutils.PartitionInfo,
+	partIdToPartUuid map[string]string, partitions []diskutils.PartitionInfo, buildDir string,
 ) (string, error) {
 	partUuid := partIdToPartUuid[configDeviceId]
 
-	partition, _, err := findPartition(imagecustomizerapi.MountIdentifierTypePartUuid, partUuid, partitions)
+	partition, _, err := findPartition(imagecustomizerapi.MountIdentifierTypePartUuid, partUuid, partitions, buildDir)
 	if err != nil {
 		return "", err
 	}
@@ -274,7 +275,7 @@ func validateVerityDependencies(imageChroot *safechroot.Chroot) error {
 func updateUkiKernelArgsForVerity(rootfsVerity imagecustomizerapi.Verity, rootHash string,
 	partIdToPartUuid map[string]string, partitions []diskutils.PartitionInfo, buildDir string,
 ) error {
-	newArgs, err := constructVerityKernelCmdlineArgs(rootfsVerity, rootHash, partIdToPartUuid, partitions)
+	newArgs, err := constructVerityKernelCmdlineArgs(rootfsVerity, rootHash, partIdToPartUuid, partitions, buildDir)
 	if err != nil {
 		return fmt.Errorf("failed to generate verity kernel arguments:\n%w", err)
 	}

toolkit/tools/pkg/imagecustomizerlib/dracututils.go

@@ -7,25 +7,58 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/microsoft/azurelinux/toolkit/tools/internal/file"
 	"github.com/microsoft/azurelinux/toolkit/tools/internal/safechroot"
 )
 
+const (
+	dracutConfigDir = "etc/dracut.conf.d"
+)
+
 func addDracutConfig(dracutConfigFile string, lines []string) error {
 	if _, err := os.Stat(dracutConfigFile); os.IsNotExist(err) {
+		// File does not exist, create and write the lines.
 		err := file.WriteLines(lines, dracutConfigFile)
 		if err != nil {
-			return fmt.Errorf("failed to write to dracut config file (%s): %w", dracutConfigFile, err)
+			return fmt.Errorf("failed to write to dracut config file (%s):\n%w", dracutConfigFile, err)
 		}
 	} else {
-		return fmt.Errorf("dracut config file (%s) already exists", dracutConfigFile)
+		// File exists, append the lines.
+		existingLines, err := file.ReadLines(dracutConfigFile)
+		if err != nil {
+			return fmt.Errorf("failed to read existing dracut config file (%s):\n%w", dracutConfigFile, err)
+		}
+
+		// Avoid duplicate lines by checking if they already exist.
+		existingLineSet := make(map[string]struct{})
+		for _, line := range existingLines {
+			existingLineSet[line] = struct{}{}
+		}
+
+		linesToAppend := []string{}
+		for _, line := range lines {
+			if _, exists := existingLineSet[line]; !exists {
+				linesToAppend = append(linesToAppend, line)
+			}
+		}
+
+		// Append only non-duplicate lines.
+		if len(linesToAppend) > 0 {
+			content := strings.Join(linesToAppend, "\n") + "\n"
+			err = file.Append(content, dracutConfigFile)
+			if err != nil {
+				return fmt.Errorf("failed to append to dracut config file (%s):\n%w", dracutConfigFile, err)
+			}
+		}
 	}
+
 	return nil
 }
 
 func addDracutModuleAndDriver(dracutModuleName string, dracutDriverName string, imageChroot *safechroot.Chroot) error {
-	dracutConfigFile := filepath.Join(imageChroot.RootDir(), "etc", "dracut.conf.d", dracutModuleName+".conf")
+	dracutConfigFile := filepath.Join(imageChroot.RootDir(), dracutConfigDir, dracutModuleName+".conf")
 	lines := []string{
 		"add_dracutmodules+=\" " + dracutModuleName + " \"",
 		"add_drivers+=\" " + dracutDriverName + " \"",
@@ -34,15 +67,15 @@ func addDracutModuleAndDriver(dracutModuleName string, dracutDriverName string,
 }
 
 func addDracutModule(dracutModuleName string, imageChroot *safechroot.Chroot) error {
-	dracutConfigFile := filepath.Join(imageChroot.RootDir(), "etc", "dracut.conf.d", dracutModuleName+".conf")
+	dracutConfigFile := filepath.Join(imageChroot.RootDir(), dracutConfigDir, dracutModuleName+".conf")
 	lines := []string{
 		"add_dracutmodules+=\" " + dracutModuleName + " \"",
 	}
 	return addDracutConfig(dracutConfigFile, lines)
 }
 
 func addDracutDriver(dracutDriverName string, imageChroot *safechroot.Chroot) error {
-	dracutConfigFile := filepath.Join(imageChroot.RootDir(), "etc", "dracut.conf.d", dracutDriverName+".conf")
+	dracutConfigFile := filepath.Join(imageChroot.RootDir(), dracutConfigDir, dracutDriverName+".conf")
 	lines := []string{
 		"add_drivers+=\" " + dracutDriverName + " \"",
 	}

toolkit/tools/pkg/imagecustomizerlib/extendedmountidentifiertype.go

@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerlib
+
+// ExtendedMountIdentifierType indicates how a partition should be identified in the fstab file.
+// This type was introduced to extend the functionality of MountIdentifierType while preserving
+// the original public API structure. MountIdentifierType is part of a public API and cannot be
+// modified to include new identifiers without breaking backward compatibility.
+// ExtendedMountIdentifierType provides additional flexibility for internal use without affecting the public API.
+type ExtendedMountIdentifierType string
+
+const (
+	// ExtendedMountIdentifierTypeUuid mounts this partition via the filesystem UUID.
+	ExtendedMountIdentifierTypeUuid ExtendedMountIdentifierType = "uuid"
+
+	// ExtendedMountIdentifierTypePartUuid mounts this partition via the GPT/MBR PARTUUID.
+	ExtendedMountIdentifierTypePartUuid ExtendedMountIdentifierType = "part-uuid"
+
+	// ExtendedMountIdentifierTypePartLabel mounts this partition via the GPT PARTLABEL.
+	ExtendedMountIdentifierTypePartLabel ExtendedMountIdentifierType = "part-label"
+
+	// ExtendedMountIdentifierTypeDev mounts this partition via a device.
+	ExtendedMountIdentifierTypeDev ExtendedMountIdentifierType = "dev"
+
+	// ExtendedMountIdentifierTypeDefault uses the default type, which is PARTUUID.
+	ExtendedMountIdentifierTypeDefault ExtendedMountIdentifierType = ""
+)

toolkit/tools/pkg/imagecustomizerlib/imagecustomizer.go

@@ -22,8 +22,9 @@ import (
 )
 
 const (
-	tmpParitionDirName    = "tmp-partition"
-	tmpEspParitionDirName = "tmp-esp-partition"
+	tmpParitionDirName      = "tmp-partition"
+	tmpEspPartitionDirName  = "tmp-esp-partition"
+	tmpBootPartitionDirName = "tmp-boot-partition"
 
 	// supported input formats
 	ImageFormatVhd      = "vhd"
@@ -372,14 +373,14 @@ func customizeOSContents(ic *ImageCustomizerParameters) error {
 		ic.config.OS = &imagecustomizerapi.OS{}
 	}
 
-	// Check if the partition is using DM_verity_hash file system type.
-	// The presence of this type indicates that dm-verity has been enabled on the base image. If dm-verity is not enabled,
-	// the verity hash device should not be assigned this type. We do not support customization on verity enabled base
-	// images at this time because such modifications would compromise the integrity and security mechanisms enforced by dm-verity.
-	err := checkDmVerityEnabled(ic.rawImageFile)
+	// Check if dm-verity is enabled on the base image.
+	verityEnabled, err := isDmVerityEnabled(ic.rawImageFile)
 	if err != nil {
 		return err
 	}
+	if verityEnabled && !ic.customizeOSPartitions {
+		return fmt.Errorf("dm-verity is enabled on the base image so partitions must be specified.")
+	}
 
 	// Customize the partitions.
 	partitionsCustomized, newRawImageFile, partIdToPartUuid, err := customizePartitions(ic.buildDirAbs,
@@ -901,7 +902,7 @@ func customizeVerityImageHelper(buildDir string, config *imagecustomizerapi.Conf
 		}
 	} else {
 		// UKI is not enabled, update grub.cfg as usual.
-		err = updateGrubConfigForVerity(rootfsVerity, rootHash, grubCfgFullPath, partIdToPartUuid, diskPartitions)
+		err = updateGrubConfigForVerity(rootfsVerity, rootHash, grubCfgFullPath, partIdToPartUuid, diskPartitions, buildDir)
 		if err != nil {
 			return "", "", "", fmt.Errorf("failed to update grub config for verity:\n%w", err)
 		}
@@ -920,34 +921,35 @@ func customizeVerityImageHelper(buildDir string, config *imagecustomizerapi.Conf
 	return rootHash, dataPartUuid, hashPartUuid, nil
 }
 
-func checkDmVerityEnabled(rawImageFile string) error {
-	logger.Log.Debugf("Check if dm-verity is enabled in base image")
-
+func isDmVerityEnabled(rawImageFile string) (bool, error) {
 	loopback, err := safeloopback.NewLoopback(rawImageFile)
 	if err != nil {
-		return fmt.Errorf("failed to check if dm-verity is enabled in base image:\n%w", err)
+		return false, fmt.Errorf("failed to create a loopback device for checking if dm-verity is enabled on the base image:\n%w", err)
 	}
 	defer loopback.Close()
 
+	verityEnabled := false
+
 	diskPartitions, err := diskutils.GetDiskPartitions(loopback.DevicePath())
 	if err != nil {
-		return err
+		return false, err
 	}
 
 	for i := range diskPartitions {
 		diskPartition := diskPartitions[i]
 
 		if diskPartition.FileSystemType == "DM_verity_hash" {
-			return fmt.Errorf("cannot customize base image that has dm-verity enabled")
+			verityEnabled = true
+			break
 		}
 	}
 
 	err = loopback.CleanClose()
 	if err != nil {
-		return fmt.Errorf("failed to check if dm-verity is enabled in base image:\n%w", err)
+		return false, fmt.Errorf("failed to cleanly close loopback device:\n%w", err)
 	}
 
-	return nil
+	return verityEnabled, nil
 }
 
 func warnOnLowFreeSpace(buildDir string, imageConnection *ImageConnection) {

toolkit/tools/pkg/imagecustomizerlib/imageutils.go

@@ -241,7 +241,7 @@ func createImageBoilerplate(targetOs targetos.TargetOs, imageConnection *ImageCo
 	}
 
 	// Read back the fstab file.
-	mountPoints, err := findMountsFromFstabFile(tmpFstabFile, diskPartitions)
+	mountPoints, err := findMountsFromFstabFile(tmpFstabFile, diskPartitions, buildDir)
 	if err != nil {
 		return nil, "", err
 	}