Image Customizer: Use safeloopback.Loopback instead of ImageConnection for split partitions (#7300)

toolkit/tools/pkg/imagecustomizerlib/extractpartitions.go

@@ -13,14 +13,11 @@ import (
 )
 
 // Extract all partitions of connected image into separate files with specified format.
-func extractPartitions(imageConnection *ImageConnection, outputImageFile string, partitionFormat string) error {
+func extractPartitions(imageLoopDevice string, outputImageFile string, partitionFormat string) error {
 
 	// Extract basename from outputImageFile. E.g. if outputImageFile is "image.qcow2", then basename is "image".
 	basename := strings.TrimSuffix(filepath.Base(outputImageFile), filepath.Ext(outputImageFile))
 
-	// Get path of loop device associated with the image.
-	imageLoopDevice := imageConnection.Loopback().DevicePath()
-
 	// Get output directory path.
 	outDir := filepath.Dir(outputImageFile)
 

toolkit/tools/pkg/imagecustomizerlib/imagecustomizer.go

@@ -13,6 +13,7 @@ import (
 	"github.com/microsoft/CBL-Mariner/toolkit/tools/imagegen/diskutils"
 	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/file"
 	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/logger"
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/safeloopback"
 	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/safemount"
 	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/shell"
 )
@@ -131,7 +132,7 @@ func CustomizeImage(buildDir string, baseConfigPath string, config *imagecustomi
 	// If outputSplitPartitionsFormat is specified, extract the partition files.
 	if outputSplitPartitionsFormat != "" {
 		logger.Log.Infof("Extracting partition files")
-		err = extractPartitionsHelper(buildDirAbs, buildImageFile, outputImageFile, outputSplitPartitionsFormat)
+		err = extractPartitionsHelper(buildImageFile, outputImageFile, outputSplitPartitionsFormat)
 		if err != nil {
 			return err
 		}
@@ -307,20 +308,20 @@ func customizeImageHelper(buildDir string, baseConfigPath string, config *imagec
 	return nil
 }
 
-func extractPartitionsHelper(buildDir string, buildImageFile string, outputImageFile string, outputSplitPartitionsFormat string) error {
-	imageConnection, err := connectToExistingImage(buildImageFile, buildDir, "imageroot")
+func extractPartitionsHelper(buildImageFile string, outputImageFile string, outputSplitPartitionsFormat string) error {
+	imageLoopback, err := safeloopback.NewLoopback(buildImageFile)
 	if err != nil {
 		return err
 	}
-	defer imageConnection.Close()
+	defer imageLoopback.Close()
 
 	// Extract the partitions as files.
-	err = extractPartitions(imageConnection, outputImageFile, outputSplitPartitionsFormat)
+	err = extractPartitions(imageLoopback.DevicePath(), outputImageFile, outputSplitPartitionsFormat)
 	if err != nil {
 		return err
 	}
 
-	err = imageConnection.CleanClose()
+	err = imageLoopback.CleanClose()
 	if err != nil {
 		return err
 	}