From 8e2b2b9bb068ddaf9c9847f71f11169235c7eac2 Mon Sep 17 00:00:00 2001
From: Achilleas Koutsou <achilleas@koutsou.net>
Date: Mon, 9 Dec 2024 16:37:49 +0100
Subject: [PATCH] osbuild: validate sfdisk stage options

DOS partition tables only support up to 4 partitions.  Don't create the
sfdisk stage options if this rule is violated.  Stage validation
failures panic, so we panic on the stage creation.  This might change in
the future (to return errors) but for now we use panics and put the
responsibility on the callers to validate their inputs if they want to
have better error control.

Test sfdisk stage validation failure
---
 pkg/osbuild/sfdisk_stage.go      | 17 +++++++++++++++++
 pkg/osbuild/sfdisk_stage_test.go | 19 +++++++++++++++++++
 2 files changed, 36 insertions(+)

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)
+	})
+}