Relocate Apple Silicon APFS recovery partition to the end of a macOS disk image for safe APFS container expansion.

RelocatePartition.go

// RelocatePartition is a standalone utility for relocating the outer
// Apple_APFS_Recovery GPT partition (RecoveryOSContainer) to the end of an
// Apple Silicon macOS disk image so that the main Apple_APFS container can be
// expanded into the freed space.
//
// Origin:
//   - Based on Cirrus Labs' recovery relocation logic from:
//     https://github.com/cirruslabs/packer-plugin-tart/blob/main/builder/tart/recoverypartition/relocate.go
//
// What this fork adds:
//   - Matches the recovery partition by both GPT name and GPT type GUID.
//   - Supports either a whole-disk device (`/dev/diskN`) or an image path.
//   - Uses `diskutil image attach --noMount` for image files, which works with
//     newer Apple disk image formats that are not directly readable as raw GPT.
//   - Avoids the macOS-incompatible partition-table re-read ioctl used by
//     `go-diskfs` after writing GPT.
//   - Detects and resumes an interrupted relocation where GPT was updated but
//     the recovery container contents were not yet copied to the new location.
//   - Prints actionable errors instead of silently exiting on failure.
//
// Typical workflow:
//  1. Run this tool on a macOS VM disk image or on an already attached
//     `/dev/diskN`.
//  2. Confirm the resulting outer GPT layout looks like:
//     `Apple_APFS_ISC`, `Apple_APFS`, free space, `Apple_APFS_Recovery`.
//  3. Expand the main APFS container with:
//     `diskutil apfs resizeContainer diskNs2 0`
//
// Build:
// 	go mod init RelocatePartition.go
// 	go mod tidy
// 	go build RelocatePartition.go
//
// Usage:
// 	./RelocatePartition path_to_disk_image
// 	./RelocatePartition /dev/diskN
//
// Notes:
//   - Always keep a backup before modifying GPT.
//   - This tool moves the outer recovery partition, not the APFS volumes inside
//     the main container.
//   - If a previous run stopped after writing GPT, rerunning this tool on the
//     same image/device will attempt to finish restoring the recovery contents.

package main

import (
	"bytes"
	"fmt"
	"github.com/diskfs/go-diskfs"
	"github.com/diskfs/go-diskfs/disk"
	"github.com/diskfs/go-diskfs/partition/gpt"
	"io"
	"os"
	"os/exec"
	"regexp"
	"strings"
)

const (
	RecoveryPartitionName = "RecoveryOSContainer"
	// Apple Silicon RecoveryOS container GUID, shown by diskutil as Apple_APFS_Recovery.
	RecoveryPartitionType = gpt.Type("52637672-7900-11AA-AA11-00306543ECAC")
	APFSContainerMagic    = "NXSB"
	APFSContainerOffset   = 32
)

var wholeDiskDevicePattern = regexp.MustCompile(`^/dev/disk\d+$`)

func isRecoveryPartition(partition *gpt.Partition) bool {
	if partition == nil {
		return false
	}

	return partition.Name == RecoveryPartitionName || partition.Type == RecoveryPartitionType
}

func findRecoveryPartition(partitions []*gpt.Partition) (*gpt.Partition, int, bool) {
	for i, partition := range partitions {
		if isRecoveryPartition(partition) {
			return partition, i, true
		}
	}

	return nil, -1, false
}

func isLastPartitionOnDisk(partitions []*gpt.Partition, partition *gpt.Partition) bool {
	if partition == nil {
		return false
	}

	for _, candidate := range partitions {
		if candidate == nil || candidate == partition {
			continue
		}

		if candidate.End > partition.End {
			return false
		}
	}

	return true
}

func relocatePartitionBounds(partition *gpt.Partition, lastSectorAvailableForPartitions uint64) {
	recoveryPartitionSizeInSectors := partition.End - partition.Start

	partition.Start = lastSectorAvailableForPartitions - recoveryPartitionSizeInSectors
	partition.End = lastSectorAvailableForPartitions
}

