From 3120ba58644ef26a5233a6323becc6d892b35e73 Mon Sep 17 00:00:00 2001
From: Jon Johnson <jon.johnson@chainguard.dev>
Date: Mon, 1 May 2023 11:58:50 -0700
Subject: [PATCH] Keep order of env in crane mutate (#1683)

---
 cmd/crane/cmd/mutate.go       | 86 ++++++++++++++++++++++++++++++-----
 cmd/crane/doc/crane_mutate.md |  2 +-
 2 files changed, 75 insertions(+), 13 deletions(-)

diff --git a/cmd/crane/cmd/mutate.go b/cmd/crane/cmd/mutate.go
index a113b4251..2eb207423 100644
--- a/cmd/crane/cmd/mutate.go
+++ b/cmd/crane/cmd/mutate.go
@@ -30,8 +30,8 @@ import (
 func NewCmdMutate(options *[]crane.Option) *cobra.Command {
 	var labels map[string]string
 	var annotations map[string]string
+	var envVars keyToValue
 	var entrypoint, cmd []string
-	var envVars map[string]string
 	var newLayers []string
 	var outFile string
 	var newRef string
@@ -163,7 +163,7 @@ func NewCmdMutate(options *[]crane.Option) *cobra.Command {
 	}
 	mutateCmd.Flags().StringToStringVarP(&annotations, "annotation", "a", nil, "New annotations to add")
 	mutateCmd.Flags().StringToStringVarP(&labels, "label", "l", nil, "New labels to add")
-	mutateCmd.Flags().StringToStringVarP(&envVars, "env", "e", nil, "New envvar to add")
+	mutateCmd.Flags().VarP(&envVars, "env", "e", "New envvar to add")
 	mutateCmd.Flags().StringSliceVar(&entrypoint, "entrypoint", nil, "New entrypoint to set")
 	mutateCmd.Flags().StringSliceVar(&cmd, "cmd", nil, "New cmd to set")
 	mutateCmd.Flags().StringVar(&newRepo, "repo", "", "Repository to push the mutated image to. If provided, push by digest to this repository.")
@@ -186,29 +186,91 @@ func validateKeyVals(kvPairs map[string]string) error {
 }
 
 // setEnvVars override envvars in a config
-func setEnvVars(cfg *v1.ConfigFile, envVars map[string]string) error {
+func setEnvVars(cfg *v1.ConfigFile, envVars keyToValue) error {
+	eMap := envVars.Map()
 	newEnv := make([]string, 0, len(cfg.Config.Env))
+	isWindows := cfg.OS == "windows"
+
+	// Keep the old values.
 	for _, old := range cfg.Config.Env {
-		split := strings.SplitN(old, "=", 2)
-		if len(split) != 2 {
+		oldKey, _, ok := strings.Cut(old, "=")
+		if !ok {
 			return fmt.Errorf("invalid key value pair in config: %s", old)
 		}
-		// keep order so override if specified again
-		oldKey := split[0]
-		if v, ok := envVars[oldKey]; ok {
-			newEnv = append(newEnv, fmt.Sprintf("%s=%s", oldKey, v))
-			delete(envVars, oldKey)
+
+		if v, ok := eMap[oldKey]; ok {
+			// Override in place to keep ordering of original env.
+			newEnv = append(newEnv, oldKey+"="+v)
+
+			// Remove this from eMap so we don't add it twice.
+			delete(eMap, oldKey)
 		} else {
 			newEnv = append(newEnv, old)
 		}
 	}
-	isWindows := cfg.OS == "windows"
-	for k, v := range envVars {
+
+	// Append the new values.
+	for _, e := range envVars.values {
+		k, v := e.key, e.value
+
+		if _, ok := eMap[k]; !ok {
+			// If we come across a value not in eMap, it means we replaced the
+			// old env in-place and deleted it from eMap, so we can skip adding.
+			continue
+		}
+
 		if isWindows {
 			k = strings.ToUpper(k)
 		}
+
 		newEnv = append(newEnv, fmt.Sprintf("%s=%s", k, v))
 	}
+
 	cfg.Config.Env = newEnv
 	return nil
 }
+
+type env struct {
+	key   string
+	value string
+}
+
+type keyToValue struct {
+	values  []env
+	changed bool
+	mapped  map[string]string
+}
+
+func (o *keyToValue) Set(val string) error {
+	before, after, ok := strings.Cut(val, "=")
+	if !ok {
+		return fmt.Errorf("%s must be formatted as key=value", val)
+	}
+
+	if !o.changed {
+		o.values = []env{}
+		o.mapped = map[string]string{}
+	}
+
+	o.values = append(o.values, env{before, after})
+	o.mapped[before] = after
+	o.changed = true
+
+	return nil
+}
+
+func (o *keyToValue) Type() string {
+	return "keyToValue"
+}
+
+func (o *keyToValue) String() string {
+	ss := make([]string, 0, len(o.values))
+	for _, e := range o.values {
+		ss = append(ss, e.key+"="+e.value)
+	}
+	return strings.Join(ss, ",")
+}
+
+func (o *keyToValue) Map() map[string]string {
+	return o.mapped
+}
diff --git a/cmd/crane/doc/crane_mutate.md b/cmd/crane/doc/crane_mutate.md
index 76e81a0be..6c484279b 100644
--- a/cmd/crane/doc/crane_mutate.md
+++ b/cmd/crane/doc/crane_mutate.md
@@ -13,7 +13,7 @@ crane mutate [flags]
       --append strings              Path to tarball to append to image
       --cmd strings                 New cmd to set
       --entrypoint strings          New entrypoint to set
-  -e, --env stringToString          New envvar to add (default [])
+  -e, --env keyToValue              New envvar to add
   -h, --help                        help for mutate
   -l, --label stringToString        New labels to add (default [])
   -o, --output string               Path to new tarball of resulting image