Addressed comment

Author: liulanze

toolkit/tools/pkg/imagecustomizerlib/imagecustomizer.go

@@ -379,7 +379,7 @@ func customizeOSContents(ic *ImageCustomizerParameters) error {
 		return err
 	}
 	if verityEnabled && !ic.customizeOSPartitions {
-		return fmt.Errorf("dm-verity is enabled on the base image. To customize a verity-enabled base image, the verity section must be reconfigured.")
+		return fmt.Errorf("dm-verity is enabled on the base image so partitions must be specified.")
 	}
 
 	// Customize the partitions.

toolkit/tools/pkg/imagecustomizerlib/partitionutils.go

@@ -347,15 +347,16 @@ func findPartition(mountIdType imagecustomizerapi.MountIdentifierType, mountId s
 func findExtendedPartition(mountIdType ExtendedMountIdentifierType, mountId string,
 	partitions []diskutils.PartitionInfo, buildDir string,
 ) (diskutils.PartitionInfo, int, error) {
-	var devPartUuid string
 	if mountIdType == ExtendedMountIdentifierTypeDev {
-		var err error
-		devPartUuid, err = convertDevToPartUuid(partitions, buildDir)
+		cmdline, err := extractKernelCmdline(partitions, buildDir)
+		if err != nil {
+			return diskutils.PartitionInfo{}, 0, err
+		}
+
+		mountIdType, mountId, err = extractVerityRootPartitionId(cmdline)
 		if err != nil {
 			return diskutils.PartitionInfo{}, 0, err
 		}
-		mountId = devPartUuid
-		mountIdType = ExtendedMountIdentifierTypePartUuid
 	}
 
 	matchedPartitionIndexes := []int(nil)
@@ -389,44 +390,6 @@ func findExtendedPartition(mountIdType ExtendedMountIdentifierType, mountId stri
 	return partition, partitionIndex, nil
 }
 
-func convertDevToPartUuid(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.PartUuid, nil
-		}
-	}
-	return "", fmt.Errorf("unable to find PARTUUID for /dev path")
-}
-
-func checkExtendedDevPartition(partition diskutils.PartitionInfo, partitions []diskutils.PartitionInfo,
-	buildDir string,
-) (bool, error) {
-	cmdline, err := extractKernelCmdline(partitions, buildDir)
-	if err != nil {
-		return false, err
-	}
-
-	mountIdType, verityPartitionId, err := extractVerityRootPartitionId(cmdline)
-	if err != nil {
-		return false, err
-	}
-
-	switch mountIdType {
-	case ExtendedMountIdentifierTypePartUuid:
-		return partition.PartUuid == verityPartitionId, nil
-	case ExtendedMountIdentifierTypePartLabel:
-		return partition.PartLabel == verityPartitionId, nil
-	case ExtendedMountIdentifierTypeUuid:
-		return partition.Uuid == verityPartitionId, nil
-	default:
-		return false, fmt.Errorf("unsupported mount identifier type: (%s)", mountIdType)
-	}
-}
-
 func extractKernelCmdline(partitions []diskutils.PartitionInfo, buildDir string) ([]grubConfigLinuxArg, error) {
 	espPartition, err := findSystemBootPartition(partitions)
 	if err != nil {
@@ -438,46 +401,31 @@ func extractKernelCmdline(partitions []diskutils.PartitionInfo, buildDir string)
 		return nil, fmt.Errorf("failed to find boot partition: %w", err)
 	}
 
-	tmpDirEsp := filepath.Join(buildDir, tmpEspPartitionDirName)
-	espPartitionMount, err := safemount.NewMount(espPartition.Path, tmpDirEsp, espPartition.FileSystemType, 0, "", true)
-	if err != nil {
-		return nil, fmt.Errorf("failed to mount ESP partition (%s):\n%w", espPartition.Path, err)
-	}
-	defer espPartitionMount.Close()
-
-	cmdline, err := extractKernelCmdlineFromUki(tmpDirEsp, buildDir)
+	cmdline, err := extractKernelCmdlineFromUki(espPartition, buildDir)
 	if err == nil {
 		return cmdline, nil
 	} else if !errors.Is(err, os.ErrNotExist) {
 		return nil, err
 	}
 
-	tmpDirBoot := filepath.Join(buildDir, tmpBootPartitionDirName)
-	bootPartitionMount, err := safemount.NewMount(bootPartition.Path, tmpDirBoot, bootPartition.FileSystemType, unix.MS_RDONLY, "", true)
-	if err != nil {
-		return nil, fmt.Errorf("failed to mount boot partition (%s):\n%w", bootPartition.Path, err)
-	}
-	defer bootPartitionMount.Close()
-
-	cmdline, err = extractKernelCmdlineFromGrub(tmpDirBoot)
+	cmdline, err = extractKernelCmdlineFromGrub(bootPartition, buildDir)
 	if err != nil {
 		return nil, fmt.Errorf("failed to extract kernel arguments from grub.cfg:\n%w", err)
 	}
 
-	err = bootPartitionMount.CleanClose()
-	if err != nil {
-		return nil, fmt.Errorf("failed to close bootPartitionMount:\n%w", err)
-	}
+	return cmdline, nil
+}
 
