Merge branch 'master' into repack-pr

pkg/visitors/dxecleaner.go

@@ -0,0 +1,162 @@
+// Copyright 2018 the LinuxBoot Authors. All rights reserved
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package visitors
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"os"
+	"os/exec"
+	"os/signal"
+	"path/filepath"
+	"syscall"
+
+	"github.com/linuxboot/fiano/pkg/guid"
+	"github.com/linuxboot/fiano/pkg/uefi"
+)
+
+// DXECleaner removes DXEs sequentially in multiple rounds. Each round, an
+// attempt is made to remove each DXE. The Test function determines if the
+// removal was successful. Additional rounds are performed until all DXEs are
+// removed.
+type DXECleaner struct {
+	// This function tests whether the firmware boots. The return values can be:
+	//
+	//     - (false, nil): The firmware was tested and failed to boot.
+	//     - (false, err): The firmware was tested and failed to boot due to err.
+	//     - (true, nil):  The firmware was tested and booted.
+	//     - (true, err):  Failed to test the firmware due to err.
+	Test func(f uefi.Firmware) (bool, error)
+
+	// List of GUIDs which were removed.
+	Removals []guid.GUID
+
+	// Logs are written to this writer.
+	W io.Writer
+}
+
+// Run wraps Visit and performs some setup and teardown tasks.
+func (v *DXECleaner) Run(f uefi.Firmware) error {
+	var printf = func(format string, a ...interface{}) {
+		if v.W != nil {
+			fmt.Fprintf(v.W, format, a...)
+		}
+	}
+
+	// Find list of DXEs.
+	find := (&Find{
+		Predicate: FindFileTypePredicate(uefi.FVFileTypeDriver),
+	})
+	if err := find.Run(f); err != nil {
+		return err
+	}
+	var dxes []guid.GUID
+	for i := range find.Matches {
+		dxes = append(dxes, find.Matches[i].(*uefi.File).Header.GUID)
+	}
+	if len(dxes) == 0 {
+		return errors.New("found no DXEs in firmware image")
+	}
+
+	// Print list of removals in a format which can be passed back into UTK.
+	defer func() {
+		printf("Summary of removed DXEs:\n")
+		if len(v.Removals) == 0 {
+			printf("  Could not remove any DXEs\n")
+		} else {
+			for _, r := range v.Removals {
+				printf("  remove %s \\\n", r)
+			}
+		}
+	}()
+
+	// Main algorithm to remove DXEs.
+	moreRoundsNeeded := true
+	for i := 0; moreRoundsNeeded; i++ {
+		printf("Beginning of round %d\n", i+1)
+		moreRoundsNeeded = false
+		for i := 0; i < len(dxes); i++ {
+			// Remove the DXE from the image.
+			printf("Trying to remove %v\n", dxes[i])
+			remove := &Remove{Predicate: FindFileGUIDPredicate(dxes[i])}
+			if err := remove.Run(f); err != nil {
+				return err
+			}
+
+			if removedSuccessfully, err := v.Test(f); err == context.Canceled {
+				printf("Canceled by user!\n")
+				return nil
+			} else if removedSuccessfully && err != nil {
+				return err
+			} else if removedSuccessfully {
+				printf("  Success!\n")
+				v.Removals = append(v.Removals, dxes[i])
+				dxes = append(dxes[:i], dxes[i+1:]...)
+				i--
+				moreRoundsNeeded = true
+			} else {
+				printf("  Failed!\n")
+				remove.Undo()
+			}
+		}
+	}
+	return nil
+}
+
+// Visit applies the DXEClearn visitor to any Firmware type.
+func (v *DXECleaner) Visit(f uefi.Firmware) error {
+	return nil
+}
+
+func init() {
+	RegisterCLI("dxecleaner", "automates removal of UEFI drivers", 1, func(args []string) (uefi.Visitor, error) {
+		// When the user enters CTRL-C, the DXECleaner should stop, but
+		// also output the current progress.
+		ctx, cancel := context.WithCancel(context.Background())
+		c := make(chan os.Signal, 1)
+		signal.Notify(c, os.Interrupt)
+		go func() {
+			<-c
+			cancel()
+		}()
+
+		return &DXECleaner{
+			Test: func(f uefi.Firmware) (bool, error) {
+				tmpDir, err := ioutil.TempDir("", "dxecleaner")
+				if err != nil {
+					return true, err
+				}
+				defer os.RemoveAll(tmpDir)
+				tmpFile := filepath.Join(tmpDir, "bios.bin")
+
+				if err := (&Save{tmpFile}).Run(f); err != nil {
+					return true, err
+				}
+				if err := exec.CommandContext(ctx, args[0], tmpFile).Run(); err != nil {
+					if _, ok := err.(*exec.ExitError); !ok {
+						return true, err
+					}
+					status, ok := err.(*exec.ExitError).Sys().(syscall.WaitStatus)
+					if !ok {
+						return true, err
+					}
+					switch status.ExitStatus() {
+					case 1:
+						return true, err
+					case 2:
+						return false, err
+					default:
+						return true, fmt.Errorf("unexpected exit status %d", status.ExitStatus())
+					}
+				}
+				return true, nil
+			},
+			W: os.Stdout,
+		}, nil
+	})
+}

