run on creation and update of firmware

Signed-off-by: karlodwyer <[email protected]>

pkg/uefi/uefi.go

@@ -147,6 +147,18 @@ func UnmarshalFirmware(b []byte) (Firmware, error) {
 // implement any parser itself, but it calls known parsers that implement the
 // Firmware interface.
 func Parse(buf []byte) (Firmware, error) {
+	f, err := parse(buf)
+	if err != nil {
+		return f, err
+	}
+	err = (&PositionUpdater{}).Run(f)
+	if err != nil {
+		return nil, fmt.Errorf("unable to update offsets: %w", err)
+	}
+	return f, nil
+}
+
+func parse(buf []byte) (Firmware, error) {
 	if _, err := FindSignature(buf); err == nil {
 		// Intel rom.
 		return NewFlashImage(buf)

pkg/visitors/update_positions.go → pkg/uefi/update_positions.go

@@ -2,14 +2,10 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-package visitors
+package uefi
 
-import (
-	"github.com/linuxboot/fiano/pkg/uefi"
-)
-
-// positionUpdater updates the position of the varisous firmwares in memory
-type positionUpdater struct {
+// PositionUpdater updates the position of the varisous firmwares in memory
+type PositionUpdater struct {
 	Scan      bool
 	Layout    bool
 	Depth     int
@@ -19,58 +15,58 @@ type positionUpdater struct {
 }
 
 // Run wraps Visit and performs some setup and teardown tasks.
-func (v *positionUpdater) Run(f uefi.Firmware) error {
+func (v *PositionUpdater) Run(f Firmware) error {
 	return f.Apply(v)
 }
 
 // Visit applies the Table visitor to any Firmware type.
-func (v *positionUpdater) Visit(f uefi.Firmware) error {
+func (v *PositionUpdater) Visit(f Firmware) error {
 	var offset uint64
 	switch f := f.(type) {
-	case *uefi.FlashImage:
+	case *FlashImage:
 		if v.Depth > 0 { // Depth <= 0 means all
 			v.Depth++
 		}
 		return v.GoDeeper(f, 0)
-	case *uefi.FirmwareVolume:
+	case *FirmwareVolume:
 		f.AbsOffSet = v.offset + f.FVOffset
 		return v.GoDeeper(f, v.offset+f.FVOffset+f.DataOffset)
-	case *uefi.File:
+	case *File:
 		f.AbsOffSet = v.curOffset
 		return v.GoDeeper(f, v.curOffset+f.DataOffset)
-	case *uefi.Section:
+	case *Section:
 		f.AbsOffSet = v.curOffset
 
 		// Reset offset to O for (compressed) section content
 		return v.GoDeeper(f, 0)
-	case *uefi.FlashDescriptor:
+	case *FlashDescriptor:
 		f.AbsOffSet = 0
 		return v.GoDeeper(f, 0)
-	case *uefi.BIOSRegion:
+	case *BIOSRegion:
 		if f.FRegion != nil {
 			offset = uint64(f.FRegion.BaseOffset())
 			f.AbsOffSet = offset
 		}
 		return v.GoDeeper(f, offset)
-	case *uefi.BIOSPadding:
+	case *BIOSPadding:
 		f.AbsOffSet = v.offset + f.Offset
 		return v.GoDeeper(f, 0)
-	case *uefi.NVarStore:
+	case *NVarStore:
 		f.AbsOffSet = v.curOffset
 		return v.GoDeeper(f, v.curOffset)
-	case *uefi.NVar:
+	case *NVar:
 		f.AbsOffSet = v.curOffset
 		return v.GoDeeper(f, v.curOffset+uint64(f.DataOffset))
-	case *uefi.MERegion:
+	case *MERegion:
 		if f.FRegion != nil {
 			offset = uint64(f.FRegion.BaseOffset())
 			f.AbsOffSet = offset
 		}
 		return v.GoDeeper(f, offset)
-	case *uefi.MEFPT:
+	case *MEFPT:
 		f.AbsOffSet = v.offset
 		return v.GoDeeper(f, 0)
-	case *uefi.RawRegion:
+	case *RawRegion:
 		if f.FRegion != nil {
 			offset = uint64(f.FRegion.BaseOffset())
 			f.AbsOffSet = offset
@@ -81,7 +77,7 @@ func (v *positionUpdater) Visit(f uefi.Firmware) error {
 	}
 }
 
-func (v *positionUpdater) GoDeeper(f uefi.Firmware, dataOffset uint64) error {
+func (v *PositionUpdater) GoDeeper(f Firmware, dataOffset uint64) error {
 	// Prepare data and print
 	length := uint64(len(f.Buf()))
 

pkg/visitors/update_positions_test.go → pkg/uefi/update_positions_test.go

@@ -2,28 +2,40 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-package visitors
+package uefi
 
 import (
+	"io/ioutil"
 	"testing"
 
 	"github.com/linuxboot/fiano/pkg/guid"
-	"github.com/linuxboot/fiano/pkg/uefi"
 )
 
 type visitor struct {
 	T *testing.T
 }
 
-func (v *visitor) Run(f uefi.Firmware) error {
+func parseImage(t *testing.T) Firmware {
+	image, err := ioutil.ReadFile("../../integration/roms/OVMF.rom")
+	if err != nil {
+		t.Fatal(err)
+	}
+	parsedRoot, err := Parse(image)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return parsedRoot
+}
+
+func (v *visitor) Run(f Firmware) error {
 	return f.Apply(v)
 }
-func (v *visitor) Visit(f uefi.Firmware) error {
+func (v *visitor) Visit(f Firmware) error {
 	var guid *guid.GUID
 	switch f := f.(type) {
-	case *uefi.File:
+	case *File:
 		guid = &f.Header.GUID
-	case *uefi.FirmwareVolume:
+	case *FirmwareVolume:
 		guid = &f.FVName
 	default:
 		return f.ApplyChildren(v)
@@ -46,7 +58,7 @@ func TestUpdatePositions(t *testing.T) {
 
 	f := parseImage(t)
 
-	updater := &positionUpdater{}
+	updater := &PositionUpdater{}
 	if err := updater.Run(f); err != nil {
 		t.Fatal(err)
 	}

pkg/visitors/assemble.go

@@ -30,7 +30,11 @@ type Assemble struct {
 
 // Run just applies the visitor.
 func (v *Assemble) Run(f uefi.Firmware) error {
-	return f.Apply(v)
+	err := f.Apply(v)
+	if err != nil {
+		return err
+	}
+	return (&uefi.PositionUpdater{}).Run(f)
 }
 
 // Visit applies the Assemble visitor to any Firmware type.