Reusable flags package

Signed-off-by: Chris Koch <[email protected]>

cmd/mkuimage/main.go

@@ -16,78 +16,16 @@ import (
 	"strings"
 
 	"github.com/dustin/go-humanize"
-	"github.com/hugelgupf/go-shlex"
 	"github.com/u-root/gobusybox/src/pkg/golang"
-	"github.com/u-root/gobusybox/src/pkg/uflag"
 	"github.com/u-root/mkuimage/uimage"
-	"github.com/u-root/mkuimage/uimage/builder"
-	"github.com/u-root/mkuimage/uimage/initramfs"
+	"github.com/u-root/mkuimage/uimage/uflags"
 	"github.com/u-root/uio/llog"
 )
 
-// multiFlag is used for flags that support multiple invocations, e.g. -files.
-type multiFlag []string
-
-func (m *multiFlag) String() string {
-	return fmt.Sprint(*m)
-}
-
-// Set implements flag.Value.Set.
-func (m *multiFlag) Set(value string) error {
-	*m = append(*m, value)
-	return nil
-}
-
 var (
 	errEmptyFilesArg = errors.New("empty argument to -files")
 )
 
-// Flags for u-root builder.
-var (
-	build, format, tmpDir, basePath, outputPath *string
-	uinitCmd, initCmd                           *string
-	defaultShell                                *string
-	useExistingInit                             *bool
-	noCommands                                  *bool
-	extraFiles                                  multiFlag
-	shellbang                                   *bool
-	// For the new "filepath only" logic.
-	urootSourceDir *string
-)
-
-func init() {
-	var sh string
-	switch golang.Default().GOOS {
-	case "plan9":
-		sh = ""
-	default:
-		sh = "elvish"
-	}
-
-	build = flag.String("build", "gbb", "u-root build format (e.g. bb/gbb or binary).")
-	format = flag.String("format", "cpio", "Archival format.")
-
-	tmpDir = flag.String("tmpdir", "", "Temporary directory to put binaries in.")
-
-	basePath = flag.String("base", "", "Base archive to add files to. By default, this is a couple of directories like /bin, /etc, etc. u-root has a default internally supplied set of files; use base=/dev/null if you don't want any base files.")
-	useExistingInit = flag.Bool("useinit", false, "Use existing init from base archive (only if --base was specified).")
-	outputPath = flag.String("o", "", "Path to output initramfs file.")
-
-	initCmd = flag.String("initcmd", "init", "Symlink target for /init. Can be an absolute path or a u-root command name. Use initcmd=\"\" if you don't want the symlink.")
-	uinitCmd = flag.String("uinitcmd", "", "Symlink target and arguments for /bin/uinit. Can be an absolute path or a u-root command name. Use uinitcmd=\"\" if you don't want the symlink. E.g. -uinitcmd=\"echo foobar\"")
-	defaultShell = flag.String("defaultsh", sh, "Default shell. Can be an absolute path or a u-root command name. Use defaultsh=\"\" if you don't want the symlink.")
-
-	noCommands = flag.Bool("nocmd", false, "Build no Go commands; initramfs only")
-
-	flag.Var(&extraFiles, "files", "Additional files, directories, and binaries (with their ldd dependencies) to add to archive. Can be specified multiple times.")
-
-	shellbang = flag.Bool("shellbang", false, "Use #! instead of symlinks for busybox")
-
-	// Flag for the new filepath only mode. This will be required to find the u-root commands and make templates work
-	// In almost every case, "." is fine.
-	urootSourceDir = flag.String("uroot-source", ".", "Path to the locally checked out u-root source tree in case commands from there are desired.")
-}
-
 // checkArgs checks for common mistakes that cause confusion.
 //  1. -files as the last argument
 //  2. -files followed by any switch, indicating a shell expansion problem
@@ -123,41 +61,42 @@ func main() {
 		log.Fatal(err)
 	}
 
