Improve bb main

src/pkg/bb/bbmain/cmd/main.go

@@ -6,6 +6,7 @@
 package main
 
 import (
+	"errors"
 	"log"
 	"os"
 	"path/filepath"
@@ -68,8 +69,25 @@ func ResolveUntilLastSymlink(p string) string {
 
 func run() {
 	name := filepath.Base(os.Args[0])
-	if err := bbmain.Run(name); err != nil {
-		log.Fatalf("%s: %v", name, err)
+	err := bbmain.Run(name)
+	if errors.Is(err, bbmain.ErrNotRegistered) {
+		if len(os.Args) > 1 {
+			os.Args = os.Args[1:]
+			err = bbmain.Run(filepath.Base(os.Args[0]))
+		}
+	}
+	if errors.Is(err, bbmain.ErrNotRegistered) {
+		log.SetFlags(0)
+		log.Printf("Failed to run command: %v", err)
+
+		log.Printf("Supported commands are:")
+		for _, cmd := range bbmain.ListCmds() {
+			log.Printf(" - %s", cmd)
+		}
+		os.Exit(1)
+	} else if err != nil {
+		log.SetFlags(0)
+		log.Fatalf("Failed to run command: %v", err)
 	}
 }
 
@@ -78,16 +96,3 @@ func main() {
 
 	run()
 }
-
-func init() {
-	m := func() {
-		if len(os.Args) == 1 {
-			log.Fatalf("Invalid busybox command: %q", os.Args)
-		}
-		// Use argv[1] as the name.
-		os.Args = os.Args[1:]
-		run()
-	}
-	bbmain.Register("bbdiagnose", bbmain.Noop, bbmain.ListCmds)
-	bbmain.RegisterDefault(bbmain.Noop, m)
-}

src/pkg/bb/bbmain/register.go

@@ -10,84 +10,25 @@ import (
 	"errors"
 	"fmt"
 	"os"
-	"path/filepath"
 	// There MUST NOT be any other dependencies here.
 	//
 	// It is preferred to copy minimal code necessary into this file, as
 	// dependency management for this main file is... hard.
 )
 
 // ErrNotRegistered is returned by Run if the given command is not registered.
-var ErrNotRegistered = errors.New("command not registered")
+var ErrNotRegistered = errors.New("command is not present in busybox")
 
 // Noop is a noop function.
 var Noop = func() {}
 
-// ListCmds lists bb commands and verifies symlinks.
-// It is by convention called when the bb command is invoked directly.
-// For every command, there should be a symlink in /bbin,
-// and for every symlink, there should be a command.
-// Occasionally, we have bugs that result in one of these
-// being false. Just running bb is an easy way to tell if something
-// in your image is messed up.
-func ListCmds() {
-	type known struct {
-		name string
-		bb   string
-	}
-	names := map[string]*known{}
-	g, err := filepath.Glob("/bbin/*")
-	if err != nil {
-		fmt.Printf("bb: unable to enumerate /bbin")
-	}
-
-	// First step is to assemble a list of all possible
-	// names, both from /bbin/* and our built in commands.
-	for _, l := range g {
-		if l == "/bbin/bb" {
-			continue
-		}
-		b := filepath.Base(l)
-		names[b] = &known{name: l}
-	}
-	for n := range bbCmds {
-		if n == "bb" {
-			continue
-		}
-		if c, ok := names[n]; ok {
-			c.bb = n
-			continue
-		}
-		names[n] = &known{bb: n}
-	}
-	// Now walk the array of structs.
-	// We don't sort as we don't want the
-	// footprint of bringing in the package.
-	// If you want it sorted, bb | sort
-	var hadError bool
-	for c, k := range names {
-		if len(k.name) == 0 || len(k.bb) == 0 {
-			hadError = true
-			fmt.Printf("%s:\t", c)
-			if k.name == "" {
-				fmt.Printf("NO SYMLINK\t")
-			} else {
-				fmt.Printf("%q\t", k.name)
-			}
-			if k.bb == "" {
-				fmt.Printf("NO COMMAND\n")
-			} else {
-				fmt.Printf("%s\n", k.bb)
-			}
-		}
-	}
-	if hadError {
-		fmt.Println("There is at least one problem. Known causes:")
-		fmt.Println("At least two initrds -- one compiled in to the kernel, a second supplied by the bootloader.")
-		fmt.Println("The initrd cpio was changed after creation or merged with another one.")
-		fmt.Println("When the initrd was created, files were inserted into /bbin by mistake.")
-		fmt.Println("Post boot, files were added to /bbin.")
+// ListCmds returns all supported commands.
+func ListCmds() []string {
+	var cmds []string
+	for c := range bbCmds {
+		cmds = append(cmds, c)
 	}
+	return cmds
 }
 
 type bbCmd struct {
@@ -128,7 +69,7 @@ func Run(name string) error {
 	} else if defaultCmd != nil {
 		cmd = defaultCmd
 	} else {
-		return ErrNotRegistered
+		return fmt.Errorf("%w: %s", ErrNotRegistered, name)
 	}
 	cmd.init()
 	cmd.main()

src/pkg/bb/bbmain_src.go

@@ -1,3 +1,3 @@
 package bb
 
-var bbMainSource = []byte("// Copyright 2018 the u-root Authors. All rights reserved\n// Use of this source code is governed by a BSD-style\n// license that can be found in the LICENSE file.\n\n// Package main is the busybox main.go template.\npackage main\n\nimport (\n\t\"log\"\n\t\"os\"\n\t\"path/filepath\"\n\n\t\"github.com/u-root/gobusybox/src/pkg/bb/bbmain\"\n\t// There MUST NOT be any other dependencies here.\n\t//\n\t// It is preferred to copy minimal code necessary into this file, as\n\t// dependency management for this main file is... hard.\n)\n\n// AbsSymlink returns an absolute path for the link from a file to a target.\nfunc AbsSymlink(originalFile, target string) string {\n\tif !filepath.IsAbs(originalFile) {\n\t\tvar err error\n\t\toriginalFile, err = filepath.Abs(originalFile)\n\t\tif err != nil {\n\t\t\t// This should not happen on Unix systems, or you're\n\t\t\t// already royally screwed.\n\t\t\tlog.Fatalf(\"could not determine absolute path for %v: %v\", originalFile, err)\n\t\t}\n\t}\n\t// Relative symlinks are resolved relative to the original file's\n\t// parent directory.\n\t//\n\t// E.g. /bin/defaultsh -> ../bbin/elvish\n\tif !filepath.IsAbs(target) {\n\t\treturn filepath.Join(filepath.Dir(originalFile), target)\n\t}\n\treturn target\n}\n\n// IsTargetSymlink returns true if a target of a symlink is also a symlink.\nfunc IsTargetSymlink(originalFile, target string) bool {\n\ts, err := os.Lstat(AbsSymlink(originalFile, target))\n\tif err != nil {\n\t\treturn false\n\t}\n\treturn (s.Mode() & os.ModeSymlink) == os.ModeSymlink\n}\n\n// ResolveUntilLastSymlink resolves until the last symlink.\n//\n// This is needed when we have a chain of symlinks and want the last\n// symlink, not the file pointed to (which is why we don't use\n// filepath.EvalSymlinks)\n//\n// I.e.\n//\n// /foo/bar -> ../baz/foo\n// /baz/foo -> bla\n//\n// ResolveUntilLastSymlink(/foo/bar) returns /baz/foo.\nfunc ResolveUntilLastSymlink(p string) string {\n\tfor target, err := os.Readlink(p); err == nil && IsTargetSymlink(p, target); target, err = os.Readlink(p) {\n\t\tp = AbsSymlink(p, target)\n\t}\n\treturn p\n}\n\nfunc run() {\n\tname := filepath.Base(os.Args[0])\n\tif err := bbmain.Run(name); err != nil {\n\t\tlog.Fatalf(\"%s: %v\", name, err)\n\t}\n}\n\nfunc main() {\n\tos.Args[0] = ResolveUntilLastSymlink(os.Args[0])\n\n\trun()\n}\n\nfunc init() {\n\tm := func() {\n\t\tif len(os.Args) == 1 {\n\t\t\tlog.Fatalf(\"Invalid busybox command: %q\", os.Args)\n\t\t}\n\t\t// Use argv[1] as the name.\n\t\tos.Args = os.Args[1:]\n\t\trun()\n\t}\n\tbbmain.Register(\"bbdiagnose\", bbmain.Noop, bbmain.ListCmds)\n\tbbmain.RegisterDefault(bbmain.Noop, m)\n}\n")
+var bbMainSource = []byte("// Copyright 2018 the u-root Authors. All rights reserved\n// Use of this source code is governed by a BSD-style\n// license that can be found in the LICENSE file.\n\n// Package main is the busybox main.go template.\npackage main\n\nimport (\n\t\"errors\"\n\t\"log\"\n\t\"os\"\n\t\"path/filepath\"\n\n\t\"github.com/u-root/gobusybox/src/pkg/bb/bbmain\"\n\t// There MUST NOT be any other dependencies here.\n\t//\n\t// It is preferred to copy minimal code necessary into this file, as\n\t// dependency management for this main file is... hard.\n)\n\n// AbsSymlink returns an absolute path for the link from a file to a target.\nfunc AbsSymlink(originalFile, target string) string {\n\tif !filepath.IsAbs(originalFile) {\n\t\tvar err error\n\t\toriginalFile, err = filepath.Abs(originalFile)\n\t\tif err != nil {\n\t\t\t// This should not happen on Unix systems, or you're\n\t\t\t// already royally screwed.\n\t\t\tlog.Fatalf(\"could not determine absolute path for %v: %v\", originalFile, err)\n\t\t}\n\t}\n\t// Relative symlinks are resolved relative to the original file's\n\t// parent directory.\n\t//\n\t// E.g. /bin/defaultsh -> ../bbin/elvish\n\tif !filepath.IsAbs(target) {\n\t\treturn filepath.Join(filepath.Dir(originalFile), target)\n\t}\n\treturn target\n}\n\n// IsTargetSymlink returns true if a target of a symlink is also a symlink.\nfunc IsTargetSymlink(originalFile, target string) bool {\n\ts, err := os.Lstat(AbsSymlink(originalFile, target))\n\tif err != nil {\n\t\treturn false\n\t}\n\treturn (s.Mode() & os.ModeSymlink) == os.ModeSymlink\n}\n\n// ResolveUntilLastSymlink resolves until the last symlink.\n//\n// This is needed when we have a chain of symlinks and want the last\n// symlink, not the file pointed to (which is why we don't use\n// filepath.EvalSymlinks)\n//\n// I.e.\n//\n// /foo/bar -> ../baz/foo\n// /baz/foo -> bla\n//\n// ResolveUntilLastSymlink(/foo/bar) returns /baz/foo.\nfunc ResolveUntilLastSymlink(p string) string {\n\tfor target, err := os.Readlink(p); err == nil && IsTargetSymlink(p, target); target, err = os.Readlink(p) {\n\t\tp = AbsSymlink(p, target)\n\t}\n\treturn p\n}\n\nfunc run() {\n\tname := filepath.Base(os.Args[0])\n\terr := bbmain.Run(name)\n\tif errors.Is(err, bbmain.ErrNotRegistered) {\n\t\tif len(os.Args) > 1 {\n\t\t\tos.Args = os.Args[1:]\n\t\t\terr = bbmain.Run(filepath.Base(os.Args[0]))\n\t\t}\n\t}\n\tif errors.Is(err, bbmain.ErrNotRegistered) {\n\t\tlog.SetFlags(0)\n\t\tlog.Printf(\"Failed to run command: %v\", err)\n\n\t\tlog.Printf(\"Supported commands are:\")\n\t\tfor _, cmd := range bbmain.ListCmds() {\n\t\t\tlog.Printf(\" - %s\", cmd)\n\t\t}\n\t\tos.Exit(1)\n\t} else if err != nil {\n\t\tlog.SetFlags(0)\n\t\tlog.Fatalf(\"Failed to run command: %v\", err)\n\t}\n}\n\nfunc main() {\n\tos.Args[0] = ResolveUntilLastSymlink(os.Args[0])\n\n\trun()\n}\n")

src/pkg/bb/bbregister_src.go

@@ -1,3 +1,3 @@
 package bb
 
-var bbRegisterSource = []byte("// Copyright 2018 the u-root Authors. All rights reserved\n// Use of this source code is governed by a BSD-style\n// license that can be found in the LICENSE file.\n\n// Package bbmain is the package imported by all rewritten busybox\n// command-packages to register themselves.\npackage bbmain\n\nimport (\n\t\"errors\"\n\t\"fmt\"\n\t\"os\"\n\t\"path/filepath\"\n\t// There MUST NOT be any other dependencies here.\n\t//\n\t// It is preferred to copy minimal code necessary into this file, as\n\t// dependency management for this main file is... hard.\n)\n\n// ErrNotRegistered is returned by Run if the given command is not registered.\nvar ErrNotRegistered = errors.New(\"command not registered\")\n\n// Noop is a noop function.\nvar Noop = func() {}\n\n// ListCmds lists bb commands and verifies symlinks.\n// It is by convention called when the bb command is invoked directly.\n// For every command, there should be a symlink in /bbin,\n// and for every symlink, there should be a command.\n// Occasionally, we have bugs that result in one of these\n// being false. Just running bb is an easy way to tell if something\n// in your image is messed up.\nfunc ListCmds() {\n\ttype known struct {\n\t\tname string\n\t\tbb   string\n\t}\n\tnames := map[string]*known{}\n\tg, err := filepath.Glob(\"/bbin/*\")\n\tif err != nil {\n\t\tfmt.Printf(\"bb: unable to enumerate /bbin\")\n\t}\n\n\t// First step is to assemble a list of all possible\n\t// names, both from /bbin/* and our built in commands.\n\tfor _, l := range g {\n\t\tif l == \"/bbin/bb\" {\n\t\t\tcontinue\n\t\t}\n\t\tb := filepath.Base(l)\n\t\tnames[b] = &known{name: l}\n\t}\n\tfor n := range bbCmds {\n\t\tif n == \"bb\" {\n\t\t\tcontinue\n\t\t}\n\t\tif c, ok := names[n]; ok {\n\t\t\tc.bb = n\n\t\t\tcontinue\n\t\t}\n\t\tnames[n] = &known{bb: n}\n\t}\n\t// Now walk the array of structs.\n\t// We don't sort as we don't want the\n\t// footprint of bringing in the package.\n\t// If you want it sorted, bb | sort\n\tvar hadError bool\n\tfor c, k := range names {\n\t\tif len(k.name) == 0 || len(k.bb) == 0 {\n\t\t\thadError = true\n\t\t\tfmt.Printf(\"%s:\\t\", c)\n\t\t\tif k.name == \"\" {\n\t\t\t\tfmt.Printf(\"NO SYMLINK\\t\")\n\t\t\t} else {\n\t\t\t\tfmt.Printf(\"%q\\t\", k.name)\n\t\t\t}\n\t\t\tif k.bb == \"\" {\n\t\t\t\tfmt.Printf(\"NO COMMAND\\n\")\n\t\t\t} else {\n\t\t\t\tfmt.Printf(\"%s\\n\", k.bb)\n\t\t\t}\n\t\t}\n\t}\n\tif hadError {\n\t\tfmt.Println(\"There is at least one problem. Known causes:\")\n\t\tfmt.Println(\"At least two initrds -- one compiled in to the kernel, a second supplied by the bootloader.\")\n\t\tfmt.Println(\"The initrd cpio was changed after creation or merged with another one.\")\n\t\tfmt.Println(\"When the initrd was created, files were inserted into /bbin by mistake.\")\n\t\tfmt.Println(\"Post boot, files were added to /bbin.\")\n\t}\n}\n\ntype bbCmd struct {\n\tinit, main func()\n}\n\nvar bbCmds = map[string]bbCmd{}\n\nvar defaultCmd *bbCmd\n\n// Register registers an init and main function for name.\nfunc Register(name string, init, main func()) {\n\tif _, ok := bbCmds[name]; ok {\n\t\tpanic(fmt.Sprintf(\"cannot register two commands with name %q\", name))\n\t}\n\tbbCmds[name] = bbCmd{\n\t\tinit: init,\n\t\tmain: main,\n\t}\n}\n\n// RegisterDefault registers a default init and main function.\nfunc RegisterDefault(init, main func()) {\n\tdefaultCmd = &bbCmd{\n\t\tinit: init,\n\t\tmain: main,\n\t}\n}\n\n// Run runs the command with the given name.\n//\n// If the command's main exits without calling os.Exit, Run will exit with exit\n// code 0.\nfunc Run(name string) error {\n\tvar cmd *bbCmd\n\tif c, ok := bbCmds[name]; ok {\n\t\tcmd = &c\n\t} else if defaultCmd != nil {\n\t\tcmd = defaultCmd\n\t} else {\n\t\treturn ErrNotRegistered\n\t}\n\tcmd.init()\n\tcmd.main()\n\tos.Exit(0)\n\t// Unreachable.\n\treturn nil\n}\n")
+var bbRegisterSource = []byte("// Copyright 2018 the u-root Authors. All rights reserved\n// Use of this source code is governed by a BSD-style\n// license that can be found in the LICENSE file.\n\n// Package bbmain is the package imported by all rewritten busybox\n// command-packages to register themselves.\npackage bbmain\n\nimport (\n\t\"errors\"\n\t\"fmt\"\n\t\"os\"\n\t// There MUST NOT be any other dependencies here.\n\t//\n\t// It is preferred to copy minimal code necessary into this file, as\n\t// dependency management for this main file is... hard.\n)\n\n// ErrNotRegistered is returned by Run if the given command is not registered.\nvar ErrNotRegistered = errors.New(\"command is not present in busybox\")\n\n// Noop is a noop function.\nvar Noop = func() {}\n\n// ListCmds returns all supported commands.\nfunc ListCmds() []string {\n\tvar cmds []string\n\tfor c := range bbCmds {\n\t\tcmds = append(cmds, c)\n\t}\n\treturn cmds\n}\n\ntype bbCmd struct {\n\tinit, main func()\n}\n\nvar bbCmds = map[string]bbCmd{}\n\nvar defaultCmd *bbCmd\n\n// Register registers an init and main function for name.\nfunc Register(name string, init, main func()) {\n\tif _, ok := bbCmds[name]; ok {\n\t\tpanic(fmt.Sprintf(\"cannot register two commands with name %q\", name))\n\t}\n\tbbCmds[name] = bbCmd{\n\t\tinit: init,\n\t\tmain: main,\n\t}\n}\n\n// RegisterDefault registers a default init and main function.\nfunc RegisterDefault(init, main func()) {\n\tdefaultCmd = &bbCmd{\n\t\tinit: init,\n\t\tmain: main,\n\t}\n}\n\n// Run runs the command with the given name.\n//\n// If the command's main exits without calling os.Exit, Run will exit with exit\n// code 0.\nfunc Run(name string) error {\n\tvar cmd *bbCmd\n\tif c, ok := bbCmds[name]; ok {\n\t\tcmd = &c\n\t} else if defaultCmd != nil {\n\t\tcmd = defaultCmd\n\t} else {\n\t\treturn fmt.Errorf(\"%w: %s\", ErrNotRegistered, name)\n\t}\n\tcmd.init()\n\tcmd.main()\n\tos.Exit(0)\n\t// Unreachable.\n\treturn nil\n}\n")

test/all_test.go

@@ -160,3 +160,56 @@ func TestMakeBB(t *testing.T) {
 		})
 	}
 }
+
+func TestBBSymlink(t *testing.T) {
+	if *makebb == "" {
+		t.Fatalf("Path to makebb is not set")
+	}
+	dir := t.TempDir()
+
+	cmdPath := "./implicitimport/cmd/loghello"
+	want := "Log Hello\n"
+
+	goEnv := "GO111MODULE=on"
+	binary := filepath.Join(dir, "bb")
+
+	// Build the bb.
+	t.Logf("Run: %s %s -o %s %s", goEnv, *makebb, binary, cmdPath)
+
+	cmd := exec.Command(*makebb, "-o", binary, cmdPath)
+	cmd.Env = append(os.Environ(), goEnv)
+	out, err := cmd.CombinedOutput()
+	if err != nil {
+		t.Logf("makebb: %s", string(out))
+		t.Fatalf("cmd: %v", err)
+	}
+
+	t.Run("symlink", func(t *testing.T) {
+		link := filepath.Join(dir, "loghello")
+		// Test that symlinking works.
+		if err := os.Symlink(binary, link); err != nil {
+			t.Fatal(err)
+		}
+
+		t.Logf("Run: %s", link)
+		out, err = exec.Command(link).CombinedOutput()
+		if err != nil {
+			t.Fatalf("cmd: %v", err)
+		}
+		if got := string(out); got != want {
+			t.Errorf("Output of %s = %v, want %v", link, got, want)
+		}
+	})
+
+	t.Run("argv1", func(t *testing.T) {
+		cmdName := "loghello"
+		t.Logf("Run: %s %s", binary, cmdName)
+		out, err = exec.Command(binary, cmdName).CombinedOutput()
+		if err != nil {
+			t.Fatalf("cmd: %v", err)
+		}
+		if got := string(out); got != want {
+			t.Errorf("Output of %s %s = %v, want %v", binary, cmdName, got, want)
+		}
+	})
+}