Fix VHD parsing - Take 2.

docs/imagecustomizer/api/cli.md

@@ -32,6 +32,11 @@ But it can also be an Azure Linux image that has been customized.
 
 Supported image file formats: vhd, vhdx, qcow2, and raw.
 
+Note: Image Customizer will reject VHD files created by `qemu-img` unless the
+`-o force_size=on` option was passed. Without this option, `qemu-img` will
+likely change the size of the disk (to a non-multiple of 1 MiB), which can cause
+problems when trying to upload the disk to Azure.
+
 If verity is enabled in the base image, then:
 
 - If the partitions are recustomized using the

docs/imagecustomizer/api/configuration/inputImage.md

@@ -27,6 +27,11 @@ But it can also be an Azure Linux image that has been customized.
 
 Supported image file formats: vhd, vhdx, qcow2, and raw.
 
+Note: Image Customizer will reject VHD files created by `qemu-img` unless the
+`-o force_size=on` option was passed. Without this option, `qemu-img` will
+likely change the size of the disk (to a non-multiple of 1 MiB), which can cause
+problems when trying to upload the disk to Azure.
+
 If verity is enabled in the base image, then:
 
 - If the partitions are recustomized using the

toolkit/tools/internal/vhdutils/main_test.go

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

toolkit/tools/internal/vhdutils/vhdutils.go

