gopkg.in/hugelgupf/u-root.v2@v2.0.0-20180831055005-3f8fdb0ce09d/pkg/uroot/builder/bb.go (about)

     1  // Copyright 2015-2017 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 builder
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"go/ast"
    11  	"go/build"
    12  	"go/format"
    13  	"go/importer"
    14  	"go/parser"
    15  	"go/token"
    16  	"go/types"
    17  	"io/ioutil"
    18  	"os"
    19  	"path"
    20  	"path/filepath"
    21  	"sort"
    22  	"strconv"
    23  	"strings"
    24  
    25  	"golang.org/x/tools/go/ast/astutil"
    26  	"golang.org/x/tools/imports"
    27  
    28  	"github.com/u-root/u-root/pkg/cpio"
    29  	"github.com/u-root/u-root/pkg/golang"
    30  	"github.com/u-root/u-root/pkg/uroot/initramfs"
    31  )
    32  
    33  // Commands to skip building in bb mode.
    34  var skip = map[string]struct{}{
    35  	"bb": {},
    36  }
    37  
    38  // BBBuilder is an implementation of Builder that compiles many Go commands
    39  // into one busybox-style binary.
    40  type BBBuilder struct{}
    41  
    42  func (BBBuilder) DefaultBinaryDir() string {
    43  	return "bbin"
    44  }
    45  
    46  // Build is an implementation of uroot.Builder.Build for a busybox-like initramfs.
    47  //
    48  // Build rewrites the source files of the packages given to create one
    49  // busybox-like binary containing all commands in opts.Packages.
    50  //
    51  // See pkg/uroot/README.md for a detailed explanation of busybox mode.
    52  func (BBBuilder) Build(af initramfs.Files, opts Opts) error {
    53  	// Build the busybox binary.
    54  	bbPath := filepath.Join(opts.TempDir, "bb")
    55  	if err := BuildBusybox(opts.Env, opts.Packages, bbPath); err != nil {
    56  		return err
    57  	}
    58  
    59  	if len(opts.BinaryDir) == 0 {
    60  		return fmt.Errorf("must specify binary directory")
    61  	}
    62  
    63  	// Build initramfs.
    64  	if err := af.AddFile(bbPath, path.Join(opts.BinaryDir, "bb")); err != nil {
    65  		return err
    66  	}
    67  
    68  	// Add symlinks for included commands to initramfs.
    69  	for _, pkg := range opts.Packages {
    70  		if _, ok := skip[path.Base(pkg)]; ok {
    71  			continue
    72  		}
    73  
    74  		// Add a symlink /bbin/{cmd} -> /bbin/bb to our initramfs.
    75  		if err := af.AddRecord(cpio.Symlink(filepath.Join(opts.BinaryDir, path.Base(pkg)), "bb")); err != nil {
    76  			return err
    77  		}
    78  	}
    79  	return nil
    80  }
    81  
    82  // BuildBusybox builds a busybox of the given Go packages.
    83  //
    84  // pkgs is a list of Go import paths. If nil is returned, binaryPath will hold
    85  // the busybox-style binary.
    86  func BuildBusybox(env golang.Environ, pkgs []string, binaryPath string) error {
    87  	urootPkg, err := env.Package("github.com/u-root/u-root")
    88  	if err != nil {
    89  		return err
    90  	}
    91  
    92  	bbDir := filepath.Join(urootPkg.Dir, "bb")
    93  	// Blow bb away before trying to re-create it.
    94  	if err := os.RemoveAll(bbDir); err != nil {
    95  		return err
    96  	}
    97  	if err := os.MkdirAll(bbDir, 0755); err != nil {
    98  		return err
    99  	}
   100  
   101  	var bbPackages []string
   102  	// Move and rewrite package files.
   103  	importer := importer.For("source", nil)
   104  	for _, pkg := range pkgs {
   105  		if _, ok := skip[path.Base(pkg)]; ok {
   106  			continue
   107  		}
   108  
   109  		pkgDir := filepath.Join(bbDir, "cmds", path.Base(pkg))
   110  		// TODO: use bbDir to derive import path below or vice versa.
   111  		if err := RewritePackage(env, pkg, pkgDir, "github.com/u-root/u-root/pkg/bb", importer); err != nil {
   112  			return err
   113  		}
   114  
   115  		bbPackages = append(bbPackages, fmt.Sprintf("github.com/u-root/u-root/bb/cmds/%s", path.Base(pkg)))
   116  	}
   117  
   118  	bb, err := NewPackageFromEnv(env, "github.com/u-root/u-root/pkg/bb/cmd", importer)
   119  	if err != nil {
   120  		return err
   121  	}
   122  	if bb == nil {
   123  		return fmt.Errorf("bb cmd template missing")
   124  	}
   125  	if len(bb.ast.Files) != 1 {
   126  		return fmt.Errorf("bb cmd template is supposed to only have one file")
   127  	}
   128  	// Create bb main.go.
   129  	if err := CreateBBMainSource(bb.fset, bb.ast, bbPackages, bbDir); err != nil {
   130  		return err
   131  	}
   132  
   133  	// Compile bb.
   134  	return env.Build("github.com/u-root/u-root/bb", binaryPath, golang.BuildOpts{})
   135  }
   136  
   137  // CreateBBMainSource creates a bb Go command that imports all given pkgs.
   138  //
   139  // p must be the bb template.
   140  //
   141  // - For each pkg in pkgs, add
   142  //     import _ "pkg"
   143  //   to astp's first file.
   144  // - Write source file out to destDir.
   145  func CreateBBMainSource(fset *token.FileSet, astp *ast.Package, pkgs []string, destDir string) error {
   146  	for _, pkg := range pkgs {
   147  		for _, sourceFile := range astp.Files {
   148  			// Add side-effect import to bb binary so init registers itself.
   149  			//
   150  			// import _ "pkg"
   151  			astutil.AddNamedImport(fset, sourceFile, "_", pkg)
   152  			break
   153  		}
   154  	}
   155  
   156  	// Write bb main binary out.
   157  	for filePath, sourceFile := range astp.Files {
   158  		path := filepath.Join(destDir, filepath.Base(filePath))
   159  		if err := writeFile(path, fset, sourceFile); err != nil {
   160  			return err
   161  		}
   162  		break
   163  	}
   164  	return nil
   165  }
   166  
   167  // Package is a Go package.
   168  //
   169  // It holds AST, type, file, and Go package information about a Go package.
   170  type Package struct {
   171  	name string
   172  
   173  	fset        *token.FileSet
   174  	ast         *ast.Package
   175  	typeInfo    types.Info
   176  	types       *types.Package
   177  	sortedFiles []*ast.File
   178  
   179  	// initCount keeps track of what the next init's index should be.
   180  	initCount uint
   181  
   182  	// init is the cmd.Init function that calls all other InitXs in the
   183  	// right order.
   184  	init *ast.FuncDecl
   185  
   186  	// initAssigns is a map of assignment expression -> InitN function call
   187  	// statement.
   188  	//
   189  	// That InitN should contain the assignment statement for the
   190  	// appropriate assignment expression.
   191  	//
   192  	// types.Info.InitOrder keeps track of Initializations by Lhs name and
   193  	// Rhs ast.Expr.  We reparent the Rhs in assignment statements in InitN
   194  	// functions, so we use the Rhs as an easy key here.
   195  	// types.Info.InitOrder + initAssigns can then easily be used to derive
   196  	// the order of Stmts in the "real" init.
   197  	//
   198  	// The key Expr must also be the AssignStmt.Rhs[0].
   199  	initAssigns map[ast.Expr]ast.Stmt
   200  }
   201  
   202  // NewPackageFromEnv finds the package identified by importPath, and gathers
   203  // AST, type, and token information.
   204  func NewPackageFromEnv(env golang.Environ, importPath string, importer types.Importer) (*Package, error) {
   205  	p, err := env.Package(importPath)
   206  	if err != nil {
   207  		return nil, err
   208  	}
   209  	return NewPackage(filepath.Base(p.Dir), p, importer)
   210  }
   211  
   212  // ParseAST parses p's package files into an AST.
   213  func ParseAST(p *build.Package) (*token.FileSet, *ast.Package, error) {
   214  	name := filepath.Base(p.Dir)
   215  	if !p.IsCommand() {
   216  		return nil, nil, fmt.Errorf("package %q is not a command and cannot be included in bb", name)
   217  	}
   218  
   219  	fset := token.NewFileSet()
   220  	pars, err := parser.ParseDir(fset, p.Dir, func(fi os.FileInfo) bool {
   221  		// Only parse Go files that match build tags of this package.
   222  		for _, name := range p.GoFiles {
   223  			if name == fi.Name() {
   224  				return true
   225  			}
   226  		}
   227  		return false
   228  	}, parser.ParseComments)
   229  	if err != nil {
   230  		return nil, nil, fmt.Errorf("failed to parse AST for pkg %q: %v", name, err)
   231  	}
   232  
   233  	astp, ok := pars[p.Name]
   234  	if !ok {
   235  		return nil, nil, fmt.Errorf("parsed files, but could not find package %q in ast: %v", p.Name, pars)
   236  	}
   237  	return fset, astp, nil
   238  }
   239  
   240  // NewPackage gathers AST, type, and token information about package p, using
   241  // the given importer to resolve dependencies.
   242  func NewPackage(name string, p *build.Package, importer types.Importer) (*Package, error) {
   243  	fset, astp, err := ParseAST(p)
   244  	if err != nil {
   245  		return nil, err
   246  	}
   247  
   248  	pp := &Package{
   249  		name: name,
   250  		fset: fset,
   251  		ast:  astp,
   252  		typeInfo: types.Info{
   253  			Types: make(map[ast.Expr]types.TypeAndValue),
   254  		},
   255  		initAssigns: make(map[ast.Expr]ast.Stmt),
   256  	}
   257  
   258  	// This Init will hold calls to all other InitXs.
   259  	pp.init = &ast.FuncDecl{
   260  		Name: ast.NewIdent("Init"),
   261  		Type: &ast.FuncType{
   262  			Params:  &ast.FieldList{},
   263  			Results: nil,
   264  		},
   265  		Body: &ast.BlockStmt{},
   266  	}
   267  
   268  	// The order of types.Info.InitOrder depends on this list of files
   269  	// always being passed to conf.Check in the same order.
   270  	filenames := make([]string, 0, len(pp.ast.Files))
   271  	for name := range pp.ast.Files {
   272  		filenames = append(filenames, name)
   273  	}
   274  	sort.Strings(filenames)
   275  
   276  	pp.sortedFiles = make([]*ast.File, 0, len(pp.ast.Files))
   277  	for _, name := range filenames {
   278  		pp.sortedFiles = append(pp.sortedFiles, pp.ast.Files[name])
   279  	}
   280  	// Type-check the package before we continue. We need types to rewrite
   281  	// some statements.
   282  	conf := types.Config{
   283  		Importer: importer,
   284  
   285  		// We only need global declarations' types.
   286  		IgnoreFuncBodies: true,
   287  	}
   288  	tpkg, err := conf.Check(p.ImportPath, pp.fset, pp.sortedFiles, &pp.typeInfo)
   289  	if err != nil {
   290  		return nil, fmt.Errorf("type checking failed: %#v: %v", importer, err)
   291  	}
   292  	pp.types = tpkg
   293  	return pp, nil
   294  }
   295  
   296  func (p *Package) nextInit(addToCallList bool) *ast.Ident {
   297  	i := ast.NewIdent(fmt.Sprintf("Init%d", p.initCount))
   298  	if addToCallList {
   299  		p.init.Body.List = append(p.init.Body.List, &ast.ExprStmt{X: &ast.CallExpr{Fun: i}})
   300  	}
   301  	p.initCount++
   302  	return i
   303  }
   304  
   305  // TODO:
   306  // - write an init name generator, in case InitN is already taken.
   307  // - also rewrite all non-Go-stdlib dependencies.
   308  func (p *Package) rewriteFile(f *ast.File) bool {
   309  	hasMain := false
   310  
   311  	// Change the package name declaration from main to the command's name.
   312  	f.Name = ast.NewIdent(p.name)
   313  
   314  	// Map of fully qualified package name -> imported alias in the file.
   315  	importAliases := make(map[string]string)
   316  	for _, impt := range f.Imports {
   317  		if impt.Name != nil {
   318  			importPath, err := strconv.Unquote(impt.Path.Value)
   319  			if err != nil {
   320  				panic(err)
   321  			}
   322  			importAliases[importPath] = impt.Name.Name
   323  		}
   324  	}
   325  
   326  	// When the types.TypeString function translates package names, it uses
   327  	// this function to map fully qualified package paths to a local alias,
   328  	// if it exists.
   329  	qualifier := func(pkg *types.Package) string {
   330  		name, ok := importAliases[pkg.Path()]
   331  		if ok {
   332  			return name
   333  		}
   334  		// When referring to self, don't use any package name.
   335  		if pkg == p.types {
   336  			return ""
   337  		}
   338  		return pkg.Name()
   339  	}
   340  
   341  	for _, decl := range f.Decls {
   342  		switch d := decl.(type) {
   343  		case *ast.GenDecl:
   344  			// We only care about vars.
   345  			if d.Tok != token.VAR {
   346  				break
   347  			}
   348  			for _, spec := range d.Specs {
   349  				s := spec.(*ast.ValueSpec)
   350  				if s.Values == nil {
   351  					continue
   352  				}
   353  
   354  				// For each assignment, create a new init
   355  				// function, and place it in the same file.
   356  				for i, name := range s.Names {
   357  					varInit := &ast.FuncDecl{
   358  						Name: p.nextInit(false),
   359  						Type: &ast.FuncType{
   360  							Params:  &ast.FieldList{},
   361  							Results: nil,
   362  						},
   363  						Body: &ast.BlockStmt{
   364  							List: []ast.Stmt{
   365  								&ast.AssignStmt{
   366  									Lhs: []ast.Expr{name},
   367  									Tok: token.ASSIGN,
   368  									Rhs: []ast.Expr{s.Values[i]},
   369  								},
   370  							},
   371  						},
   372  					}
   373  					// Add a call to the new init func to
   374  					// this map, so they can be added to
   375  					// Init0() in the correct init order
   376  					// later.
   377  					p.initAssigns[s.Values[i]] = &ast.ExprStmt{X: &ast.CallExpr{Fun: varInit.Name}}
   378  					f.Decls = append(f.Decls, varInit)
   379  				}
   380  
   381  				// Add the type of the expression to the global
   382  				// declaration instead.
   383  				if s.Type == nil {
   384  					typ := p.typeInfo.Types[s.Values[0]]
   385  					s.Type = ast.NewIdent(types.TypeString(typ.Type, qualifier))
   386  				}
   387  				s.Values = nil
   388  			}
   389  
   390  		case *ast.FuncDecl:
   391  			if d.Recv == nil && d.Name.Name == "main" {
   392  				d.Name.Name = "Main"
   393  				hasMain = true
   394  			}
   395  			if d.Recv == nil && d.Name.Name == "init" {
   396  				d.Name = p.nextInit(true)
   397  			}
   398  		}
   399  	}
   400  
   401  	// Now we change any import names attached to package declarations. We
   402  	// just upcase it for now; it makes it easy to look in bbsh for things
   403  	// we changed, e.g. grep -r bbsh Import is useful.
   404  	for _, cg := range f.Comments {
   405  		for _, c := range cg.List {
   406  			if strings.HasPrefix(c.Text, "// import") {
   407  				c.Text = "// Import" + c.Text[9:]
   408  			}
   409  		}
   410  	}
   411  	return hasMain
   412  }
   413  
   414  // RewritePackage rewrites pkgPath to be bb-mode compatible, where destDir is
   415  // the file system destination of the written files and bbImportPath is the Go
   416  // import path of the bb package to register with.
   417  func RewritePackage(env golang.Environ, pkgPath, destDir, bbImportPath string, importer types.Importer) error {
   418  	p, err := NewPackageFromEnv(env, pkgPath, importer)
   419  	if err != nil {
   420  		return err
   421  	}
   422  	if p == nil {
   423  		return nil
   424  	}
   425  	return p.Rewrite(destDir, bbImportPath)
   426  }
   427  
   428  // Rewrite rewrites p into destDir as a bb package using bbImportPath for the
   429  // bb implementation.
   430  func (p *Package) Rewrite(destDir, bbImportPath string) error {
   431  	if err := os.MkdirAll(destDir, 0755); err != nil {
   432  		return err
   433  	}
   434  
   435  	// This init holds all variable initializations.
   436  	//
   437  	// func Init0() {}
   438  	varInit := &ast.FuncDecl{
   439  		Name: p.nextInit(true),
   440  		Type: &ast.FuncType{
   441  			Params:  &ast.FieldList{},
   442  			Results: nil,
   443  		},
   444  		Body: &ast.BlockStmt{},
   445  	}
   446  
   447  	var mainFile *ast.File
   448  	for _, sourceFile := range p.sortedFiles {
   449  		if hasMainFile := p.rewriteFile(sourceFile); hasMainFile {
   450  			mainFile = sourceFile
   451  		}
   452  	}
   453  	if mainFile == nil {
   454  		return os.RemoveAll(destDir)
   455  	}
   456  
   457  	// Add variable initializations to Init0 in the right order.
   458  	for _, initStmt := range p.typeInfo.InitOrder {
   459  		a, ok := p.initAssigns[initStmt.Rhs]
   460  		if !ok {
   461  			return fmt.Errorf("couldn't find init assignment %s", initStmt)
   462  		}
   463  		varInit.Body.List = append(varInit.Body.List, a)
   464  	}
   465  
   466  	// import bb "bbImportPath"
   467  	astutil.AddNamedImport(p.fset, mainFile, "bb", bbImportPath)
   468  
   469  	// func init() {
   470  	//   bb.Register(p.name, Init, Main)
   471  	// }
   472  	bbRegisterInit := &ast.FuncDecl{
   473  		Name: ast.NewIdent("init"),
   474  		Type: &ast.FuncType{},
   475  		Body: &ast.BlockStmt{
   476  			List: []ast.Stmt{
   477  				&ast.ExprStmt{X: &ast.CallExpr{
   478  					Fun: ast.NewIdent("bb.Register"),
   479  					Args: []ast.Expr{
   480  						// name=
   481  						&ast.BasicLit{
   482  							Kind:  token.STRING,
   483  							Value: strconv.Quote(p.name),
   484  						},
   485  						// init=
   486  						ast.NewIdent("Init"),
   487  						// main=
   488  						ast.NewIdent("Main"),
   489  					},
   490  				}},
   491  			},
   492  		},
   493  	}
   494  
   495  	// We could add these statements to any of the package files. We choose
   496  	// the one that contains Main() to guarantee reproducibility of the
   497  	// same bbsh binary.
   498  	mainFile.Decls = append(mainFile.Decls, varInit, p.init, bbRegisterInit)
   499  
   500  	// Write all files out.
   501  	for filePath, sourceFile := range p.ast.Files {
   502  		path := filepath.Join(destDir, filepath.Base(filePath))
   503  		if err := writeFile(path, p.fset, sourceFile); err != nil {
   504  			return err
   505  		}
   506  	}
   507  	return nil
   508  }
   509  
   510  func writeFile(path string, fset *token.FileSet, f *ast.File) error {
   511  	var buf bytes.Buffer
   512  	if err := format.Node(&buf, fset, f); err != nil {
   513  		return fmt.Errorf("error formatting Go file %q: %v", path, err)
   514  	}
   515  	return writeGoFile(path, buf.Bytes())
   516  }
   517  
   518  func writeGoFile(path string, code []byte) error {
   519  	// Format the file. Do not fix up imports, as we only moved code around
   520  	// within files.
   521  	opts := imports.Options{
   522  		Comments:   true,
   523  		TabIndent:  true,
   524  		TabWidth:   8,
   525  		FormatOnly: true,
   526  	}
   527  	code, err := imports.Process("commandline", code, &opts)
   528  	if err != nil {
   529  		return fmt.Errorf("bad parse while processing imports %q: %v", path, err)
   530  	}
   531  
   532  	if err := ioutil.WriteFile(path, code, 0644); err != nil {
   533  		return fmt.Errorf("error writing Go file to %q: %v", path, err)
   534  	}
   535  	return nil
   536  }