pkg/visitors/dxecleaner_test.go

@@ -0,0 +1,59 @@
+// Copyright 2018 the LinuxBoot Authors. All rights reserved
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package visitors
+
+import (
+	"os"
+	"testing"
+
+	"github.com/linuxboot/fiano/pkg/guid"
+	"github.com/linuxboot/fiano/pkg/uefi"
+)
+
+var (
+	testDXE1 = guid.MustParse("93B80004-9FB3-11D4-9A3A-0090273FC14D")
+	testDXE2 = guid.MustParse("4B28E4C7-FF36-4E10-93CF-A82159E777C5")
+	testDXE3 = guid.MustParse("C8339973-A563-4561-B858-D8476F9DEFC4")
+	testDXE4 = guid.MustParse("378D7B65-8DA9-4773-B6E4-A47826A833E1")
+	testDXE5 = guid.MustParse("33CB97AF-6C33-4C42-986B-07581FA366D4")
+)
+
+func contains(t *testing.T, f uefi.Firmware, g *guid.GUID) bool {
+	return len(find(t, f, g)) > 0
+}
+
+func TestDXECleaner(t *testing.T) {
+	// This test to see if an image "boots" by looking for fake dependencies
+	// between the DXEs.
+	testDXEDependencies := func(f uefi.Firmware) (bool, error) {
+		// Dependencies
+		return contains(t, f, testDXE5) &&
+			(!contains(t, f, testDXE5) || contains(t, f, testDXE4)) &&
+			(!contains(t, f, testDXE2) || contains(t, f, testDXE1)) &&
+			(!contains(t, f, testDXE3) || contains(t, f, testDXE2)), nil
+	}
+
+	// Parse image and run the visitor.
+	f := parseImage(t)
+	dxeCleaner := DXECleaner{
+		Test: testDXEDependencies,
+		W:    os.Stdout,
+	}
+	if err := dxeCleaner.Run(f); err != nil {
+		t.Fatal(err)
+	}
+
+	// Check that the correct DXEs remain.
+	for _, d := range []*guid.GUID{testDXE1, testDXE2, testDXE3} {
+		if contains(t, f, d) {
+			t.Errorf("expected %v to be deleted", d)
+		}
+	}
+	for _, d := range []*guid.GUID{testDXE4, testDXE5} {
+		if !contains(t, f, d) {
+			t.Errorf("expected %v to be remain", d)
+		}
+	}
+}

pkg/visitors/find.go

@@ -94,6 +94,16 @@ func FindFileGUIDPredicate(r guid.GUID) func(f uefi.Firmware) bool {
 	}
 }
 
+// FindFileTypePredicate is a generic predicate for searching file types only.
+func FindFileTypePredicate(t uefi.FVFileType) func(f uefi.Firmware) bool {
+	return func(f uefi.Firmware) bool {
+		if f, ok := f.(*uefi.File); ok {
+			return f.Header.Type == t
+		}
+		return false
+	}
+}
+
 // FindFilePredicate is a generic predicate for searching files and UI sections only.
 func FindFilePredicate(r string) (func(f uefi.Firmware) bool, error) {
 	searchRE, err := regexp.Compile("^" + r + "$")

pkg/visitors/remove.go

@@ -16,6 +16,8 @@ type Remove struct {
 
 	// Output
 	Matches []uefi.Firmware
+	// Calling this function undoes the removals performed by this visitor.
+	Undo func()
 }
 
 // Run wraps Visit and performs some setup and teardown tasks.
@@ -40,6 +42,8 @@ func (v *Remove) Visit(f uefi.Firmware) error {
 		for i := 0; i < len(f.Files); i++ {
 			for _, m := range v.Matches {
 				if f.Files[i] == m {
+					originalList := append([]*uefi.File{}, f.Files...)
+
 					m := m.(*uefi.File)
 					if v.Pad || m.Header.Type == uefi.FVFileTypePEIM {
 						// Create a new pad file of the exact same size
@@ -51,6 +55,15 @@ func (v *Remove) Visit(f uefi.Firmware) error {
 					} else {
 						f.Files = append(f.Files[:i], f.Files[i+1:]...)
 					}
+
+					// Creates a stack of undoes in case there are multiple FVs.
+					prev := v.Undo
+					v.Undo = func() {
+						f.Files = originalList
+						if prev != nil {
+							prev()
+						}
+					}
 				}
 			}
 		}