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 }