github.com/linuxboot/fiano@v1.2.0/pkg/visitors/replacepe32_test.go (about)

     1  // Copyright 2018 the LinuxBoot Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package visitors
     6  
     7  import (
     8  	"reflect"
     9  	"testing"
    10  
    11  	"github.com/linuxboot/fiano/pkg/uefi"
    12  )
    13  
    14  func TestReplacePE32(t *testing.T) {
    15  	f := parseImage(t)
    16  
    17  	// Apply the visitor.
    18  	replace := &ReplacePE32{
    19  		Predicate: FindFileGUIDPredicate(*testGUID),
    20  		NewPE32:   []byte("MZbanana"),
    21  	}
    22  	if err := replace.Run(f); err != nil {
    23  		t.Fatal(err)
    24  	}
    25  
    26  	// We expect one match.
    27  	if len(replace.Matches) != 1 {
    28  		t.Fatalf("got %d matches; expected 1", len(replace.Matches))
    29  	}
    30  
    31  	// Find the section and make sure it contains the expected data.
    32  	results := find(t, f, testGUID)
    33  	if len(results) != 1 {
    34  		t.Fatalf("got %d matches; expected 1", len(results))
    35  	}
    36  	want := []byte{0x0c, 0x00, 0x00, byte(uefi.SectionTypePE32), 'M', 'Z', 'b', 'a', 'n', 'a', 'n', 'a'}
    37  	file, ok := results[0].(*uefi.File)
    38  	if !ok {
    39  		t.Fatalf("did not match a file, got type :%T", file)
    40  	}
    41  	got := file.Sections[0].Buf()
    42  	if !reflect.DeepEqual(want, got) {
    43  		t.Fatalf("want %v; got %v", want, got)
    44  	}
    45  }
    46  
    47  func TestErrors(t *testing.T) {
    48  	f := parseImage(t)
    49  
    50  	var tests = []struct {
    51  		name    string
    52  		newPE32 []byte
    53  		match   string
    54  		err     string
    55  	}{
    56  		{"No Matches", []byte("MZbanana"), "no-match-string",
    57  			"no matches found for replacement"},
    58  		{"Multiple Matches", []byte("MZbanana"), ".*",
    59  			"multiple matches found! There can be only one. Use find to list all matches"},
    60  		{"Not PE32", []byte("banana"), ".*",
    61  			"supplied binary is not a valid pe32 image"},
    62  	}
    63  	for _, test := range tests {
    64  		t.Run(test.name, func(t *testing.T) {
    65  			// Apply the visitor.
    66  			pred, err := FindFilePredicate(test.match)
    67  			if err != nil {
    68  				t.Fatal(err)
    69  			}
    70  			replace := &ReplacePE32{
    71  				Predicate: pred,
    72  				NewPE32:   test.newPE32,
    73  			}
    74  			err = replace.Run(f)
    75  			if err == nil {
    76  				t.Fatalf("Expected Error (%v), got nil", test.err)
    77  			} else if err.Error() != test.err {
    78  				t.Fatalf("Mismatched Error: Expected %v, got %v", test.err, err.Error())
    79  			}
    80  		})
    81  	}
    82  }