ensure GPT partition entries are at correct index

disk/disk_test.go

@@ -124,7 +124,7 @@ func TestPartition(t *testing.T) {
 		partitionEnd := partitionSize + partitionStart - 1
 		table := &gpt.Table{
 			Partitions: []*gpt.Partition{
-				{Start: partitionStart, End: partitionEnd, Type: gpt.EFISystemPartition, Name: "EFI System"},
+				{Index: 1, Start: partitionStart, End: partitionEnd, Type: gpt.EFISystemPartition, Name: "EFI System"},
 			},
 			LogicalSectorSize: sectorSize,
 		}

partition/gpt/table.go

@@ -225,18 +225,34 @@ func (t *Table) generateProtectiveMBR() []byte {
 func (t *Table) toPartitionArrayBytes() ([]byte, error) {
 	blocksize := uint64(t.LogicalSectorSize)
 
-	// go through the partitions, make sure Start/End/Size are correct, and each has a GUID
+	// go through the partitions, make sure Start/End/Size are correct, and each has a GUID.
+	// In addition, the Partition slice could be in order, or not, e.g. it might have partitions 1,3,4,7 only, yet the
+	// slice will have 4 positions. Or it could be out of order. So we need to write out the partition entries in
+	// order, and fill in blanks.
+	partMap := make(map[int]*Partition)
 	for i, part := range t.Partitions {
 		err := part.initEntry(blocksize)
 		if err != nil {
 			return nil, fmt.Errorf("could not initialize partition %d correctly: %v", i, err)
 		}
+		if part.Index < 1 || part.Index > t.partitionArraySize {
+			return nil, fmt.Errorf("partition %d has invalid index %d for partition array size %d", i, part.Index, t.partitionArraySize)
+		}
+		if _, exists := partMap[part.Index]; exists {
+			return nil, fmt.Errorf("duplicate partition index %d found", part.Index)
+		}
+		partMap[part.Index] = part
 	}
 
 	// generate the partition bytes
 	partSize := t.partitionEntrySize * uint32(t.partitionArraySize)
 	bpart := make([]byte, partSize)
-	for i, p := range t.Partitions {
+	for i := 0; i < t.partitionArraySize; i++ {
+		p, ok := partMap[i+1]
+		if !ok {
+			// unused partition
+			continue
+		}
 		// write the primary partition entry
 		b2, err := p.toBytes()
 		if err != nil {

partition/gpt/table_internal_test.go

@@ -1,6 +1,7 @@
 package gpt
 
 import (
+	"bytes"
 	"crypto/rand"
 	"fmt"
 	"io"
@@ -46,6 +47,186 @@ func GetValidTable() *Table {
 	return &table
 }
 
+func TestToPartitionArrayBytes(t *testing.T) {
+	t.Run("empty partition array", func(t *testing.T) {
+		table := GetValidTable()
+		table.Partitions = []*Partition{}
+		b, err := table.toPartitionArrayBytes()
+		if err != nil {
+			t.Errorf("unexpected error: %v", err)
+		}
+		expectedSize := int(table.partitionEntrySize) * table.partitionArraySize
+		if len(b) != expectedSize {
+			t.Errorf("unexpected byte slice size %d, expected %d", len(b), expectedSize)
+		}
+		for i := 0; i < table.partitionArraySize; i++ {
+			offset := i * int(table.partitionEntrySize)
+			entryBytes := b[offset : offset+int(table.partitionEntrySize)]
+			emptyEntry := make([]byte, table.partitionEntrySize)
+			if !bytes.Equal(entryBytes, emptyEntry) {
+				t.Errorf("expected empty partition entry at index %d, got %v", i, entryBytes)
+			}
+		}
+	})
+	t.Run("null partition array", func(t *testing.T) {
+		table := GetValidTable()
+		table.Partitions = nil
+		b, err := table.toPartitionArrayBytes()
+		if err != nil {
+			t.Errorf("unexpected error: %v", err)
+		}
+		expectedSize := int(table.partitionEntrySize) * table.partitionArraySize
+		if len(b) != expectedSize {
+			t.Errorf("unexpected byte slice size %d, expected %d", len(b), expectedSize)
+		}
+		for i := 0; i < table.partitionArraySize; i++ {
+			offset := i * int(table.partitionEntrySize)
+			entryBytes := b[offset : offset+int(table.partitionEntrySize)]
+			emptyEntry := make([]byte, table.partitionEntrySize)
+			if !bytes.Equal(entryBytes, emptyEntry) {
+				t.Errorf("expected empty partition entry at index %d, got %v", i, entryBytes)
+			}
+		}
+	})
+	t.Run("slice with duplicate indexes", func(t *testing.T) {
+		table := GetValidTable()
+		table.Partitions = []*Partition{
+			{
+				Index: 1,
+				Start: 4000,
+				End:   5000,
+				Size:  (5000 - 4000 + 1) * 512,
+				Name:  "Duplicate Index 1",
+			},
+			{
+				Index: 1,
+				Start: 6000,
+				End:   7000,
+				Size:  (7000 - 6000 + 1) * 512,
+				Name:  "Duplicate Index 1 Again",
+			},
+		}
+		_, err := table.toPartitionArrayBytes()
+		if err == nil {
+			t.Error("expected error due to duplicate partition indexes, got nil")
+		}
+	})
+	t.Run("negative index", func(t *testing.T) {
+		table := GetValidTable()
+		table.Partitions = []*Partition{
+			{
+				Index: -1,
+				Start: 4000,
+				End:   5000,
+				Size:  (5000 - 4000 + 1) * 512,
+				Name:  "Index -1",
+			},
+		}
+		_, err := table.toPartitionArrayBytes()
+		if err == nil {
+			t.Error("expected error due to negative partition index, got nil")
+		}
+	})
+	t.Run("zero index", func(t *testing.T) {
+		table := GetValidTable()
+		table.Partitions = []*Partition{
+			{
+				Index: 0,
+				Start: 4000,
+				End:   5000,
+				Size:  (5000 - 4000 + 1) * 512,
+				Name:  "Index 0",
+			},
+		}
+		_, err := table.toPartitionArrayBytes()
+		if err == nil {
+			t.Error("expected error due to zero partition index, got nil")
+		}
+	})
+	t.Run("index too large", func(t *testing.T) {
+		table := GetValidTable()
+		table.Partitions = []*Partition{
+			{
+				Index: table.partitionArraySize + 1,
+				Start: 4000,
+				End:   5000,
+				Size:  (5000 - 4000 + 1) * 512,
+				Name:  "Index Too Large",
+			},
+		}
+		_, err := table.toPartitionArrayBytes()
+		if err == nil {
+			t.Error("expected error due to partition index too large, got nil")
+		}
+	})
+	t.Run("normal slice in order", func(t *testing.T) {
+		table := GetValidTable()
+		table.Partitions = []*Partition{
+			{
+				Index: 1,
+				Start: 4000,
+				End:   5000,
+				Size:  (5000 - 4000 + 1) * 512,
+				Name:  "Index 1",
+				Type:  LinuxFilesystem,
+			},
+			{
+				Index: 2,
+				Start: 6000,
+				End:   7000,
+				Size:  (7000 - 6000 + 1) * 512,
+				Name:  "Index 2",
+				Type:  LinuxFilesystem,
+			},
+			{
+				Index: 3,
+				Start: 8000,
+				End:   9000,
+				Size:  (9000 - 8000 + 1) * 512,
+				Name:  "Index 3",
+				Type:  LinuxFilesystem,
+			},
+		}
+		_, err := table.toPartitionArrayBytes()
+		if err != nil {
+			t.Errorf("unexpected error: %v", err)
+		}
+	})
+	t.Run("normal slice out of order", func(t *testing.T) {
+		table := GetValidTable()
+		table.Partitions = []*Partition{
+			{
+				Index: 3,
+				Start: 8000,
+				End:   9000,
+				Size:  (9000 - 8000 + 1) * 512,
+				Name:  "Index 3",
+				Type:  LinuxFilesystem,
+			},
+			{
+				Index: 8,
+				Start: 4000,
+				End:   5000,
+				Size:  (5000 - 4000 + 1) * 512,
+				Name:  "Index 8",
+				Type:  LinuxFilesystem,
+			},
+			{
+				Index: 2,
+				Start: 6000,
+				End:   7000,
+				Size:  (7000 - 6000 + 1) * 512,
+				Name:  "Index 2",
+				Type:  LinuxFilesystem,
+			},
+		}
+		_, err := table.toPartitionArrayBytes()
+		if err != nil {
+			t.Errorf("unexpected error: %v", err)
+		}
+	})
+}
+
 func TestTableFromBytes(t *testing.T) {
 	t.Run("Short byte slice", func(t *testing.T) {
 		b := make([]byte, 512+512-1)
@@ -270,6 +451,7 @@ func TestRepairVerify(t *testing.T) {
 	table := &Table{
 		Partitions: []*Partition{
 			{
+				Index: 1,
 				Start: 2048,
 				End:   sizeBefore,
 				Type:  LinuxFilesystem,

partition/gpt/table_test.go

@@ -282,7 +282,7 @@ func TestTableWrite(t *testing.T) {
 		name := "EFI System Tester"
 		table := &gpt.Table{
 			Partitions: []*gpt.Partition{
-				{Start: partitionStart, End: partitionEnd, Type: gpt.EFISystemPartition, Name: name},
+				{Index: 1, Start: partitionStart, End: partitionEnd, Type: gpt.EFISystemPartition, Name: name},
 			},
 			LogicalSectorSize: sectorSize,
 			ProtectiveMBR:     true,