-	gbbOpts := &golang.BuildOpts{}
-	gbbOpts.RegisterFlags(flag.CommandLine)
-	// Register an alias for -go-no-strip for backwards compatibility.
-	flag.CommandLine.BoolVar(&gbbOpts.NoStrip, "no-strip", false, "Build unstripped binaries")
+	var sh string
+	if golang.Default().GOOS != "plan9" {
+		sh = "gosh"
+	}
 
-	env := golang.Default()
-	env.RegisterFlags(flag.CommandLine)
-	tags := (*uflag.Strings)(&env.BuildTags)
-	flag.CommandLine.Var(tags, "tags", "Go build tags -- repeat the flag for multiple values")
+	env := golang.Default(golang.DisableCGO())
+	f := &uflags.Flags{
+		Commands: uflags.CommandFlags{
+			Env:       env,
+			Builder:   "bb",
+			BuildOpts: &golang.BuildOpts{},
+		},
+		Init:          "init",
+		Shell:         sh,
+		ArchiveFormat: "cpio",
+		OutputFile:    defaultFile(env),
+	}
+	f.RegisterFlags(flag.CommandLine)
 
 	l := llog.Default()
 	l.RegisterVerboseFlag(flag.CommandLine, "v", slog.LevelDebug)
 	flag.Parse()
 
-	if usrc := os.Getenv("UROOT_SOURCE"); usrc != "" && *urootSourceDir == "" {
-		*urootSourceDir = usrc
-	}
-
-	if env.CgoEnabled {
-		l.Infof("Disabling CGO for u-root...")
-		env.CgoEnabled = false
-	}
 	l.Infof("Build environment: %s", env)
 	if env.GOOS != "linux" {
 		l.Warnf("GOOS is not linux. Did you mean to set GOOS=linux?")
 	}
 
 	// Main is in a separate functions so defers run on return.
