Add --preserve-numbers to keep partition numbers

cmd/resizer/main.go

@@ -16,6 +16,7 @@ var rootCmd = func() *cobra.Command {
 		growPartitions  []string
 		fixErrors       bool
 		dryRun          bool
+		preserveNumbers bool
 	)
 	cmd := &cobra.Command{
 		Use:   "resizer",
@@ -71,7 +72,7 @@ var rootCmd = func() *cobra.Command {
 			if len(args) > 0 {
 				disk = args[0]
 			}
-			if err := resizer.Run(disk, &shrinkPartitionParsed, growPartitionsParsed, fixErrors, dryRun); err != nil {
+			if err := resizer.Run(disk, &shrinkPartitionParsed, growPartitionsParsed, fixErrors, dryRun, preserveNumbers); err != nil {
 				log.Fatalf("Resize operation failed: %v", err)
 			}
 		},
@@ -80,6 +81,7 @@ var rootCmd = func() *cobra.Command {
 	cmd.Flags().StringSliceVar(&growPartitions, "grow-partition", []string{}, "Partitions to grow, along with their desired sizes, in format identifier:partition:size, see help (e.g. name:sda1:20G or label:EFI System:100M)")
 	cmd.Flags().BoolVar(&dryRun, "dry-run", false, "If set, will only simulate the resize operations without making any changes")
 	cmd.Flags().BoolVar(&fixErrors, "fix-errors", false, "If set, will attempt to fix any ext4 filesystem errors found during fsck before shrinking")
+	cmd.Flags().BoolVar(&preserveNumbers, "preserve-numbers", false, "If set, a grown partition that is relocated is renumbered back to its original partition number, so labels keep their original partition numbers (e.g. /dev/sda2)")
 	return cmd
 }
 

resize.go

@@ -16,8 +16,11 @@ type copyData struct {
 	err   error
 }
 
-// resize performs the actual resize operations on the given disk
-func resize(d *disk.Disk, resizes []partitionResizeTarget, fixErrors bool) error {
+// resize performs the actual resize operations on the given disk.
+// When preserveNumbers is set, a relocated partition is renumbered back to its
+// original partition number after the copy, so that consumers referencing a
+// partition by number (e.g. boot loaders) continue to find it.
+func resize(d *disk.Disk, resizes []partitionResizeTarget, fixErrors, preserveNumbers bool) error {
 	// do any shrinks first
 	// this is idempotent. If I have a 500MB partition with a 500MB filesystem,
 	// and shrink it to 400MB. If I stop, and then run it again, it will just say
@@ -54,9 +57,16 @@ func resize(d *disk.Disk, resizes []partitionResizeTarget, fixErrors bool) error
 		return err
 	}
 
-	// remove the old partitions
-	if err := removePartitions(d, resizes); err != nil {
-		return err
+	// remove the old partitions, optionally renumbering the relocated partitions
+	// back to their original partition numbers
+	if preserveNumbers {
+		if err := renumberPartitions(d, resizes); err != nil {
+			return err
+		}
+	} else {
+		if err := removePartitions(d, resizes); err != nil {
+			return err
+		}
 	}
 
 	return nil
@@ -219,6 +229,67 @@ func removePartitions(d *disk.Disk, resizes []partitionResizeTarget) error {
 	return nil
 }
 
+// renumberPartitions removes the original partitions and reassigns each relocated
+// target partition's GPT slot index to the original partition's number, so the
+// resized partition keeps the same partition number it had before. This is the
+// preserve-numbers counterpart to removePartitions, and must run after the data
+// has been copied and the identities swapped onto the target partitions.
+//
+// The renumbered entry stays at its new on-disk offset, so the resulting GPT entries
+// are no longer in disk-offset order. That is permitted by the GPT specification and
+// is invisible to consumers that locate a partition by its number (e.g. a boot loader
+// referencing (hd0,gptN)); only tools that expect entries sorted by offset will note it.
+func renumberPartitions(d *disk.Disk, resizes []partitionResizeTarget) error {
+	tableRaw, err := d.GetPartitionTable()
+	if err != nil {
+		return err
+	}
+	table, ok := tableRaw.(*gpt.Table)
+	if !ok {
+		return fmt.Errorf("unsupported partition table type, only GPT is supported")
+	}
+	// map partition number -> position in the slice, captured before any mutation so
+	// that the lookups below are unaffected by the index reassignments we make.
+	indexToPosition := make(map[int]int)
+	for i, p := range table.Partitions {
+		indexToPosition[p.Index] = i
+	}
+	// slice positions of the original partitions to drop, keyed by position rather
+	// than partition number: once we reassign a target's Index to the original number,
+	// keying removal on Index would also match (and wrongly drop) the renumbered target.
+	removePositions := make(map[int]bool)
+	for _, r := range resizes {
+		if r.original.number == r.target.number {
+			log.Printf("partition %d %s: no change in partition number, no need to renumber", r.original.number, r.original.label)
+			continue
+		}
+		origPos, ok := indexToPosition[r.original.number]
+		if !ok {
+			return fmt.Errorf("original partition %d not found in partition table", r.original.number)
+		}
+		targetPos, ok := indexToPosition[r.target.number]
+		if !ok {
+			return fmt.Errorf("target partition %d not found in partition table", r.target.number)
+		}
+		log.Printf("renumbering partition %d -> %d (label %s) and removing original slot", r.target.number, r.original.number, r.original.label)
+		table.Partitions[targetPos].Index = r.original.number
+		removePositions[origPos] = true
+	}
+	// rebuild the slice, dropping the vacated original slots so their numbers are free
+	partitions := make([]*gpt.Partition, 0, len(table.Partitions))
+	for i, p := range table.Partitions {
+		if removePositions[i] {
+			continue
+		}
+		partitions = append(partitions, p)
+	}
+	table.Partitions = partitions
+	if err := d.Partition(table); err != nil {
+		return fmt.Errorf("failed to write renumbered partition table: %v", err)
+	}
+	return nil
+}
+
 // swapPartitions swaps the labels, Type GUIDs, and UUIDs of the original and target partitions,
 // as well as any attributes flags.
 func swapPartitions(d *disk.Disk, resizes []partitionResizeTarget) error {

resize_test.go

@@ -236,6 +236,111 @@ func TestRemovePartitions(t *testing.T) {
 	}
 }
 
+func TestRenumberPartitions(t *testing.T) {
+	// Model the state after createPartitions+copy+swap for two grown partitions:
+	// originals 2 and 3 (now carrying throwaway identities) plus their relocated
+	// copies in slots 5 and 6 (carrying the real labels). renumberPartitions must
+	// drop the originals and move the relocated copies back to numbers 2 and 3.
+	workDir := t.TempDir()
+	f, err := os.CreateTemp(workDir, "disk.img")
+	if err != nil {
+		t.Fatalf("failed to create temp disk image: %v", err)
+	}
+	if err := os.Truncate(f.Name(), 1*GB); err != nil {
+		t.Fatalf("failed to truncate disk image: %v", err)
+	}
+	defer func() { _ = f.Close() }()
+
+	backend := file.New(f, false)
+	d, err := diskfs.OpenBackend(backend, diskfs.WithOpenMode(diskfs.ReadWrite))
+	if err != nil {
+		t.Fatalf("failed to open disk: %v", err)
+	}
+	var offset uint64 = 2048
+	table := &gpt.Table{
+		Partitions: []*gpt.Partition{
+			{Index: 1, Start: offset, Size: 36 * MB, Type: gpt.LinuxFilesystem, Name: "part1"},
+			// originals, to be dropped
+			{Index: 2, Start: offset + 36*MB, Size: 100 * MB, Type: gpt.LinuxFilesystem, Name: "IMGA_old"},
+			{Index: 3, Start: offset + 36*MB + 100*MB, Size: 50 * MB, Type: gpt.LinuxFilesystem, Name: "DATA_old"},
+			{Index: 4, Start: offset + 36*MB + 100*MB + 50*MB, Size: 36 * MB, Type: gpt.LinuxFilesystem, Name: "part4"},
+			// relocated copies carrying the real identities, to be renumbered to 2 and 3
+			{Index: 5, Start: offset + 36*MB + 100*MB + 50*MB + 36*MB, Size: 300 * MB, Type: gpt.LinuxFilesystem, Name: "IMGA"},
+			{Index: 6, Start: offset + 36*MB + 100*MB + 50*MB + 36*MB + 300*MB, Size: 200 * MB, Type: gpt.LinuxFilesystem, Name: "DATA"},
+		},
+	}
+	if err := d.Partition(table); err != nil {
+		t.Fatalf("failed to write partition table: %v", err)
+	}
+	// capture the relocated copies' on-disk starts so we can confirm renumbering keeps
+	// the data in place and only changes the slot number
+	wantStart := make(map[string]uint64)
+	for _, p := range table.Partitions {
+		wantStart[p.Name] = p.Start
+	}
+
+	resizes := []partitionResizeTarget{
+		{
+			original: partitionData{number: 2, label: "IMGA"},
+			target:   partitionData{number: 5, label: "IMGA"},
+		},
+		{
+			original: partitionData{number: 3, label: "DATA"},
+			target:   partitionData{number: 6, label: "DATA"},
+		},
+	}
+
+	if err := renumberPartitions(d, resizes); err != nil {
+		t.Fatalf("renumberPartitions failed: %v", err)
+	}
+
+	tableRaw, err := d.GetPartitionTable()
+	if err != nil {
+		t.Fatalf("failed to get partition table: %v", err)
+	}
+	newTable, ok := tableRaw.(*gpt.Table)
+	if !ok {
+		t.Fatalf("unsupported partition table type, only GPT is supported")
+	}
+
+	byIndex := make(map[int]*gpt.Partition)
+	for _, p := range newTable.Partitions {
+		if p.Type == gpt.Unused {
+			continue
+		}
+		byIndex[p.Index] = p
+	}
+
+	// originals are gone, relocated copies now own numbers 2 and 3
+	wantByNumber := map[int]string{1: "part1", 2: "IMGA", 3: "DATA", 4: "part4"}
+	if len(byIndex) != len(wantByNumber) {
+		t.Fatalf("expected %d partitions after renumber, got %d", len(wantByNumber), len(byIndex))
+	}
+	for number, name := range wantByNumber {
+		p, ok := byIndex[number]
+		if !ok {
+			t.Fatalf("expected partition number %d (%s) to exist after renumber", number, name)
+		}
+		if p.Name != name {
+			t.Errorf("partition %d: expected label %q, got %q", number, name, p.Name)
+		}
+	}
+	// the renumbered partitions must keep the relocated copies' on-disk location
+	if got := byIndex[2].Start; got != wantStart["IMGA"] {
+		t.Errorf("renumbered IMGA start moved: expected %d, got %d", wantStart["IMGA"], got)
+	}
+	if got := byIndex[3].Start; got != wantStart["DATA"] {
+		t.Errorf("renumbered DATA start moved: expected %d, got %d", wantStart["DATA"], got)
+	}
+	// the old high-numbered slots must be free
+	if _, ok := byIndex[5]; ok {
+		t.Errorf("slot 5 should have been vacated by renumbering")
+	}
+	if _, ok := byIndex[6]; ok {
+		t.Errorf("slot 6 should have been vacated by renumbering")
+	}
+}
+
 func TestCopyFilesystems(t *testing.T) {
 	// create a duplicate disk with a partition with the specified filesystem type
 	tmpdir := t.TempDir()

run.go

@@ -15,7 +15,9 @@ import (
 // if it has an identifiable ext4 filesystem to shrink, and there is enough space to shrink it.
 // It always will try to run e2fsck before shrinking. By default, it will not fix any found errors, in which case it will
 // error out if any filesystem errors are found. If fixErrors is true, it will attempt to fix any found errors.
-func Run(disk string, shrinkPartition *PartitionIdentifier, growPartitions []PartitionChange, fixErrors, dryRun bool) error {
+// If preserveNumbers is true, any partition that is relocated while growing is renumbered back to its original
+// partition number once the data has been copied, so its partition number (e.g. /dev/sda2) is unchanged by the resize.
+func Run(disk string, shrinkPartition *PartitionIdentifier, growPartitions []PartitionChange, fixErrors, dryRun, preserveNumbers bool) error {
 	// we always work solely with partition UUIDs internally, so convert any other identifiers to UUIDs
 	// see if a disk was specified
 	// no disk specified, try to discover
@@ -75,5 +77,5 @@ func Run(disk string, shrinkPartition *PartitionIdentifier, growPartitions []Par
 		return nil
 	}
 	log.Printf("Will perform resizes %+v", resizes)
-	return resize(d, resizes, fixErrors)
+	return resize(d, resizes, fixErrors, preserveNumbers)
 }

run_test.go

@@ -16,6 +16,18 @@ const (
 )
 
 func TestRun(t *testing.T) {
+	for _, preserveNumbers := range []bool{false, true} {
+		name := "renumber"
+		if preserveNumbers {
+			name = "preserveNumbers"
+		}
+		t.Run(name, func(t *testing.T) {
+			runResizeTest(t, preserveNumbers)
+		})
+	}
+}
+
+func runResizeTest(t *testing.T, preserveNumbers bool) {
 	tmpDir := t.TempDir()
 	tmpFile := filepath.Join(tmpDir, "diskfull.img")
 	if err := testCopyFile(diskfullImg, tmpFile); err != nil {
@@ -38,10 +50,14 @@ func TestRun(t *testing.T) {
 	}
 	table0 := tableRaw0.(*gpt.Table)
 	var origShrinkSize int64
+	origNumber := map[string]int{}
 	for _, p := range table0.Partitions {
+		if p.Type == gpt.Unused {
+			continue
+		}
+		origNumber[p.Name] = int(p.Index)
 		if p.Name == "shrinker" {
 			origShrinkSize = int64(p.GetSize())
-			break
 		}
 	}
 	if origShrinkSize == 0 {
@@ -54,7 +70,7 @@ func TestRun(t *testing.T) {
 		NewPartitionChange(IdentifierByLabel, "partb", 2*GB),
 		NewPartitionChange(IdentifierByLabel, "ESP", 1*GB),
 	}
-	if err := Run(tmpFile, &shrink, growList, false, false); err != nil {
+	if err := Run(tmpFile, &shrink, growList, false, false, preserveNumbers); err != nil {
 		t.Fatalf("Run failed: %v", err)
 	}
 
@@ -114,6 +130,11 @@ func TestRun(t *testing.T) {
 		default:
 			t.Errorf("unexpected active partition %q", name)
 		}
+		// with preserveNumbers, every partition (including the relocated grown ones)
+		// must keep the number it had before the resize
+		if preserveNumbers && int(p.Index) != origNumber[name] {
+			t.Errorf("%s partition number = %d, want %d (preserveNumbers)", name, p.Index, origNumber[name])
+		}
 		seen[name] = true
 	}
 	for _, n := range []string{"shrinker", "parta", "partb", "ESP"} {