Address comments 2nd round.

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,7 +166,7 @@ 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, buildDir)
 		if err != nil {
@@ -177,10 +176,10 @@ func fixPartitionUuidsInFstabFile(partitions []diskutils.PartitionInfo, newUuids
 		// 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/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/extendedmountidentifiertype_test.go

@@ -0,0 +1,33 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerlib
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestExtendedMountIdentifierTypeIsValid(t *testing.T) {
+	validCases := []ExtendedMountIdentifierType{
+		ExtendedMountIdentifierTypeUuid,
+		ExtendedMountIdentifierTypePartUuid,
+		ExtendedMountIdentifierTypePartLabel,
+		ExtendedMountIdentifierTypeDev,
+		ExtendedMountIdentifierTypeDefault,
+	}
+
+	for _, validCase := range validCases {
+		t.Run(string(validCase), func(t *testing.T) {
+			err := validCase.IsValid()
+			assert.NoError(t, err)
+		})
+	}
+
+	// Test invalid case
+	invalidCase := ExtendedMountIdentifierType("invalid-type")
+	err := invalidCase.IsValid()
+	assert.Error(t, err)
+	assert.ErrorContains(t, err, "invalid value (invalid-type)")
+}

toolkit/tools/pkg/imagecustomizerlib/imagecustomizer.go

@@ -363,6 +363,15 @@ func customizeOSContents(ic *ImageCustomizerParameters) error {
 		ic.config.OS = &imagecustomizerapi.OS{}
 	}
 
+	// Check if dm-verity is enabled on the base image.
+	verityEnabled, err := isDmVerityEnabled(ic.rawImageFile)
+	if err != nil {
+		return err
+	}
+	if verityEnabled && len(ic.config.Storage.Verity) == 0 {
+		return fmt.Errorf("dm-verity is enabled on the base image. To customize a verity-enabled base image, the verity section must be reconfigured.")
+	}
+
 	// Customize the partitions.
 	partitionsCustomized, newRawImageFile, partIdToPartUuid, err := customizePartitions(ic.buildDirAbs,
 		ic.configPath, ic.config, ic.rawImageFile)
@@ -863,7 +872,7 @@ func customizeVerityImageHelper(buildDir string, baseConfigPath string, config *
 		}
 	} 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)
 		}
@@ -882,34 +891,34 @@ func customizeVerityImageHelper(buildDir string, baseConfigPath string, config *
 	return nil
 }
 
