Add support for custom partition types/GUIDs in blueprints

pkg/blueprint/disk_customizations.go

@@ -6,9 +6,11 @@ import (
 	"errors"
 	"fmt"
 	"path/filepath"
+	"regexp"
 	"slices"
 	"strings"
 
+	"github.com/google/uuid"
 	"github.com/osbuild/images/pkg/datasizes"
 	"github.com/osbuild/images/pkg/pathpolicy"
 )
@@ -64,6 +66,15 @@ type PartitionCustomization struct {
 	// (optional, defaults depend on payload and mountpoints).
 	MinSize uint64 `json:"minsize" toml:"minsize"`
 
+	// The partition type GUID for GPT partitions. For DOS partitions, this
+	// field can be used to set the partition type (e.g. "swap").
+	// If not set, the code will make the best guess based on the mountpoint or
+	// the payload type.
+	// Examples:
+	// 3B8F8425-20E0-4F3B-907F-1A25A76F98E8 (/srv on GPT)
+	// 06 (FAT16 on DOS)
+	GUID string `json:"guid,omitempty" toml:"guid,omitempty"`
+
 	BtrfsVolumeCustomization
 
 	VGCustomization
@@ -156,6 +167,7 @@ func (v *PartitionCustomization) UnmarshalJSON(data []byte) error {
 	var typeSniffer struct {
 		Type    string `json:"type"`
 		MinSize any    `json:"minsize"`
+		GUID    string `json:"guid"`
 	}
 	if err := json.Unmarshal(data, &typeSniffer); err != nil {
 		return fmt.Errorf("%s %w", errPrefix, err)
@@ -184,6 +196,7 @@ func (v *PartitionCustomization) UnmarshalJSON(data []byte) error {
 	}
 
 	v.Type = partType
+	v.GUID = typeSniffer.GUID
 
 	if typeSniffer.MinSize == nil {
 		return fmt.Errorf("minsize is required")
@@ -203,10 +216,11 @@ func (v *PartitionCustomization) UnmarshalJSON(data []byte) error {
 // the type is "plain", none of the fields for btrfs or lvm are used.
 func decodePlain(v *PartitionCustomization, data []byte) error {
 	var plain struct {
-		// Type and minsize are handled by the caller. These are added here to
+		// Type, minsize and guid are handled by the caller. These are added here to
 		// satisfy "DisallowUnknownFields" when decoding.
 		Type    string `json:"type"`
 		MinSize any    `json:"minsize"`
+		GUID    string `json:"guid"`
 		FilesystemTypedCustomization
 	}
 
@@ -226,10 +240,11 @@ func decodePlain(v *PartitionCustomization, data []byte) error {
 // the type is btrfs, none of the fields for plain or lvm are used.
 func decodeBtrfs(v *PartitionCustomization, data []byte) error {
 	var btrfs struct {
-		// Type and minsize are handled by the caller. These are added here to
+		// Type, minsize and guid are handled by the caller. These are added here to
 		// satisfy "DisallowUnknownFields" when decoding.
 		Type    string `json:"type"`
 		MinSize any    `json:"minsize"`
+		GUID    string `json:"guid"`
 		BtrfsVolumeCustomization
 	}
 
@@ -249,10 +264,11 @@ func decodeBtrfs(v *PartitionCustomization, data []byte) error {
 // is lvm, none of the fields for plain or btrfs are used.
 func decodeLVM(v *PartitionCustomization, data []byte) error {
 	var vg struct {
-		// Type and minsize are handled by the caller. These are added here to
+		// Type, minsize and guid are handled by the caller. These are added here to
 		// satisfy "DisallowUnknownFields" when decoding.
 		Type    string `json:"type"`
 		MinSize any    `json:"minsize"`
+		GUID    string `json:"guid"`
 		VGCustomization
 	}
 
@@ -367,6 +383,9 @@ func (p *DiskCustomization) Validate() error {
 	vgnames := make(map[string]bool)
 	var errs []error
 	for _, part := range p.Partitions {
+		if err := part.validateGeneric(); err != nil {
+			errs = append(errs, err)
+		}
 		switch part.Type {
 		case "plain", "":
 			errs = append(errs, part.validatePlain(mountpoints))
@@ -471,6 +490,32 @@ var validPlainFSTypes = []string{
 	"xfs",
 }
 
+// exactly 2 hex digits
+var validDosPartitionType = regexp.MustCompile(`^[0-9a-fA-F]{2}$`)
+
+// validateGeneric checks the partition validity regardless of its type.
+// Currently, it only checks the GUID field.
+func (p *PartitionCustomization) validateGeneric() error {
+	// Empty GUID is fine, the code will auto-generate it later.
+	if p.GUID == "" {
+		return nil
+	}
+
+	// We don't know the partition table type yet, so check:
+
+	// 1) the partition GUID is either a valid UUID (for GPT)
+	if _, err := uuid.Parse(p.GUID); err == nil {
+		return nil
+	}
+
+	// 2) or a valid DOS partition type (for MBR)
+	if validDosPartitionType.MatchString(p.GUID) {
+		return nil
+	}
+
+	return fmt.Errorf("invalid partition GUID: %q (use UUIDs for GPT partition tables, or 2-digit hex numbers for DOS partition tables)", p.GUID)
+}
+
 func (p *PartitionCustomization) validatePlain(mountpoints map[string]bool) error {
 	if p.FSType == "swap" {
 		// make sure the mountpoint is empty and return

pkg/blueprint/disk_customizations_test.go

@@ -970,6 +970,48 @@ func TestPartitioningValidation(t *testing.T) {
 			},
 			expectedMsg: `invalid partitioning customizations: "dos" partition table type only supports up to 4 partitions: got 6`,
 		},
+		"happy-partition-guids": {
+			partitioning: &blueprint.DiskCustomization{
+				Partitions: []blueprint.PartitionCustomization{
+					{
+						GUID: "12345678-1234-1234-1234-1234567890ab",
+						FilesystemTypedCustomization: blueprint.FilesystemTypedCustomization{
+							FSType:     "ext4",
+							Mountpoint: "/gpt",
+						},
+					},
+					{
+						GUID: "ef",
+						FilesystemTypedCustomization: blueprint.FilesystemTypedCustomization{
+							FSType:     "ext4",
+							Mountpoint: "/dos",
+						},
+					},
+				},
+			},
+		},
+		"unhappy-partition-guids": {
+			partitioning: &blueprint.DiskCustomization{
+				Partitions: []blueprint.PartitionCustomization{
+
+					{
+						GUID: "12345678-uuid-1234-1234-1234567890ab",
+						FilesystemTypedCustomization: blueprint.FilesystemTypedCustomization{
+							FSType:     "ext4",
+							Mountpoint: "/gpt",
+						},
+					},
+					{
+						GUID: "0x52",
+						FilesystemTypedCustomization: blueprint.FilesystemTypedCustomization{
+							FSType:     "ext4",
+							Mountpoint: "/dos",
+						},
+					},
+				},
+			},
+			expectedMsg: "invalid partitioning customizations:\ninvalid partition GUID: \"12345678-uuid-1234-1234-1234567890ab\" (use UUIDs for GPT partition tables, or 2-digit hex numbers for DOS partition tables)\ninvalid partition GUID: \"0x52\" (use UUIDs for GPT partition tables, or 2-digit hex numbers for DOS partition tables)",
+		},
 	}
 
 	for name := range testCases {
@@ -1187,13 +1229,15 @@ func TestPartitionCustomizationUnmarshalJSON(t *testing.T) {
 			input: `{
 				"type": "plain",
 				"minsize": "1 GiB",
+				"guid": "12345678-1234-1234-1234-1234567890ab",
 				"mountpoint": "/",
 				"label": "root",
 				"fs_type": "xfs"
 			}`,
 			expected: &blueprint.PartitionCustomization{
 				Type:    "plain",
 				MinSize: 1 * datasizes.GiB,
+				GUID:    "12345678-1234-1234-1234-1234567890ab",
 				FilesystemTypedCustomization: blueprint.FilesystemTypedCustomization{
 					Mountpoint: "/",
 					Label:      "root",
@@ -1223,6 +1267,7 @@ func TestPartitionCustomizationUnmarshalJSON(t *testing.T) {
 			input: `{
 				"type": "btrfs",
 				"minsize": "10 GiB",
+				"guid": "12345678-1234-1234-1234-1234567890ab",
 				"subvolumes": [
 					{
 						"name": "subvols/root",
@@ -1237,6 +1282,7 @@ func TestPartitionCustomizationUnmarshalJSON(t *testing.T) {
 			expected: &blueprint.PartitionCustomization{
 				Type:    "btrfs",
 				MinSize: 10 * datasizes.GiB,
+				GUID:    "12345678-1234-1234-1234-1234567890ab",
 				BtrfsVolumeCustomization: blueprint.BtrfsVolumeCustomization{
 					Subvolumes: []blueprint.BtrfsSubvolumeCustomization{
 						{
@@ -1288,6 +1334,7 @@ func TestPartitionCustomizationUnmarshalJSON(t *testing.T) {
 				"type": "lvm",
 				"name": "myvg",
 				"minsize": "99 GiB",
+				"guid": "12345678-1234-1234-1234-1234567890ab",
 				"logical_volumes": [
 					{
 						"name": "homelv",
@@ -1308,6 +1355,7 @@ func TestPartitionCustomizationUnmarshalJSON(t *testing.T) {
 			expected: &blueprint.PartitionCustomization{
 				Type:    "lvm",
 				MinSize: 99 * datasizes.GiB,
+				GUID:    "12345678-1234-1234-1234-1234567890ab",
 				VGCustomization: blueprint.VGCustomization{
 					Name: "myvg",
 					LogicalVolumes: []blueprint.LVCustomization{
@@ -1399,6 +1447,14 @@ func TestPartitionCustomizationUnmarshalJSON(t *testing.T) {
 			}`,
 			errorMsg: "JSON unmarshal: error decoding minsize for partition: cannot be negative",
 		},
+		"guid-not-string": {
+			input: `{
+				"minsize": "10 GiB",
+				"mountpoint": "/",
+				"guid": 12345678
+			}`,
+			errorMsg: "JSON unmarshal: json: cannot unmarshal number into Go struct field .guid of type string",
+		},
 		"wrong-type/btrfs-with-lvm": {
 			input: `{
 				"type": "btrfs",

pkg/disk/disk.go

@@ -24,6 +24,7 @@ import (
 	"io"
 	"math/rand"
 	"reflect"
+	"regexp"
 	"strings"
 
 	"slices"
@@ -120,6 +121,26 @@ func getPartitionTypeIDfor(ptType PartitionTableType, partTypeName string) (stri
 	return id, nil
 }
 
+// exactly 2 hex digits
+var validDosPartitionType = regexp.MustCompile(`^[0-9a-fA-F]{2}$`)
+
+func validatePartitionTypeID(ptType PartitionTableType, partTypeName string) error {
+	switch ptType {
+	case PT_DOS:
+		if !validDosPartitionType.MatchString(partTypeName) {
+			return fmt.Errorf("invalid dos partition type ID: %s", partTypeName)
+		}
+	case PT_GPT:
+		if _, err := uuid.Parse(partTypeName); err != nil {
+			return fmt.Errorf("invalid gpt partition type GUID: %s", partTypeName)
+		}
+	default:
+		return fmt.Errorf("unknown or unsupported partition table enum: %d", ptType)
+	}
+
+	return nil
+}
+
 // FSType is the filesystem type enum.
 //
 // There should always be one value for each filesystem type supported by

pkg/disk/partition_table.go

@@ -1313,24 +1313,35 @@ func addPlainPartition(pt *PartitionTable, partition blueprint.PartitionCustomiz
 		return fmt.Errorf("error creating partition with mountpoint %q: %w", partition.Mountpoint, err)
 	}
 
-	// all user-defined partitions are data partitions except boot and swap
-	var typeName string
-	switch {
-	case partition.Mountpoint == "/boot":
-		typeName = "boot"
-	case fstype == "swap":
-		typeName = "swap"
-	default:
-		typeName = "data"
-	}
+	partType := partition.GUID
 
-	partType, err := getPartitionTypeIDfor(pt.Type, typeName)
-	if err != nil {
-		return fmt.Errorf("error getting partition type ID for %q: %w", partition.Mountpoint, err)
+	if partType != "" {
+		if err := validatePartitionTypeID(pt.Type, partType); err != nil {
+			return fmt.Errorf("error validating partition type ID for %q: %w", partition.Mountpoint, err)
+		}
+	} else {
+		// if the partition type is not specified, determine it based on the
+		// mountpoint and the partition type
+
+		// all user-defined partitions are data partitions except boot and swap
+		var typeName string
+		switch {
+		case partition.Mountpoint == "/boot":
+			typeName = "boot"
+		case fstype == "swap":
+			typeName = "swap"
+		default:
+			typeName = "data"
+		}
+
+		partType, err = getPartitionTypeIDfor(pt.Type, typeName)
+		if err != nil {
+			return fmt.Errorf("error getting partition type ID for %q: %w", partition.Mountpoint, err)
+		}
 	}
 
 	var payload PayloadEntity
-	switch typeName {
+	switch fstype {
 	case "swap":
 		payload = &Swap{
 			Label:        partition.Label,
@@ -1408,10 +1419,19 @@ func addLVMPartition(pt *PartitionTable, partition blueprint.PartitionCustomizat
 	}
 
 	// create partition for volume group
-	partType, err := getPartitionTypeIDfor(pt.Type, "lvm")
-	if err != nil {
-		return fmt.Errorf("error creating lvm partition %q: %w", vgname, err)
+	partType := partition.GUID
+	if partType != "" {
+		if err := validatePartitionTypeID(pt.Type, partType); err != nil {
+			return fmt.Errorf("error validating partition type ID for %q: %w", vgname, err)
+		}
+	} else {
+		var err error
+		partType, err = getPartitionTypeIDfor(pt.Type, "lvm")
+		if err != nil {
+			return fmt.Errorf("error creating lvm partition %q: %w", vgname, err)
+		}
 	}
+
 	newpart := Partition{
 		Type:     partType,
 		Size:     partition.MinSize,
@@ -1437,9 +1457,17 @@ func addBtrfsPartition(pt *PartitionTable, partition blueprint.PartitionCustomiz
 	}
 
 	// create partition for btrfs volume
-	partType, err := getPartitionTypeIDfor(pt.Type, "data")
-	if err != nil {
-		return fmt.Errorf("error creating btrfs partition: %w", err)
+	partType := partition.GUID
+	if partType != "" {
+		if err := validatePartitionTypeID(pt.Type, partType); err != nil {
+			return fmt.Errorf("error validating partition type ID for btrfs: %w", err)
+		}
+	} else {
+		var err error
+		partType, err = getPartitionTypeIDfor(pt.Type, "data")
+		if err != nil {
+			return fmt.Errorf("error creating btrfs partition: %w", err)
+		}
 	}
 	newpart := Partition{
 		Type:     partType,