diff --git a/pkg/osbuild/sfdisk_stage.go b/pkg/osbuild/sfdisk_stage.go index d5a73daedb..6fcad97846 100644 --- a/pkg/osbuild/sfdisk_stage.go +++ b/pkg/osbuild/sfdisk_stage.go @@ -1,5 +1,11 @@ package osbuild +import ( + "fmt" + + "github.com/osbuild/images/pkg/disk" +) + // Partition a target using sfdisk(8) type SfdiskStageOptions struct { @@ -36,7 +42,18 @@ type SfdiskPartition struct { UUID string `json:"uuid,omitempty"` } +func (o SfdiskStageOptions) validate() error { + if o.Label == disk.PT_DOS.String() && len(o.Partitions) > 4 { + return fmt.Errorf("sfdisk stage creation failed: \"dos\" partition table only supports up to 4 partitions: got %d", len(o.Partitions)) + } + return nil +} + func NewSfdiskStage(options *SfdiskStageOptions, device *Device) *Stage { + if err := options.validate(); err != nil { + panic(err) + } + return &Stage{ Type: "org.osbuild.sfdisk", Options: options, diff --git a/pkg/osbuild/sfdisk_stage_test.go b/pkg/osbuild/sfdisk_stage_test.go index b6a659426a..e7cf19c1b6 100644 --- a/pkg/osbuild/sfdisk_stage_test.go +++ b/pkg/osbuild/sfdisk_stage_test.go @@ -41,3 +41,22 @@ func TestNewSfdiskStage(t *testing.T) { actualStageDOS := NewSfdiskStage(&options, device) assert.Equal(t, expectedStage, actualStageDOS) } + +func TestNewSfdiskStageInvalid(t *testing.T) { + + partition := SfdiskPartition{ + // doesn't really matter + } + + options := SfdiskStageOptions{ + Label: "dos", + UUID: "D209C89E-EA5E-4FBD-B161-B461CCE297E0", + Partitions: []SfdiskPartition{partition, partition, partition, partition, partition}, // 5 partitions + } + + device := NewLoopbackDevice(&LoopbackDeviceOptions{Filename: "disk.raw"}) + + assert.Panics(t, func() { + NewSfdiskStage(&options, device) + }) +}