-	if err := Main(l, env, gbbOpts); err != nil {
+	if err := Main(l, f); err != nil {
 		l.Errorf("Build error: %v", err)
 		return
 	}
 
-	if stat, err := os.Stat(*outputPath); err == nil {
-		l.Infof("Successfully built %q (size %d bytes -- %s).", *outputPath, stat.Size(), humanize.IBytes(uint64(stat.Size())))
+	if stat, err := os.Stat(f.OutputFile); err == nil {
+		l.Infof("Successfully built %q (size %d bytes -- %s).", f.OutputFile, stat.Size(), humanize.IBytes(uint64(stat.Size())))
 	}
 }
 
@@ -176,26 +115,6 @@ func isRecommendedVersion(v string) bool {
 	return false
 }
 
-func getReader(format string, path string) initramfs.ReadOpener {
-	switch format {
-	case "cpio":
-		return &initramfs.CPIOFile{Path: path}
-	default:
-		return nil
-	}
-}
-
-func getWriter(format string, path string) initramfs.WriteOpener {
-	switch format {
-	case "cpio":
-		return &initramfs.CPIOFile{Path: path}
-	case "dir":
-		return &initramfs.Dir{Path: path}
-	default:
-		return nil
-	}
-}
-
 func defaultFile(env *golang.Environ) string {
 	if len(env.GOOS) == 0 || len(env.GOARCH) == 0 {
 		return "/tmp/initramfs.cpio"
@@ -205,7 +124,8 @@ func defaultFile(env *golang.Environ) string {
 
 // Main is a separate function so defers are run on return, which they wouldn't
 // on exit.
-func Main(l *llog.Logger, env *golang.Environ, buildOpts *golang.BuildOpts) error {
+func Main(l *llog.Logger, f *uflags.Flags) error {
+	env := f.Commands.Env
 	v, err := env.Version()
 	if err != nil {
 		l.Infof("Could not get environment's Go version, using runtime's version: %v", err)
@@ -219,73 +139,32 @@ func Main(l *llog.Logger, env *golang.Environ, buildOpts *golang.BuildOpts) erro
 			v, recommendedVersions, recommendedVersions[0])
 	}
 
-	if *outputPath == "" && *format == "cpio" {
-		*outputPath = defaultFile(env)
-	}
-	output := getWriter(*format, *outputPath)
-
-	var base initramfs.ReadOpener
-	base = &initramfs.Archive{Archive: uimage.DefaultRamfs()}
-	if *basePath != "" {
-		base = getReader(*format, *basePath)
-	}
-
-	tempDir := *tmpDir
-	if tempDir == "" {
+	if f.TempDir == "" {
 		var err error
-		tempDir, err = os.MkdirTemp("", "u-root")
+		f.TempDir, err = os.MkdirTemp("", "u-root")
 		if err != nil {
 			return err
 		}
-		defer os.RemoveAll(tempDir)
-	} else if _, err := os.Stat(tempDir); os.IsNotExist(err) {
-		if err := os.MkdirAll(tempDir, 0o755); err != nil {
-			return fmt.Errorf("temporary directory %q did not exist; tried to mkdir but failed: %v", tempDir, err)
+		if f.KeepTempDir {
+			defer func() {
+				l.Infof("Keeping temp dir %s", f.TempDir)
+			}()
+		} else {
+			defer os.RemoveAll(f.TempDir)
 		}
-	}
-
-	var c []uimage.Commands
-	if !*noCommands {
-		var b builder.Builder
-		switch *build {
-		case "bb", "gbb":
-			b = builder.GBBBuilder{ShellBang: *shellbang}
-		case "binary":
-			b = builder.BinaryBuilder{}
-		default:
-			return fmt.Errorf("could not find builder %q", *build)
+	} else if _, err := os.Stat(f.TempDir); os.IsNotExist(err) {
+		if err := os.MkdirAll(f.TempDir, 0o755); err != nil {
+			return fmt.Errorf("temporary directory %q did not exist; tried to mkdir but failed: %v", f.TempDir, err)
 		}
-
-		pkgs := flag.Args()
-		if len(pkgs) == 0 {
-			pkgs = []string{"github.com/u-root/u-root/cmds/core/*"}
-		}
-
-		c = append(c, uimage.Commands{
-			Builder:   b,
-			Packages:  pkgs,
-			BuildOpts: buildOpts,
-		})
 	}
 
-	opts := uimage.Opts{
-		Env:             env,
-		Commands:        c,
-		UrootSource:     *urootSourceDir,
-		TempDir:         tempDir,
-		ExtraFiles:      extraFiles,
-		OutputFile:      output,
-		BaseArchive:     base,
-		UseExistingInit: *useExistingInit,
-		InitCmd:         *initCmd,
-		DefaultShell:    *defaultShell,
+	// Set default output.
+	m := []uimage.Modifier{
+		uimage.WithCPIOOutput(defaultFile(env)),
 	}
-	uinitArgs := shlex.Split(*uinitCmd)
-	if len(uinitArgs) > 0 {
-		opts.UinitCmd = uinitArgs[0]
-	}
-	if len(uinitArgs) > 1 {
-		opts.UinitArgs = uinitArgs[1:]
+	more, err := f.Modifiers(flag.Args()...)
+	if err != nil {
+		return err
 	}
-	return uimage.CreateInitramfs(l, opts)
+	return uimage.Create(l, append(m, more...)...)
 }

uimage/builder/builder.go

@@ -16,10 +16,10 @@ import (
 
 var (
 	// Busybox is a shared GBBBuilder instance.
-	Busybox = GBBBuilder{}
+	Busybox = &GBBBuilder{}
 
 	// Binary is a shared BinaryBuilder instance.
-	Binary = BinaryBuilder{}
+	Binary = &BinaryBuilder{}
 )
 
 // Possible build errors.

uimage/builder/gbb.go

@@ -48,7 +48,7 @@ func (GBBBuilder) DefaultBinaryDir() string {
 }
 
 // Build is an implementation of Builder.Build for a busybox-like initramfs.
-func (b GBBBuilder) Build(l *llog.Logger, af *initramfs.Files, opts Opts) error {
+func (b *GBBBuilder) Build(l *llog.Logger, af *initramfs.Files, opts Opts) error {
 	// Build the busybox binary.
 	if len(opts.TempDir) == 0 {
 		return ErrTempDirMissing