Mariner Image Customizer boilerplate (#5982)

.github/workflows/go-test-coverage.yml

@@ -25,6 +25,11 @@ jobs:
     - name: Check out code into the Go module directory
       uses: actions/checkout@v4
 
+    - name: Install prerequisites
+      run: |
+        sudo apt-get update
+        sudo apt -y install qemu-utils
+
     - name: Check for bad go formatting
       run: |
         pushd toolkit

toolkit/tools/go.mod

@@ -18,6 +18,7 @@ require (
 	golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f
 	gonum.org/v1/gonum v0.11.0
 	gopkg.in/alecthomas/kingpin.v2 v2.2.6
+	gopkg.in/yaml.v3 v3.0.0
 )
 
 require (
@@ -35,5 +36,4 @@ require (
 	golang.org/x/exp v0.0.0-20191002040644-a1355ae1e2c3 // indirect
 	golang.org/x/text v0.3.8 // indirect
 	gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f // indirect
-	gopkg.in/yaml.v3 v3.0.0 // indirect
 )

toolkit/tools/imagecustomizer/main.go

@@ -0,0 +1,64 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package main
+
+import (
+	"log"
+	"os"
+
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/exe"
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/logger"
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/timestamp"
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/pkg/imagecustomizerlib"
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/pkg/profile"
+	"gopkg.in/alecthomas/kingpin.v2"
+)
+
+var (
+	app = kingpin.New("imagecustomizer", "Customizes a pre-built CBL-Mariner image")
+
+	buildDir          = app.Flag("build-dir", "Directory to run build out of.").Required().String()
+	imageFile         = app.Flag("image-file", "Path of the base CBL-Mariner image which the customization will be applied to.").Required().String()
+	outputImageFile   = app.Flag("output-image-file", "Path to write the customized image to.").Required().String()
+	outputImageFormat = app.Flag("output-image-format", "Format of output image. Supported: vhd, vhdx, qcow2, raw.").Required().Enum("vhd", "vhdx", "qcow2", "raw")
+	configFile        = app.Flag("config-file", "Path of the image customization config file.").Required().String()
+	logFile           = exe.LogFileFlag(app)
+	logLevel          = exe.LogLevelFlag(app)
+	profFlags         = exe.SetupProfileFlags(app)
+	timestampFile     = app.Flag("timestamp-file", "File that stores timestamps for this program.").String()
+)
+
+func main() {
+	var err error
+
+	kingpin.MustParse(app.Parse(os.Args[1:]))
+
+	logger.InitBestEffort(*logFile, *logLevel)
+
+	prof, err := profile.StartProfiling(profFlags)
+	if err != nil {
+		logger.Log.Warnf("Could not start profiling: %s", err)
+	}
+	defer prof.StopProfiler()
+
+	timestamp.BeginTiming("imagecustomizer", *timestampFile)
+	defer timestamp.CompleteTiming()
+
+	err = customizeImage()
+	if err != nil {
+		log.Fatalf("image customization failed: %v", err)
+	}
+}
+
+func customizeImage() error {
+	var err error
+
+	err = imagecustomizerlib.CustomizeImageWithConfigFile(*buildDir, *configFile, *imageFile,
+		*outputImageFile, *outputImageFormat)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

toolkit/tools/imagecustomizerapi/README.md

@@ -0,0 +1,16 @@
+# Mariner Image Customizer API
+
+The Mariner image customizer (imgcustomizer) will be released as a standalone tool and
+will provide strong backwards compatibility guarantees (after the first official
+release).
+This is contrast to the Mariner toolkit's new image config, which isn't officially
+released and doesn't provide any backwards compatibility guarantees.
+
+While currently the new image config and imgcustomizer config are very similar, in the
+future there is the possibility they will diverge.
+
+## Known differences
+
+- For the new image config, `AdditionalFiles`' source files are relative to the working
+  directory.
+  Whereas, for imgcustomizer, the source files are relative to the config file.

toolkit/tools/imagecustomizerapi/fileconfig.go

@@ -0,0 +1,109 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+//
+
+package imagecustomizerapi
+
+import (
+	"fmt"
+
+	"gopkg.in/yaml.v3"
+)
+
+// DestinationFileConfigList is a list of destination files where the source file will be copied to in the final image.
+// This type exists to allow a custom marshaller to be attached to it.
+type FileConfigList []FileConfig
+
+// FileConfig specifies options for how a file is copied in the target OS.
+type FileConfig struct {
+	// The file path in the target OS that the file will be copied to.
+	Path string `yaml:"Path"`
+
+	// The file permissions to set on the file.
+	Permissions *FilePermissions `yaml:"Permissions"`
+}
+
+var (
+	DefaultFileConfig = FileConfig{
+		Path:        "",
+		Permissions: nil,
+	}
+)
+
+func (l *FileConfigList) IsValid() (err error) {
+	if len(*l) <= 0 {
+		return fmt.Errorf("list is empty")
+	}
+
+	for i, fileConfig := range *l {
+		err = fileConfig.IsValid()
+		if err != nil {
+			return fmt.Errorf("invalid FileConfig at index %d: %w", i, err)
+		}
+	}
+
+	return nil
+}
+
+func (l *FileConfigList) UnmarshalYAML(value *yaml.Node) error {
+	var err error
+
+	// Try to parse as a single value.
+	var fileConfig FileConfig
+	err = value.Decode(&fileConfig)
+	if err == nil {
+		*l = FileConfigList{fileConfig}
+		return nil
+	}
+
+	// Try to parse as a list.
+	type IntermediateTypeFileConfigList FileConfigList
+	err = value.Decode((*IntermediateTypeFileConfigList)(l))
+	if err != nil {
+		return fmt.Errorf("failed to parse FileConfigList: %w", err)
+	}
+
+	return nil
+}
+
+func (f *FileConfig) IsValid() (err error) {
+	// Path
+	if f.Path == "" {
+		return fmt.Errorf("invalid Path value: empty string")
+	}
+
+	// Permissions
+	if f.Permissions != nil {
+		err = f.Permissions.IsValid()
+		if err != nil {
+			return fmt.Errorf("invalid Permissions value: %w", err)
+		}
+	}
+
+	return nil
+}
+
+func (f *FileConfig) UnmarshalYAML(value *yaml.Node) error {
+	var err error
+
+	if value.Kind == yaml.ScalarNode {
+		// Parse as a string.
+		*f = FileConfig{
+			Path:        value.Value,
+			Permissions: nil,
+		}
+		return nil
+	}
+
+	// Parse as a struct.
+	*f = DefaultFileConfig
+
+	type IntermediateTypeFileConfig FileConfig
+	err = value.Decode((*IntermediateTypeFileConfig)(f))
+	if err != nil {
+		return fmt.Errorf("failed to parse FileConfig: %w", err)
+	}
+
+	return nil
+}

toolkit/tools/imagecustomizerapi/fileconfig_test.go

@@ -0,0 +1,57 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerapi
+
+import (
+	"testing"
+
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/ptrutils"
+)
+
+func TestParseFileConfigValidString(t *testing.T) {
+	testValidYamlValue(t, "\"/a.txt\"", &FileConfigList{{Path: "/a.txt"}})
+}
+
+func TestParseFileConfigValidStringInArray(t *testing.T) {
+	testValidYamlValue(t, "[ \"/a.txt\" ]", &FileConfigList{{Path: "/a.txt"}})
+}
+
+func TestParseFileConfigValidBasicStruct(t *testing.T) {
+	testValidYamlValue(t, "{ \"Path\": \"/b.txt\" }", &FileConfigList{{Path: "/b.txt"}})
+}
+
+func TestParseFileConfigValidFullStruct(t *testing.T) {
+	testValidYamlValue(t, "{ \"Path\": \"/b.txt\", \"Permissions\": \"770\" }",
+		&FileConfigList{{Path: "/b.txt", Permissions: ptrutils.PtrTo(FilePermissions(0o770))}},
+	)
+}
+
+func TestParseFileConfigValidMixedArray(t *testing.T) {
+	testValidYamlValue(t, "[ { \"Path\": \"/b.txt\" }, \"/c.txt\" ]",
+		&FileConfigList{
+			{Path: "/b.txt"},
+			{Path: "/c.txt"},
+		},
+	)
+}
+
+func TestParseFileConfigInvalidEmptyArray(t *testing.T) {
+	// Empty array.
+	testInvalidYamlValue[*FileConfigList](t, "[ ]")
+}
+
+func TestParseFileConfigInvalidArrayArray(t *testing.T) {
+	// Empty array.
+	testInvalidYamlValue[*FileConfigList](t, "[ [ ] ]")
+}
+
+func TestParseFileConfigInvalidEmptyString(t *testing.T) {
+	// Empty string.
+	testInvalidYamlValue[*FileConfigList](t, "\"\"")
+}
+
+func TestParseFileConfigInvalidFilePermissions(t *testing.T) {
+	// Empty string.
+	testInvalidYamlValue[*FileConfigList](t, "{ \"Path\": \"/b.txt\", \"Permissions\": \"7777\" }")
+}

toolkit/tools/imagecustomizerapi/filepermissions.go

@@ -0,0 +1,48 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerapi
+
+import (
+	"fmt"
+	"os"
+	"strconv"
+
+	"gopkg.in/yaml.v3"
+)
+
+// The file permissions to set on the file.
+//
+// Accepted formats:
+//
+// - Octal string (e.g. "660")
+type FilePermissions os.FileMode
+
+func (p *FilePermissions) IsValid() error {
+	// Check if there are set bits outside of the permissions bits.
+	if *p & ^FilePermissions(os.ModePerm) != 0 {
+		return fmt.Errorf("0o%o contains non-permission bits", *p)
+	}
+
+	return nil
+}
+
+func (p *FilePermissions) UnmarshalYAML(value *yaml.Node) error {
+	var err error
+
+	// Try to parse as a string.
+	var strValue string
+	err = value.Decode(&strValue)
+	if err != nil {
+		return fmt.Errorf("failed to parse FilePermissions: %w", err)
+	}
+
+	// Try to parse the string as an octal number.
+	fileModeUint, err := strconv.ParseUint(strValue, 8, 32)
+	if err != nil {
+		return fmt.Errorf("failed to parse FilePermissions: %w", err)
+	}
+
+	*p = (FilePermissions)(fileModeUint)
+	return nil
+}

toolkit/tools/imagecustomizerapi/filepermissions_test.go

@@ -0,0 +1,37 @@
+// Copyright Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerapi
+
+import (
+	"testing"
+
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/ptrutils"
+)
+
+func TestParseFilePermissionsValid1(t *testing.T) {
+	testValidYamlValue(t, "\"777\"", ptrutils.PtrTo(FilePermissions(0o777)))
+}
+
+func TestParseFilePermissionsValid2(t *testing.T) {
+	testValidYamlValue(t, "\"000\"", ptrutils.PtrTo(FilePermissions(0)))
+}
+
+func TestParseFilePermissionsValid3(t *testing.T) {
+	testValidYamlValue(t, "\"0\"", ptrutils.PtrTo(FilePermissions(0)))
+}
+
+func TestParseFilePermissionsInvalidOutOfRange(t *testing.T) {
+	// Value out of range.
+	testInvalidYamlValue[*FilePermissions](t, "\"1000\"")
+}
+
+func TestParseFilePermissionsInvalidType(t *testing.T) {
+	// Array value.
+	testInvalidYamlValue[*FilePermissions](t, "[]")
+}
+
+func TestParseFilePermissionsInvalidNotOctal(t *testing.T) {
+	// Not an octal value.
+	testInvalidYamlValue[*FilePermissions](t, "\"999\"")
+}

toolkit/tools/imagecustomizerapi/systemconfig.go

@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerapi
+
+import (
+	"fmt"
+)
+
+// SystemConfig defines how each system present on the image is supposed to be configured.
+type SystemConfig struct {
+	AdditionalFiles map[string]FileConfigList `yaml:"AdditionalFiles"`
+}
+
+func (s *SystemConfig) IsValid() error {
+	var err error
+
+	for sourcePath, fileConfigList := range s.AdditionalFiles {
+		err = fileConfigList.IsValid()
+		if err != nil {
+			return fmt.Errorf("invalid file configs for (%s): %w", sourcePath, err)
+		}
+	}
+
+	return nil
+}

toolkit/tools/imagecustomizerapi/systemconfig_test.go

@@ -0,0 +1,16 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerapi
+
+import (
+	"testing"
+)
+
+func TestSystemConfigValidEmpty(t *testing.T) {
+	testValidYamlValue[*SystemConfig](t, "{ }", &SystemConfig{})
+}
+
+func TestSystemConfigInvalidAdditionalFiles(t *testing.T) {
+	testInvalidYamlValue[*SystemConfig](t, "{ \"AdditionalFiles\": { \"a.txt\": [] } }")
+}

toolkit/tools/imagecustomizerapi/utils.go

@@ -0,0 +1,46 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerapi
+
+import (
+	"os"
+
+	"gopkg.in/yaml.v3"
+)
+
+type HasIsValid interface {
+	IsValid() error
+}
+
+func UnmarshalYamlFile[ValueType HasIsValid](yamlFilePath string, value ValueType) error {
+	var err error
+
+	yamlFile, err := os.ReadFile(yamlFilePath)
+	if err != nil {
+		return err
+	}
+
+	err = UnmarshalYaml(yamlFile, value)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func UnmarshalYaml[ValueType HasIsValid](yamlData []byte, value ValueType) error {
+	var err error
+
+	err = yaml.Unmarshal(yamlData, value)
+	if err != nil {
+		return err
+	}
+
+	err = value.IsValid()
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

toolkit/tools/imagecustomizerapi/utils_test.go

@@ -0,0 +1,32 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerapi
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func testValidYamlValue[DataType HasIsValid](t *testing.T, yamlString string, expectedValue DataType) {
+	value := makeValue[DataType]()
+
+	err := UnmarshalYaml([]byte(yamlString), value)
+	assert.NoError(t, err)
+	assert.Equal(t, expectedValue, value)
+}
+
+func testInvalidYamlValue[DataType HasIsValid](t *testing.T, yamlString string) {
+	value := makeValue[DataType]()
+	err := UnmarshalYaml([]byte(yamlString), value)
+	assert.Errorf(t, err, "value: %v", value)
+}
+
+func makeValue[DataType any]() DataType {
+	// When DataType is a pointer, there is no built-in way to create a new value
+	// of the underlying type. So, use reflection to do this.
+	var placeholder DataType
+	return reflect.New(reflect.TypeOf(placeholder).Elem()).Interface().(DataType)
+}

toolkit/tools/imagegen/diskutils/diskutils.go

@@ -321,6 +321,17 @@ func DetachLoopbackDevice(diskDevPath string) (err error) {
 	return
 }
 
+// WaitForDevicesToSettle waits for all udev events to be processed on the system.
+// This can be used to wait for partitions to be discovered after mounting a disk.
+func WaitForDevicesToSettle() error {
+	logger.Log.Debugf("Waiting for devices to settle")
+	_, _, err := shell.Execute("udevadm", "settle")
+	if err != nil {
+		return fmt.Errorf("failed to wait for devices to settle: %w", err)
+	}
+	return nil
+}
+
 // CreatePartitions creates partitions on the specified disk according to the disk config
 func CreatePartitions(diskDevPath string, disk configuration.Disk, rootEncryption configuration.RootEncryption, readOnlyRootConfig configuration.ReadOnlyVerityRoot) (partDevPathMap map[string]string, partIDToFsTypeMap map[string]string, encryptedRoot EncryptedRootDevice, readOnlyRoot VerityDevice, err error) {
 	const timeoutInSeconds = "5"

toolkit/tools/internal/safechroot/safechroot.go

@@ -43,7 +43,8 @@ type MountPoint struct {
 	flags  uintptr
 	data   string
 
-	isMounted bool
+	isMounted           bool
+	mountBeforeDefaults bool
 }
 
 // Chroot represents a Chroot environment with automatic synchronization protections
@@ -98,6 +99,18 @@ func NewMountPoint(source, target, fstype string, flags uintptr, data string) (m
 	}
 }
 
+// NewPreDefaultsMountPoint creates a new MountPoint struct to be created by a Chroot but before the default mount points.
+func NewPreDefaultsMountPoint(source, target, fstype string, flags uintptr, data string) (mountPoint *MountPoint) {
+	return &MountPoint{
+		source:              source,
+		target:              target,
+		fstype:              fstype,
+		flags:               flags,
+		data:                data,
+		mountBeforeDefaults: true,
+	}
+}
+
 // NewOverlayMountPoint creates a new MountPoint struct and extra directories slice configured for a given overlay
 func NewOverlayMountPoint(chrootDir, source, target, lowerDir, upperDir, workDir string) (mountPoint *MountPoint, extaDirsNeeds []string) {
 	const (
@@ -230,7 +243,21 @@ func (c *Chroot) Initialize(tarPath string, extraDirectories []string, extraMoun
 	// mount is only supported in regular pipeline
 	if buildpipeline.IsRegularBuild() {
 		// Create kernel mountpoints
-		allMountPoints := append(defaultMountPoints(), extraMountPoints...)
+		allMountPoints := []*MountPoint{}
+
+		for _, mountPoint := range extraMountPoints {
+			if mountPoint.mountBeforeDefaults {
+				allMountPoints = append(allMountPoints, mountPoint)
+			}
+		}
+
+		allMountPoints = append(allMountPoints, defaultMountPoints()...)
+
+		for _, mountPoint := range extraMountPoints {
+			if !mountPoint.mountBeforeDefaults {
+				allMountPoints = append(allMountPoints, mountPoint)
+			}
+		}
 
 		// Mount with the original unsorted order. Assumes the order of mounts is important.
 		err = c.createMountPoints(allMountPoints)
@@ -269,6 +296,7 @@ func (c *Chroot) AddFiles(filesToCopy ...FileToCopy) (err error) {
 		} else {
 			err = file.Copy(f.Src, dest)
 		}
+
 		if err != nil {
 			logger.Log.Errorf("Error provisioning worker with '%s'", f.Src)
 			return
@@ -561,7 +589,7 @@ func (c *Chroot) createMountPoints(allMountPoints []*MountPoint) (err error) {
 
 		err = unix.Mount(mountPoint.source, fullPath, mountPoint.fstype, mountPoint.flags, mountPoint.data)
 		if err != nil {
-			logger.Log.Errorf("Mount failed on (%s). Error: %s", fullPath, err)
+			logger.Log.Errorf("Mount of (%s) to (%s) failed. Error: %s", mountPoint.source, fullPath, err)
 			return
 		}
 

toolkit/tools/pkg/imagecustomizerlib/customizeutils.go

@@ -0,0 +1,44 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerlib
+
+import (
+	"io/fs"
+	"path/filepath"
+
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/imagecustomizerapi"
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/safechroot"
+)
+
+func doCustomizations(baseConfigPath string, config *imagecustomizerapi.SystemConfig, imageChroot *safechroot.Chroot) error {
+	var err error
+
+	err = copyAdditionalFiles(baseConfigPath, config.AdditionalFiles, imageChroot)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func copyAdditionalFiles(baseConfigPath string, additionalFiles map[string]imagecustomizerapi.FileConfigList, imageChroot *safechroot.Chroot) error {
+	var err error
+
+	for sourceFile, fileConfigs := range additionalFiles {
+		for _, fileConfig := range fileConfigs {
+			fileToCopy := safechroot.FileToCopy{
+				Src:         filepath.Join(baseConfigPath, sourceFile),
+				Dest:        fileConfig.Path,
+				Permissions: (*fs.FileMode)(fileConfig.Permissions),
+			}
+
+			err = imageChroot.AddFiles(fileToCopy)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}

toolkit/tools/pkg/imagecustomizerlib/customizeutils_test.go

@@ -0,0 +1,70 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerlib
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/imagecustomizerapi"
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/ptrutils"
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/safechroot"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCopyAdditionalFiles(t *testing.T) {
+	proposedDir := filepath.Join(tmpDir, "chroot", "TestCopyAdditionalFiles")
+	chroot := safechroot.NewChroot(proposedDir, false)
+	baseConfigPath := testDir
+
+	err := chroot.Initialize("", []string{}, []*safechroot.MountPoint{})
+	assert.NoError(t, err)
+	defer chroot.Close(false)
+
+	copy_2_filemode := os.FileMode(0o777)
+
+	err = copyAdditionalFiles(baseConfigPath, map[string]imagecustomizerapi.FileConfigList{
+		"files/a.txt": {
+			{Path: "/a_copy_1.txt"},
+			{Path: "/a_copy_2.txt", Permissions: ptrutils.PtrTo(imagecustomizerapi.FilePermissions(copy_2_filemode))},
+		},
+	}, chroot)
+	assert.NoError(t, err)
+
+	orig_path := filepath.Join(baseConfigPath, "files/a.txt")
+	copy_1_path := filepath.Join(chroot.RootDir(), "a_copy_1.txt")
+	copy_2_path := filepath.Join(chroot.RootDir(), "a_copy_2.txt")
+
+	// Make sure the files exist.
+	orig_stat, err := os.Stat(orig_path)
+	assert.NoError(t, err)
+
+	copy_1_stat, err := os.Stat(copy_1_path)
+	assert.NoError(t, err)
+
+	copy_2_stat, err := os.Stat(copy_2_path)
+	assert.NoError(t, err)
+
+	// Make sure the filemode of the original file is different from the target filemode,
+	// as otherwise it would defeat the purpose of the test.
+	assert.NotEqual(t, copy_2_filemode, orig_stat.Mode()&os.ModePerm)
+
+	// Make sure the file permissions are the expected values.
+	assert.Equal(t, orig_stat.Mode()&os.ModePerm, copy_1_stat.Mode()&os.ModePerm)
+	assert.Equal(t, copy_2_filemode, copy_2_stat.Mode()&os.ModePerm)
+
+	// Make sure the files' contents are correct.
+	orig_contents, err := os.ReadFile(orig_path)
+	assert.NoError(t, err)
+
+	copy_1_contents, err := os.ReadFile(copy_1_path)
+	assert.NoError(t, err)
+
+	copy_2_contents, err := os.ReadFile(copy_2_path)
+	assert.NoError(t, err)
+
+	assert.Equal(t, orig_contents, copy_1_contents)
+	assert.Equal(t, orig_contents, copy_2_contents)
+}

toolkit/tools/pkg/imagecustomizerlib/imagecustomizer.go

@@ -0,0 +1,171 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerlib
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/imagecustomizerapi"
+	"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/safechroot"
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/shell"
+)
+
+func CustomizeImageWithConfigFile(buildDir string, configFile string, imageFile string,
+	outputImageFile string, outputImageFormat string,
+) error {
+	var err error
+
+	var config imagecustomizerapi.SystemConfig
+	err = imagecustomizerapi.UnmarshalYamlFile(configFile, &config)
+	if err != nil {
+		return err
+	}
+
+	baseConfigPath, _ := filepath.Split(configFile)
+
+	err = CustomizeImage(buildDir, baseConfigPath, &config, imageFile, outputImageFile, outputImageFormat)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func CustomizeImage(buildDir string, baseConfigPath string, config *imagecustomizerapi.SystemConfig, imageFile string,
+	outputImageFile string, outputImageFormat string,
+) error {
+	var err error
+
+	// Validate 'outputImageFormat' value.
+	qemuOutputImageFormat, err := toQemuImageFormat(outputImageFormat)
+	if err != nil {
+		return err
+	}
+
+	// Normalize 'buildDir' path.
+	buildDirAbs, err := filepath.Abs(buildDir)
+	if err != nil {
+		return err
+	}
+
+	// Create 'buildDir' directory.
+	err = os.MkdirAll(buildDirAbs, os.ModePerm)
+	if err != nil {
+		return err
+	}
+
+	// Validate config.
+	err = validateConfig(baseConfigPath, config)
+	if err != nil {
+		return fmt.Errorf("invalid image config: %w", err)
+	}
+
+	// Convert image file to raw format, so that a kernel loop device can be used to make changes to the image.
+	buildImageFile := filepath.Join(buildDirAbs, "image.raw")
+
+	_, _, err = shell.Execute("qemu-img", "convert", "-O", "raw", imageFile, buildImageFile)
+	if err != nil {
+		return fmt.Errorf("failed to convert image file to raw format: %w", err)
+	}
+
+	// Customize the raw image file.
+	err = customizeImageHelper(buildDirAbs, baseConfigPath, config, buildImageFile)
+	if err != nil {
+		return err
+	}
+
+	// Create final output image file.
+	_, _, err = shell.Execute("qemu-img", "convert", "-O", qemuOutputImageFormat, buildImageFile, outputImageFile)
+	if err != nil {
+		return fmt.Errorf("failed to convert image file to format: %s: %w", outputImageFormat, err)
+	}
+
+	return nil
+}
+
+func toQemuImageFormat(imageFormat string) (string, error) {
+	switch imageFormat {
+	case "vhd":
+		return "vpc", nil
+
+	case "vhdx", "raw", "qcow2":
+		return imageFormat, nil
+
+	default:
+		return "", fmt.Errorf("unsupported image format (supported: vhd, vhdx, raw, qcow2): %s", imageFormat)
+	}
+}
+
+func validateConfig(baseConfigPath string, config *imagecustomizerapi.SystemConfig) error {
+	for sourceFile := range config.AdditionalFiles {
+		sourceFileFullPath := filepath.Join(baseConfigPath, sourceFile)
+		isFile, err := file.IsFile(sourceFileFullPath)
+		if err != nil {
+			return fmt.Errorf("invalid AdditionalFiles source file (%s): %w", sourceFile, err)
+		}
+
+		if !isFile {
+			return fmt.Errorf("invalid AdditionalFiles source file (%s): not a file", sourceFile)
+		}
+	}
+
+	return nil
+}
+
+func customizeImageHelper(buildDir string, baseConfigPath string, config *imagecustomizerapi.SystemConfig,
+	buildImageFile string,
+) error {
+	// Mount the raw disk image file.
+	diskDevPath, err := diskutils.SetupLoopbackDevice(buildImageFile)
+	if err != nil {
+		return fmt.Errorf("failed to mount raw disk (%s) as a loopback device: %w", buildImageFile, err)
+	}
+	defer diskutils.DetachLoopbackDevice(diskDevPath)
+
+	// Wait for the partitions to show up.
+	err = diskutils.WaitForDevicesToSettle()
+	if err != nil {
+		return err
+	}
+
+	// Look for all the partitions on the image.
+	newMountDirectories, mountPoints, err := findPartitions(diskDevPath)
+	if err != nil {
+		return err
+	}
+
+	// Create chroot environment.
+	imageChrootDir := filepath.Join(buildDir, "imageroot")
+
+	imageChroot := safechroot.NewChroot(imageChrootDir, false)
+	err = imageChroot.Initialize("", newMountDirectories, mountPoints)
+	if err != nil {
+		return err
+	}
+	defer imageChroot.Close(false)
+
+	// Do the actual customizations.
+	err = doCustomizations(baseConfigPath, config, imageChroot)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func findPartitions(diskDevice string) ([]string, []*safechroot.MountPoint, error) {
+	newMountDirectories := []string{}
+
+	// TODO: Dynamically find partitions instead of hardcoding the mappings.
+	mountPoints := []*safechroot.MountPoint{
+		safechroot.NewPreDefaultsMountPoint(fmt.Sprintf("%sp2", diskDevice), "/", "ext4", 0, ""),
+		safechroot.NewMountPoint(fmt.Sprintf("%sp1", diskDevice), "/boot", "vfat", 0, ""),
+	}
+
+	return newMountDirectories, mountPoints, nil
+}

toolkit/tools/pkg/imagecustomizerlib/imagecustomizer_test.go

@@ -0,0 +1,212 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerlib
+
+import (
+	"bytes"
+	"fmt"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/imagecustomizerapi"
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/imagegen/configuration"
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/imagegen/diskutils"
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/safechroot"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCustomizeImageEmptyConfig(t *testing.T) {
+	var err error
+
+	buildDir := filepath.Join(tmpDir, "TestCustomizeImageEmptyConfig")
+	outImageFilePath := filepath.Join(buildDir, "image.vhd")
+
+	// Create empty disk.
+	diskFilePath, err := createEmptyDisk(buildDir)
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	// Customize image.
+	err = CustomizeImage(buildDir, buildDir, &imagecustomizerapi.SystemConfig{}, diskFilePath, outImageFilePath, "vhd")
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	// Check output file type.
+	checkFileType(t, outImageFilePath, "vhd")
+}
+
+func TestCustomizeImageCopyFiles(t *testing.T) {
+	var err error
+
+	buildDir := filepath.Join(tmpDir, "TestCustomizeImageCopyFiles")
+	configFile := filepath.Join(testDir, "addfiles-config.yaml")
+	outImageFilePath := filepath.Join(buildDir, "image.qcow2")
+
+	// Create empty disk.
+	diskFilePath, err := createEmptyDisk(buildDir)
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	// Customize image.
+	err = CustomizeImageWithConfigFile(buildDir, configFile, diskFilePath, outImageFilePath, "raw")
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	// Check output file type.
+	checkFileType(t, outImageFilePath, "raw")
+
+	// Mount the output disk image so that its contents can be checked.
+	diskDevPath, err := diskutils.SetupLoopbackDevice(outImageFilePath)
+	if !assert.NoError(t, err) {
+		return
+	}
+	defer diskutils.DetachLoopbackDevice(diskDevPath)
+
+	newMountDirectories, mountPoints := emptyDiskPartitions(diskDevPath)
+
+	imageChroot := safechroot.NewChroot(filepath.Join(buildDir, "imageroot"), false)
+	err = imageChroot.Initialize("", newMountDirectories, mountPoints)
+	if !assert.NoError(t, err) {
+		return
+	}
+	defer imageChroot.Close(false)
+
+	// Check the contents of the copied file.
+	file_contents, err := os.ReadFile(filepath.Join(imageChroot.RootDir(), "a.txt"))
+	assert.NoError(t, err)
+	assert.Equal(t, "abcdefg\n", string(file_contents))
+}
+
+func TestValidateConfigValidAdditionalFiles(t *testing.T) {
+	var err error
+
+	err = validateConfig(testDir, &imagecustomizerapi.SystemConfig{
+		AdditionalFiles: map[string]imagecustomizerapi.FileConfigList{
+			"files/a.txt": {{Path: "/a.txt"}},
+		},
+	})
+	assert.NoError(t, err)
+}
+
+func TestValidateConfigMissingAdditionalFiles(t *testing.T) {
+	var err error
+
+	err = validateConfig(testDir, &imagecustomizerapi.SystemConfig{
+		AdditionalFiles: map[string]imagecustomizerapi.FileConfigList{
+			"files/missing_a.txt": {{Path: "/a.txt"}},
+		},
+	})
+	assert.Error(t, err)
+}
+
+func TestValidateConfigdditionalFilesIsDir(t *testing.T) {
+	var err error
+
+	err = validateConfig(testDir, &imagecustomizerapi.SystemConfig{
+		AdditionalFiles: map[string]imagecustomizerapi.FileConfigList{
+			"files": {{Path: "/a.txt"}},
+		},
+	})
+	assert.Error(t, err)
+}
+
+func createEmptyDisk(buildDir string) (string, error) {
+	var err error
+
+	err = os.MkdirAll(buildDir, os.ModePerm)
+	if err != nil {
+		return "", fmt.Errorf("failed to make build directory (%s): %w", buildDir, err)
+	}
+
+	// Use a prototypical Mariner image partition config.
+	diskConfig := configuration.Disk{
+		PartitionTableType: configuration.PartitionTableTypeGpt,
+		MaxSize:            4096,
+		Partitions: []configuration.Partition{
+			{
+				ID:     "boot",
+				Flags:  []configuration.PartitionFlag{"esp", "boot"},
+				Start:  1,
+				End:    9,
+				FsType: "fat32",
+			},
+			{
+				ID:     "rootfs",
+				Start:  9,
+				End:    0,
+				FsType: "ext4",
+			},
+		},
+	}
+
+	// Create raw disk image file.
+	rawDisk, err := diskutils.CreateEmptyDisk(buildDir, "disk.raw", diskConfig)
+	if err != nil {
+		return "", fmt.Errorf("failed to create empty disk file in (%s): %w", buildDir, err)
+	}
+
+	// Mount raw disk image file.
+	diskDevPath, err := diskutils.SetupLoopbackDevice(rawDisk)
+	if err != nil {
+		return "", fmt.Errorf("failed to mount raw disk (%s) as a loopback device: %w", rawDisk, err)
+	}
+	defer diskutils.DetachLoopbackDevice(diskDevPath)
+
+	// Set up partitions.
+	_, _, _, _, err = diskutils.CreatePartitions(diskDevPath, diskConfig,
+		configuration.RootEncryption{}, configuration.ReadOnlyVerityRoot{})
+	if err != nil {
+		return "", fmt.Errorf("failed to create partitions on disk (%s): %w", diskDevPath, err)
+	}
+
+	return rawDisk, nil
+}
+
+func emptyDiskPartitions(diskDevPath string) ([]string, []*safechroot.MountPoint) {
+	newMountDirectories := []string{}
+	mountPoints := []*safechroot.MountPoint{
+		safechroot.NewPreDefaultsMountPoint(fmt.Sprintf("%sp2", diskDevPath), "/", "ext4", 0, ""),
+		safechroot.NewMountPoint(fmt.Sprintf("%sp1", diskDevPath), "/boot", "vfat", 0, ""),
+	}
+	return newMountDirectories, mountPoints
+}
+
+func checkFileType(t *testing.T, filePath string, expectedFileType string) {
+	fileType, err := getImageFileType(filePath)
+	assert.NoError(t, err)
+	assert.Equal(t, expectedFileType, fileType)
+}
+
+func getImageFileType(filePath string) (string, error) {
+	file, err := os.OpenFile(filePath, os.O_RDONLY, 0)
+	if err != nil {
+		return "", err
+	}
+	defer file.Close()
+
+	firstBytes := make([]byte, 512)
+	readByteCount, err := file.Read(firstBytes)
+	if err != nil {
+		return "", err
+	}
+
+	switch {
+	case readByteCount >= 8 && bytes.Equal(firstBytes[:8], []byte("conectix")):
+		return "vhd", nil
+
+	case readByteCount >= 8 && bytes.Equal(firstBytes[:8], []byte("vhdxfile")):
+		return "vhdx", nil
+
+	// Check for the MBR signature (which exists even on GPT formatted drives).
+	case readByteCount >= 512 && bytes.Equal(firstBytes[510:512], []byte{0x55, 0xAA}):
+		return "raw", nil
+	}
+
+	return "", fmt.Errorf("Unknown file type")
+}

toolkit/tools/pkg/imagecustomizerlib/main_test.go

@@ -0,0 +1,46 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerlib
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/microsoft/CBL-Mariner/toolkit/tools/internal/logger"
+)
+
+var (
+	testDir    string
+	tmpDir     string
+	workingDir string
+)
+
+func TestMain(m *testing.M) {
+	var err error
+
+	logger.InitStderrLog()
+
+	workingDir, err = os.Getwd()
+	if err != nil {
+		logger.Log.Panicf("Failed to get working directory, error: %s", err)
+	}
+
+	testDir = filepath.Join(workingDir, "testdata")
+	tmpDir = filepath.Join(workingDir, "_tmp")
+
+	err = os.MkdirAll(tmpDir, os.ModePerm)
+	if err != nil {
+		logger.Log.Panicf("Failed to create tmp directory, error: %s", err)
+	}
+
+	retVal := m.Run()
+
+	err = os.RemoveAll(tmpDir)
+	if err != nil {
+		logger.Log.Warnf("Failed to cleanup tmp dir (%s). Error: %s", tmpDir, err)
+	}
+
+	os.Exit(retVal)
+}

toolkit/tools/pkg/imagecustomizerlib/testdata/addfiles-config.yaml

@@ -0,0 +1,2 @@
+AdditionalFiles:
+  files/a.txt: /a.txt

toolkit/tools/pkg/imagecustomizerlib/testdata/files/a.txt

@@ -0,0 +1 @@
+abcdefg