Keep order of env in crane mutate (#1683)

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
+}