@@ -0,0 +1,137 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package vhdutils
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"io"
+	"os"
+)
+
+const (
+	VhdFileFooterSize = 512
+	VhdFileSignature  = "conectix"
+	VhdFileVersion    = 0x00010000
+)
+
+type VhdFileFooter struct {
+	Cookie             [8]byte
+	Features           uint32
+	FileFormatVersion  uint32
+	DataOffset         uint64
+	TimeStamp          uint32
+	CreatorApplication [4]byte
+	CreatorVersion     [4]byte
+	CreatorHostOS      [4]byte
+	OriginalSize       uint64
+	CurrentSize        uint64
+	Cylinder           uint16
+	Heads              uint8
+	SectorsPerCylinder uint8
+	DiskType           uint32
+	Checksum           [4]byte
+	UniqueId           [16]byte
+	SavedState         uint8
+	Reserved           [427]byte
+}
+
+var (
+	ErrVhdFileTooSmall       = errors.New("file is too small to be a VHD")
+	ErrVhdWrongFileSignature = errors.New("footer does not have correct VHD file signature")
+	ErrVhdWrongFileVersion   = errors.New("VHD footer has unsupported file format version")
+)
+
+type VhdFileSizeCalcType int
+
+const (
+	// File is not a VHD file.
+	VhdFileSizeCalcTypeNone VhdFileSizeCalcType = iota
+	// VHD uses "Current Size" field to calculate the disk size.
+	VhdFileSizeCalcTypeCurrentSize
+	// VHD uses "Disk Geometry" fields to calculate the disk size.
+	VhdFileSizeCalcTypeDiskGeometry
+)
+
+func GetVhdFileSizeCalcType(filename string) (VhdFileSizeCalcType, error) {
+	footer, err := ParseVhdFileFooter(filename)
+	if errors.Is(err, ErrVhdFileTooSmall) || errors.Is(err, ErrVhdWrongFileSignature) {
+		// Not a VHD file.
+		return VhdFileSizeCalcTypeNone, nil
+	}
+	if err != nil {
+		return VhdFileSizeCalcTypeNone, err
+	}
+
+	creatorApplication := string(footer.CreatorApplication[:])
+
+	//   There are actually two different ways of calculating the disk size of a VHD file. The old method, which is
+	// used by Microsoft Virtual PC, uses the VHD's footer's 'Disk Geometry' (cylinder, heads, and sectors per
+	// track/cylinder) fields. Using 'Disk Geometry' limits what file sizes are possible. So, Hyper-V uses only uses the
+	// the 'Current Size' field, which allows it to accept disks of any size.
+	//   Microsoft Virtual PC is pretty dead at this point. So, it is fairly safe to assume that almost all VHD files
+	// use the Hyper-V format. Unfortunately, qemu-img still defaults to using 'Disk Geometry' when a user requests a
+	// VHD (i.e. 'vpc') image. Image Customizer knows to pass the 'force_size=on' arg to qemu-img so that it uses
+	// 'Current Size'. But users might not know that they need to do this when using qemu-img manually.
+	//   Fortunately, qemu-img is nice enough to use different values of the 'Creator Application' field depending on
+	// the value of 'force_size'. Specifically, "qemu" for 'Disk Geometry' and "qem2 " for 'Current Size'. This can be
+	// used to determine which type of VHD we are dealing with.
+	//   qemu-img uses the 'Creator Application' field internally to determine what type of VHD it is dealing with.
+	// However, if it sees a 'Creator Application' value it doesn't recognize, it will assume it uses 'Disk Geometry'.
+	// Whereas, nowadays it is more likely for a VHD to use 'Current Size'.
+	switch creatorApplication {
+	case "vpc ", "vs  ", "qemu":
+		return VhdFileSizeCalcTypeDiskGeometry, nil
+
+	default:
+		return VhdFileSizeCalcTypeCurrentSize, nil
+	}
+}
+
+func ParseVhdFileFooter(filename string) (VhdFileFooter, error) {
+	fd, err := os.Open(filename)
+	if err != nil {
+		return VhdFileFooter{}, err
+	}
+	defer fd.Close()
+
+	stat, err := fd.Stat()
+	if err != nil {
+		return VhdFileFooter{}, err
+	}
+
+	if stat.Size() < VhdFileFooterSize {
+		return VhdFileFooter{}, ErrVhdFileTooSmall
+	}
+
+	_, err = fd.Seek(-VhdFileFooterSize, io.SeekEnd)
+	if err != nil {
+		return VhdFileFooter{}, err
+	}
+
+	footerBytes := [VhdFileFooterSize]byte{}
+	_, err = fd.Read([]byte(footerBytes[:]))
+	if err != nil {
+		return VhdFileFooter{}, err
+	}
+
+	footerReader := bytes.NewReader(footerBytes[:])
+
+	var footer VhdFileFooter
+	err = binary.Read(footerReader, binary.BigEndian, &footer)
+	if err != nil {
+		return VhdFileFooter{}, err
+	}
+
+	if string(footer.Cookie[:]) != VhdFileSignature {
+		return VhdFileFooter{}, ErrVhdWrongFileSignature
+	}
+
+	if footer.FileFormatVersion != VhdFileVersion {
+		return VhdFileFooter{}, ErrVhdWrongFileVersion
+	}
+
+	return footer, nil
+}

toolkit/tools/internal/vhdutils/vhdutils_test.go

