Skip to content

Commit

Permalink
Add initial COSI support (#40)
Browse files Browse the repository at this point in the history
<!-- Description: Please provide a summary of the changes and the
motivation behind them. -->
PR summary:

This is a phase 1 COSI support, verity images support is not included.
Enable COSI output, which involves steps like converting vhdx image to
raw, extract partition images, generate metadata, and package everything
into the final COSI tarball.

some local testing
with command `sudo ./imagecustomizer --log-level debug --build-dir
./build --image-file core-3.0.20241213.vhdx --output-image-file
./output-image-test.cosi --output-image-format cosi --config-file
../pkg/imagecustomizerlib/testdata/newpartitionsuuids-config.yaml`


![image](https://github.com/user-attachments/assets/3f092aee-2996-460a-ad56-a1b4c1f98966)
 
---

### **Checklist**
- [ ] Tests added/updated
- [ ] Documentation updated (if needed)
- [ ] Code conforms to style guidelines
  • Loading branch information
elainezhao96 authored Dec 20, 2024
1 parent 11ad05f commit 30fe585
Show file tree
Hide file tree
Showing 6 changed files with 283 additions and 37 deletions.
2 changes: 1 addition & 1 deletion docs/imagecustomizer/api/cli.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
2 changes: 1 addition & 1 deletion toolkit/tools/imagecustomizer/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,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()
Expand Down
192 changes: 192 additions & 0 deletions toolkit/tools/pkg/imagecustomizerlib/cosicommon.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,192 @@
package imagecustomizerlib

import (
"archive/tar"
"crypto/sha512"
_ "embed"
"encoding/json"
"fmt"
"io"
"os"
"path"
"path/filepath"
"runtime"

"github.com/microsoft/azurelinux/toolkit/tools/internal/logger"
"github.com/microsoft/azurelinux/toolkit/tools/internal/safeloopback"
)

type ImageBuildData struct {
Source string
KnownInfo outputPartitionMetadata
Metadata *Image
}

func convertToCosi(ic *ImageCustomizerParameters) error {
logger.Log.Infof("Extracting partition files")
outputDir := filepath.Join(ic.buildDir, "cosiimages")
err := os.MkdirAll(outputDir, os.ModePerm)
if err != nil {
return fmt.Errorf("failed to create folder %s:\n%w", outputDir, err)
}

imageLoopback, err := safeloopback.NewLoopback(ic.rawImageFile)
if err != nil {
return err
}
defer imageLoopback.Close()

partitionMetadataOutput, err := extractPartitions(imageLoopback.DevicePath(), outputDir, ic.outputImageBase, "raw-zst", ic.imageUuid)
if err != nil {
return err
}

err = buildCosiFile(outputDir, ic.outputImageFile, partitionMetadataOutput, ic.imageUuidStr)
if err != nil {
return fmt.Errorf("failed to build COSI:\n%w", err)
}

logger.Log.Infof("Successfully converted to COSI: %s", ic.outputImageFile)

err = imageLoopback.CleanClose()
if err != nil {
return err
}

return nil
}

func buildCosiFile(sourceDir string, outputFile string, expectedImages []outputPartitionMetadata, imageUuidStr string) error {
metadata := MetadataJson{
Version: "1.0",
OsArch: runtime.GOARCH,
Id: imageUuidStr,
Images: make([]Image, len(expectedImages)),
}

if len(expectedImages) == 0 {
return fmt.Errorf("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.PartitionFilename),
Metadata: metadata,
KnownInfo: image,
}

metadata.Image.Path = path.Join("images", image.PartitionFilename)
metadata.PartType = image.PartitionTypeUuid
metadata.MountPoint = image.Mountpoint
metadata.FsType = image.FileSystemType
metadata.FsUuid = image.Uuid
metadata.UncompressedSize = image.UncompressedSize
}

// Populate metadata for each image
for _, data := range imageData {
logger.Log.Infof("Processing image %s", data.Source)
err := populateMetadata(data)
if err != nil {
return fmt.Errorf("failed to populate metadata for %s:\n%w", data.Source, err)
}

logger.Log.Infof("Populated metadata for image %s", data.Source)
}

// Marshal metadata.json
metadataJson, err := json.MarshalIndent(metadata, "", " ")
if err != nil {
return fmt.Errorf("failed to marshal metadata:\n%w", err)
}

// Create COSI file
cosiFile, err := os.Create(outputFile)
if err != nil {
return fmt.Errorf("failed to create COSI file:\n%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 := addToCosi(data, tw); err != nil {
return fmt.Errorf("failed to add %s to COSI:\n%w", data.Source, err)
}
}

logger.Log.Infof("Finished building COSI: %s", outputFile)
return nil
}

func addToCosi(data ImageBuildData, tw *tar.Writer) error {
imageFile, err := os.Open(data.Source)
if err != nil {
return fmt.Errorf("failed to open image file:\n%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:\n%w", err)
}

_, err = io.Copy(tw, imageFile)
if err != nil {
return fmt.Errorf("failed to write image to COSI:\n%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 populateMetadata(data ImageBuildData) error {
stat, err := os.Stat(data.Source)
if err != nil {
return fmt.Errorf("filed to stat %s:\n%w", data.Source, err)
}
if stat.IsDir() {
return 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 fmt.Errorf("failed to calculate sha384 of %s:\n%w", data.Source, err)
}
data.Metadata.Image.Sha384 = sha384
return nil
}
31 changes: 31 additions & 0 deletions toolkit/tools/pkg/imagecustomizerlib/cosimetadata.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
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 string `json:"partType"`
Verity *Verity `json:"verity"`
UncompressedSize uint64 `json:"uncompressedSize"`
}

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"`
}
52 changes: 28 additions & 24 deletions toolkit/tools/pkg/imagecustomizerlib/extractpartitions.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,15 @@ import (
)

type outputPartitionMetadata struct {
PartitionNum int `json:"partitionnum"` // Example: 1
PartitionFilename string `json:"filename"` // Example: image_1.raw.zst
PartLabel string `json:"partlabel"` // Example: boot
FileSystemType string `json:"fstype"` // Example: vfat
PartitionTypeUuid string `json:"parttype"` // Example: c12a7328-f81f-11d2-ba4b-00a0c93ec93b
Uuid string `json:"uuid"` // Example: 4BD9-3A78
PartUuid string `json:"partuuid"` // Example: 7b1367a6-5845-43f2-99b1-a742d873f590
Mountpoint string `json:"mountpoint"` // Example: /mnt/os/boot
PartitionNum int `json:"partitionnum"` // Example: 1
PartitionFilename string `json:"filename"` // Example: image_1.raw.zst
PartLabel string `json:"partlabel"` // Example: boot
FileSystemType string `json:"fstype"` // Example: vfat
PartitionTypeUuid string `json:"parttype"` // Example: c12a7328-f81f-11d2-ba4b-00a0c93ec93b
Uuid string `json:"uuid"` // Example: 4BD9-3A78
PartUuid string `json:"partuuid"` // Example: 7b1367a6-5845-43f2-99b1-a742d873f590
Mountpoint string `json:"mountpoint"` // Example: /mnt/os/boot
UncompressedSize uint64 `json:"uncompressedsize"` // Example: 104857600
}

const (
Expand All @@ -32,12 +33,12 @@ const (
)

// Extract all partitions of connected image into separate files with specified format.
func extractPartitions(imageLoopDevice string, outDir string, basename string, partitionFormat string, imageUuid [UuidSize]byte) error {
func extractPartitions(imageLoopDevice string, outDir string, basename string, partitionFormat string, imageUuid [UuidSize]byte) ([]outputPartitionMetadata, error) {

// Get partition info
diskPartitions, err := diskutils.GetDiskPartitions(imageLoopDevice)
if err != nil {
return err
return nil, err
}

// Stores the output partition metadata that will be written to JSON file
Expand All @@ -51,55 +52,58 @@ func extractPartitions(imageLoopDevice string, outDir string, basename string, p

partitionNum, err := getPartitionNum(partition.Path)
if err != nil {
return err
return nil, err
}

partitionFilename := basename + "_" + strconv.Itoa(partitionNum)
rawFilename := partitionFilename + ".raw"

partitionFilepath, err := copyBlockDeviceToFile(outDir, partition.Path, rawFilename)
if err != nil {
return err
return nil, err
}

partitionFullFilePath, err := filepath.Abs(partitionFilepath)
if err != nil {
return fmt.Errorf("failed to get absolute path (%s):\n%w", partitionFilepath, err)
return nil, fmt.Errorf("failed to get absolute path (%s):\n%w", partitionFilepath, err)
}

// Sanity check the partition file.
err = checkFileSystemFile(partition.FileSystemType, partitionFullFilePath)
if err != nil {
return fmt.Errorf("failed to check file system integrity (%s):\n%w", partitionFilepath, err)
return nil, fmt.Errorf("failed to check file system integrity (%s):\n%w", partitionFilepath, err)
}

// Get uncompressed size for raw files
var uncompressedPartitionFileSize uint64
stat, err := os.Stat(partitionFullFilePath)
if err != nil {
return nil, fmt.Errorf("failed to stat raw file %s: %w", partitionFilepath, err)
}
uncompressedPartitionFileSize = uint64(stat.Size())

switch partitionFormat {
case "raw":
// Do nothing for "raw" case
case "raw-zst":
partitionFilepath, err = extractRawZstPartition(partitionFilepath, imageUuid, partitionFilename, outDir)
if err != nil {
return err
return nil, err
}
default:
return fmt.Errorf("unsupported partition format (supported: raw, raw-zst): %s", partitionFormat)
return nil, fmt.Errorf("unsupported partition format (supported: raw, raw-zst): %s", partitionFormat)
}

partitionMetadata, err := constructOutputPartitionMetadata(partition, partitionNum, partitionFilepath)
if err != nil {
return fmt.Errorf("failed to construct partition metadata:\n%w", err)
return nil, fmt.Errorf("failed to construct partition metadata:\n%w", err)
}
partitionMetadata.UncompressedSize = uncompressedPartitionFileSize
partitionMetadataOutput = append(partitionMetadataOutput, partitionMetadata)
logger.Log.Infof("Partition file created: %s", partitionFilepath)
}

// Write partition metadata JSON to a file
jsonFilename := basename + "_partition_metadata.json"
err = writePartitionMetadataJson(outDir, jsonFilename, &partitionMetadataOutput)
if err != nil {
return fmt.Errorf("failed to write partition metadata json:\n%w", err)
}
return nil
return partitionMetadataOutput, nil
}

// Extract raw-zst partition.
Expand Down
Loading

0 comments on commit 30fe585

Please sign in to comment.