Add initial COSI support

docs/imagecustomizer/cli.md

@@ -31,7 +31,7 @@ The file path to write the final customized image to.
 
 The image format of the the final customized image.
 
-Options: vhd, vhd-fixed, vhdx, qcow2, raw, and iso.
+Options: vhd, vhd-fixed, vhdx, qcow2, raw, iso, and cosi.
 
 At least one of `--output-image-format` and `--output-split-partitions-format` is
 required.

toolkit/tools/imagecustomizer/main.go

@@ -21,7 +21,7 @@ var (
 	buildDir                    = app.Flag("build-dir", "Directory to run build out of.").Required().String()
 	imageFile                   = app.Flag("image-file", "Path of the base Azure Linux 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, iso.").Enum("vhd", "vhd-fixed", "vhdx", "qcow2", "raw", "iso")
+	outputImageFormat           = app.Flag("output-image-format", "Format of output image. Supported: vhd, vhdx, qcow2, raw, iso, cosi.").Enum("vhd", "vhd-fixed", "vhdx", "qcow2", "raw", "iso", "cosi")
 	outputSplitPartitionsFormat = app.Flag("output-split-partitions-format", "Format of partition files. Supported: raw, raw-zst").Enum("raw", "raw-zst")
 	configFile                  = app.Flag("config-file", "Path of the image customization config file.").Required().String()
 	rpmSources                  = app.Flag("rpm-source", "Path to a RPM repo config file or a directory containing RPMs.").Strings()

toolkit/tools/pkg/imagecustomizerlib/cosicommon.go

@@ -0,0 +1,199 @@
+package imagecustomizerlib
+
+import (
+	"archive/tar"
+	"crypto/sha512"
+	_ "embed"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"os"
+	"path"
+
+	"github.com/google/uuid"
+	log "github.com/sirupsen/logrus"
+)
+
+type ImageBuildData struct {
+	Source    string
+	KnownInfo ExpectedImage
+	Metadata  *Image
+}
+
+type ExpectedImage struct {
+	Name                string
+	PartType            PartitionType
+	MountPoint          string
+	FsType              string
+	FsUuid              string
+	OsReleasePath       *string
+	GrubCfgPath         *string
+	ContainsRpmDatabase bool
+}
+
+func (ex ExpectedImage) ShouldMount() bool {
+	return ex.OsReleasePath != nil || ex.GrubCfgPath != nil || ex.ContainsRpmDatabase
+}
+
+type ExtractedImageData struct {
+	OsRelease string
+	GrubCfg   string
+}
+
+func buildCosiFile(sourceDir string, outputFile string, expectedImages []ExpectedImage) error {
+	metadata := MetadataJson{
+		Version: "1.0",
+		OsArch:  "x86_64",
+		Id:      uuid.New().String(),
+		Images:  make([]Image, len(expectedImages)),
+	}
+
+	if len(expectedImages) == 0 {
+		return errors.New("no images to build")
+	}
+
+	// Create an interim metadata struct to combine the known data with the metadata
+	imageData := make([]ImageBuildData, len(expectedImages))
+	for i, image := range expectedImages {
+		metadata := &metadata.Images[i]
+		imageData[i] = ImageBuildData{
+			Source:    path.Join(sourceDir, image.Name),
+			Metadata:  metadata,
+			KnownInfo: image,
+		}
+
+		metadata.Image.Path = path.Join("images", image.Name)
+		metadata.PartType = image.PartType
+		metadata.MountPoint = image.MountPoint
+		metadata.FsType = image.FsType
+		metadata.FsUuid = image.FsUuid
+	}
+
+	// Populate metadata for each image
+	for _, data := range imageData {
+		log.WithField("image", data.Source).Info("Processing image...")
+		extracted, err := data.populateMetadata()
+		if err != nil {
+			return fmt.Errorf("failed to populate metadata for %s: %w", data.Source, err)
+		}
+
+		log.WithField("image", data.Source).Info("Populated metadata for image.")
+
+		metadata.OsRelease = extracted.OsRelease
+	}
+
+	// Marshal metadata.json
+	metadataJson, err := json.MarshalIndent(metadata, "", "  ")
+	if err != nil {
+		return fmt.Errorf("failed to marshal metadata: %w", err)
+	}
+
+	// Create COSI file
+	cosiFile, err := os.Create(outputFile)
+	if err != nil {
+		return fmt.Errorf("failed to create COSI file: %w", err)
+	}
+	defer cosiFile.Close()
+
+	tw := tar.NewWriter(cosiFile)
+	defer tw.Close()
+
+	tw.WriteHeader(&tar.Header{
+		Typeflag: tar.TypeReg,
+		Name:     "metadata.json",
+		Size:     int64(len(metadataJson)),
+		Mode:     0o400,
+		Format:   tar.FormatPAX,
+	})
+	tw.Write(metadataJson)
+
+	for _, data := range imageData {
+		if err := data.addToCosi(tw); err != nil {
+			return fmt.Errorf("failed to add %s to COSI: %w", data.Source, err)
+		}
+	}
+
+	log.Infof("Finished building COSI: %s", outputFile)
+	return nil
+}
+
+func (data *ImageBuildData) addToCosi(tw *tar.Writer) error {
+	imageFile, err := os.Open(data.Source)
+	if err != nil {
+		return fmt.Errorf("failed to open image file: %w", err)
+	}
+	defer imageFile.Close()
+
+	err = tw.WriteHeader(&tar.Header{
+		Typeflag: tar.TypeReg,
+		Name:     data.Metadata.Image.Path,
+		Size:     int64(data.Metadata.Image.CompressedSize),
+		Mode:     0o400,
+		Format:   tar.FormatPAX,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to write tar header: %w", err)
+	}
+
+	_, err = io.Copy(tw, imageFile)
+	if err != nil {
+		return fmt.Errorf("failed to write image to COSI: %w", err)
+	}
+
+	return nil
+}
+
+func sha384sum(path string) (string, error) {
+	sha384 := sha512.New384()
+	file, err := os.Open(path)
+	if err != nil {
+		return "", err
+	}
+	defer file.Close()
+
+	if _, err := io.Copy(sha384, file); err != nil {
+		return "", err
+	}
+	return fmt.Sprintf("%x", sha384.Sum(nil)), nil
+}
+
+func (data *ImageBuildData) populateMetadata() (*ExtractedImageData, error) {
+	stat, err := os.Stat(data.Source)
+	if err != nil {
+		return nil, fmt.Errorf("filed to stat %s: %w", data.Source, err)
+	}
+	if stat.IsDir() {
+		return nil, fmt.Errorf("%s is a directory", data.Source)
+	}
+	data.Metadata.Image.CompressedSize = uint64(stat.Size())
+
+	// Calculate the sha384 of the image
+	sha384, err := sha384sum(data.Source)
+	if err != nil {
+		return nil, fmt.Errorf("failed to calculate sha384 of %s: %w", data.Source, err)
+	}
+	data.Metadata.Image.Sha384 = sha384
+
+	tmpFile, err := os.Open(data.Source)
+	if err != nil {
+		return nil, fmt.Errorf("failed to open uncompressed file %s: %w", data.Source, err)
+	}
+
+	defer tmpFile.Close()
+
+	stat, err = tmpFile.Stat()
+	if err != nil {
+		return nil, fmt.Errorf("failed to stat decompressed image: %w", err)
+	}
+
+	data.Metadata.Image.UncompressedSize = uint64(stat.Size())
+	var extractedData ExtractedImageData
+
+	// If this image doesn't need to be mounted, we're done
+	if !data.KnownInfo.ShouldMount() {
+		return &extractedData, nil
+	}
+
+	return &extractedData, nil
+}

toolkit/tools/pkg/imagecustomizerlib/cosimetadata.go

@@ -0,0 +1,62 @@
+package imagecustomizerlib
+
+type MetadataJson struct {
+	Version   string  `json:"version"`
+	OsArch    string  `json:"osArch"`
+	Images    []Image `json:"images"`
+	OsRelease string  `json:"osRelease"`
+	Id        string  `json:"id"`
+}
+
+type Image struct {
+	Image      ImageFile     `json:"image"`
+	MountPoint string        `json:"mountPoint"`
+	FsType     string        `json:"fsType"`
+	FsUuid     string        `json:"fsUuid"`
+	PartType   PartitionType `json:"partType"`
+	Verity     *Verity       `json:"verity"`
+}
+
+type Verity struct {
+	Image    ImageFile `json:"image"`
+	Roothash string    `json:"roothash"`
+}
+
+type ImageFile struct {
+	Path             string `json:"path"`
+	CompressedSize   uint64 `json:"compressedSize"`
+	UncompressedSize uint64 `json:"uncompressedSize"`
+	Sha384           string `json:"sha384"`
+}
+
+type PartitionType string
+
+const (
+	PartitionTypeEsp                PartitionType = "c12a7328-f81f-11d2-ba4b-00a0c93ec93b"
+	PartitionTypeXbootldr           PartitionType = "bc13c2ff-59e6-4262-a352-b275fd6f7172"
+	PartitionTypeSwap               PartitionType = "0657fd6d-a4ab-43c4-84e5-0933c84b4f4f"
+	PartitionTypeHome               PartitionType = "933ac7e1-2eb4-4f13-b844-0e14e2aef915"
+	PartitionTypeSrv                PartitionType = "3b8f8425-20e0-4f3b-907f-1a25a76f98e8"
+	PartitionTypeVar                PartitionType = "4d21b016-b534-45c2-a9fb-5c16e091fd2d"
+	PartitionTypeTmp                PartitionType = "7ec6f557-3bc5-4aca-b293-16ef5df639d1"
+	PartitionTypeLinuxGeneric       PartitionType = "0fc63daf-8483-4772-8e79-3d69d8477de4"
+	PartitionTypeRootAmd64          PartitionType = "4f68bce3-e8cd-4db1-96e7-fbcaf984b709"
+	PartitionTypeRootAmd64Verity    PartitionType = "2c7357ed-ebd2-46d9-aec1-23d437ec2bf5"
+	PartitionTypeRootAmd64VeritySig PartitionType = "41092b05-9fc8-4523-994f-2def0408b176"
+	PartitionTypeUsrAmd64           PartitionType = "8484680c-9521-48c6-9c11-b0720656f69e"
+	PartitionTypeUsrAmd64Verity     PartitionType = "77ff5f63-e7b6-4633-acf4-1565b864c0e6"
+	PartitionTypeUsrAmd64VeritySig  PartitionType = "e7bb33fb-06cf-4e81-8273-e543b413e2e2"
+	PartitionTypeRootArm64          PartitionType = "b921b045-1df0-41c3-af44-4c6f280d3fae"
+	PartitionTypeRootArm64Verity    PartitionType = "df3300ce-d69f-4c92-978c-9bfb0f38d820"
+	PartitionTypeRootArm64VeritySig PartitionType = "6db69de6-29f4-4758-a7a5-962190f00ce3"
+	PartitionTypeUsrArm64           PartitionType = "b0e01050-ee5f-4390-949a-9101b17104e9"
+	PartitionTypeUsrArm64Verity     PartitionType = "6e11a4e7-fbca-4ded-b9e9-e1a512bb664e"
+	PartitionTypeUsrArm64VeritySig  PartitionType = "c23ce4ff-44bd-4b00-b2d4-b41b3419e02a"
+
+	PartitionTypeRoot          PartitionType = PartitionTypeRootAmd64
+	PartitionTypeRootVerity    PartitionType = PartitionTypeRootAmd64Verity
+	PartitionTypeRootVeritySig PartitionType = PartitionTypeRootAmd64VeritySig
+	PartitionTypeUsr           PartitionType = PartitionTypeUsrAmd64
+	PartitionTypeUsrVerity     PartitionType = PartitionTypeUsrAmd64Verity
+	PartitionTypeUsrVeritySig  PartitionType = PartitionTypeUsrAmd64VeritySig
+)