@@ -0,0 +1,73 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package vhdutils
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/microsoft/azurelinux/toolkit/tools/internal/file"
+	"github.com/microsoft/azurelinux/toolkit/tools/internal/shell"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetVhdFileSizeCalcTypeDiskGeometryDynamic(t *testing.T) {
+	testGetVhdFileSizeCalcTypeHelper(t, "TestGetVhdFileSizeCalcTypeDiskGeometryDynamic", VhdFileSizeCalcTypeDiskGeometry,
+		[]string{"-f", "vpc", "-o", "force_size=off,subformat=dynamic"})
+}
+
+func TestGetVhdFileSizeCalcTypeDiskGeometryFixed(t *testing.T) {
+	testGetVhdFileSizeCalcTypeHelper(t, "TestGetVhdFileSizeCalcTypeDiskGeometryFixed", VhdFileSizeCalcTypeDiskGeometry,
+		[]string{"-f", "vpc", "-o", "force_size=off,subformat=fixed"})
+}
+
+func TestGetVhdFileSizeCalcTypeCurrentSizeDynamic(t *testing.T) {
+	testGetVhdFileSizeCalcTypeHelper(t, "TestGetVhdFileSizeCalcTypeCurrentSizeDynamic", VhdFileSizeCalcTypeCurrentSize,
+		[]string{"-f", "vpc", "-o", "force_size=on,subformat=dynamic"})
+}
+
+func TestGetVhdFileSizeCalcTypeCurrentSizeFixed(t *testing.T) {
+	testGetVhdFileSizeCalcTypeHelper(t, "TestGetVhdFileSizeCalcTypeCurrentSizeFixed", VhdFileSizeCalcTypeCurrentSize,
+		[]string{"-f", "vpc", "-o", "force_size=on,subformat=fixed"})
+}
+
+func TestGetVhdFileSizeCalcTypeNoneVhdx(t *testing.T) {
+	testGetVhdFileSizeCalcTypeHelper(t, "TestGetVhdFileSizeCalcTypeNoneVhdx", VhdFileSizeCalcTypeNone,
+		[]string{"-f", "vhdx"})
+}
+
+func TestGetVhdFileSizeCalcTypeNoneQcow2(t *testing.T) {
+	testGetVhdFileSizeCalcTypeHelper(t, "TestGetVhdFileSizeCalcTypeNoneQcow2", VhdFileSizeCalcTypeNone,
+		[]string{"-f", "qcow2"})
+}
+
+func TestGetVhdFileSizeCalcTypeNoneRaw(t *testing.T) {
+	testGetVhdFileSizeCalcTypeHelper(t, "TestGetVhdFileSizeCalcTypeNoneRaw", VhdFileSizeCalcTypeNone,
+		[]string{"-f", "raw"})
+}
+
+func testGetVhdFileSizeCalcTypeHelper(t *testing.T, testName string, expectedVhdFileSizeCalcType VhdFileSizeCalcType, qemuImgArgs []string) {
+	qemuimgExists, err := file.CommandExists("qemu-img")
+	assert.NoError(t, err)
+	if !qemuimgExists {
+		t.Skip("The 'qemu-img' command is not available")
+	}
+
+	testTempDir := filepath.Join(testsTempDir, testName)
+	testVhdFile := filepath.Join(testTempDir, "test.vhd")
+
+	err = os.MkdirAll(testTempDir, os.ModePerm)
+	assert.NoError(t, err)
+
+	args := []string{"create", testVhdFile, "1M"}
+	args = append(args, qemuImgArgs...)
+
+	err = shell.ExecuteLive(true, "qemu-img", args...)
+	assert.NoError(t, err)
+
+	vhdType, err := GetVhdFileSizeCalcType(testVhdFile)
+	assert.NoError(t, err)
+	assert.Equal(t, expectedVhdFileSizeCalcType, vhdType)
+}

toolkit/tools/pkg/imagecustomizerlib/artifactsinputoutput.go

