github.com/oweisse/u-root@v0.0.0-20181109060735-d005ad25fef1/pkg/boot/package_test.go (about)

     1  // Copyright 2018 the u-root 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 boot
     6  
     7  import (
     8  	"crypto/rand"
     9  	"crypto/rsa"
    10  	"log"
    11  	"reflect"
    12  	"strings"
    13  	"testing"
    14  
    15  	"github.com/u-root/u-root/pkg/cpio"
    16  )
    17  
    18  type mockOSImage struct {
    19  	packErr error
    20  }
    21  
    22  var _ OSImage = &mockOSImage{}
    23  
    24  func newMockImage(*cpio.Archive) (OSImage, error) {
    25  	return &mockOSImage{}, nil
    26  }
    27  
    28  func (mockOSImage) ExecutionInfo(log *log.Logger)     {}
    29  func (mockOSImage) Execute() error                    { return nil }
    30  func (m mockOSImage) Pack(sw cpio.RecordWriter) error { return m.packErr }
    31  
    32  func packageEqual(p1, p2 *Package) bool {
    33  	li1, ok := p1.OSImage.(*LinuxImage)
    34  	if !ok {
    35  		return false
    36  	}
    37  	li2, ok := p2.OSImage.(*LinuxImage)
    38  	if !ok {
    39  		return false
    40  	}
    41  	return imageEqual(li1, li2) && reflect.DeepEqual(p1.Metadata, p2.Metadata)
    42  }
    43  
    44  func TestBootPackage(t *testing.T) {
    45  	privateKey, err := rsa.GenerateKey(rand.Reader, 4096)
    46  	if err != nil {
    47  		t.Errorf("GenerateKey() = %v", err)
    48  	}
    49  
    50  	for _, tt := range []struct {
    51  		pkg       *Package
    52  		packErr   error
    53  		unpackErr error
    54  		signer    *rsa.PrivateKey
    55  		verifier  *rsa.PublicKey
    56  	}{
    57  		{
    58  			pkg: &Package{
    59  				OSImage: &LinuxImage{
    60  					Kernel:  strings.NewReader("lana"),
    61  					Initrd:  strings.NewReader("mcnulty"),
    62  					Cmdline: "foo=bar",
    63  				},
    64  				Metadata: map[string]string{
    65  					"stuff": "fooasdf",
    66  				},
    67  			},
    68  			signer:   privateKey,
    69  			verifier: &privateKey.PublicKey,
    70  			packErr:  nil,
    71  		},
    72  		{
    73  			pkg: &Package{
    74  				OSImage: &LinuxImage{
    75  					Kernel:  strings.NewReader("lana"),
    76  					Initrd:  strings.NewReader("mcnulty"),
    77  					Cmdline: "foo=bar",
    78  				},
    79  				Metadata: map[string]string{
    80  					"stuff": "fooasdf",
    81  				},
    82  			},
    83  		},
    84  		{
    85  			pkg: &Package{
    86  				OSImage: &LinuxImage{
    87  					Kernel:  strings.NewReader("lana"),
    88  					Initrd:  strings.NewReader("mcnulty"),
    89  					Cmdline: "foo=bar",
    90  				},
    91  				Metadata: map[string]string{
    92  					"stuff": "fooasdf",
    93  				},
    94  			},
    95  			verifier:  &privateKey.PublicKey,
    96  			unpackErr: rsa.ErrVerification,
    97  		},
    98  		{
    99  			pkg: &Package{
   100  				OSImage: &LinuxImage{
   101  					Kernel:  strings.NewReader("lana"),
   102  					Initrd:  nil,
   103  					Cmdline: "foo=bar",
   104  				},
   105  				Metadata: map[string]string{},
   106  			},
   107  			signer:   privateKey,
   108  			verifier: &privateKey.PublicKey,
   109  			packErr:  nil,
   110  		},
   111  		{
   112  			pkg: &Package{
   113  				OSImage: &LinuxImage{
   114  					Kernel:  strings.NewReader("lana"),
   115  					Initrd:  nil,
   116  					Cmdline: "",
   117  				},
   118  				Metadata: map[string]string{},
   119  			},
   120  			signer:   privateKey,
   121  			verifier: &privateKey.PublicKey,
   122  			packErr:  nil,
   123  		},
   124  		{
   125  			pkg: &Package{
   126  				OSImage: &LinuxImage{
   127  					Kernel:  strings.NewReader("lana"),
   128  					Initrd:  nil,
   129  					Cmdline: "",
   130  				},
   131  				Metadata: map[string]string{
   132  					"abc":      "def",
   133  					"abcd/foo": "haha",
   134  				},
   135  			},
   136  			signer:   privateKey,
   137  			verifier: &privateKey.PublicKey,
   138  			packErr:  nil,
   139  		},
   140  		{
   141  			pkg: &Package{
   142  				OSImage: &LinuxImage{
   143  					Kernel:  strings.NewReader("lana"),
   144  					Initrd:  nil,
   145  					Cmdline: "",
   146  				},
   147  				Metadata: map[string]string{
   148  					"abc":     "def",
   149  					"abc/foo": "haha",
   150  				},
   151  			},
   152  			signer:   privateKey,
   153  			verifier: &privateKey.PublicKey,
   154  			packErr:  nil,
   155  		},
   156  	} {
   157  		a := cpio.InMemArchive()
   158  		if err := tt.pkg.Pack(a, tt.signer); err != tt.packErr {
   159  			t.Errorf("Pack(%v) = %v, want %v", tt.pkg, err, tt.packErr)
   160  		}
   161  
   162  		var p2 Package
   163  		if err := (&p2).Unpack(a.Reader(), tt.verifier); err != tt.unpackErr {
   164  			t.Errorf("Unpack() = %v, want %v", err, tt.unpackErr)
   165  		} else if err == nil {
   166  			if !packageEqual(tt.pkg, &p2) {
   167  				t.Errorf("packages are not equal: got %v\nwant %v", p2, tt.pkg)
   168  			}
   169  		}
   170  	}
   171  }