mkctr: add --env flag, fix volume support

This adds support for setting environment variables and fixes two problems
with volume support (merged in #27):

* if no volumes were specified, it added a bogus "" (empty string) volume
  which made docker fail to run the container
* if no args were given, volumes were omitted

This also pulls out some common code into a new withPlatformPrefix helper,
and removes some log spam by silently skipping over "unknown" OS layers
in the source image when discovering what OSes to build for.

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>

Author: bradfitz

mkctr.go

@@ -15,6 +15,7 @@ import (
 	"io"
 	"io/fs"
 	"log"
+	"maps"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -42,6 +43,14 @@ func withPrefix(f logf, prefix string) logf {
 	}
 }
 
+func withPlatformPrefix(f logf, p v1.Platform) logf {
+	var variantSlash string
+	if v := p.Variant; v != "" {
+		variantSlash = "/" + v
+	}
+	return withPrefix(f, fmt.Sprintf("%v/%v%s: ", p.OS, p.Architecture, variantSlash))
+}
+
 // parseFiles parses a comma-separated list of colon-separated pairs
 // into a map of filePathOnDisk -> filePathInContainer.
 func parseFiles(s string) (map[string]string, error) {
@@ -87,6 +96,7 @@ type buildParams struct {
 	verbose     bool
 	annotations map[string]string // OCI image annotations
 	volumes     map[string]struct{}
+	envVars     []string // Environment variables to add to image config
 }
 
 func main() {
@@ -107,6 +117,7 @@ func main() {
 		Annotations must be comma separated key=value pairs, i.e key1=val1,key2=val2. For a single image manifest annotations will get added to the image manifest.
 		For an image index (a multi-platform manifest list) annotations will get added to each image manifest as well as the image index.
 		Annotations with empty values are not supported.`)
+		envArg = flag.String("env", "", "comma-separated list of environment variables in KEY=value form to add to the image config")
 	)
 	flag.Parse()
 	if *tagArg == "" {
@@ -139,11 +150,13 @@ func main() {
 		log.Fatal("at least one of --files or --gopaths must be set")
 	}
 	var vols map[string]struct{}
-	for vol := range strings.SplitSeq(*volumes, ",") {
-		if vols == nil {
-			vols = make(map[string]struct{})
+	if *volumes != "" {
+		for vol := range strings.SplitSeq(*volumes, ",") {
+			if vols == nil {
+				vols = make(map[string]struct{})
+			}
+			vols[strings.TrimSpace(vol)] = struct{}{}
 		}
-		vols[strings.TrimSpace(vol)] = struct{}{}
 	}
 
 	bp := &buildParams{
@@ -159,6 +172,7 @@ func main() {
 		goarch:      strings.Split(*goarch, ","),
 		annotations: parseAnnotations(*annotations),
 		volumes:     vols,
+		envVars:     parseEnv(*envArg),
 	}
 
 	if err := fetchAndBuild(bp); err != nil {
@@ -250,11 +264,15 @@ func fetchAndBuild(bp *buildParams) error {
 		if err := bp.verifyPlatform(p); err != nil {
 			return err
 		}
-		logf := withPrefix(logf, fmt.Sprintf("%v/%v: ", p.OS, p.Architecture))
+		logf := withPlatformPrefix(logf, p)
 		img, err := createImageForBase(bp, logf, baseImage, p)
 		if err != nil {
 			return err
 		}
+		img, err = applyEnvVars(img, bp.envVars)
+		if err != nil {
+			return err
+		}
 		if !bp.publish {
 			logf("not pushing")
 			return nil
@@ -292,10 +310,13 @@ func fetchAndBuild(bp *buildParams) error {
 	var adds []mutate.IndexAddendum
 	// Try to build images for all supported platforms.
 	for _, id := range im.Manifests {
-		logf := withPrefix(logf, fmt.Sprintf("%v/%v: ", id.Platform.OS, id.Platform.Architecture))
 		if id.Platform == nil {
 			return fmt.Errorf("unknown platform for image: %v", bp.baseImage)
 		}
+		if id.Platform.OS == "unknown" {
+			continue
+		}
+		logf := withPlatformPrefix(logf, *id.Platform)
 		if err := bp.verifyPlatform(*id.Platform); err != nil {
 			logf("skipping: %v", err)
 			continue
@@ -314,15 +335,31 @@ func fetchAndBuild(bp *buildParams) error {
 		// Ensure that any provided OCI annotations are added to each OCI image manifest.
 		img = mutate.Annotations(img, bp.annotations).(v1.Image)
 
+		img, err = applyEnvVars(img, bp.envVars)
+		if err != nil {
+			return err
+		}
+
+		if bp.volumes != nil {
+			img, err = mutateConfig(img, func(c *v1.Config) error {
+				c.Volumes = bp.volumes
+				return nil
+			})
+			if err != nil {
+				return err
+			}
+		}
+
 		if args := flag.Args(); len(args) > 0 {
-			img, err = mutate.Config(img, v1.Config{
-				Cmd:     args,
-				Volumes: bp.volumes,
+			img, err = mutateConfig(img, func(c *v1.Config) error {
+				c.Cmd = args
+				return nil
 			})
 			if err != nil {
 				return err
 			}
 		}
+
 		d, err := img.Digest()
 		if err != nil {
 			return err
@@ -634,3 +671,74 @@ func parseAnnotations(s string) map[string]string {
 	}
 	return annotations
 }
+
+// parseEnv accepts a string with comma separated KEY=value pairs of environment variables
+// and returns them as a slice of "KEY=value" strings.
+func parseEnv(s string) []string {
+	if len(s) == 0 {
+		return nil
+	}
+	var envVars []string
+	for env := range strings.SplitSeq(s, ",") {
+		env = strings.TrimSpace(env)
+		if len(env) == 0 {
+			continue
+		}
+		if !strings.Contains(env, "=") {
+			continue
+		}
+		envVars = append(envVars, env)
+	}
+	return envVars
+}
+
+// applyEnvVars applies environment variables to an image config, merging with existing env vars.
+// New env vars override existing ones with the same key.
+func applyEnvVars(img v1.Image, newEnvVars []string) (v1.Image, error) {
+	if len(newEnvVars) == 0 {
+		return img, nil
+	}
+	config, err := img.ConfigFile()
+	if err != nil {
+		return nil, fmt.Errorf("error getting config: %w", err)
+	}
+
+	envMap := make(map[string]string)
+	for _, kv := range config.Config.Env {
+		if k, v, ok := strings.Cut(kv, "="); ok {
+			envMap[k] = v
+		}
+	}
+
+	for _, env := range newEnvVars {
+		if k, v, ok := strings.Cut(env, "="); ok {
+			envMap[k] = v
+		}
+	}
+
+	// Apply the merged env vars
+	return mutateConfig(img, func(c *v1.Config) error {
+		c.Env = nil
+		for _, k := range slices.Sorted(maps.Keys(envMap)) {
+			c.Env = append(c.Env, k+"="+envMap[k])
+		}
+		return nil
+	})
+}
+
+// mutateConfig returns img with its config mutated by f.
+//
+// The pointer given to f is a deep copy of the existing config,
+// so any fields untouched by f will be preserved.
+func mutateConfig(img v1.Image, f func(*v1.Config) error) (v1.Image, error) {
+	config, err := img.ConfigFile()
+	if err != nil {
+		return nil, fmt.Errorf("error getting config: %w", err)
+	}
+
+	confCopy := config.DeepCopy()
+	if err := f(&confCopy.Config); err != nil {
+		return nil, err
+	}
+	return mutate.Config(img, confCopy.Config)
+}