-func checkDmVerityEnabled(rawImageFile string) error {
+func isDmVerityEnabled(rawImageFile string) (bool, error) {
 	logger.Log.Debugf("Check if dm-verity is enabled in base image")
 
 	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 check if dm-verity is enabled in base image:\n%w", err)
 	}
 	defer loopback.Close()
 
 	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")
+			return true, nil
 		}
 	}
 
 	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 false, nil
 }
 
 func warnOnLowFreeSpace(buildDir string, imageConnection *ImageConnection) {

toolkit/tools/pkg/imagecustomizerlib/partitionutils.go

@@ -308,27 +308,10 @@ func findSourcePartition(source string, partitions []diskutils.PartitionInfo,
 
 func findSourcePartitionHelper(source string, partitions []diskutils.PartitionInfo,
 	buildDir string,
-) (imagecustomizerapi.MountIdentifierType, diskutils.PartitionInfo, int, error) {
-	mountIdTypeExt, partition, partitionIndex, err := findExtendedSourcePartitionHelper(source, partitions, buildDir)
+) (ExtendedMountIdentifierType, diskutils.PartitionInfo, int, error) {
+	mountIdType, partition, partitionIndex, err := findExtendedSourcePartitionHelper(source, partitions, buildDir)
 	if err != nil {
-		return imagecustomizerapi.MountIdentifierTypeDefault, diskutils.PartitionInfo{}, 0, err
-	}
-
-	// Convert ExtendedMountIdentifierType to MountIdentifierType
-	var mountIdType imagecustomizerapi.MountIdentifierType
-	switch mountIdTypeExt {
-	case ExtendedMountIdentifierTypeUuid:
-		mountIdType = imagecustomizerapi.MountIdentifierTypeUuid
-	case ExtendedMountIdentifierTypePartUuid:
-		mountIdType = imagecustomizerapi.MountIdentifierTypePartUuid
-	case ExtendedMountIdentifierTypePartLabel:
-		mountIdType = imagecustomizerapi.MountIdentifierTypePartLabel
-	case ExtendedMountIdentifierTypeDev:
-		// Dev type is not supported in MountIdentifierType.
-		// Return the partition, but set MountIdentifierTypeDefault as a placeholder.
-		return imagecustomizerapi.MountIdentifierTypeDefault, partition, partitionIndex, nil
-	default:
-		return imagecustomizerapi.MountIdentifierTypeDefault, diskutils.PartitionInfo{}, 0, fmt.Errorf("unsupported identifier type: %v", mountIdTypeExt)
+		return ExtendedMountIdentifierTypeDefault, diskutils.PartitionInfo{}, 0, err
 	}
 
 	return mountIdType, partition, partitionIndex, nil
@@ -379,6 +362,18 @@ func findPartition(mountIdType imagecustomizerapi.MountIdentifierType, mountId s
 func findExtendedPartition(mountIdType ExtendedMountIdentifierType, mountId string,
 	partitions []diskutils.PartitionInfo, buildDir string,
 ) (diskutils.PartitionInfo, int, error) {
+	var devUuid string
+	if mountIdType == ExtendedMountIdentifierTypeDev {
+		var err error
+		devUuid, err = convertDevToUuid(partitions, buildDir)
+		if err != nil {
+			return diskutils.PartitionInfo{}, 0, err
+		}
+		// Replace the original mountId with the UUID type.
+		mountId = devUuid
+		mountIdType = ExtendedMountIdentifierTypeUuid
+	}
+
 	matchedPartitionIndexes := []int(nil)
 	for i, partition := range partitions {
 		matches := false
@@ -389,14 +384,6 @@ func findExtendedPartition(mountIdType ExtendedMountIdentifierType, mountId stri
 			matches = partition.PartUuid == mountId
 		case ExtendedMountIdentifierTypePartLabel:
 			matches = partition.PartLabel == mountId
-		case ExtendedMountIdentifierTypeDev:
-			matches, err := findExtendedDevPartition(partition, partitions, buildDir)
-			if err != nil {
-				return diskutils.PartitionInfo{}, 0, err
-			}
-			if matches {
-				matchedPartitionIndexes = append(matchedPartitionIndexes, i)
-			}
 		}
 		if matches {
 			matchedPartitionIndexes = append(matchedPartitionIndexes, i)
@@ -418,7 +405,25 @@ func findExtendedPartition(mountIdType ExtendedMountIdentifierType, mountId stri
 	return partition, partitionIndex, nil
 }
 
-func findExtendedDevPartition(partition diskutils.PartitionInfo, partitions []diskutils.PartitionInfo,
+func convertDevToUuid(partitions []diskutils.PartitionInfo, buildDir string) (string, error) {
+	for _, partition := range partitions {
+		matches, err := checkExtendedDevPartition(partition, partitions, buildDir)
+		if err != nil {
+			return "", err
+		}
+		if matches {
+			return partition.Uuid, nil
+		}
+	}
+	return "", fmt.Errorf("unable to find UUID for /dev path")
+}
+
+// checkExtendedDevPartition checks whether a given partition is associated with a verity-enabled root partition.
+// This function is specific to the case where the partition type is "dev" (device path, e.g., /dev/mapper/root).
+// It is an extension of the findExtendedPartition logic and specializes in handling verity partitions,
+// supporting both UKI and non-UKI configurations by validating against the kernel cmdline arguments or grub configuration.
+// Returns true if the partition matches the expected verity configuration, otherwise false.
+func checkExtendedDevPartition(partition diskutils.PartitionInfo, partitions []diskutils.PartitionInfo,
 	buildDir string,
 ) (bool, error) {
 	espPartition, err := findSystemBootPartition(partitions)
@@ -431,17 +436,9 @@ func findExtendedDevPartition(partition diskutils.PartitionInfo, partitions []di
 		return false, err
 	}
 
-	// Temporarily mount the boot partition so that the grub config file can be read.
-	tmpDirBoot := filepath.Join(buildDir, tmpBootPartitionDirName)
-	bootPartitionMount, err := safemount.NewMount(bootPartition.Path, tmpDirBoot, bootPartition.FileSystemType, 0, "", true)
-	if err != nil {
-		return false, fmt.Errorf("failed to mount boot partition (%s):\n%w", bootPartition.Path, err)
-	}
-	defer bootPartitionMount.Close()
-
-	// Mount ESP partition to check for UKI.
+	// Temporarily mount ESP partition to check for UKIs.
 	tmpDirEsp := filepath.Join(buildDir, tmpEspPartitionDirName)
-	espPartitionMount, err := safemount.NewMount(espPartition.Path, tmpDirEsp, espPartition.FileSystemType, 0, "", true)
+	espPartitionMount, err := safemount.NewMount(espPartition.Path, tmpDirEsp, espPartition.FileSystemType, unix.MS_RDONLY, "", true)
 	if err != nil {
 		return false, fmt.Errorf("failed to mount ESP partition (%s):\n%w", espPartition.Path, err)
 	}
@@ -469,6 +466,14 @@ func findExtendedDevPartition(partition diskutils.PartitionInfo, partitions []di
 
 		return checkVerityRootPartUUID(partition, string(cmdlineContent))
 	} else {
+		// Temporarily mount the boot partition so that the grub config file can be read.
+		tmpDirBoot := filepath.Join(buildDir, tmpBootPartitionDirName)
+		bootPartitionMount, err := safemount.NewMount(bootPartition.Path, tmpDirBoot, bootPartition.FileSystemType, unix.MS_RDONLY, "", true)
+		if err != nil {
+			return false, fmt.Errorf("failed to mount boot partition (%s):\n%w", bootPartition.Path, err)
+		}
+		defer bootPartitionMount.Close()
+
 		grubCfgPath := filepath.Join(tmpDirBoot, "/grub2/grub.cfg")
 		kernelToArgs, err := extractKernelToArgsFromGrub(grubCfgPath)
 		if err != nil {
@@ -480,14 +485,14 @@ func findExtendedDevPartition(partition diskutils.PartitionInfo, partitions []di
 				return matches, err
 			}
 		}
-	}
 
-	err = espPartitionMount.CleanClose()
-	if err != nil {
-		return false, fmt.Errorf("failed to close bootPartitionMount: %w", err)
+		err = bootPartitionMount.CleanClose()
+		if err != nil {
+			return false, fmt.Errorf("failed to close bootPartitionMount: %w", err)
+		}
 	}
 
-	err = bootPartitionMount.CleanClose()
+	err = espPartitionMount.CleanClose()
 	if err != nil {
 		return false, fmt.Errorf("failed to close bootPartitionMount: %w", err)
 	}