-	err = espPartitionMount.CleanClose()
+func extractKernelCmdlineFromUki(espPartition *diskutils.PartitionInfo,
+	buildDir string,
+) ([]grubConfigLinuxArg, error) {
+	tmpDirEsp := filepath.Join(buildDir, tmpEspPartitionDirName)
+	espPartitionMount, err := safemount.NewMount(espPartition.Path, tmpDirEsp, espPartition.FileSystemType, 0, "", true)
 	if err != nil {
-		return nil, fmt.Errorf("failed to close espPartitionMount:\n%w", err)
+		return nil, fmt.Errorf("failed to mount ESP partition (%s):\n%w", espPartition.Path, err)
 	}
+	defer espPartitionMount.Close()
 
-	return cmdline, nil
-}
-
-func extractKernelCmdlineFromUki(tmpDirEsp, buildDir string) ([]grubConfigLinuxArg, error) {
 	espLinuxPath := filepath.Join(tmpDirEsp, UkiOutputDir)
 	ukiFiles, err := filepath.Glob(filepath.Join(espLinuxPath, "vmlinuz-*.efi"))
 	if err != nil {
@@ -509,10 +457,24 @@ func extractKernelCmdlineFromUki(tmpDirEsp, buildDir string) ([]grubConfigLinuxA
 		return nil, fmt.Errorf("failed to parse kernel command-line from UKI: %w", err)
 	}
 
+	err = espPartitionMount.CleanClose()
+	if err != nil {
+		return nil, fmt.Errorf("failed to close espPartitionMount:\n%w", err)
+	}
+
 	return args, nil
 }
 
-func extractKernelCmdlineFromGrub(tmpDirBoot string) ([]grubConfigLinuxArg, error) {
+func extractKernelCmdlineFromGrub(bootPartition *diskutils.PartitionInfo,
+	buildDir string,
+) ([]grubConfigLinuxArg, error) {
+	tmpDirBoot := filepath.Join(buildDir, tmpBootPartitionDirName)
+	bootPartitionMount, err := safemount.NewMount(bootPartition.Path, tmpDirBoot, bootPartition.FileSystemType, unix.MS_RDONLY, "", true)
+	if err != nil {
+		return nil, fmt.Errorf("failed to mount boot partition (%s):\n%w", bootPartition.Path, err)
+	}
+	defer bootPartitionMount.Close()
+
 	grubCfgPath := filepath.Join(tmpDirBoot, DefaultGrubCfgPath)
 	grubCfgContent, err := os.ReadFile(grubCfgPath)
 	if err != nil {
@@ -524,6 +486,11 @@ func extractKernelCmdlineFromGrub(tmpDirBoot string) ([]grubConfigLinuxArg, erro
 		return nil, fmt.Errorf("failed to extract kernel command-line arguments: %w", err)
 	}
 
+	err = bootPartitionMount.CleanClose()
+	if err != nil {
+		return nil, fmt.Errorf("failed to close bootPartitionMount:\n%w", err)
+	}
+
 	return args, nil
 }