From c9385719c202bf002f4e785696545aa3b036c7d0 Mon Sep 17 00:00:00 2001 From: Qinqi Qu Date: Mon, 30 Oct 2023 17:09:28 +0800 Subject: [PATCH] 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 Signed-off-by: Qinqi Qu --- diskfs_test.go | 106 +++++++++++++++++++++++++++++++++-- partition/gpt/table.go | 122 ++++++++++++++++++++++++++++++++++++----- partition/mbr/table.go | 10 ++++ partition/table.go | 2 + 4 files changed, 220 insertions(+), 20 deletions(-) diff --git a/diskfs_test.go b/diskfs_test.go index ca4c3790..e2a3b120 100644 --- 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,20 @@ 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 +69,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 index bbea768d..df19d0ee 100644 --- 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,12 +390,7 @@ 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, @@ -409,9 +398,80 @@ func tableFromBytes(b []byte, logicalBlockSize, physicalBlockSize int) (*Table, 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 teh 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,35 @@ 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 index 631018a3..d19554c5 100644 --- 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 index 29cf0046..5e1e84b3 100644 --- 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 }