From 2ec0c3cc0ff5ff705ed5c910ca1bcd5d93c7b102 Mon Sep 17 00:00:00 2001
From: Andrey Smirnov <smirnov.andrey@gmail.com>
Date: Wed, 4 Aug 2021 23:51:56 +0300
Subject: [PATCH] fix: preserve the PMBR bootable flag when opening GPT
 partition

Otherwise we set the flag when creating the partition table, but it gets
reset after `Open()` and `Write()`.

Signed-off-by: Andrey Smirnov <smirnov.andrey@gmail.com>
---
 blockdevice/partition/gpt/gpt.go      | 22 ++++++++++------
 blockdevice/partition/gpt/gpt_test.go | 37 +++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/blockdevice/partition/gpt/gpt.go b/blockdevice/partition/gpt/gpt.go
index b66e20d..06e1ad5 100644
--- a/blockdevice/partition/gpt/gpt.go
+++ b/blockdevice/partition/gpt/gpt.go
@@ -5,7 +5,6 @@
 package gpt
 
 import (
-	"bytes"
 	"encoding/binary"
 	"errors"
 	"fmt"
@@ -50,17 +49,23 @@ type GPT struct {
 
 // Open attempts to open a partition table on f.
 func Open(f *os.File) (g *GPT, err error) {
-	buf := make([]byte, 1)
+	buf := make([]byte, 16)
 
 	// PMBR protective entry starts at 446. The partition type is at offset
 	// 4 from the start of the PMBR protective entry.
-	_, err = f.ReadAt(buf, 450)
+	var n int
+
+	n, err = f.ReadAt(buf, 446)
 	if err != nil {
 		return nil, err
 	}
 
+	if n != len(buf) {
+		return nil, fmt.Errorf("incomplete read: %d != %d", n, len(buf))
+	}
+
 	// For GPT, the partition type should be 0xEE (EFI GPT).
-	if bytes.Equal(buf, []byte{0xEE}) {
+	if buf[4] == 0xEE {
 		l, err := lba.NewLBA(f)
 		if err != nil {
 			return nil, err
@@ -71,10 +76,11 @@ func Open(f *os.File) (g *GPT, err error) {
 		h := &Header{Buffer: b, LBA: l}
 
 		g = &GPT{
-			f: f,
-			l: l,
-			h: h,
-			e: &Partitions{h: h, devname: f.Name()},
+			f:               f,
+			l:               l,
+			h:               h,
+			e:               &Partitions{h: h, devname: f.Name()},
+			markMBRBootable: buf[0] == 0x80,
 		}
 
 		return g, nil
diff --git a/blockdevice/partition/gpt/gpt_test.go b/blockdevice/partition/gpt/gpt_test.go
index 9b5f635..0a17c7d 100644
--- a/blockdevice/partition/gpt/gpt_test.go
+++ b/blockdevice/partition/gpt/gpt_test.go
@@ -49,6 +49,8 @@ func (suite *GPTSuite) TestReset() {
 
 	suite.Require().NoError(g.Write())
 
+	suite.checkMBRFlag(0x00)
+
 	// re-read the partition table
 	g, err = gpt.Open(suite.Dev)
 	suite.Require().NoError(err)
@@ -61,6 +63,8 @@ func (suite *GPTSuite) TestReset() {
 
 	suite.Require().NoError(g.Write())
 
+	suite.checkMBRFlag(0x00)
+
 	// re-read the partition table
 	g, err = gpt.Open(suite.Dev)
 	suite.Require().NoError(err)
@@ -424,6 +428,39 @@ func (suite *GPTSuite) TestPartitionGUUID() {
 	suite.Assert().NotEqual(g.Header().GUUID.String(), "00000000-0000-0000-0000-000000000000")
 }
 
+func (suite *GPTSuite) TestMarkPMBRBootable() {
+	g, err := gpt.New(suite.Dev, gpt.WithMarkMBRBootable(true))
+	suite.Require().NoError(err)
+
+	_, err = g.Add(1048576, gpt.WithPartitionName("boot"))
+	suite.Require().NoError(err)
+
+	_, err = g.Add(1048576, gpt.WithPartitionName("efi"))
+	suite.Require().NoError(err)
+
+	suite.Require().NoError(g.Write())
+
+	suite.checkMBRFlag(0x80)
+
+	// re-read the partition table
+	g, err = gpt.Open(suite.Dev)
+	suite.Require().NoError(err)
+
+	suite.Require().NoError(g.Read())
+	suite.Require().NoError(g.Write())
+
+	suite.checkMBRFlag(0x80)
+}
+
+func (suite *GPTSuite) checkMBRFlag(flag byte) {
+	buf := make([]byte, 1)
+
+	_, err := suite.Dev.ReadAt(buf, 446)
+	suite.Require().NoError(err)
+
+	suite.Assert().Equal(flag, buf[0])
+}
+
 func TestGPTSuite(t *testing.T) {
 	if os.Getuid() != 0 {
 		t.Skip("can't run the test as non-root")