osbase/gpt: fix bugs in AddPartition and add tests

AddPartition was very buggy: Many of the new tests fail on the old 
implementation. For example, in empty-fill, it fails to create a 
partition because it calculates the maximum free space without 
considering alignment. In haveone-basic, it creates overlapping 
partitions.

Change-Id: I4ab9ea833a72f694b5f5116ba084b923190c0bd2
Reviewed-on: https://review.monogon.dev/c/monogon/+/3347
Reviewed-by: Lorenz Brun <lorenz@monogon.tech>
Tested-by: Jenkins CI
diff --git a/osbase/gpt/gpt.go b/osbase/gpt/gpt.go
index ca8566a..8b1fc4d 100644
--- a/osbase/gpt/gpt.go
+++ b/osbase/gpt/gpt.go
@@ -10,7 +10,7 @@
 	"errors"
 	"fmt"
 	"hash/crc32"
-	"sort"
+	"slices"
 	"strings"
 	"unicode/utf16"
 
@@ -21,7 +21,8 @@
 )
 
 var gptSignature = [8]byte{'E', 'F', 'I', ' ', 'P', 'A', 'R', 'T'}
-var gptRevision uint32 = 0x00010000 // First 2 bytes major, second 2 bytes minor
+
+const gptRevision uint32 = 0x00010000 // First 2 bytes major, second 2 bytes minor
 
 // See UEFI Specification 2.9 Table 5-5
 type header struct {
@@ -190,14 +191,15 @@
 // value. By default, these are aligned to 1MiB.
 // Only use this flag if you are certain you need it, it can cause quite severe
 // performance degradation under certain conditions.
-func WithAlignment(alignmenet int64) AddOption {
+func WithAlignment(alignment int64) AddOption {
 	return func(options *addOptions) {
-		options.alignment = alignmenet
+		options.alignment = alignment
 	}
 }
 
 // AddPartition takes a pointer to a partition and adds it, placing it into
 // the first (or last using WithPreferEnd) continuous free space which fits it.
+// If size is -1, the partition will fill the largest free space.
 // It writes the placement information (FirstBlock, LastBlock) back to p.
 // By default, AddPartition aligns FirstBlock to 1MiB boundaries, but this can
 // be overridden using WithAlignment.
@@ -212,9 +214,13 @@
 	for _, o := range options {
 		o(&opts)
 	}
+	if opts.alignment <= 0 {
+		return fmt.Errorf("alignment (%d bytes) must be positive", opts.alignment)
+	}
 	if opts.alignment%blockSize != 0 {
 		return fmt.Errorf("requested alignment (%d bytes) is not an integer multiple of the block size (%d), unable to align", opts.alignment, blockSize)
 	}
+	alignBlocks := opts.alignment / blockSize
 	if p.ID == uuid.Nil {
 		p.ID = uuid.New()
 	}
@@ -225,71 +231,61 @@
 	}
 	if opts.preferEnd {
 		// Reverse fs slice to start iteration at the end
-		for i, j := 0, len(fs)-1; i < j; i, j = i+1, j-1 {
-			fs[i], fs[j] = fs[j], fs[i]
-		}
+		slices.Reverse(fs)
 	}
-	// Number of blocks the partition should occupy, rounded up.
-	blocks := (size + blockSize - 1) / blockSize
-	if size == -1 {
-		var largestFreeSpace int64
-		for _, freeInt := range fs {
-			intSz := freeInt[1] - freeInt[0]
-			if intSz > largestFreeSpace {
-				largestFreeSpace = intSz
-			}
-		}
-		blocks = largestFreeSpace
-	}
-	var maxFreeBlocks int64
+	var largestFreeSpace int64
 	for _, freeInt := range fs {
-		start := freeInt[0]
-		end := freeInt[1]
-		freeBlocks := end - start
-		// Align start properly
-		alignTo := opts.alignment / blockSize
-		// Go doesn't implement the euclidean modulus, thus this construction
-		// is necessary.
-		paddingBlocks := ((alignTo - start) % alignTo) % alignTo
-		freeBlocks -= paddingBlocks
-		start += paddingBlocks
-		if maxFreeBlocks < freeBlocks {
-			maxFreeBlocks = freeBlocks
-		}
-		if freeBlocks >= blocks {
-			if !opts.preferEnd {
-				p.FirstBlock = uint64(start)
-				p.LastBlock = uint64(start + blocks - 1)
-			} else {
-				// Realign FirstBlock. This will always succeed as
-				// there is enough space to align to the start.
-				moveLeft := (end - blocks - 1) % (opts.alignment / blockSize)
-				p.FirstBlock = uint64(end - (blocks + 1 + moveLeft))
-				p.LastBlock = uint64(end - (2 + moveLeft))
-			}
-			newPartPos := -1
-			if !opts.keepEmptyEntries {
-				for i, part := range gpt.Partitions {
-					if part.IsUnused() {
-						newPartPos = i
-						break
-					}
-				}
-			}
-			if newPartPos == -1 {
-				gpt.Partitions = append(gpt.Partitions, p)
-			} else {
-				gpt.Partitions[newPartPos] = p
-			}
-			p.Section, err = blockdev.NewSection(gpt.b, int64(p.FirstBlock), int64(p.LastBlock)+1)
-			if err != nil {
-				return fmt.Errorf("failed to create blockdev Section for partition: %w", err)
-			}
-			return nil
+		startAligned := (freeInt[0] + alignBlocks - 1) / alignBlocks * alignBlocks
+		freeBlocks := freeInt[1] - startAligned
+		if freeBlocks > largestFreeSpace {
+			largestFreeSpace = freeBlocks
 		}
 	}
 
-	return fmt.Errorf("no space for partition of %d blocks, largest continuous free space after alignment is %d blocks", blocks, maxFreeBlocks)
+	// Number of blocks the partition should occupy, rounded up.
+	blocks := (size + blockSize - 1) / blockSize
+	if size == -1 {
+		if largestFreeSpace == 0 {
+			return errors.New("no free space")
+		}
+		blocks = largestFreeSpace
+	} else if size <= 0 {
+		return fmt.Errorf("partition size (%d bytes) must be positive or the special value -1", size)
+	}
+	for _, freeInt := range fs {
+		startAligned := (freeInt[0] + alignBlocks - 1) / alignBlocks * alignBlocks
+		freeBlocks := freeInt[1] - startAligned
+		if freeBlocks < blocks {
+			continue
+		}
+		if opts.preferEnd {
+			startAligned += (freeBlocks - blocks) / alignBlocks * alignBlocks
+		}
+		p.FirstBlock = uint64(startAligned)
+		p.LastBlock = uint64(startAligned + blocks - 1)
+
+		newPartPos := -1
+		if !opts.keepEmptyEntries {
+			for i, part := range gpt.Partitions {
+				if part.IsUnused() {
+					newPartPos = i
+					break
+				}
+			}
+		}
+		if newPartPos == -1 {
+			gpt.Partitions = append(gpt.Partitions, p)
+		} else {
+			gpt.Partitions[newPartPos] = p
+		}
+		p.Section, err = blockdev.NewSection(gpt.b, int64(p.FirstBlock), int64(p.LastBlock)+1)
+		if err != nil {
+			return fmt.Errorf("failed to create blockdev Section for partition: %w", err)
+		}
+		return nil
+	}
+
+	return fmt.Errorf("no space for partition of %d blocks, largest continuous free space after alignment is %d blocks", blocks, largestFreeSpace)
 }
 
 // FirstUsableBlock returns the first usable (i.e. a partition can start there)
@@ -364,8 +360,8 @@
 	// Sort both sets of blocks independently in ascending order. Note that it
 	// is now no longer possible to extract the original intervals. Integers
 	// have no identity thus it doesn't matter if the sort is stable or not.
-	sort.Slice(startBlocks, func(i, j int) bool { return startBlocks[i] < startBlocks[j] })
-	sort.Slice(endBlocks, func(i, j int) bool { return endBlocks[i] < endBlocks[j] })
+	slices.Sort(startBlocks)
+	slices.Sort(endBlocks)
 
 	var freeSpaces [][2]int64
 
diff --git a/osbase/gpt/gpt_test.go b/osbase/gpt/gpt_test.go
index e662e69..6691e05 100644
--- a/osbase/gpt/gpt_test.go
+++ b/osbase/gpt/gpt_test.go
@@ -5,6 +5,7 @@
 	"crypto/sha256"
 	"io"
 	"os"
+	"strings"
 	"testing"
 
 	"github.com/google/uuid"
@@ -77,6 +78,179 @@
 	}
 }
 
