Fix recoverypartition.Delete() and add a test (#145)

* Fix recoverypartition.Delete() and add a test

* recoverypartition.Delete(): fix typo and reword the error message

Co-authored-by: Tor Arne Vestbø <[email protected]>

---------

Co-authored-by: Tor Arne Vestbø <[email protected]>

builder/tart/recoverypartition/delete.go

@@ -10,6 +10,7 @@ import (
 )
 
 func Delete(diskImagePath string, ui packer.Ui, state multistep.StateBag) error {
+	// Open the disk image and read its partition table
 	disk, err := diskfs.Open(diskImagePath)
 	if err != nil {
 		return fmt.Errorf("failed to open the disk image: %w", err)
@@ -19,8 +20,9 @@ func Delete(diskImagePath string, ui packer.Ui, state multistep.StateBag) error
 
 	partitionTable, err := disk.GetPartitionTable()
 	if err != nil {
+		// Disk may not be initialized with a partition table yet
+		// when running on a freshly created Linux VMs, for example
 		if err.Error() == "unknown disk partition type" {
-			// Disk may not be initialized with a partition table yet
 			return nil
 		}
 
@@ -29,29 +31,44 @@ func Delete(diskImagePath string, ui packer.Ui, state multistep.StateBag) error
 
 	gptTable := partitionTable.(*gpt.Table)
 
-	for i, partition := range gptTable.Partitions {
+	recoveryPartitionIdx := -1
+
+	for idx, partition := range gptTable.Partitions {
 		if partition.Name != Name {
 			continue
 		}
 
-		ui.Say("Found recovery partition. Let's remove it to save space...")
-
-		// there are max 128 partitions and we probably on the third one
-		// the rest are just empty structs so let's reuse them
-		gptTable.Partitions[i] = gptTable.Partitions[len(gptTable.Partitions)-1]
-
-		if err = disk.Partition(gptTable); err != nil {
-			return fmt.Errorf("failed to write the new partition table: %w", err)
+		if recoveryPartitionIdx != -1 {
+			return fmt.Errorf("found a recovery partition at GPT entry %d, but there's another recovery "+
+				"partition at GPT entry %d, refusing to proceed", idx+1, recoveryPartitionIdx+1)
 		}
 
-		ui.Say("Successfully updated partitions...")
+		recoveryPartitionIdx = idx
+	}
 
-		state.Put(statekey.DiskChanged, true)
+	if recoveryPartitionIdx == -1 {
+		ui.Say("No recovery partition was found, assuming that it was already deleted.")
 
 		return nil
 	}
 
-	ui.Say("No recovery partition was found, assuming that it was already deleted.")
+	if recoveryPartitionIdx != len(gptTable.Partitions)-1 {
+		return fmt.Errorf("found a recovery partition at GPT entry %d, but it's "+
+			"not the last partition on the disk, refusing to proceed", recoveryPartitionIdx+1)
+	}
+
+	ui.Say(fmt.Sprintf("Found a recovery partition at GPT entry %d, let's remove it "+
+		"to save space and allow for resizing the main partition...", recoveryPartitionIdx+1))
+
+	gptTable.Partitions = gptTable.Partitions[:recoveryPartitionIdx]
+
+	if err = disk.Partition(gptTable); err != nil {
+		return fmt.Errorf("failed to write the new partition table: %w", err)
+	}
+
+	ui.Say("Successfully updated partitions!")
+
+	state.Put(statekey.DiskChanged, true)
 
 	return nil
 }

builder/tart/recoverypartition/delete_test.go

@@ -0,0 +1,68 @@
+package recoverypartition_test
+
+import (
+	"github.com/diskfs/go-diskfs"
+	"github.com/diskfs/go-diskfs/partition/gpt"
+	"github.com/hashicorp/packer-plugin-sdk/multistep"
+	"github.com/hashicorp/packer-plugin-sdk/packer"
+	"github.com/stretchr/testify/require"
+	"os"
+	"packer-plugin-tart/builder/tart/recoverypartition"
+	"path/filepath"
+	"testing"
+)
+
+func TestDelete(t *testing.T) {
+	// Create a disk
+	const diskSizeBytes = 1 * 1024 * 1024
+
+	diskPath := filepath.Join(t.TempDir(), "disk.img")
+
+	diskFile, err := os.Create(diskPath)
+	require.NoError(t, err)
+	require.NoError(t, diskFile.Truncate(diskSizeBytes))
+	require.NoError(t, diskFile.Close())
+
+	// Partition our disk as GPT with a macOS recovery partition
+	const sectorSizeBytes = 512
+	const partitionSizeSectors = 5
+	const partitionSizeBytes = partitionSizeSectors * sectorSizeBytes
+
+	firstPartition := &gpt.Partition{
+		Start: 34,
+		Size:  partitionSizeBytes,
+		Type:  gpt.AppleAPFS,
+		Name:  "Doesn't matter",
+	}
+	secondPartition := &gpt.Partition{
+		Start: 34 + partitionSizeSectors,
+		Size:  partitionSizeBytes,
+		Type:  gpt.AppleAPFS,
+		Name:  recoverypartition.Name,
+	}
+	gptTable := &gpt.Table{
+		LogicalSectorSize:  sectorSizeBytes,
+		PhysicalSectorSize: sectorSizeBytes,
+		Partitions: []*gpt.Partition{
+			firstPartition,
+			secondPartition,
+		},
+	}
+	oldDisk, err := diskfs.Open(diskPath)
+	require.NoError(t, err)
+	require.NoError(t, oldDisk.Partition(gptTable))
+
+	// Delete the recovery partition
+	require.NoError(t, recoverypartition.Delete(diskPath, packer.TestUi(t), &multistep.BasicStateBag{}))
+
+	// Ensure that the recovery partition was deleted
+	disk, err := diskfs.Open(diskPath)
+	require.NoError(t, err)
+
+	partitionTable, err := disk.GetPartitionTable()
+	require.NoError(t, err)
+
+	partitions := partitionTable.(*gpt.Table).Partitions
+	require.Len(t, partitions, 1)
+	require.Equal(t, "Doesn't matter", partitions[0].Name)
+}