disk: grow root partition while re-layouting

Instead of growing the root partition in `CreatePartitionTable` do it
as part of `updatePartitionStartPointOffsets`.

internal/disk/customizations.go

@@ -48,12 +48,7 @@ func CreatePartitionTable(
 
 	// start point for all of the arches is
 	// 2048 sectors.
-	var start uint64 = table.updatePartitionStartPointOffsets(2048, imageSize)
-
-	// treat the root partition as a special case
-	// by setting the size dynamically
-	rootPartition := table.RootPartition()
-	rootPartition.Size = (table.BytesToSectors(imageSize) - start - 100)
+	table.updatePartitionStartPointOffsets(2048, imageSize)
 
 	// Generate new UUIDs for filesystems and partitions
 	table.GenerateUUIDs(rng)

internal/disk/disk.go

@@ -316,6 +316,7 @@ func (pt *PartitionTable) GenerateUUIDs(rng *rand.Rand) {
 // Dynamically calculate and update the start point for each of the existing
 // partitions. Adjusts the overall size of image to either the supplied
 // value in `size` or to the sum of all partitions if that is lager.
+// Will grow the root partition if there is any empty space.
 // Returns the updated start point.
 func (pt *PartitionTable) updatePartitionStartPointOffsets(start, size uint64) uint64 {
 	var rootIdx = -1
@@ -342,6 +343,13 @@ func (pt *PartitionTable) updatePartitionStartPointOffsets(start, size uint64) u
 		pt.Size = size
 	}
 
+	// If there is space left in the partition table, grow root
+	root.Size = pt.BytesToSectors(pt.Size) - root.Start
+
+	// Finally we shrink the last partition, i.e. the root partition,
+	// to leave space for the secondary GPT header.
+	root.Size -= 100
+
 	return start
 }