func sectorOffset(sector uint64, logicalBlockSize int64) int64 {
	return int64(sector) * logicalBlockSize
}

func partitionSizeInSectors(partition *gpt.Partition) uint64 {
	if partition == nil {
		return 0
	}

	return partition.End - partition.Start + 1
}

func writeGPTWithoutKernelReread(disk *disk.Disk, table *gpt.Table) error {
	writableBackend, err := disk.Backend.Writable()
	if err != nil {
		return fmt.Errorf("failed to get writable disk backend: %w", err)
	}

	if err := table.Write(writableBackend, disk.Size); err != nil {
		return fmt.Errorf("failed to write partition table: %w", err)
	}

	disk.Table = table

	return nil
}

func lastUsableSector(disk *disk.Disk, table *gpt.Table) uint64 {
	lastSectorAvailableForPartitions := table.LastDataSector()
	if lastSectorAvailableForPartitions != 0 {
		return lastSectorAvailableForPartitions
	}

	// Fallback for unexpected GPT metadata.
	return uint64((disk.Size / disk.LogicalBlocksize) - 34)
}

func isWholeDiskDevice(path string) bool {
	return wholeDiskDevicePattern.MatchString(path)
}

func parseAttachedWholeDiskDevice(output string) (string, error) {
	for _, line := range strings.Split(output, "\n") {
		fields := strings.Fields(line)
		if len(fields) == 0 {
			continue
		}

		candidate := fields[0]
		if isWholeDiskDevice(candidate) {
			return candidate, nil
		}
	}

	return "", fmt.Errorf("failed to find an attached whole-disk device in output:\n%s", output)
}

func attachImageAsDevice(imagePath string) (string, func(), error) {
	cmd := exec.Command("diskutil", "image", "attach", "--noMount", imagePath)
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		return "", nil, fmt.Errorf("failed to attach disk image: %w: %s", err, strings.TrimSpace(stderr.String()))
	}

	devicePath, err := parseAttachedWholeDiskDevice(stdout.String())
	if err != nil {
		return "", nil, err
	}

	cleanup := func() {
		if err := detachImageDevice(devicePath); err != nil {
			fmt.Fprintf(os.Stderr, "Warning: failed to detach %s: %v\n", devicePath, err)
		}
	}

	return devicePath, cleanup, nil
}

func detachImageDevice(devicePath string) error {
	cmd := exec.Command("diskutil", "eject", devicePath)
	output, err := cmd.CombinedOutput()
	if err == nil {
		return nil
	}

	unmountCmd := exec.Command("diskutil", "unmountDisk", "force", devicePath)
	unmountOutput, unmountErr := unmountCmd.CombinedOutput()
	if unmountErr != nil {
		return fmt.Errorf(
			"failed to eject %s: %w: %s; failed to force-unmount: %v: %s",
			devicePath,
			err,
			strings.TrimSpace(string(output)),
			unmountErr,
			strings.TrimSpace(string(unmountOutput)),
		)
	}

	cmd = exec.Command("diskutil", "eject", devicePath)
	output, err = cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("%w: %s", err, strings.TrimSpace(string(output)))
	}

	return nil
}

func resolveDiskPath(path string) (string, func(), error) {
	if isWholeDiskDevice(path) {
		return path, func() {}, nil
	}

	return attachImageAsDevice(path)
}

func inferPreviousPartitionEnd(partitions []*gpt.Partition, recovery *gpt.Partition) (uint64, bool) {
	var (
		maxEnd uint64
		found  bool
	)

	for _, partition := range partitions {
		if partition == nil || partition == recovery {
			continue
		}

		if partition.End < recovery.Start && (!found || partition.End > maxEnd) {
			maxEnd = partition.End
			found = true
		}
	}

	return maxEnd, found
}

