From 613ad0b8629fee490f90f69f5e49d0cebdb8a04d Mon Sep 17 00:00:00 2001
From: Gianluca Zuccarelli <gzuccare@redhat.com>
Date: Thu, 19 Aug 2021 15:17:55 +0100
Subject: [PATCH] disk: refactor partition table size & start points

---
 internal/disk/customizations.go | 36 +++++++++++--------------
 internal/disk/disk.go           | 11 ++++++++
 internal/disk/disk_test.go      | 47 +++++++++++++++++++++++++++++++++
 3 files changed, 74 insertions(+), 20 deletions(-)
 create mode 100644 internal/disk/disk_test.go

diff --git a/internal/disk/customizations.go b/internal/disk/customizations.go
index b77f7bd77..a1cc8ea0d 100644
--- a/internal/disk/customizations.go
+++ b/internal/disk/customizations.go
@@ -6,7 +6,6 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/osbuild/osbuild-composer/internal/blueprint"
-	"github.com/osbuild/osbuild-composer/internal/distro"
 )
 
 const (
@@ -31,42 +30,41 @@ func CreatePartitionTable(
 	rng *rand.Rand,
 ) PartitionTable {
 
-	basePartitionTable.Size = imageSize
-	partitions := []Partition{}
-
 	if bootPartition := basePartitionTable.BootPartition(); bootPartition != nil {
 		// the boot partition UUID needs to be set since this
 		// needs to be randomly generated
 		bootPartition.Filesystem.UUID = uuid.Must(newRandomUUIDFromReader(rng)).String()
 	}
 
+	for _, m := range mountpoints {
+		if m.Mountpoint != "/" {
+			partitionSize := uint64(m.MinSize) / sectorSize
+			partition := basePartitionTable.createPartition(m.Mountpoint, partitionSize, rng)
+			basePartitionTable.Partitions = append(basePartitionTable.Partitions, partition)
+		}
+	}
+
+	if tableSize := basePartitionTable.getPartitionTableSize(); imageSize < tableSize {
+		imageSize = tableSize
+	}
+
+	basePartitionTable.Size = imageSize
+
 	// start point for all of the arches is
 	// 2048 sectors.
 	var start uint64 = basePartitionTable.updatePartitionStartPointOffsets(2048)
 
-	for _, m := range mountpoints {
-		if m.Mountpoint != "/" {
-			partitionSize := uint64(m.MinSize) / sectorSize
-			partition := createPartition(m.Mountpoint, partitionSize, start, archName, rng)
-			partitions = append(partitions, partition)
-			start += uint64(m.MinSize / sectorSize)
-		}
-	}
-
 	// treat the root partition as a special case
 	// by setting the size dynamically
 	rootPartition := basePartitionTable.RootPartition()
-	rootPartition.Start = start
 	rootPartition.Size = ((imageSize / sectorSize) - start - 100)
 	rootPartition.Filesystem.UUID = uuid.Must(newRandomUUIDFromReader(rng)).String()
-
 	basePartitionTable.updateRootPartition(*rootPartition)
-	basePartitionTable.Partitions = append(basePartitionTable.Partitions, partitions...)
 
 	return basePartitionTable
 }
 
-func createPartition(mountpoint string, size uint64, start uint64, archName string, rng *rand.Rand) Partition {
+func (pt *PartitionTable) createPartition(mountpoint string, size uint64, rng *rand.Rand) Partition {
 	filesystem := Filesystem{
 		Type:         "xfs",
 		UUID:         uuid.Must(newRandomUUIDFromReader(rng)).String(),
@@ -75,15 +73,13 @@ func createPartition(mountpoint string, size uint64, start uint64, archName stri
 		FSTabFreq:    0,
 		FSTabPassNo:  0,
 	}
-	if archName == distro.Ppc64leArchName || archName == distro.S390xArchName {
+	if pt.Type != "gpt" {
 		return Partition{
-			Start:      start,
 			Size:       size,
 			Filesystem: &filesystem,
 		}
 	}
 	return Partition{
-		Start:      start,
 		Size:       size,
 		Type:       FilesystemDataGUID,
 		UUID:       uuid.Must(newRandomUUIDFromReader(rng)).String(),
diff --git a/internal/disk/disk.go b/internal/disk/disk.go
index 548043523..50689c466 100644
--- a/internal/disk/disk.go
+++ b/internal/disk/disk.go
@@ -174,14 +174,17 @@ func (pt PartitionTable) RootPartitionIndex() int {
 // for each of the existing partitions
 // return the updated start point
 func (pt *PartitionTable) updatePartitionStartPointOffsets(start uint64) uint64 {
+	var rootIdx = -1
 	for i := range pt.Partitions {
 		partition := &pt.Partitions[i]
 		if partition.Filesystem != nil && partition.Filesystem.Mountpoint == "/" {
+			rootIdx = i
 			continue
 		}
 		partition.Start = start
 		start += partition.Size
 	}
+	pt.Partitions[rootIdx].Start = start
 	return start
 }
 
@@ -189,6 +192,14 @@ func (pt *PartitionTable) updateRootPartition(rootPartition Partition) {
 	pt.Partitions[pt.RootPartitionIndex()] = rootPartition
 }
 
+func (pt *PartitionTable) getPartitionTableSize() uint64 {
+	var size uint64
+	for _, p := range pt.Partitions {
+		size += p.Size
+	}
+	return size
+}
+
 // Converts Partition to osbuild.QEMUPartition that encodes the same partition.
 func (p Partition) QEMUPartition() osbuild.QEMUPartition {
 	var fs *osbuild.QEMUFilesystem
diff --git a/internal/disk/disk_test.go b/internal/disk/disk_test.go
new file mode 100644
index 000000000..02f9fbb02
--- /dev/null
+++ b/internal/disk/disk_test.go
@@ -0,0 +1,47 @@
+package disk_test
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/osbuild/osbuild-composer/internal/blueprint"
+	"github.com/osbuild/osbuild-composer/internal/disk"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDisk_DynamicallyResizePartitionTable(t *testing.T) {
+	mountpoints := []blueprint.FilesystemCustomization{
+		{
+			MinSize:    2147483648,
+			Mountpoint: "/usr",
+		},
+	}
+	pt := disk.PartitionTable{
+		UUID: "D209C89E-EA5E-4FBD-B161-B461CCE297E0",
+		Type: "gpt",
+		Partitions: []disk.Partition{
+			{
+				Size:     2048,
+				Bootable: true,
+				Type:     disk.BIOSBootPartitionGUID,
+				UUID:     disk.BIOSBootPartitionUUID,
+			},
+			{
+				Type: disk.FilesystemDataGUID,
+				UUID: disk.RootPartitionUUID,
+				Filesystem: &disk.Filesystem{
+					Type:         "xfs",
+					Label:        "root",
+					Mountpoint:   "/",
+					FSTabOptions: "defaults",
+					FSTabFreq:    0,
+					FSTabPassNo:  0,
+				},
+			},
+		},
+	}
+	var expectedSize uint64 = 2147483648
+	rng := rand.New(rand.NewSource(0))
+	pt = disk.CreatePartitionTable(mountpoints, 1024, pt, rng)
+	assert.GreaterOrEqual(t, expectedSize, pt.Size)
+}