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,22 @@ import (
 	"github.com/u-root/mkuimage/uimage/templates"
 )
 
+type optionalStringVar struct {
+	s **string
+}
+
+func (o optionalStringVar) Set(s string) error {
+	*o.s = &s
+	return nil
+}
+
+func (o *optionalStringVar) String() string {
+	if *o.s == nil {
+		return ""
+	}
+	return string(**o.s)
+}
+
 // CommandFlags are flags related to Go commands to be built by mkuimage.
 type CommandFlags struct {
 	NoCommands bool
@@ -79,14 +95,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 +122,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 +158,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.")