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