+func TestAddPartition(t *testing.T) {
+	if os.Getenv("IN_KTEST") == "true" {
+		t.Skip("In ktest")
+	}
+	cases := []struct {
+		name        string
+		parts       []*Partition
+		addSize     int64
+		addOptions  []AddOption
+		expectErr   string
+		expectParts []*Partition
+	}{
+		{
+			name:        "empty-basic",
+			addSize:     9 * 512,
+			expectParts: []*Partition{{Name: "added", FirstBlock: 2048, LastBlock: 2048 + 9 - 1}},
+		},
+		{
+			name:        "empty-fill",
+			addSize:     -1,
+			expectParts: []*Partition{{Name: "added", FirstBlock: 2048, LastBlock: 5*2048 - 16384/512 - 2}},
+		},
+		{
+			name:        "empty-end",
+			addSize:     9 * 512,
+			addOptions:  []AddOption{WithPreferEnd()},
+			expectParts: []*Partition{{Name: "added", FirstBlock: 4 * 2048, LastBlock: 4*2048 + 9 - 1}},
+		},
+		{
+			name:       "empty-align0",
+			addSize:    9 * 512,
+			addOptions: []AddOption{WithAlignment(0)},
+			expectErr:  "must be positive",
+		},
+		{
+			name:        "empty-align512",
+			addSize:     9 * 512,
+			addOptions:  []AddOption{WithAlignment(512)},
+			expectParts: []*Partition{{Name: "added", FirstBlock: 2 + 16384/512, LastBlock: 2 + 16384/512 + 9 - 1}},
+		},
+		{
+			name:      "empty-zero-size",
+			addSize:   0,
+			expectErr: "must be positive",
+		},
+		{
+			name:        "full-fill",
+			parts:       []*Partition{{Name: "full", FirstBlock: 2048, LastBlock: 5*2048 - 16384/512 - 2}},
+			addSize:     -1,
+			expectErr:   "no free space",
+			expectParts: []*Partition{{Name: "full", FirstBlock: 2048, LastBlock: 5*2048 - 16384/512 - 2}},
+		},
+		{
+			name:    "haveone-basic",
+			parts:   []*Partition{{Name: "one", FirstBlock: 2048, LastBlock: 2048 + 5}},
+			addSize: 9 * 512,
+			expectParts: []*Partition{
+				{Name: "one", FirstBlock: 2048, LastBlock: 2048 + 5},
+				{Name: "added", FirstBlock: 2 * 2048, LastBlock: 2*2048 + 9 - 1},
+			},
+		},
+		{
+			name:    "havemiddle-basic",
+			parts:   []*Partition{{Name: "middle", FirstBlock: 2 * 2048, LastBlock: 2*2048 + 5}},
+			addSize: 9 * 512,
+			expectParts: []*Partition{
+				{Name: "middle", FirstBlock: 2 * 2048, LastBlock: 2*2048 + 5},
+				{Name: "added", FirstBlock: 2048, LastBlock: 2048 + 9 - 1},
+			},
+		},
+		{
+			name:       "havemiddle-end",
+			parts:      []*Partition{{Name: "middle", FirstBlock: 2 * 2048, LastBlock: 2*2048 + 5}},
+			addSize:    9 * 512,
+			addOptions: []AddOption{WithPreferEnd()},
+			expectParts: []*Partition{
+				{Name: "middle", FirstBlock: 2 * 2048, LastBlock: 2*2048 + 5},
+				{Name: "added", FirstBlock: 4 * 2048, LastBlock: 4*2048 + 9 - 1},
+			},
+		},
+		{
+			name:       "end-aligned",
+			parts:      []*Partition{{Name: "endalign", FirstBlock: 4 * 2048, LastBlock: 4*2048 + 8}},
+			addSize:    2048 * 512,
+			addOptions: []AddOption{WithPreferEnd()},
+			expectParts: []*Partition{
+				{Name: "endalign", FirstBlock: 4 * 2048, LastBlock: 4*2048 + 8},
+				{Name: "added", FirstBlock: 3 * 2048, LastBlock: 4*2048 - 1},
+			},
+		},
+		{
+			name: "emptyslots-basic",
+			parts: []*Partition{
+				{Name: "one", FirstBlock: 2048, LastBlock: 2048},
+				nil, nil,
+				{Name: "two", FirstBlock: 2048 + 1, LastBlock: 2048 + 1},
+			},
+			addSize: 9 * 512,
+			expectParts: []*Partition{
+				{Name: "one", FirstBlock: 2048, LastBlock: 2048},
+				{Name: "added", FirstBlock: 2 * 2048, LastBlock: 2*2048 + 9 - 1},
+				nil,
+				{Name: "two", FirstBlock: 2048 + 1, LastBlock: 2048 + 1},
+			},
+		},
+		{
+			name: "emptyslots-keep",
+			parts: []*Partition{
+				{Name: "one", FirstBlock: 2048, LastBlock: 2048},
+				nil, nil,
+				{Name: "two", FirstBlock: 2048 + 1, LastBlock: 2048 + 1},
+			},
+			addSize:    9 * 512,
+			addOptions: []AddOption{WithKeepEmptyEntries()},
+			expectParts: []*Partition{
+				{Name: "one", FirstBlock: 2048, LastBlock: 2048},
+				nil, nil,
+				{Name: "two", FirstBlock: 2048 + 1, LastBlock: 2048 + 1},
+				{Name: "added", FirstBlock: 2 * 2048, LastBlock: 2*2048 + 9 - 1},
+			},
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			for _, part := range c.parts {
+				if part != nil {
+					part.Type = PartitionTypeEFISystem
+				}
+			}
+			addPart := &Partition{Name: "added", Type: PartitionTypeEFISystem}
+			d := blockdev.MustNewMemory(512, 5*2048) // 5MiB
+			g, err := New(d)
+			if err != nil {
+				panic(err)
+			}
+			g.Partitions = c.parts
+			err = g.AddPartition(addPart, c.addSize, c.addOptions...)
+			if (err == nil) != (c.expectErr == "") || (err != nil && !strings.Contains(err.Error(), c.expectErr)) {
+				t.Errorf("expected %q, got %v", c.expectErr, err)
+			}
+			_, overlap, err := g.GetFreeSpaces()
+			if err != nil {
+				t.Fatal(err)
+			}
+			if overlap {
+				t.Errorf("partitions overlap")
+			}
+			if len(g.Partitions) != len(c.expectParts) {
+				t.Fatalf("expected %+v, got %+v", c.expectParts, g.Partitions)
+			}
+			for i := range g.Partitions {
+				gotPart, wantPart := g.Partitions[i], c.expectParts[i]
+				if (gotPart == nil) != (wantPart == nil) {
+					t.Errorf("partition %d: got %+v, expected %+v", i, gotPart, wantPart)
+				}
+				if gotPart == nil || wantPart == nil {
+					continue
+				}
+				if gotPart.Name != wantPart.Name {
+					t.Errorf("partition %d: got Name %q, expected %q", i, gotPart.Name, wantPart.Name)
+				}
+				if gotPart.FirstBlock != wantPart.FirstBlock {
+					t.Errorf("partition %d: got FirstBlock %d, expected %d", i, gotPart.FirstBlock, wantPart.FirstBlock)
+				}
+				if gotPart.LastBlock != wantPart.LastBlock {
+					t.Errorf("partition %d: got LastBlock %d, expected %d", i, gotPart.LastBlock, wantPart.LastBlock)
+				}
+			}
+		})
+	}
+}
+
 func TestRoundTrip(t *testing.T) {
 	if os.Getenv("IN_KTEST") == "true" {
 		t.Skip("In ktest")