address comments

Author: elainezhao1

toolkit/tools/pkg/imagecustomizerlib/cosicommon.go

@@ -36,12 +36,14 @@ func convertToCosi(ic *ImageCustomizerParameters) error {
 	}
 	defer imageLoopback.Close()
 
-	partitionMetadataOutput, err := extractPartitions(imageLoopback.DevicePath(), outputDir, ic.outputImageBase, "raw-zst", ic.imageUuid)
+	partitionMetadataOutput, err := extractPartitions(imageLoopback.DevicePath(), outputDir, ic.outputImageBase,
+		"raw-zst", ic.imageUuid)
 	if err != nil {
 		return err
 	}
 
-	err = buildCosiFile(outputDir, ic.outputImageFile, partitionMetadataOutput, ic.verityMetadata, ic.imageUuidStr)
+	err = buildCosiFile(outputDir, ic.outputImageFile, partitionMetadataOutput, ic.verityMetadata, ic.partUuidToMountPath,
+		ic.imageUuidStr)
 	if err != nil {
 		return fmt.Errorf("failed to build COSI file:\n%w", err)
 	}
@@ -56,8 +58,9 @@ func convertToCosi(ic *ImageCustomizerParameters) error {
 	return nil
 }
 
-func buildCosiFile(sourceDir string, outputFile string, partitions []outputPartitionMetadata, verityMetadata []verityDeviceMetadata, imageUuidStr string) error {
-
+func buildCosiFile(sourceDir string, outputFile string, partitions []outputPartitionMetadata,
+	verityMetadata []verityDeviceMetadata, partUuidToMountPath map[string]string, imageUuidStr string,
+) error {
 	// Pre-compute a map for quick lookup of partition metadata by UUID
 	partUuidToMetadata := make(map[string]outputPartitionMetadata)
 	for _, partition := range partitions {
@@ -84,7 +87,7 @@ func buildCosiFile(sourceDir string, outputFile string, partitions []outputParti
 				UncompressedSize: partition.UncompressedSize,
 			},
 			PartType:   partition.PartitionTypeUuid,
-			MountPoint: partition.Mountpoint,
+			MountPoint: partUuidToMountPath[partition.PartUuid],
 			FsType:     partition.FileSystemType,
 			FsUuid:     partition.Uuid,
 		}

toolkit/tools/pkg/imagecustomizerlib/customizepartitionsfilecopy.go

@@ -16,7 +16,7 @@ import (
 func customizePartitionsUsingFileCopy(buildDir string, baseConfigPath string, config *imagecustomizerapi.Config,
 	buildImageFile string, newBuildImageFile string,
 ) (map[string]string, error) {
-	existingImageConnection, err := connectToExistingImage(buildImageFile, buildDir, "imageroot", false)
+	existingImageConnection, _, err := connectToExistingImage(buildImageFile, buildDir, "imageroot", false)
 	if err != nil {
 		return nil, err
 	}

toolkit/tools/pkg/imagecustomizerlib/imagecustomizer.go

@@ -84,6 +84,8 @@ type ImageCustomizerParameters struct {
 	imageUuidStr string
 
 	verityMetadata []verityDeviceMetadata
+
+	partUuidToMountPath map[string]string
 }
 
 type verityDeviceMetadata struct {
@@ -388,12 +390,15 @@ func customizeOSContents(ic *ImageCustomizerParameters) error {
 	ic.rawImageFile = newRawImageFile
 
 	// Customize the raw image file.
-	err = customizeImageHelper(ic.buildDirAbs, ic.configPath, ic.config, ic.rawImageFile, ic.rpmsSources,
+	partUuidToMountPath, err := customizeImageHelper(ic.buildDirAbs, ic.configPath, ic.config, ic.rawImageFile, ic.rpmsSources,
 		ic.useBaseImageRpmRepos, partitionsCustomized, ic.imageUuidStr)
 	if err != nil {
 		return err
 	}
 
+	// Set partition to mountpath mapping for COSI.
+	ic.partUuidToMountPath = partUuidToMountPath
+
 	// Shrink the filesystems.
 	if ic.enableShrinkFilesystems {
 		err = shrinkFilesystemsHelper(ic.rawImageFile, ic.config.Storage.Verity, partIdToPartUuid)
@@ -713,12 +718,12 @@ func validatePackageLists(baseConfigPath string, config *imagecustomizerapi.OS,
 func customizeImageHelper(buildDir string, baseConfigPath string, config *imagecustomizerapi.Config,
 	rawImageFile string, rpmsSources []string, useBaseImageRpmRepos bool, partitionsCustomized bool,
 	imageUuidStr string,
-) error {
+) (map[string]string, error) {
 	logger.Log.Debugf("Customizing OS")
 
-	imageConnection, err := connectToExistingImage(rawImageFile, buildDir, "imageroot", true)
+	imageConnection, partUuidToMountPath, err := connectToExistingImage(rawImageFile, buildDir, "imageroot", true)
 	if err != nil {
-		return err
+		return nil, err
 	}
 	defer imageConnection.Close()
 
@@ -738,15 +743,15 @@ func customizeImageHelper(buildDir string, baseConfigPath string, config *imagec
 	warnOnLowFreeSpace(buildDir, imageConnection)
 
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	err = imageConnection.CleanClose()
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	return nil
+	return partUuidToMountPath, nil
 }
 
 func extractPartitionsHelper(rawImageFile string, outputDir string, outputBasename string, outputSplitPartitionsFormat string, imageUuid [UuidSize]byte) error {

toolkit/tools/pkg/imagecustomizerlib/imageutils.go

@@ -21,41 +21,49 @@ import (
 type installOSFunc func(imageChroot *safechroot.Chroot) error
 
 func connectToExistingImage(imageFilePath string, buildDir string, chrootDirName string, includeDefaultMounts bool,
-) (*ImageConnection, error) {
+) (*ImageConnection, map[string]string, error) {
 	imageConnection := NewImageConnection()
 
-	err := connectToExistingImageHelper(imageConnection, imageFilePath, buildDir, chrootDirName, includeDefaultMounts)
+	partUuidToMountPath, err := connectToExistingImageHelper(imageConnection, imageFilePath, buildDir, chrootDirName, includeDefaultMounts)
 	if err != nil {
 		imageConnection.Close()
-		return nil, err
+		return nil, nil, err
 	}
-	return imageConnection, nil
+	return imageConnection, partUuidToMountPath, nil
 }
 
 func connectToExistingImageHelper(imageConnection *ImageConnection, imageFilePath string,
 	buildDir string, chrootDirName string, includeDefaultMounts bool,
-) error {
+) (map[string]string, error) {
 	// Connect to image file using loopback device.
 	err := imageConnection.ConnectLoopback(imageFilePath)
 	if err != nil {
-		return err
+		return nil, err
+	}
+
+	partitions, err := diskutils.GetDiskPartitions(imageConnection.Loopback().DevicePath())
+	if err != nil {
+		return nil, err
 	}
 
 	// Look for all the partitions on the image.
-	mountPoints, err := findPartitions(buildDir, imageConnection.Loopback().DevicePath())
+	mountPoints, err := findPartitions(buildDir, partitions)
 	if err != nil {
-		return fmt.Errorf("failed to find disk partitions:\n%w", err)
+		return nil, fmt.Errorf("failed to find disk partitions:\n%w", err)
 	}
 
+	// Create mapping from partition UUID to mount path
+	partUuidToMountPath := createPartUuidToMountPathMap(partitions, mountPoints)
+
 	// Create chroot environment.
 	imageChrootDir := filepath.Join(buildDir, chrootDirName)
 
 	err = imageConnection.ConnectChroot(imageChrootDir, false, []string(nil), mountPoints, includeDefaultMounts)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	return nil
+	return partUuidToMountPath, nil
 }
 
 func createNewImage(targetOs targetos.TargetOs, filename string, diskConfig imagecustomizerapi.Disk,
@@ -265,3 +273,17 @@ func createPartIdToPartUuidMap(partIDToDevPathMap map[string]string, diskPartiti
 
 	return partIdToPartUuid, nil
 }
+
+func createPartUuidToMountPathMap(partitions []diskutils.PartitionInfo, mountPoints []*safechroot.MountPoint,
+) map[string]string {
+	partUuidToMountPath := make(map[string]string)
+	for _, mountPoint := range mountPoints {
+		for _, partition := range partitions {
+			if partition.Path == mountPoint.GetSource() {
+				partUuidToMountPath[partition.PartUuid] = mountPoint.GetTarget()
+				break
+			}
+		}
+	}
+	return partUuidToMountPath
+}

toolkit/tools/pkg/imagecustomizerlib/liveosisobuilder.go

@@ -1071,7 +1071,7 @@ func (b *LiveOSIsoBuilder) prepareArtifactsFromFullImage(inputSavedConfigsFilePa
 	logger.Log.Infof("Preparing iso artifacts")
 
 	logger.Log.Debugf("Connecting to raw image (%s)", rawImageFile)
-	rawImageConnection, err := connectToExistingImage(rawImageFile, b.workingDirs.isoBuildDir, "readonly-rootfs-mount", false /*includeDefaultMounts*/)
+	rawImageConnection, _, err := connectToExistingImage(rawImageFile, b.workingDirs.isoBuildDir, "readonly-rootfs-mount", false /*includeDefaultMounts*/)
 	if err != nil {
 		return err
 	}

toolkit/tools/pkg/imagecustomizerlib/partitionutils.go

@@ -30,14 +30,9 @@ var (
 	partitionNumberRegex = regexp.MustCompile(`^/dev/loop\d+p(\d+)$`)
 )
 
-func findPartitions(buildDir string, diskDevice string) ([]*safechroot.MountPoint, error) {
+func findPartitions(buildDir string, diskPartitions []diskutils.PartitionInfo) ([]*safechroot.MountPoint, error) {
 	var err error
 
-	diskPartitions, err := diskutils.GetDiskPartitions(diskDevice)
-	if err != nil {
-		return nil, err
-	}
-
 	rootfsPartition, err := findRootfsPartition(diskPartitions, buildDir)
 	if err != nil {
 		return nil, fmt.Errorf("failed to find rootfs partition:\n%w", err)