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 }