Merge pull request #235 from paul-rodriguez/resizeGpt

Allow resizing a GPT

partition/gpt/partition.go

@@ -296,3 +296,9 @@ func (p *Partition) Equal(o *Partition) bool {
 func (p *Partition) UUID() string {
 	return p.GUID
 }
+
+// Expand increases the size of the partition by a number of sectors
+func (p *Partition) Expand(sectors uint64) {
+	p.End += sectors
+	p.Size += sectors * uint64(p.logicalSectorSize)
+}

partition/gpt/table.go

@@ -20,6 +20,7 @@ const (
 	// just defaults
 	physicalSectorSize = 512
 	logicalSectorSize  = 512
+	gptHeaderSector    = 1
 )
 
 // Table represents a partition table to be applied to a disk or read from a disk
@@ -640,3 +641,31 @@ func (t *Table) Repair(diskSize uint64) error {
 
 	return nil
 }
+
+// TotalSize returns the total size of the GPT in bytes.
+//
+// This is counted from the start of the MBR to the end of the secondary
+// header.
+func (t *Table) TotalSize() uint64 {
+	return (t.secondaryHeader + gptHeaderSector) * uint64(t.LogicalSectorSize)
+}
+
+func (t *Table) LastDataSector() uint64 {
+	return t.lastDataSector
+}
+
+// Resize changes the size of the GPT.
+//
+// The size argument is in bytes and must be a multiple of the logical sector
+// size.
+// Use this function in case a storage device is not the same as the total
+// size of its GPT.
+func (t *Table) Resize(size uint64) {
+	// how many sectors on the disk?
+	diskSectors := size / uint64(t.LogicalSectorSize)
+	// how many sectors used for partition entries?
+	partSectors := uint64(t.partitionArraySize) * uint64(t.partitionEntrySize) / uint64(t.LogicalSectorSize)
+
+	t.secondaryHeader = diskSectors - 1
+	t.lastDataSector = t.secondaryHeader - 1 - partSectors
+}

partition/gpt/table_test.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"path/filepath"
 	"regexp"
 	"strconv"
 	"strings"
@@ -465,3 +466,75 @@ func TestWritePartitionContents(t *testing.T) {
 		t.Log(b2)
 	}
 }
+
+func TestResizeTableAndExpandPartition(t *testing.T) {
+	const newSize = 11 * 1024 * 1024
+
+	blkFile, err := os.Open(gptFile)
+	if err != nil {
+		t.Errorf("cannot open file: %v", err)
+	}
+	defer blkFile.Close()
+	table, err := gpt.Read(blkFile, 512, 512)
+	if err != nil {
+		t.Errorf("cannot read gpt: %v", err)
+	}
+	table.Resize(newSize)
+
+	tmpDir := t.TempDir()
+	tmpImgPath := filepath.Join(tmpDir, "gpt.img")
+	tmpImgFile, err := os.Create(tmpImgPath)
+	if err != nil {
+		t.Errorf("cannot create output image file: %v", err)
+	}
+	err = os.Truncate(tmpImgPath, newSize)
+	if err != nil {
+		t.Errorf("cannot truncate file: %v", err)
+	}
+
+	table.Partitions[0].Expand(100)
+	err = table.Write(tmpImgFile, newSize)
+	if err != nil {
+		t.Errorf("cannot write table back: %v", err)
+	}
+	newTable, err := gpt.Read(tmpImgFile, 512, 512)
+	if err != nil {
+		t.Errorf("cannot read table back: %v", err)
+	}
+	if newTable.Partitions[0].Start != 2048 {
+		t.Fail()
+	}
+	if newTable.Partitions[0].End != 3148 {
+		t.Fail()
+	}
+	if newTable.Partitions[0].Size != 563712 {
+		t.Fail()
+	}
+}
+
+func TestResize(t *testing.T) {
+	const newSize = 1024 * 1024
+	table := gpt.GetValidTable()
+	table.Resize(newSize)
+	resultSize := table.TotalSize()
+	if resultSize != newSize {
+		t.Fail()
+	}
+	resultLastDataSector := table.LastDataSector()
+	expectedLastDataSector := uint64((newSize / 512) - 34)
+	if resultLastDataSector != expectedLastDataSector {
+		t.Fail()
+	}
+}
+
+func TestExpandPartition(t *testing.T) {
+	table := gpt.GetValidTable()
+	part := table.Partitions[0]
+	part.Expand(100)
+	if part.End != 3148 {
+		t.Fail()
+	}
+	if part.Size != (3148-2048+1)*512 {
+		t.Fail()
+	}
+}