refactor: rewrite GPT serialize/deserialize functions

Remove reads of GPT header by bytes, read whole header in memory. Use simple versions of functions to encode/decode fields without allocating any memory. Simplify things by removing functions and removing failure scenarios. Signed-off-by: Andrey Smirnov <andrey.smirnov@talos-systems.com>
siderolabs · Dec 14, 2021 · 6928ee4 · 6928ee4
1 parent 0c7e429
commit 6928ee4
Show file tree

Hide file tree

Showing 5 changed files with 137 additions and 814 deletions.
diff --git a/blockdevice/lba/lba.go b/blockdevice/lba/lba.go
@@ -5,50 +5,9 @@
 package lba
 
 import (
-	"fmt"
 	"os"
 )
 
-// Buffer is an in-memory buffer for writing to byte slices in units of LBA.
-type Buffer struct {
-	lba *LBA
-	b   []byte
-}
-
-// NewBuffer intializes and returns a `Buffer`.
-func NewBuffer(lba *LBA, b []byte) *Buffer {
-	return &Buffer{lba: lba, b: b}
-}
-
-// Read reads from a `Buffer`.
-func (buf *Buffer) Read(off, length int64) (b []byte, err error) {
-	b = make([]byte, length)
-
-	n := copy(b, buf.b[off:off+length])
-
-	if n != len(buf.b[off:off+length]) {
-		return nil, fmt.Errorf("expected to write %d bytes, read %d", len(b), n)
-	}
-
-	return b, nil
-}
-
-// Write writes to a `Buffer`.
-func (buf *Buffer) Write(b []byte, off int64) (err error) {
-	n := copy(buf.b[off:off+int64(len(b))], b)
-
-	if n != len(b) {
-		return fmt.Errorf("expected to write %d bytes, wrote %d", len(b), n)
-	}
-
-	return nil
-}
-
-// Bytes returns the buffer bytes.
-func (buf *Buffer) Bytes() []byte {
-	return buf.b
-}
-
 // LBA represents logical block addressing.
 //
 //nolint:govet

diff --git a/blockdevice/partition/gpt/gpt.go b/blockdevice/partition/gpt/gpt.go
@@ -72,9 +72,11 @@ func Open(f *os.File) (g *GPT, err error) {
 			return nil, err
 		}
 
-		b := lba.NewBuffer(l, make([]byte, l.LogicalBlockSize))
+		h := &Header{LBA: l}
 
-		h := &Header{Buffer: b, LBA: l}
+		if err = h.verifySignature(); err != nil {
+			return nil, ErrPartitionTableDoesNotExist
+		}
 
 		g = &GPT{
 			f:               f,
@@ -84,10 +86,6 @@ func Open(f *os.File) (g *GPT, err error) {
 			markMBRBootable: buf[0] == 0x80,
 		}
 
-		if err = h.DeserializeSignature(); err != nil {
-			return nil, ErrPartitionTableDoesNotExist
-		}
-
 		return g, nil
 	}
 
@@ -106,9 +104,7 @@ func New(f *os.File, setters ...Option) (g *GPT, err error) {
 		return nil, err
 	}
 
-	b := lba.NewBuffer(l, make([]byte, l.LogicalBlockSize))
-
-	h := &Header{Buffer: b, LBA: l}
+	h := &Header{LBA: l}
 
 	h.Signature = MagicEFIPart
 	h.Revision = binary.LittleEndian.Uint32([]byte{0x00, 0x00, 0x01, 0x00})

diff --git a/blockdevice/partition/gpt/gpt_test.go b/blockdevice/partition/gpt/gpt_test.go
@@ -7,6 +7,7 @@ package gpt_test
 import (
 	"encoding/binary"
 	"os"
+	"os/exec"
 	"testing"
 
 	"github.com/stretchr/testify/suite"
@@ -73,6 +74,8 @@ func (suite *GPTSuite) TestReset() {
 	suite.Require().NoError(g.Read())
 
 	suite.Assert().Empty(g.Partitions().Items())
+
+	suite.validatePartitions()
 }
 
 func (suite *GPTSuite) TestPartitionAdd() {
@@ -129,6 +132,8 @@ func (suite *GPTSuite) TestPartitionAdd() {
 	suite.Require().NoError(g.Read())
 
 	assertPartitions(g.Partitions())
+
+	suite.validatePartitions()
 }
 
 func (suite *GPTSuite) TestRepairResize() {
@@ -223,6 +228,8 @@ func (suite *GPTSuite) TestRepairResize() {
 	suite.Require().NoError(g.Read())
 
 	assertPartitions(g.Partitions())
+
+	suite.validatePartitions()
 }
 
 func (suite *GPTSuite) TestPartitionAddOutOfSpace() {
@@ -319,6 +326,8 @@ func (suite *GPTSuite) TestPartitionDelete() {
 	partSystem := partitions[2]
 	suite.Assert().EqualValues(3, partSystem.Number)
 	suite.Assert().EqualValues((size-bootSize-efiSize-grubSize-configSize)/blockSize-headReserved-tailReserved, partSystem.Length())
+
+	suite.validatePartitions()
 }
 
 func (suite *GPTSuite) TestPartitionInsertAt() {
@@ -409,6 +418,8 @@ func (suite *GPTSuite) TestPartitionInsertAt() {
 	suite.Assert().EqualValues(headReserved+(oldBootSize+configSize+efiSize)/blockSize, partSystem.FirstLBA)
 
 	suite.Require().NoError(g.Write())
+
+	suite.validatePartitions()
 }
 
 func (suite *GPTSuite) TestPartitionInsertOffsetAndResize() {
@@ -481,6 +492,8 @@ func (suite *GPTSuite) TestPartitionInsertOffsetAndResize() {
 	suite.Assert().EqualValues(configStart/blockSize, int(partConfig.FirstLBA))
 
 	suite.Require().NoError(g.Write())
+
+	suite.validatePartitions()
 }
 
 func (suite *GPTSuite) TestPartitionGUUID() {
@@ -523,6 +536,8 @@ func (suite *GPTSuite) TestMarkPMBRBootable() {
 	suite.Require().NoError(g.Write())
 
 	suite.validatePMBR(0x80)
+
+	suite.validatePartitions()
 }
 
 func (suite *GPTSuite) validatePMBR(flag byte) {
@@ -544,6 +559,20 @@ func (suite *GPTSuite) validatePMBR(flag byte) {
 	suite.Assert().Equal([]byte{0x55, 0xaa}, buf[510:512]) // boot signature
 }
 
+func (suite *GPTSuite) validatePartitions() {
+	for _, tool := range []string{"fdisk", "gdisk"} {
+		if toolPath, _ := exec.LookPath(tool); toolPath != "" { //nolint:errcheck
+			cmd := exec.Command(toolPath, "-l", suite.Dev.Name())
+			cmd.Stdout = os.Stdout
+			cmd.Stderr = os.Stderr
+
+			suite.Assert().NoError(cmd.Run())
+		} else {
+			suite.T().Logf("%s is not available", tool)
+		}
+	}
+}
+
 func TestGPTSuite(t *testing.T) {
 	if os.Getuid() != 0 {
 		t.Skip("can't run the test as non-root")