partition: add repair capabilities to GPT

1. add Repair and Verify interfaces in type table 2. add test to test opening the GPT image Signed-off-by: thebsdbox <[email protected]> Signed-off-by: Qinqi Qu <[email protected]>
diskfs · Oct 30, 2023 · 77a4994 · 77a4994
1 parent 94310d8
commit 77a4994
Show file tree

Hide file tree

Showing 4 changed files with 216 additions and 20 deletions.
diff --git a/diskfs_test.go b/diskfs_test.go
@@ -15,7 +15,7 @@ import (
 
 const oneMB = 10 * 1024 * 1024
 
-func tmpDisk(source string) (*os.File, error) {
+func tmpDisk(source string, padding int64) (*os.File, error) {
 	filename := "disk_test"
 	f, err := os.CreateTemp("", filename)
 	if err != nil {
@@ -35,8 +35,18 @@ func tmpDisk(source string) (*os.File, error) {
 		if err != nil {
 			return nil, fmt.Errorf("Failed to write contents of %s to %s: %v", source, filename, err)
 		}
-		if written != len(b) {
-			return nil, fmt.Errorf("wrote only %d bytes of %s to %s instead of %d", written, source, filename, len(b))
+		if padding != 0 {
+			data := make([]byte, padding) // Initialize an empty byte slice
+			writtenPadding, err := f.Write(data)
+			written += writtenPadding
+			if err != nil {
+				return nil, fmt.Errorf("Failed to write contents of %s to %s: %v", source, filename, err)
+			}
+			if written != len(b)+len(data) {
+				return nil, fmt.Errorf("Wrote only %d bytes of %s to %s instead of %d", written, source, filename, len(b))
+			}
+		} else if written != len(b) {
+			return nil, fmt.Errorf("Wrote only %d bytes of %s to %s instead of %d", written, source, filename, len(b))
 		}
 	}
 
@@ -57,8 +67,92 @@ func checkDiskfsErrs(t *testing.T, msg string, err, tterr error, d, ttd *disk.Di
 	}
 }
 
-func TestOpen(t *testing.T) {
-	f, err := tmpDisk("./partition/mbr/testdata/mbr.img")
+func TestGPTOpen(t *testing.T) {
+	f, err := tmpDisk("./partition/gpt/testdata/gpt.img", 0)
+	if err != nil {
+		t.Fatalf("Error creating new temporary disk: %v", err)
+	}
+	defer f.Close()
+	path := f.Name()
+	defer os.Remove(path)
+	fileInfo, err := f.Stat()
+	if err != nil {
+		t.Fatalf("Unable to stat temporary file %s: %v", path, err)
+	}
+	size := fileInfo.Size()
+
+	// Create a padded file, where "disk" has additional space after what the GPT table
+	// thinks should be the final sector
+	fPadded, err := tmpDisk("./partition/gpt/testdata/gpt.img", 1024*1024)
+	if err != nil {
+		t.Fatalf("Error creating new temporary disk: %v", err)
+	}
+	defer fPadded.Close()
+	filePadded := fPadded.Name()
+	defer os.Remove(path)
+	filePaddedInfo, err := fPadded.Stat()
+	if err != nil {
+		t.Fatalf("Unable to stat temporary file %s: %v", path, err)
+	}
+	filePaddedSize := filePaddedInfo.Size()
+
+	tests := []struct {
+		path string
+		disk *disk.Disk
+		err  error
+	}{
+		{"", nil, fmt.Errorf("must pass device name")},
+		{"/tmp/foo/bar/232323/23/2322/disk.img", nil, fmt.Errorf("")},
+		{path, &disk.Disk{Type: disk.File, LogicalBlocksize: 512, PhysicalBlocksize: 512, Size: size}, nil},
+		{filePadded, &disk.Disk{Type: disk.File, LogicalBlocksize: 512, PhysicalBlocksize: 512, Size: filePaddedSize}, nil},
+	}
+
+	for _, tt := range tests {
+		d, err := diskfs.Open(tt.path)
+		msg := fmt.Sprintf("Open(%s)", tt.path)
+		checkDiskfsErrs(t, msg, err, tt.err, d, tt.disk)
+		if d != nil {
+			table, err := d.GetPartitionTable()
+			if err != nil {
+				t.Errorf("%s: mismatched errors, actual %v expected %v", msg, err, tt.err)
+			}
+
+			// Verify will compare the GPT table to the disk and attempt to read the secondary header if possible
+			err = table.Verify(d.File, uint64(tt.disk.Size))
+			if err != nil {
+				// We log this as it's epected to be an error
+				t.Logf("%s: mismatched errors, actual %v expected %v", msg, err, tt.err)
+			}
+
+			// Will correct the internal structures of the primary GPT table
+			err = table.Repair(uint64(tt.disk.Size))
+			if err != nil {
+				t.Errorf("%s: mismatched errors, actual %v expected %v", msg, err, tt.err)
+			}
+
+			// Update both tables on disk
+			err = table.Write(d.File, tt.disk.Size)
+			if err != nil {
+				t.Errorf("%s: mismatched errors, actual %v expected %v", msg, err, tt.err)
+			}
+
+			// Check that things are as expected.
+			err = table.Verify(d.File, uint64(tt.disk.Size))
+			if err != nil {
+				t.Errorf("%s: mismatched errors, actual %v expected %v", msg, err, tt.err)
+			}
+		}
+	}
+
+	for i, tt := range tests {
+		d, err := diskfs.Open(tt.path, diskfs.WithOpenMode(diskfs.ReadOnly))
+		msg := fmt.Sprintf("%d: Open(%s)", i, tt.path)
+		checkDiskfsErrs(t, msg, err, tt.err, d, tt.disk)
+	}
+}
+
+func TestMBROpen(t *testing.T) {
+	f, err := tmpDisk("./partition/mbr/testdata/mbr.img", 1024*1024)
 	if err != nil {
 		t.Fatalf("error creating new temporary disk: %v", err)
 	}

diff --git a/partition/gpt/table.go b/partition/gpt/table.go
@@ -347,15 +347,9 @@ func readPartitionArrayBytes(b []byte, entrySize, logicalSectorSize, physicalSec
 	return parts, nil
 }
 
-// tableFromBytes read a partition table from a byte slice
-func tableFromBytes(b []byte, logicalBlockSize, physicalBlockSize int) (*Table, error) {
-	// minimum size - gpt entries + header + LBA0 for (protective) MBR
-	if len(b) < logicalBlockSize*2 {
-		return nil, fmt.Errorf("data for partition was %d bytes instead of expected minimum %d", len(b), logicalBlockSize*2)
-	}
-
-	// GPT starts at LBA1
-	gpt := b[logicalBlockSize:]
+// readGptHeader reads the GPT header from the given byte slice
+func readGptHeader(b []byte) (*Table, error) {
+	gpt := b
 	// start with fixed headers
 	efiSignature := gpt[0:8]
 	efiRevision := gpt[8:12]
@@ -396,22 +390,88 @@ func tableFromBytes(b []byte, logicalBlockSize, physicalBlockSize int) (*Table,
 		return nil, fmt.Errorf("invalid EFI Header Checksum, expected %v, got %v", checksum, efiHeaderCrc)
 	}
 
-	// potential protective MBR is at LBA0
-	hasProtectiveMBR := readProtectiveMBR(b[:logicalBlockSize], uint32(secondaryHeader))
-
 	table := Table{
-		LogicalSectorSize:      logicalBlockSize,
-		PhysicalSectorSize:     physicalBlockSize,
 		partitionEntrySize:     partitionEntrySize,
 		primaryHeader:          primaryHeader,
 		secondaryHeader:        secondaryHeader,
 		firstDataSector:        firstDataSector,
 		lastDataSector:         lastDataSector,
 		partitionArraySize:     int(partitionEntryCount),
 		partitionFirstLBA:      partitionEntryFirstLBA,
-		ProtectiveMBR:          hasProtectiveMBR,
 		GUID:                   strings.ToUpper(diskGUID.String()),
 		partitionEntryChecksum: partitionEntryChecksum,
+	}
+
+	return &table, nil
+}
+
+// tableHeaderFromBytes read a partition table from a byte slice, mainly used to validate the secondary header
+func tableHeaderFromBytes(b []byte, logicalBlockSize, physicalBlockSize int, skipMBR bool) (*Table, error) {
+	// minimum size - gpt entries + header + LBA0 for (protective) MBR
+	minSize := logicalBlockSize
+	if len(b) < minSize {
+		return nil, fmt.Errorf("data for partition was %d bytes instead of expected minimum %d", len(b), minSize)
+	}
+	gpt := b
+	if skipMBR {
+		gpt = b[logicalBlockSize:]
+	}
+
+	header, err := readGptHeader(gpt)
+	if err != nil {
+		return nil, err
+	}
+
+	// potential protective MBR is at LBA0
+	hasProtectiveMBR := readProtectiveMBR(b[:logicalBlockSize], uint32(header.secondaryHeader))
+
+	table := Table{
+		LogicalSectorSize:  logicalBlockSize,
+		PhysicalSectorSize: physicalBlockSize,
+		partitionEntrySize: header.partitionEntrySize,
+		primaryHeader:      header.primaryHeader,
+		secondaryHeader:    header.secondaryHeader,
+		firstDataSector:    header.firstDataSector,
+		lastDataSector:     header.lastDataSector,
+		partitionArraySize: header.partitionArraySize,
+		ProtectiveMBR:      hasProtectiveMBR,
+		GUID:               header.GUID,
+		initialized:        true,
+	}
+	return &table, nil
+}
+
+// tableFromBytes read a partition table from a byte slice
+func tableFromBytes(b []byte, logicalBlockSize, physicalBlockSize int) (*Table, error) {
+	// minimum size - gpt entries + header + LBA0 for (protective) MBR
+	if len(b) < logicalBlockSize*2 {
+		return nil, fmt.Errorf("data for partition was %d bytes instead of expected minimum %d", len(b), logicalBlockSize*2)
+	}
+
+	// GPT starts at LBA1
+	gpt := b[logicalBlockSize:]
+
+	header, err := readGptHeader(gpt)
+	if err != nil {
+		return nil, err
+	}
+
+	// potential protective MBR is at LBA0
+	hasProtectiveMBR := readProtectiveMBR(b[:logicalBlockSize], uint32(header.secondaryHeader))
+
+	table := Table{
+		LogicalSectorSize:      logicalBlockSize,
+		PhysicalSectorSize:     physicalBlockSize,
+		partitionEntrySize:     header.partitionEntrySize,
+		primaryHeader:          header.primaryHeader,
+		secondaryHeader:        header.secondaryHeader,
+		firstDataSector:        header.firstDataSector,
+		lastDataSector:         header.lastDataSector,
+		partitionArraySize:     header.partitionArraySize,
+		partitionFirstLBA:      header.partitionFirstLBA,
+		ProtectiveMBR:          hasProtectiveMBR,
+		GUID:                   header.GUID,
+		partitionEntryChecksum: header.partitionEntryChecksum,
 		initialized:            true,
 	}
 
@@ -550,3 +610,33 @@ func (t *Table) GetPartitions() []part.Partition {
 	}
 	return parts
 }
+
+// Verify will attempt to evaluate the headers
+func (t *Table) Verify(f util.File, diskSize uint64) error {
+	// Determine the size of disk that GPT expects
+	expectedDiskSize := (t.secondaryHeader + 1) * uint64(t.LogicalSectorSize)
+	if diskSize != expectedDiskSize {
+		return fmt.Errorf("secondary Header is not at end of the disk, expected =>  %d / actual => %d", expectedDiskSize, diskSize)
+	}
+	b := make([]byte, t.LogicalSectorSize)
+	seekAddress := int64(t.secondaryHeader) * int64(t.LogicalSectorSize)
+	_, err := f.ReadAt(b, seekAddress)
+	if err != nil {
+		return fmt.Errorf("error reading GPT from file at %d / disksize %d : %v", seekAddress, diskSize, err)
+	}
+	secondaryTable, err := tableHeaderFromBytes(b, t.LogicalSectorSize, t.PhysicalSectorSize, false)
+	if err != nil {
+		return fmt.Errorf("error reading GPT from file at %d / disksize %d : %v", seekAddress, diskSize, err)
+	}
+	if t.firstDataSector != secondaryTable.firstDataSector {
+		return fmt.Errorf("error comparing GPT headers expected =>  %d / actual => %d", t.firstDataSector, secondaryTable.firstDataSector)
+	}
+	return nil
+}
+
+// Repair will attempt to evaluate the headers fix the header location and re-write the primary and secondary header
+func (t *Table) Repair(diskSize uint64) error {
+	t.secondaryHeader = (diskSize / uint64(t.LogicalSectorSize)) - 1
+
+	return nil
+}
diff --git a/partition/mbr/table.go b/partition/mbr/table.go
@@ -176,3 +176,13 @@ func (t *Table) GetPartitions() []part.Partition {
 	}
 	return parts
 }
+
+// Verify will attempt to evaluate the headers
+func (t *Table) Verify(f util.File, diskSize uint64) error {
+	return nil
+}
+
+// Repair will attempt to repair a broken Master Boot Record
+func (t *Table) Repair(diskSize uint64) error {
+	return nil
+}
diff --git a/partition/table.go b/partition/table.go
@@ -10,4 +10,6 @@ type Table interface {
 	Type() string
 	Write(util.File, int64) error
 	GetPartitions() []part.Partition
+	Repair(diskSize uint64) error
+	Verify(f util.File, diskSize uint64) error
 }