github.com/rck/u-root@v0.0.0-20180106144920-7eb602e381bb/pkg/uroot/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 uroot
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"go/ast"
    11  	"go/build"
    12  	"go/format"
    13  	"go/parser"
    14  	"go/token"
    15  	"io/ioutil"
    16  	"log"
    17  	"os"
    18  	"path/filepath"
    19  	"text/template"
    20  
    21  	"golang.org/x/tools/imports"
    22  
    23  	"github.com/u-root/u-root/pkg/cpio"
    24  	"github.com/u-root/u-root/pkg/golang"
    25  )
    26  
    27  // Per-package templates.
    28  const (
    29  	cmdFunc = `package main
    30  
    31  import "github.com/u-root/u-root/bbsh/cmds/{{.CmdName}}"
    32  
    33  func _forkbuiltin_{{.CmdName}}(c *Command) (err error) {
    34  	{{.CmdName}}.Main()
    35  	return
    36  }
    37  
    38  func {{.CmdName}}Init() {
    39  	addForkBuiltIn("{{.CmdName}}", _forkbuiltin_{{.CmdName}})
    40  	{{.Init}}
    41  }
    42  `
    43  
    44  	bbsetupGo = `package {{.CmdName}}
    45  
    46  import "flag"
    47  
    48  var {{.CmdName}}flag = flag.NewFlagSet("{{.CmdName}}", flag.ExitOnError)
    49  `
    50  )
    51  
    52  // init.go
    53  const initGo = `package main
    54  
    55  import (
    56  	"flag"
    57  	"fmt"
    58  	"log"
    59  	"os"
    60  	"os/exec"
    61  	"path/filepath"
    62  	"strings"
    63          "syscall"
    64  
    65  	"github.com/u-root/u-root/pkg/uroot/util"
    66  )
    67  
    68  func usage () {
    69  	n := filepath.Base(os.Args[0])
    70  	fmt.Fprintf(os.Stderr, "Usage: %s:\n", n)
    71  	flag.VisitAll(func(f *flag.Flag) {
    72  		if !strings.HasPrefix(f.Name, n+".") {
    73  			return
    74  		}
    75  		fmt.Fprintf(os.Stderr, "\tFlag %s: '%s', Default %v, Value %v\n", f.Name[len(n)+1:], f.Usage, f.Value, f.DefValue)
    76  	})
    77  }
    78  
    79  func init() {
    80  	flag.Usage = usage
    81  	// This getpid adds a bit of cost to each invocation (not much really)
    82  	// but it allows us to merge init and sh. The 600K we save is worth it.
    83  	// Figure out which init to run. We must always do this.
    84  
    85  	// log.Printf("init: os is %v, initMap %v", filepath.Base(os.Args[0]), initMap)
    86  	// we use filepath.Base in case they type something like ./cmd
    87  	if f, ok := initMap[filepath.Base(os.Args[0])]; ok {
    88  		//log.Printf("run the Init function for %v: run %v", os.Args[0], f)
    89  		f()
    90  	}
    91  
    92  	if os.Args[0] != "/init" {
    93  		//log.Printf("Skipping root file system setup since we are not /init")
    94  		return
    95  	}
    96  	if os.Getpid() != 1 {
    97  		//log.Printf("Skipping root file system setup since /init is not pid 1")
    98  		return
    99  	}
   100  	util.Rootfs()
   101  
   102          // spawn the first shell. We had been running the shell as pid 1
   103          // but that makes control tty stuff messy. We think.
   104          cloneFlags := uintptr(0)
   105  	for _, v := range []string{"/inito", "/bbin/uinit", "/bbin/rush"} {
   106  		cmd := exec.Command(v)
   107  		cmd.Stdin = os.Stdin
   108  		cmd.Stderr = os.Stderr
   109  		cmd.Stdout = os.Stdout
   110  		var fd int
   111  		cmd.SysProcAttr = &syscall.SysProcAttr{Ctty: fd, Setctty: true, Setsid: true, Cloneflags: cloneFlags}
   112  		log.Printf("Run %v", cmd)
   113  		if err := cmd.Run(); err != nil {
   114  			log.Print(err)
   115  		}
   116  		// only the first init needs its own PID space.
   117  		cloneFlags = 0
   118  	}
   119  
   120  	// This will drop us into a rush prompt, since this is the init for rush.
   121  	// That's a nice fallback for when everything goes wrong. 
   122  	return
   123  }
   124  `
   125  
   126  // Commands to skip building in bb mode. init and rush should be obvious
   127  // builtin and script we skip as we have no toolchain in this mode.
   128  var skip = map[string]struct{}{
   129  	"builtin": struct{}{},
   130  	"init":    struct{}{},
   131  	"rush":    struct{}{},
   132  	"script":  struct{}{},
   133  }
   134  
   135  type bbBuilder struct {
   136  	opts    BuildOpts
   137  	bbshDir string
   138  	initMap string
   139  	af      ArchiveFiles
   140  }
   141  
   142  // BBBuild is an implementation of Build for the busybox-like u-root initramfs.
   143  //
   144  // BBBuild rewrites the source files of the packages given to create one
   145  // busybox-like binary containing all commands in `opts.Packages`.
   146  func BBBuild(opts BuildOpts) (ArchiveFiles, error) {
   147  	urootPkg, err := opts.Env.Package("github.com/u-root/u-root")
   148  	if err != nil {
   149  		return ArchiveFiles{}, err
   150  	}
   151  
   152  	bbshDir := filepath.Join(urootPkg.Dir, "bbsh")
   153  	// Blow bbsh away before trying to re-create it.
   154  	if err := os.RemoveAll(bbshDir); err != nil {
   155  		return ArchiveFiles{}, err
   156  	}
   157  	if err := os.MkdirAll(bbshDir, 0755); err != nil {
   158  		return ArchiveFiles{}, err
   159  	}
   160  
   161  	builder := &bbBuilder{
   162  		opts:    opts,
   163  		bbshDir: bbshDir,
   164  		initMap: "package main\n\nvar initMap = map[string]func() {",
   165  		af:      NewArchiveFiles(),
   166  	}
   167  
   168  	// Move and rewrite package files.
   169  	for _, pkg := range opts.Packages {
   170  		if _, ok := skip[filepath.Base(pkg)]; ok {
   171  			continue
   172  		}
   173  
   174  		if err := builder.moveCommand(pkg); err != nil {
   175  			return ArchiveFiles{}, err
   176  		}
   177  	}
   178  
   179  	// Create init.go.
   180  	if err := ioutil.WriteFile(filepath.Join(builder.bbshDir, "init.go"), []byte(initGo), 0644); err != nil {
   181  		return ArchiveFiles{}, err
   182  	}
   183  
   184  	// Move rush shell over.
   185  	p, err := opts.Env.Package("github.com/u-root/u-root/cmds/rush")
   186  	if err != nil {
   187  		return ArchiveFiles{}, err
   188  	}
   189  
   190  	if err := filepath.Walk(p.Dir, func(name string, fi os.FileInfo, err error) error {
   191  		if err != nil || fi.IsDir() {
   192  			return err
   193  		}
   194  		b, err := ioutil.ReadFile(name)
   195  		if err != nil {
   196  			return err
   197  		}
   198  		return ioutil.WriteFile(filepath.Join(builder.bbshDir, fi.Name()), b, 0644)
   199  	}); err != nil {
   200  		return ArchiveFiles{}, err
   201  	}
   202  
   203  	// Map init functions.
   204  	builder.initMap += "\n}"
   205  	if err := ioutil.WriteFile(filepath.Join(builder.bbshDir, "initmap.go"), []byte(builder.initMap), 0644); err != nil {
   206  		return ArchiveFiles{}, err
   207  	}
   208  
   209  	// Compile rush + commands to /bbin/rush.
   210  	rushPath := filepath.Join(opts.TempDir, "rush")
   211  	if err := opts.Env.Build("github.com/u-root/u-root/bbsh", rushPath, golang.BuildOpts{}); err != nil {
   212  		return ArchiveFiles{}, err
   213  	}
   214  	if err := builder.af.AddFile(rushPath, "bbin/rush"); err != nil {
   215  		return ArchiveFiles{}, err
   216  	}
   217  
   218  	// Symlink from /init to rush.
   219  	if err := builder.af.AddRecord(cpio.Symlink("init", "/bbin/rush")); err != nil {
   220  		return ArchiveFiles{}, err
   221  	}
   222  	return builder.af, nil
   223  }
   224  
   225  type CommandTemplate struct {
   226  	Gopath  string
   227  	CmdName string
   228  	Init    string
   229  }
   230  
   231  type Package struct {
   232  	name string
   233  	pkg  *build.Package
   234  	fset *token.FileSet
   235  	ast  *ast.Package
   236  
   237  	initCount uint
   238  	init      string
   239  }
   240  
   241  func (p *Package) CommandTemplate() CommandTemplate {
   242  	return CommandTemplate{
   243  		Gopath:  p.pkg.Root,
   244  		CmdName: p.name,
   245  		Init:    p.init,
   246  	}
   247  }
   248  
   249  func (p *Package) rewriteFile(opts BuildOpts, f *ast.File) bool {
   250  	// Inspect the AST and change all instances of main()
   251  	var pos token.Position
   252  	hasMain := false
   253  	ast.Inspect(f, func(n ast.Node) bool {
   254  		switch x := n.(type) {
   255  		// This is rather gross. Arguably, so is the way that Go has
   256  		// embedded build information in comments ... this import
   257  		// comment attachment to a package came in 1.4, a few years
   258  		// ago, and it only just bit us with one file in upspin. So we
   259  		// go with gross.
   260  		case *ast.Ident:
   261  			// We assume the first identifier is the package id.
   262  			if !pos.IsValid() {
   263  				pos = p.fset.Position(x.Pos())
   264  			}
   265  
   266  		case *ast.File:
   267  			x.Name.Name = p.name
   268  
   269  		case *ast.FuncDecl:
   270  			if x.Name.Name == "main" {
   271  				x.Name.Name = fmt.Sprintf("Main")
   272  				// Append a return.
   273  				x.Body.List = append(x.Body.List, &ast.ReturnStmt{})
   274  				hasMain = true
   275  			}
   276  			if x.Recv == nil && x.Name.Name == "init" {
   277  				x.Name.Name = fmt.Sprintf("Init%d", p.initCount)
   278  				p.init += fmt.Sprintf("%s.Init%d()\n", p.name, p.initCount)
   279  				p.initCount++
   280  			}
   281  
   282  		// Rewrite use of the flag package.
   283  		//
   284  		// The flag package uses a global variable to contain all
   285  		// flags. This works poorly with the busybox mode, as flags may
   286  		// conflict, so as part of turning commands into packages, we
   287  		// rewrite their use of flags to use a package-private FlagSet.
   288  		//
   289  		// bbsetup.go contains a var for the package flagset with
   290  		// params (packagename, os.ExitOnError).
   291  		//
   292  		// We rewrite arguments for calls to flag.Parse from () to
   293  		// (os.Args[1:]). We rewrite all other uses of 'flag.' to
   294  		// '"commandname"+flag.'.
   295  		case *ast.CallExpr:
   296  			switch s := x.Fun.(type) {
   297  			case *ast.SelectorExpr:
   298  				switch i := s.X.(type) {
   299  				case *ast.Ident:
   300  					if i.Name != "flag" {
   301  						break
   302  					}
   303  					switch s.Sel.Name {
   304  					case "Parse":
   305  						i.Name = p.name + "flag"
   306  						//debug("Found a call to flag.Parse")
   307  						x.Args = []ast.Expr{&ast.Ident{Name: "os.Args[1:]"}}
   308  					case "Flag":
   309  					default:
   310  						i.Name = p.name + "flag"
   311  					}
   312  				}
   313  			}
   314  
   315  		}
   316  		return true
   317  	})
   318  
   319  	// Now we change any import names attached to package declarations.  We
   320  	// just upcase it for now; it makes it easy to look in bbsh for things
   321  	// we changed, e.g. grep -r bbsh Import is useful.
   322  	for _, cg := range f.Comments {
   323  		for _, c := range cg.List {
   324  			l := p.fset.Position(c.Pos()).Line
   325  			if l != pos.Line {
   326  				continue
   327  			}
   328  			if c.Text[0:9] == "// import" {
   329  				c.Text = "// Import" + c.Text[9:]
   330  			}
   331  		}
   332  	}
   333  	return hasMain
   334  }
   335  
   336  func writeFile(path string, fset *token.FileSet, f *ast.File) error {
   337  	var buf bytes.Buffer
   338  	if err := format.Node(&buf, fset, f); err != nil {
   339  		return fmt.Errorf("error formating: %v", err)
   340  	}
   341  	return writeGoFile(path, buf.Bytes())
   342  }
   343  
   344  func writeGoFile(path string, code []byte) error {
   345  	// Fix up imports.
   346  	opts := imports.Options{
   347  		Fragment:  true,
   348  		AllErrors: true,
   349  		Comments:  true,
   350  		TabIndent: true,
   351  		TabWidth:  8,
   352  	}
   353  	fullCode, err := imports.Process("commandline", code, &opts)
   354  	if err != nil {
   355  		return fmt.Errorf("bad parse %q: %v", string(code), err)
   356  	}
   357  
   358  	if err := ioutil.WriteFile(path, fullCode, 0644); err != nil {
   359  		return fmt.Errorf("error writing to %q: %v", path, err)
   360  	}
   361  	return nil
   362  }
   363  
   364  func (p *Package) writeTemplate(path string, text string) error {
   365  	var b bytes.Buffer
   366  	t := template.Must(template.New("uroot").Parse(text))
   367  	if err := t.Execute(&b, p.CommandTemplate()); err != nil {
   368  		return fmt.Errorf("spec %v: %v", text, err)
   369  	}
   370  
   371  	return writeGoFile(path, b.Bytes())
   372  }
   373  
   374  func getPackage(opts BuildOpts, importPath string) (*Package, error) {
   375  	p, err := opts.Env.Package(importPath)
   376  	if err != nil {
   377  		return nil, err
   378  	}
   379  
   380  	name := filepath.Base(p.Dir)
   381  	if !p.IsCommand() {
   382  		return nil, fmt.Errorf("package %q is not a command and cannot be included in bb", name)
   383  	}
   384  
   385  	fset := token.NewFileSet()
   386  	pars, err := parser.ParseDir(fset, p.Dir, nil, parser.ParseComments)
   387  	if err != nil {
   388  		log.Printf("can't parsedir %q: %v", p.Dir, err)
   389  		return nil, nil
   390  	}
   391  
   392  	return &Package{
   393  		pkg:  p,
   394  		fset: fset,
   395  		ast:  pars[p.Name],
   396  		name: name,
   397  	}, nil
   398  }
   399  
   400  func (b *bbBuilder) moveCommand(pkgPath string) error {
   401  	p, err := getPackage(b.opts, pkgPath)
   402  	if err != nil {
   403  		return err
   404  	}
   405  	if p == nil {
   406  		return nil
   407  	}
   408  
   409  	pkgDir := filepath.Join(b.bbshDir, "cmds", p.name)
   410  	if err := os.MkdirAll(pkgDir, 0755); err != nil {
   411  		return err
   412  	}
   413  
   414  	var hasMain bool
   415  	for filePath, sourceFile := range p.ast.Files {
   416  		if hasMainFile := p.rewriteFile(b.opts, sourceFile); hasMainFile {
   417  			hasMain = true
   418  		}
   419  
   420  		path := filepath.Join(pkgDir, filepath.Base(filePath))
   421  		if err := writeFile(path, p.fset, sourceFile); err != nil {
   422  			return err
   423  		}
   424  	}
   425  
   426  	if !hasMain {
   427  		return os.RemoveAll(pkgDir)
   428  	}
   429  
   430  	bbsetupPath := filepath.Join(b.bbshDir, "cmds", p.name, "bbsetup.go")
   431  	if err := p.writeTemplate(bbsetupPath, bbsetupGo); err != nil {
   432  		return err
   433  	}
   434  
   435  	cmdPath := filepath.Join(b.bbshDir, fmt.Sprintf("cmd_%s.go", p.name))
   436  	if err := p.writeTemplate(cmdPath, cmdFunc); err != nil {
   437  		return err
   438  	}
   439  
   440  	b.initMap += "\n\t\"" + p.name + "\": " + p.name + "Init,"
   441  
   442  	// Add a symlink to our bbsh.
   443  	return b.af.AddRecord(cpio.Symlink(filepath.Join("bbin", p.name), "/bbin/rush"))
   444  }