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_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")