func inferOriginalRecoveryStart(partitions []*gpt.Partition, recovery *gpt.Partition) (uint64, bool) {
	previousPartitionEnd, ok := inferPreviousPartitionEnd(partitions, recovery)
	if !ok {
		return 0, false
	}

	return previousPartitionEnd + 1, true
}

func partitionHasAPFSContainerMagic(diskPath string, off int64) (bool, error) {
	diskFile, err := os.Open(diskPath)
	if err != nil {
		return false, err
	}
	defer diskFile.Close()

	magic := make([]byte, len(APFSContainerMagic))
	if _, err := diskFile.ReadAt(magic, off+APFSContainerOffset); err != nil {
		return false, err
	}

	return string(magic) == APFSContainerMagic, nil
}

func copyRangeViaTempFile(diskPath string, srcOff int64, dstOff int64, n int64) error {
	tmpFile, err := os.CreateTemp("", "")
	if err != nil {
		return fmt.Errorf("failed to create a temporary file for storing copied partition contents: %w", err)
	}
	defer os.Remove(tmpFile.Name())

	if err := dumpPartition(diskPath, tmpFile.Name(), srcOff, n); err != nil {
		return err
	}

	return restorePartition(diskPath, tmpFile.Name(), dstOff, n)
}

func recoveryAlreadyAtEnd(partition *gpt.Partition, lastSectorAvailableForPartitions uint64) bool {
	return partition.End >= lastSectorAvailableForPartitions
}

func resumeRelocationIfNeeded(diskPath string, logicalBlockSize int64, partitions []*gpt.Partition, recoveryPartition *gpt.Partition, lastSectorAvailableForPartitions uint64) (bool, error) {
	if !recoveryAlreadyAtEnd(recoveryPartition, lastSectorAvailableForPartitions) {
		return false, nil
	}

	originalRecoveryStart, ok := inferOriginalRecoveryStart(partitions, recoveryPartition)
	if !ok {
		return false, nil
	}

	recoveryPartitionSizeInSectors := partitionSizeInSectors(recoveryPartition)
	originalRecoveryEnd := originalRecoveryStart + recoveryPartitionSizeInSectors - 1
	if originalRecoveryEnd >= recoveryPartition.Start {
		return false, nil
	}

	sourceHasMagic, err := partitionHasAPFSContainerMagic(diskPath, sectorOffset(originalRecoveryStart, logicalBlockSize))
	if err != nil {
		return false, fmt.Errorf("failed to inspect original recovery location: %w", err)
	}

	destinationHasMagic, err := partitionHasAPFSContainerMagic(diskPath, recoveryPartition.GetStart())
	if err != nil {
		return false, fmt.Errorf("failed to inspect relocated recovery location: %w", err)
	}

	if destinationHasMagic {
		fmt.Println("Nothing to relocate: recovery partition already appears restored at the end of disk.")
		return true, nil
	}

	if !sourceHasMagic {
		return false, nil
	}

	fmt.Println("Detected an interrupted relocation. Restoring recovery contents from the original location...")

	if err := copyRangeViaTempFile(
		diskPath,
		sectorOffset(originalRecoveryStart, logicalBlockSize),
		recoveryPartition.GetStart(),
		recoveryPartition.GetSize(),
	); err != nil {
		return false, fmt.Errorf("failed to resume interrupted recovery relocation: %w", err)
	}

	fmt.Println("Recovery partition contents restored.")

	return true, nil
}