@@ -269,10 +269,9 @@ func InjectFiles(ctx context.Context, buildDir string, baseConfigPath string, in
 	if err != nil {
 		return err
 	}
-	inputImageFormat := strings.TrimLeft(filepath.Ext(inputImageFile), ".")
 	rawImageFile := filepath.Join(buildDirAbs, BaseImageName)
 
-	detectedImageFormat, err := convertImageToRaw(inputImageFile, inputImageFormat, rawImageFile)
+	detectedImageFormat, err := convertImageToRaw(inputImageFile, rawImageFile)
 	if err != nil {
 		return err
 	}

toolkit/tools/pkg/imagecustomizerlib/imagecustomizer.go

@@ -24,6 +24,7 @@ import (
 	"github.com/microsoft/azurelinux/toolkit/tools/internal/safeloopback"
 	"github.com/microsoft/azurelinux/toolkit/tools/internal/safemount"
 	"github.com/microsoft/azurelinux/toolkit/tools/internal/shell"
+	"github.com/microsoft/azurelinux/toolkit/tools/internal/vhdutils"
 	"github.com/sirupsen/logrus"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
@@ -401,7 +402,7 @@ func convertInputImageToWriteableFormat(ctx context.Context, ic *ImageCustomizer
 	} else {
 		logger.Log.Infof("Creating raw base image: %s", ic.rawImageFile)
 
-		_, err := convertImageToRaw(ic.inputImageFile, ic.inputImageFormat, ic.rawImageFile)
+		_, err := convertImageToRaw(ic.inputImageFile, ic.rawImageFile)
 		if err != nil {
 			return nil, err
 		}
@@ -410,9 +411,7 @@ func convertInputImageToWriteableFormat(ctx context.Context, ic *ImageCustomizer
 	}
 }
 
-func convertImageToRaw(inputImageFile string, inputImageFormat string,
-	rawImageFile string,
-) (imagecustomizerapi.ImageFormatType, error) {
+func convertImageToRaw(inputImageFile string, rawImageFile string) (imagecustomizerapi.ImageFormatType, error) {
 	imageInfo, err := GetImageFileInfo(inputImageFile)
 	if err != nil {
 		return "", fmt.Errorf("failed to detect input image (%s) format:\n%w", inputImageFile, err)
@@ -421,24 +420,26 @@ func convertImageToRaw(inputImageFile string, inputImageFormat string,
 	detectedImageFormat := imageInfo.Format
 	sourceArg := fmt.Sprintf("file.filename=%s", qemuImgEscapeOptionValue(inputImageFile))
 
-	// The fixed-size VHD format is just a raw disk file with small metadata footer appended to the end. Unfortunatley,
-	// that footer doesn't contain a file signature (i.e. "magic number"). So, qemu-img can't correctly detect this
-	// format and instead reports fixed-size VHDs as raw images. So, use the filename extension as a hint.
-	if inputImageFormat == "vhd" && detectedImageFormat == "raw" {
-		// Force qemu-img to treat the file as a VHD.
-		detectedImageFormat = "vpc"
-	}
-
-	if detectedImageFormat == "vpc" {
-		// There are actually two different ways of calculating the disk size of a VHD file. The old method, which is
-		// used by Microsoft Virtual PC, uses the VHD's footer's "Disk Geometry" (cylinder, heads, and sectors per
-		// track/cylinder) fields. Whereas, the new method, which is used by Hyper-V, simply uses the VHD's footer's
-		// "Current Size" field. The qemu-img tool does try to correctly detect which one is being used by looking at
-		// the footer's "Creator Application" field. But if the tool that created the VHD uses a name that qemu-img
-		// doesn't recognize, then the heuristic can pick the wrong one. This seems to be the case for VHDs downloaded
-		// from Azure. For the Image Customizer tool, it is pretty safe to assume all VHDs use the Hyper-V format.
-		// So, force qemu-img to use that format.
-		sourceArg += ",driver=vpc,force_size_calc=current_size"
+	if detectedImageFormat == "raw" || detectedImageFormat == "vpc" {
+		// The fixed-size VHD format is just a raw disk file with small metadata footer appended to the end.
+		// Unfortunatley, qemu-img doesn't look at the VHD footer when detecting file formats. So, it reports
+		// fixed-sized VHDs as raw disk images. So, manually detect if a raw image is a VHD.
+		vhdFileType, err := vhdutils.GetVhdFileSizeCalcType(inputImageFile)
+		if err != nil {
+			return "", err
+		}
+
+		switch vhdFileType {
+		case vhdutils.VhdFileSizeCalcTypeDiskGeometry:
+			return "", fmt.Errorf("rejecting VHD file that uses 'Disk Geometry' based size:\npass '-o force_size=on' to qemu-img when outputting as 'vpc' (i.e. VHD)")
+
+		case vhdutils.VhdFileSizeCalcTypeCurrentSize:
+			sourceArg += ",driver=vpc,force_size_calc=current_size"
+			detectedImageFormat = "vpc"
+
+		default:
+			// Not a VHD file.
+		}
 	}
 
 	err = shell.ExecuteLiveWithErr(1, "qemu-img", "convert", "-O", "raw", "--image-opts", sourceArg, rawImageFile)