fix: kernel partition sync when overwriting GPT

partitioning/gpt/gpt.go

@@ -62,6 +62,9 @@ type Table struct {
 
 	alignment  uint64
 	sectorSize uint
+
+	// newTable is true if the table completely overwrites existing partitions.
+	newTable bool
 }
 
 // Partition is a single partition entry in GPT.
@@ -128,6 +131,7 @@ func New(dev Device, opts ...Option) (*Table, error) {
 		dev:      dev,
 		options:  options,
 		diskGUID: diskGUID,
+		newTable: true,
 	}
 
 	t.init(lastLBA)
@@ -267,6 +271,7 @@ func (t *Table) init(lastLBA uint64) {
 // Clear the partition table.
 func (t *Table) Clear() {
 	t.entries = nil
+	t.newTable = true
 }
 
 // Compact the partition table by removing empty entries.
@@ -600,6 +605,52 @@ func (t *Table) writePMBR() error {
 }
 
 func (t *Table) syncKernel() error {
+	if t.newTable {
+		return t.syncKernelComplete()
+	}
+
+	return t.syncKernelIncremental()
+}
+
+// syncKernelComplete synchronizes the kernel partition table with the current table by overwriting the whole table.
+//
+// It is incompatible with mounted partitions.
+func (t *Table) syncKernelComplete() error {
+	kernelPartitionNum, err := t.dev.GetKernelLastPartitionNum()
+	if err != nil {
+		return fmt.Errorf("failed to get kernel last partition number: %w", err)
+	}
+
+	// delete all kernel partitions
+	for no := 1; no <= kernelPartitionNum; no++ {
+		if err := t.dev.KernelPartitionDelete(no); err != nil && !errors.Is(err, unix.ENXIO) {
+			return fmt.Errorf("failed to delete partition %d: %w", no, err)
+		}
+	}
+
+	// re-create all partitions
+	for no := 1; no <= len(t.entries); no++ {
+		myEntry := t.entries[no-1]
+
+		if myEntry == nil {
+			continue
+		}
+
+		if err := t.dev.KernelPartitionAdd(no,
+			myEntry.FirstLBA*uint64(t.sectorSize),
+			(myEntry.LastLBA-myEntry.FirstLBA+1)*uint64(t.sectorSize),
+		); err != nil {
+			return fmt.Errorf("failed to add partition %d: %w", no, err)
+		}
+	}
+
+	return nil
+}
+
+// syncKernelIncremental synchronizes the kernel partition table with the current table by doing minimal changes.
+//
+// It allows to change live partition layout, while some partitions are mounted.
+func (t *Table) syncKernelIncremental() error {
 	kernelPartitionNum, err := t.dev.GetKernelLastPartitionNum()
 	if err != nil {
 		return fmt.Errorf("failed to get kernel last partition number: %w", err)

partitioning/gpt/gpt_test.go

@@ -295,6 +295,75 @@ func TestGPT(t *testing.T) {
 	}
 }
 
+func TestGPTOverwrite(t *testing.T) {
+	if os.Geteuid() != 0 {
+		t.Skip("test requires root privileges")
+	}
+
+	if hostname, _ := os.Hostname(); hostname == "buildkitsandbox" { //nolint: errcheck
+		t.Skip("test not supported under buildkit as partition devices are not propagated from /dev")
+	}
+
+	partType1 := uuid.MustParse("C12A7328-F81F-11D2-BA4B-00A0C93EC93B")
+	partType2 := uuid.MustParse("E6D6D379-F507-44C2-A23C-238F2A3DF928")
+
+	// create a partition table, and then overwrite it with a new one with incompatible layout
+	tmpDir := t.TempDir()
+
+	rawImage := filepath.Join(tmpDir, "image.raw")
+
+	f, err := os.Create(rawImage)
+	require.NoError(t, err)
+
+	require.NoError(t, f.Truncate(int64(3*GiB)))
+	require.NoError(t, f.Close())
+
+	loDev := losetupAttachHelper(t, rawImage, false)
+
+	t.Cleanup(func() {
+		assert.NoError(t, loDev.Detach())
+	})
+
+	disk, err := os.OpenFile(loDev.Path(), os.O_RDWR, 0)
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		assert.NoError(t, disk.Close())
+	})
+
+	blkdev := block.NewFromFile(disk)
+
+	gptdev, err := gpt.DeviceFromBlockDevice(blkdev)
+	require.NoError(t, err)
+
+	table, err := gpt.New(gptdev)
+	require.NoError(t, err)
+
+	// allocate 2 1G partitions first
+	require.NoError(t, allocateError(table.AllocatePartition(100*MiB, "1G", partType1)))
+	require.NoError(t, allocateError(table.AllocatePartition(1*GiB, "2G", partType2)))
+
+	require.NoError(t, table.Write())
+
+	assert.FileExists(t, loDev.Path()+"p1")
+	assert.FileExists(t, loDev.Path()+"p2")
+
+	// now attempt to overwrite the partition table with a new one with different layout
+	table2, err := gpt.New(gptdev)
+	require.NoError(t, err)
+
+	// allocate new partitions first
+	require.NoError(t, allocateError(table2.AllocatePartition(600*MiB, "1P", partType1)))
+	require.NoError(t, allocateError(table2.AllocatePartition(600*MiB, "2P", partType2)))
+	require.NoError(t, allocateError(table2.AllocatePartition(600*MiB, "3P", partType2)))
+
+	require.NoError(t, table2.Write())
+
+	assert.FileExists(t, loDev.Path()+"p1")
+	assert.FileExists(t, loDev.Path()+"p2")
+	assert.FileExists(t, loDev.Path()+"p3")
+}
+
 func losetupAttachHelper(t *testing.T, rawImage string, readonly bool) losetup.Device {
 	t.Helper()