func Relocate(path string) error {
	diskPath, cleanup, err := resolveDiskPath(path)
	if err != nil {
		return err
	}
	defer cleanup()

	// Open the disk image and read its partition table
	disk, err := diskfs.Open(diskPath)
	if err != nil {
		return fmt.Errorf("failed to open the disk: %w", err)
	}
	defer disk.Close()

	partitionTable, err := disk.GetPartitionTable()
	if err != nil {
		return fmt.Errorf("failed to get the partition table: %w", err)
	}

	// We only support relocating a recovery partition on a GPT table
	gptTable, ok := partitionTable.(*gpt.Table)
	if !ok {
		return fmt.Errorf("expected a \"gpt\" partition table, got %q", partitionTable.Type())
	}

	// Find the recovery partition
	recoveryPartition, recoveryPartitionIndex, ok := findRecoveryPartition(gptTable.Partitions)
	if !ok {
		fmt.Println("Nothing to relocate: no recovery partition found.")

		return nil
	}

	// Only relocate if this partition is already the last one by on-disk position.
	if !isLastPartitionOnDisk(gptTable.Partitions, recoveryPartition) {
		return fmt.Errorf("cannot relocate the recovery partition since it's not the last partition "+
			"on disk (matched entry %d)", recoveryPartitionIndex)
	}

	lastSectorAvailableForPartitions := lastUsableSector(disk, gptTable)

	resumed, err := resumeRelocationIfNeeded(diskPath, disk.LogicalBlocksize, gptTable.Partitions, recoveryPartition, lastSectorAvailableForPartitions)
	if err != nil {
		return err
	}
	if resumed {
		return nil
	}

	// Perhaps the recovery partition already resides at the last sector available for partitions?
	if recoveryAlreadyAtEnd(recoveryPartition, lastSectorAvailableForPartitions) {
		fmt.Println("Nothing to relocate: recovery partition already ends at the last sector available to partitions.")

		return nil
	}

	fmt.Println("Dumping recovery partition contents...")

	tmpFile, err := os.CreateTemp("", "")
	if err != nil {
		return fmt.Errorf("failed to create a temporary file for storing "+
			"the recovery partition contents: %w", err)
	}
	defer os.Remove(tmpFile.Name())

	if err := dumpPartition(diskPath, tmpFile.Name(), recoveryPartition.GetStart(), recoveryPartition.GetSize()); err != nil {
		return fmt.Errorf("failed to dump the recovery partition contents: %w", err)
	}

	fmt.Println("Re-partitioning the disk to adjust the recovery partition bounds...")

	relocatePartitionBounds(recoveryPartition, lastSectorAvailableForPartitions)

	// Write GPT directly: go-diskfs tries a Linux-specific re-read ioctl even on macOS.
	if err := writeGPTWithoutKernelReread(disk, gptTable); err != nil {
		return fmt.Errorf("failed to write the new partition table: %w", err)
	}

	fmt.Println("Restoring recovery partition contents...")

	if err := restorePartition(diskPath, tmpFile.Name(), recoveryPartition.GetStart(), recoveryPartition.GetSize()); err != nil {
		return fmt.Errorf("failed to restore the recovery partition contents: %w", err)
	}

	return nil
}

func dumpPartition(diskFilePath string, partitionFilePath string, off int64, n int64) error {
	diskFile, err := os.Open(diskFilePath)
	if err != nil {
		return err
	}
	defer diskFile.Close()

	partitionFile, err := os.Create(partitionFilePath)
	if err != nil {
		return err
	}
	defer partitionFile.Close()

	partitionContentsReader := io.NewSectionReader(diskFile, off, n)

	if _, err := io.Copy(partitionFile, partitionContentsReader); err != nil {
		return err
	}

	return nil
}

func restorePartition(diskFilePath string, partitionFilePath string, off int64, n int64) error {
	diskFile, err := os.OpenFile(diskFilePath, os.O_RDWR, 0600)
	if err != nil {
		return err
	}

	contentsFile, err := os.Open(partitionFilePath)
	if err != nil {
		return err
	}
	defer contentsFile.Close()

	partitionContentsWriter := io.NewOffsetWriter(diskFile, off)

	if _, err := io.CopyN(partitionContentsWriter, contentsFile, n); err != nil {
		return err
	}

	return diskFile.Close()
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage:", os.Args[0], "path_to_disk_image")
		os.Exit(1)
	}
	diskImageFilePath := os.Args[1]
	if err := Relocate(diskImageFilePath); err != nil {
		fmt.Fprintln(os.Stderr, "Error:", err)
		os.Exit(1)
	}
}