From 6be8c5a28250b027e066005d14349ebf7eecac74 Mon Sep 17 00:00:00 2001
From: Adam Bouhenguel <adam@bouhenguel.com>
Date: Wed, 16 May 2018 16:44:12 -0700
Subject: [PATCH] Support for replacing image names in env vars

Specifically env vars that end with _IMAGE
---
 pkg/skaffold/deploy/kubectl/images.go | 117 ++++++++++++++++++++++----
 1 file changed, 101 insertions(+), 16 deletions(-)

diff --git a/pkg/skaffold/deploy/kubectl/images.go b/pkg/skaffold/deploy/kubectl/images.go
index c1d75f5db5b..04272192b38 100644
--- a/pkg/skaffold/deploy/kubectl/images.go
+++ b/pkg/skaffold/deploy/kubectl/images.go
@@ -17,6 +17,8 @@ limitations under the License.
 package kubectl
 
 import (
+	"strings"
+
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 
@@ -29,47 +31,63 @@ var warner Warner = &logrusWarner{}
 
 // ReplaceImages replaces image names in a list of manifests.
 func (l *ManifestList) ReplaceImages(builds []build.Artifact) (ManifestList, error) {
-	replacer := newImageReplacer(builds)
+	source := newImageNameSource(builds)
+	replacer := &multiReplacer{
+		any: []Replacer{
+			&imageFieldReplacer{
+				source: source,
+			},
+			&imageEnvReplacer{
+				source: source,
+			},
+		},
+	}
 
 	updated, err := l.Visit(replacer)
 	if err != nil {
 		return nil, errors.Wrap(err, "replacing images")
 	}
 
-	replacer.Check()
+	source.Check()
 	logrus.Debugln("manifests with tagged images", updated.String())
 
 	return updated, nil
 }
 
-type imageReplacer struct {
+type ReplacerSource interface {
+	NewValue(image string) (bool, string)
+}
+
+type imageNameSource struct {
 	tagsByImageName map[string]string
 	found           map[string]bool
 }
 
-func newImageReplacer(builds []build.Artifact) *imageReplacer {
+func newImageNameSource(builds []build.Artifact) *imageNameSource {
 	tagsByImageName := make(map[string]string)
 	for _, build := range builds {
 		tagsByImageName[build.ImageName] = build.Tag
 	}
 
-	return &imageReplacer{
+	return &imageNameSource{
 		tagsByImageName: tagsByImageName,
 		found:           make(map[string]bool),
 	}
 }
 
-func (r *imageReplacer) Matches(key string) bool {
-	return key == "image"
+func (r *imageNameSource) Check() {
+	for imageName := range r.tagsByImageName {
+		if !r.found[imageName] {
+			warner.Warnf("image [%s] is not used by the deployment", imageName)
+		}
+	}
 }
 
-func (r *imageReplacer) NewValue(key string, old interface{}) (bool, interface{}) {
-	image := old.(string)
-
+func (r *imageNameSource) NewValue(image string) (bool, string) {
 	parsed, err := docker.ParseReference(image)
 	if err != nil {
 		warner.Warnf("Couldn't parse image: %s", image)
-		return false, nil
+		return false, image
 	}
 
 	if tag, present := r.tagsByImageName[parsed.BaseName]; present {
@@ -83,13 +101,80 @@ func (r *imageReplacer) NewValue(key string, old interface{}) (bool, interface{}
 		}
 	}
 
-	return false, nil
+	return false, image
 }
 
-func (r *imageReplacer) Check() {
-	for imageName := range r.tagsByImageName {
-		if !r.found[imageName] {
-			warner.Warnf("image [%s] is not used by the deployment", imageName)
+type imageFieldReplacer struct {
+	source ReplacerSource
+}
+
+func (r *imageFieldReplacer) Matches(key string) bool {
+	return key == "image"
+}
+
+func (r *imageFieldReplacer) NewValue(key string, old interface{}) (bool, interface{}) {
+	return r.source.NewValue(old.(string))
+}
+
+type imageEnvReplacer struct {
+	source ReplacerSource
+}
+
+func (r *imageEnvReplacer) Matches(key string) bool {
+	return key == "env"
+}
+
+func (r *imageEnvReplacer) NewValue(key string, old interface{}) (bool, interface{}) {
+	if old.([]interface{}) == nil {
+		return false, old
+	}
+
+	replaced := false
+	for _, elt := range old.([]interface{}) {
+		if elt.(map[interface{}]interface{}) == nil {
+			break
+		}
+
+		envVarEntry := elt.(map[interface{}]interface{})
+		envVarName := envVarEntry["name"].(string)
+		if envVarName == "" {
+			break
+		}
+
+		if strings.HasSuffix(envVarName, "_IMAGE") {
+			image := envVarEntry["value"].(string)
+			if present, new := r.source.NewValue(image); present {
+				envVarEntry["value"] = new
+				replaced = true
+			}
+		}
+	}
+
+	return replaced, old
+}
+
+type multiReplacer struct {
+	any []Replacer
+}
+
+func (r *multiReplacer) matchingReplacer(key string) Replacer {
+	for _, replacer := range r.any {
+		if replacer.Matches(key) {
+			return replacer
 		}
 	}
+	return nil
+}
+
+func (r *multiReplacer) Matches(key string) bool {
+	return r.matchingReplacer(key) != nil
+}
+
+func (r *multiReplacer) NewValue(key string, old interface{}) (bool, interface{}) {
+	replacer := r.matchingReplacer(key)
+	if replacer == nil {
+		return false, old
+	}
+
+	return replacer.NewValue(key, old)
 }