Wire templates into mkuimage

mkuimage defaults for init & shell are removed. Templates can be used instead. Signed-off-by: Chris Koch <[email protected]>
u-root · Feb 23, 2024 · 2765692 · 2765692
1 parent 04c451a
commit 2765692
Show file tree

Hide file tree

Showing 10 changed files with 376 additions and 30 deletions.
diff --git a/cmd/mkuimage/main.go b/cmd/mkuimage/main.go
@@ -55,38 +55,33 @@ func checkArgs(args ...string) error {
 }
 
 func main() {
+	log.SetFlags(log.Ltime)
 	if err := checkArgs(os.Args...); err != nil {
 		log.Fatal(err)
 	}
 
-	var sh string
-	if golang.Default().GOOS != "plan9" {
-		sh = "gosh"
-	}
-
 	env := golang.Default(golang.DisableCGO())
 	f := &mkuimage.Flags{
-		Commands: mkuimage.CommandFlags{
-			Builder:   "bb",
-			BuildOpts: &golang.BuildOpts{},
-		},
-		Init:          "init",
-		Shell:         sh,
+		Commands:      mkuimage.CommandFlags{Builder: "bb"},
 		ArchiveFormat: "cpio",
 		OutputFile:    defaultFile(env),
 	}
 	f.RegisterFlags(flag.CommandLine)
 
 	l := llog.Default()
 	l.RegisterVerboseFlag(flag.CommandLine, "v", slog.LevelDebug)
+
+	tf := &mkuimage.TemplateFlags{}
+	tf.RegisterFlags(flag.CommandLine)
 	flag.Parse()
 
+	// Set defaults.
 	m := []uimage.Modifier{
 		uimage.WithReplaceEnv(env),
 		uimage.WithBaseArchive(uimage.DefaultRamfs()),
 		uimage.WithCPIOOutput(defaultFile(env)),
 	}
-	if err := mkuimage.CreateUimage(l, m, f, flag.Args()); err != nil {
+	if err := mkuimage.CreateUimage(l, m, tf, f, flag.Args()); err != nil {
 		l.Errorf("mkuimage error: %v", err)
 		os.Exit(1)
 	}

diff --git a/cmd/mkuimage/main_test.go b/cmd/mkuimage/main_test.go
@@ -233,15 +233,6 @@ func TestUrootCmdline(t *testing.T) {
 			},
 			exitCode: 1,
 		},
-		{
-			name: "build invalid",
-			args: []string{
-				"-build=source",
-				"github.com/u-root/u-root/cmds/core/init",
-				"github.com/u-root/u-root/cmds/core/echo",
-			},
-			exitCode: 1,
-		},
 		{
 			name: "arch invalid preserves temp dir",
 			env:  []string{"GOARCH=doesnotexist"},
@@ -263,6 +254,56 @@ func TestUrootCmdline(t *testing.T) {
 			exitCode:   1,
 			wantOutput: dirExists(tempDir),
 		},
+		{
+			name: "template config",
+			args: []string{"-config-file=./testdata/test-config.yaml", "-v", "-config=coreconf"},
+			validators: []itest.ArchiveValidator{
+				itest.HasRecord{R: cpio.CharDev("dev/tty", 0o666, 5, 0)},
+				itest.HasFile{Path: "bbin/bb"},
+				itest.HasRecord{R: cpio.Symlink("bbin/echo", "bb")},
+				itest.HasRecord{R: cpio.Symlink("bbin/ip", "bb")},
+				itest.HasRecord{R: cpio.Symlink("bbin/init", "bb")},
+				itest.HasRecord{R: cpio.Symlink("init", "bbin/init")},
+				itest.HasRecord{R: cpio.Symlink("bin/sh", "../bbin/echo")},
+				itest.HasRecord{R: cpio.Symlink("bin/uinit", "../bbin/echo")},
+				itest.HasRecord{R: cpio.Symlink("bin/defaultsh", "../bbin/echo")},
+				itest.HasContent{
+					Path:    "etc/uinit.flags",
+					Content: "\"script.sh\"",
+				},
+			},
+		},
+		{
+			name: "template command",
+			args: []string{"-config-file=./testdata/test-config.yaml", "-v", "core"},
+			validators: []itest.ArchiveValidator{
+				itest.HasRecord{R: cpio.CharDev("dev/tty", 0o666, 5, 0)},
+				itest.HasFile{Path: "bbin/bb"},
+				itest.HasRecord{R: cpio.Symlink("bbin/echo", "bb")},
+				itest.HasRecord{R: cpio.Symlink("bbin/ip", "bb")},
+				itest.HasRecord{R: cpio.Symlink("bbin/init", "bb")},
+			},
+		},
+		{
+			name:     "template config not found",
+			args:     []string{"-config-file=./testdata/test-config.yaml", "-v", "-config=foobar"},
+			exitCode: 1,
+		},
+		{
+			name:     "builder not found",
+			args:     []string{"-v", "build=source"},
+			exitCode: 1,
+		},
+		{
+			name:     "template file not found",
+			args:     []string{"-v", "-config-file=./testdata/doesnotexist"},
+			exitCode: 1,
+		},
+		{
+			name:     "config not found with no default template",
+			args:     []string{"-v", "-config=foo"},
+			exitCode: 1,
+		},
 	}
 
 	for _, tt := range append(noCmdTests, bareTests...) {

diff --git a/cmd/mkuimage/testdata/test-config.yaml b/cmd/mkuimage/testdata/test-config.yaml
@@ -0,0 +1,37 @@
+commands:
+  core:
+    - github.com/u-root/u-root/cmds/core/ip
+    - github.com/u-root/u-root/cmds/core/init
+    - github.com/u-root/u-root/cmds/core/echo
+
+  minimal:
+    - github.com/u-root/u-root/cmds/core/ls
+    - github.com/u-root/u-root/cmds/core/init
+
+  plan9:
+    - github.com/u-root/u-root/cmds/core/ls
+    - github.com/u-root/u-root/cmds/core/init
+    - github.com/u-root/u-root/cmds/core/cat
+
+configs:
+  plan9:
+    goarch: amd64
+    goos: plan9
+    build_tags: [grpcnotrace]
+    files:
+      - /bin/bash
+    init: init
+    uinit: cat script.sh
+    shell: cat
+    commands:
+      - builder: bb
+        commands: [plan9]
+
+  coreconf:
+    build_tags: [grpcnotrace]
+    init: init
+    uinit: echo script.sh
+    shell: echo
+    commands:
+      - builder: bb
+        commands: [core, minimal]
diff --git a/go.mod b/go.mod
@@ -4,6 +4,7 @@ go 1.21
 
 require (
 	github.com/dustin/go-humanize v1.0.1
+	github.com/google/go-cmp v0.5.9
 	github.com/hugelgupf/go-shlex v0.0.0-20200702092117-c80c9d0918fa
 	github.com/u-root/gobusybox/src v0.0.0-20240218001334-a32c1883bffa
 	github.com/u-root/u-root v0.12.0

diff --git a/uimage/mkuimage/cmd.go b/uimage/mkuimage/cmd.go
@@ -13,6 +13,7 @@ import (
 
 	"github.com/u-root/mkuimage/uimage"
 	"github.com/u-root/mkuimage/uimage/builder"
+	"github.com/u-root/mkuimage/uimage/templates"
 	"github.com/u-root/uio/llog"
 )
 
@@ -31,8 +32,37 @@ func isRecommendedVersion(v string) bool {
 	return false
 }
 
-// CreateUimage creates a uimage from the given base modifiers and flags.
-func CreateUimage(l *llog.Logger, base []uimage.Modifier, f *Flags, args []string) error {
+func uimageOpts(l *llog.Logger, m []uimage.Modifier, tpl *templates.Templates, f *Flags, conf string, cmds []string) (*uimage.Opts, error) {
+	// Evaluate template first -- template settings may always be
+	// appended/overridden by further flag-based settings.
+	if conf != "" {
+		mods, err := tpl.Uimage(conf)
+		if err != nil {
+			return nil, err
+		}
+		l.Debugf("Config: %#v", tpl.Configs[conf])
+		m = append(m, mods...)
+	}
+
+	// Expand command templates.
+	if tpl != nil {
+		cmds = tpl.CommandsFor(cmds...)
+	}
+
+	more, err := f.Modifiers(cmds...)
+	if err != nil {
+		return nil, err
+	}
+	return uimage.OptionsFor(append(m, more...)...)
+}
+
+// CreateUimage creates a uimage with the given base modifiers and flags, using args as the list of commands.
+func CreateUimage(l *llog.Logger, base []uimage.Modifier, tf *TemplateFlags, f *Flags, args []string) error {
+	tpl, err := tf.Get()
+	if err != nil {
+		return fmt.Errorf("failed to get template: %w", err)
+	}
+
 	keepTempDir := f.KeepTempDir
 	if f.TempDir == "" {
 		var err error
@@ -53,13 +83,7 @@ func CreateUimage(l *llog.Logger, base []uimage.Modifier, f *Flags, args []strin
 		}
 	}
 
-	// Set defaults.
-	more, err := f.Modifiers(args...)
-	if err != nil {
-		return err
-	}
-
-	opts, err := uimage.OptionsFor(append(base, more...)...)
+	opts, err := uimageOpts(l, base, tpl, f, tf.Config, args)
 	if err != nil {
 		return err
 	}
@@ -70,6 +94,7 @@ func CreateUimage(l *llog.Logger, base []uimage.Modifier, f *Flags, args []strin
 	if env.GOOS != "linux" {
 		l.Warnf("GOOS is not linux. Did you mean to set GOOS=linux?")
 	}
+
 	v, err := env.Version()
 	if err != nil {
 		l.Infof("Could not get environment's Go version, using runtime's version: %v", err)