flags: if tempdir/init/uinit/shell is not specified, default value ma…

…y come from template or base main file

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

uimage/mkuimage/cmd.go

@@ -64,22 +64,22 @@ func CreateUimage(l *llog.Logger, base []uimage.Modifier, tf *TemplateFlags, f *
  }
 
  keepTempDir := f.KeepTempDir
- if f.TempDir == "" {
- var err error
- f.TempDir, err = os.MkdirTemp("", "u-root")
+ if f.TempDir == nil {
+ tempDir, err := os.MkdirTemp("", "u-root")
  if err != nil {
  return err
  }
+ f.TempDir = &tempDir
  defer func() {
  if keepTempDir {
- l.Infof("Keeping temp dir %s", f.TempDir)
+ l.Infof("Keeping temp dir %s", tempDir)
  } else {
- os.RemoveAll(f.TempDir)
+ os.RemoveAll(tempDir)
  }
  }()
- } 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)
+ } 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)
  }
  }
 

uimage/mkuimage/cmd_test.go

@@ -39,8 +39,8 @@ func TestOpts(t *testing.T) {
  f: &Flags{
  Commands: CommandFlags{Builder: "bb"},
  ArchiveFormat: "cpio",
- Init: "init",
- Uinit: "gosh script.sh",
+ Init: String("init"),
+ Uinit: String("gosh script.sh"),
  OutputFile: "./foo.cpio",
  Files: []string{"/bin/bash"},
  },
@@ -80,7 +80,7 @@ func TestOpts(t *testing.T) {
  GOOS: "plan9",
  GOARCH: "amd64",
  BuildTags: []string{"grpcnotrace"},
- Uinit: "gosh script.sh",
+ Uinit: String("gosh script.sh"),
  Files: []string{"foobar"},
  Commands: []templates.Command{
  {
@@ -103,8 +103,8 @@ func TestOpts(t *testing.T) {
  f: &Flags{
  Commands: CommandFlags{Builder: "bb"},
  ArchiveFormat: "cpio",
- Init: "init",
- Uinit: "cat",
+ Init: String("init"),
+ Uinit: String("cat"),
  OutputFile: "./foo.cpio",
  Files: []string{"/bin/bash"},
  },

uimage/mkuimage/uflags.go

@@ -17,6 +17,23 @@ import (
  "github.com/u-root/mkuimage/uimage/templates"
 )
 
+type optionalStringVar struct {
+ s **string
+}
+
+// Set implements flag.Value.Set.
+func (o optionalStringVar) Set(s string) error {
+ *o.s = &s
+ return nil
+}
+
+func (o *optionalStringVar) String() string {
+ if o == nil || o.s == nil || *(o.s) == nil {
+ return ""
+ }
+ return **(o.s)
+}
+
 // CommandFlags are flags related to Go commands to be built by mkuimage.
 type CommandFlags struct {
  NoCommands bool
@@ -79,14 +96,19 @@ func (c *CommandFlags) Modifiers(packages ...string) ([]uimage.Modifier, error)
  }
 }
 
+// String can be used to fill in values for Init, Uinit, and Shell.
+func String(s string) *string {
+ return &s
+}
+
 // Flags are mkuimage command-line flags.
 type Flags struct {
- TempDir string
+ TempDir *string
  KeepTempDir bool
 
- Init string
- Uinit string
- Shell string
+ Init *string
+ Uinit *string
+ Shell *string
 
  Files []string
 
@@ -101,15 +123,25 @@ type Flags struct {
 // Modifiers return uimage modifiers created from the flags.
 func (f *Flags) Modifiers(packages ...string) ([]uimage.Modifier, error) {
  m := []uimage.Modifier{
- uimage.WithTempDir(f.TempDir),
- uimage.WithInit(f.Init),
- uimage.WithUinitCommand(f.Uinit),
- uimage.WithShell(f.Shell),
  uimage.WithFiles(f.Files...),
- // ArchiveFormat does not determine this, as only CPIO is supported.
- uimage.WithBaseFile(f.BaseArchive),
  uimage.WithExistingInit(f.UseExistingInit),
  }
+ if f.TempDir != nil {
+ m = append(m, uimage.WithTempDir(*f.TempDir))
+ }
+ if f.BaseArchive != "" {
+ // ArchiveFormat does not determine this, as only CPIO is supported.
+ m = append(m, uimage.WithBaseFile(f.BaseArchive))
+ }
+ if f.Init != nil {
+ m = append(m, uimage.WithInit(*f.Init))
+ }
+ if f.Uinit != nil {
+ m = append(m, uimage.WithUinitCommand(*f.Uinit))
+ }
+ if f.Shell != nil {
+ m = append(m, uimage.WithShell(*f.Shell))
+ }
  switch f.ArchiveFormat {
  case "cpio":
  m = append(m, uimage.WithCPIOOutput(f.OutputFile))
@@ -127,12 +159,12 @@ func (f *Flags) Modifiers(packages ...string) ([]uimage.Modifier, error) {
 
 // RegisterFlags registers flags.
 func (f *Flags) RegisterFlags(fs *flag.FlagSet) {
- fs.StringVar(&f.TempDir, "tmp-dir", "", "Temporary directory to build binary and archive in. Deleted after build if --keep-tmp-dir is not set.")
+ fs.Var(&optionalStringVar{&f.TempDir}, "tmp-dir", "Temporary directory to build binary and archive in. Deleted after build if --keep-tmp-dir is not set.")
  fs.BoolVar(&f.KeepTempDir, "keep-tmp-dir", f.KeepTempDir, "Keep temporary directory after build")
 
- fs.StringVar(&f.Init, "initcmd", f.Init, "Symlink target for /init. Can be an absolute path or a Go command name. Use initcmd=\"\" if you don't want the symlink.")
- fs.StringVar(&f.Uinit, "uinitcmd", f.Uinit, "Symlink target and arguments for /bin/uinit. Can be an absolute path or a Go command name, followed by command-line args. Use uinitcmd=\"\" if you don't want the symlink. E.g. -uinitcmd=\"echo foobar\"")
- fs.StringVar(&f.Shell, "defaultsh", f.Shell, "Default shell. Can be an absolute path or a Go command name. Use defaultsh=\"\" if you don't want the symlink.")
+ fs.Var(&optionalStringVar{&f.Init}, "initcmd", "Symlink target for /init. Can be an absolute path or a Go command name. Use initcmd=\"\" if you don't want the symlink.")
+ fs.Var(&optionalStringVar{&f.Uinit}, "uinitcmd", "Symlink target and arguments for /bin/uinit. Can be an absolute path or a Go command name, followed by command-line args. Use uinitcmd=\"\" if you don't want the symlink. E.g. -uinitcmd=\"echo foobar\"")
+ fs.Var(&optionalStringVar{&f.Shell}, "defaultsh", "Default shell. Can be an absolute path or a Go command name. Use defaultsh=\"\" if you don't want the symlink.")
 
  fs.Var((*uflag.Strings)(&f.Files), "files", "Additional files, directories, and binaries (with their ldd dependencies) to add to archive. Can be specified multiple times.")
 

uimage/mkuimage/uflags_test.go

@@ -8,7 +8,14 @@ import (
  "errors"
  "flag"
  "os"
+ "reflect"
  "testing"
+
+ "github.com/google/go-cmp/cmp"
+ "github.com/google/go-cmp/cmp/cmpopts"
+ "github.com/u-root/gobusybox/src/pkg/golang"
+ "github.com/u-root/mkuimage/uimage"
+ "github.com/u-root/mkuimage/uimage/builder"
 )
 
 func TestFlagErrors(t *testing.T) {
@@ -36,3 +43,114 @@ func TestFlagErrors(t *testing.T) {
  }
  }
 }
+
+func TestFlags(t *testing.T) {
+ for _, tt := range []struct {
+ input []string
+ want *Flags
+ }{
+ {
+ input: []string{"-build=bb", "-initcmd=foo"},
+ want: &Flags{
+ Init: String("foo"),
+ Commands: CommandFlags{
+ Builder: "bb",
+ BuildOpts: &golang.BuildOpts{},
+ },
+ },
+ },
+ {
+ input: []string{"-build=bb"},
+ want: &Flags{
+ Commands: CommandFlags{
+ Builder: "bb",
+ BuildOpts: &golang.BuildOpts{},
+ },
+ },
+ },
+ {
+ input: []string{"-build=bb", "-initcmd=foo", "-uinitcmd=foo bar", "-tmp-dir=bla", "-defaultsh=gosh"},
+ want: &Flags{
+ Init: String("foo"),
+ Uinit: String("foo bar"),
+ TempDir: String("bla"),
+ Shell: String("gosh"),
+ Commands: CommandFlags{
+ Builder: "bb",
+ BuildOpts: &golang.BuildOpts{},
+ },
+ },
+ },
+ {
+ input: []string{"-build=bb", "-initcmd=", "-uinitcmd=", "-tmp-dir=", "-defaultsh="},
+ want: &Flags{
+ Init: String(""),
+ Uinit: String(""),
+ TempDir: String(""),
+ Shell: String(""),
+ Commands: CommandFlags{
+ Builder: "bb",
+ BuildOpts: &golang.BuildOpts{},
+ },
+ },
+ },
+ } {
+ fs := flag.NewFlagSet("test", flag.ContinueOnError)
+ f := &Flags{}
+ f.RegisterFlags(fs)
+ if err := fs.Parse(tt.input); err != nil {
+ t.Fatal(err)
+ }
+ if !reflect.DeepEqual(f, tt.want) {
+ t.Errorf("Parse = %+v, want %+v", f, tt.want)
+ }
+ }
+}
+
+func TestFlagModifiers(t *testing.T) {
+ for _, tt := range []struct {
+ input []string
+ base []uimage.Modifier
+ want *uimage.Opts
+ cmds []string
+ }{
+ {
+ // Override modifier defaults with empty.
+ input: []string{"-build=bb", "-format=cpio", "-initcmd=", "-uinitcmd=", "-defaultsh=", "-tmp-dir="},
+ base: []uimage.Modifier{
+ uimage.WithInit("foo"),
+ uimage.WithUinit("foo bar"),
+ uimage.WithTempDir("foo"),
+ uimage.WithShell("foo"),
+ },
+ cmds: []string{"foo"},
+ want: &uimage.Opts{
+ Env: golang.Default(),
+ Commands: []uimage.Commands{
+ {
+ Builder: &builder.GBBBuilder{},
+ BuildOpts: &golang.BuildOpts{},
+ },
+ },
+ },
+ },
+ } {
+ fs := flag.NewFlagSet("test", flag.ContinueOnError)
+ f := &Flags{}
+ f.RegisterFlags(fs)
+ if err := fs.Parse(tt.input); err != nil {
+ t.Fatal(err)
+ }
+ mods, err := f.Modifiers()
+ if err != nil {
+ t.Errorf("Modifiers = %v", err)
+ }
+ opts, err := uimage.OptionsFor(append(tt.base, mods...)...)
+ if err != nil {
+ t.Errorf("Options = %v", err)
+ }
+ if diff := cmp.Diff(opts, tt.want, cmpopts.IgnoreFields(uimage.Opts{}, "BaseArchive")); diff != "" {
+ t.Errorf("opts (-got, +want) = %v", diff)
+ }
+ }
+}

uimage/uimage.go

@@ -518,13 +518,12 @@ func WithBaseArchive(archive *cpio.Archive) Modifier {
 // If this is empty, no uinit symlink will be created, but a user may
 // still specify a command called uinit or include a /bin/uinit file.
 func WithUinitCommand(cmd string) Modifier {
- if cmd == "" {
- return nil
- }
  return func(opts *Opts) error {
  args := shlex.Split(cmd)
  if len(args) > 0 {
  opts.UinitCmd = args[0]
+ } else {
+ opts.UinitCmd = ""
  }
  if len(args) > 1 {
  opts.UinitArgs = args[1:]
@@ -544,9 +543,6 @@ func WithUinitCommand(cmd string) Modifier {
 // and append arguments from both the kernel command-line
 // (uroot.uinitargs) as well as those specified in cmd.
 func WithUinit(arg0 string, args ...string) Modifier {
- if arg0 == "" && len(args) == 0 {
- return nil
- }
  return func(opts *Opts) error {
  opts.UinitCmd = arg0
  opts.UinitArgs = args