From f311ba8d4f0cc6b5c307ef4cd5255e41c735eda9 Mon Sep 17 00:00:00 2001
From: Yujun Zhang <zhangyujun@gmail.com>
Date: Sat, 9 Mar 2019 12:04:18 +0800
Subject: [PATCH] Support custom config for image transformer

---
 .../config/defaultconfig/defaultconfig.go     |  4 +-
 .../defaultconfig/{image.go => images.go}     | 14 ++--
 pkg/transformers/config/transformerconfig.go  |  2 +-
 pkg/transformers/image.go                     | 67 ++++---------------
 4 files changed, 26 insertions(+), 61 deletions(-)
 rename pkg/transformers/config/defaultconfig/{image.go => images.go} (73%)

diff --git a/pkg/transformers/config/defaultconfig/defaultconfig.go b/pkg/transformers/config/defaultconfig/defaultconfig.go
index 10fe504ac..d34d9a934 100644
--- a/pkg/transformers/config/defaultconfig/defaultconfig.go
+++ b/pkg/transformers/config/defaultconfig/defaultconfig.go
@@ -31,7 +31,7 @@ func GetDefaultFieldSpecs() []byte {
 		[]byte(namespaceFieldSpecs),
 		[]byte(varReferenceFieldSpecs),
 		[]byte(nameReferenceFieldSpecs),
-		[]byte(imageFieldSpecs),
+		[]byte(imagesFieldSpecs),
 	}
 	return bytes.Join(configData, []byte("\n"))
 }
@@ -46,6 +46,6 @@ func GetDefaultFieldSpecsAsMap() map[string]string {
 	result["namespace"] = namespaceFieldSpecs
 	result["varreference"] = varReferenceFieldSpecs
 	result["namereference"] = nameReferenceFieldSpecs
-	result["image"] = imageFieldSpecs
+	result["images"] = imagesFieldSpecs
 	return result
 }
diff --git a/pkg/transformers/config/defaultconfig/image.go b/pkg/transformers/config/defaultconfig/images.go
similarity index 73%
rename from pkg/transformers/config/defaultconfig/image.go
rename to pkg/transformers/config/defaultconfig/images.go
index d571c46aa..f67521183 100644
--- a/pkg/transformers/config/defaultconfig/image.go
+++ b/pkg/transformers/config/defaultconfig/images.go
@@ -17,9 +17,15 @@ limitations under the License.
 package defaultconfig
 
 const (
-	imageFieldSpecs = `
-image:
-- path: containers
-- path: initContainers
+	imagesFieldSpecs = `
+images:
+- kind: Pod
+  path: spec/initContainers
+- kind: Pod
+  path: spec/containers
+- kind: Deployment
+  path: spec/template/spec/initContainers
+- kind: Deployment
+  path: spec/template/spec/containers
 `
 )
diff --git a/pkg/transformers/config/transformerconfig.go b/pkg/transformers/config/transformerconfig.go
index 7ba048524..a00c6da66 100644
--- a/pkg/transformers/config/transformerconfig.go
+++ b/pkg/transformers/config/transformerconfig.go
@@ -34,7 +34,7 @@ type TransformerConfig struct {
 	CommonAnnotations fsSlice  `json:"commonAnnotations,omitempty" yaml:"commonAnnotations,omitempty"`
 	NameReference     nbrSlice `json:"nameReference,omitempty" yaml:"nameReference,omitempty"`
 	VarReference      fsSlice  `json:"varReference,omitempty" yaml:"varReference,omitempty"`
-	Images            fsSlice  `json:"image,omitempty" yaml:"image,omitempty"`
+	Images            fsSlice  `json:"images,omitempty" yaml:"images,omitempty"`
 }
 
 // MakeEmptyConfig returns an empty TransformerConfig object
diff --git a/pkg/transformers/image.go b/pkg/transformers/image.go
index 5ca3b86a1..c60f0eeb8 100644
--- a/pkg/transformers/image.go
+++ b/pkg/transformers/image.go
@@ -40,47 +40,30 @@ func NewImageTransformer(slice []image.Image, fs []config.FieldSpec) (Transforme
 }
 
 // Transform finds the matching images and replaces name, tag and/or digest
-func (pt *imageTransformer) Transform(resources resmap.ResMap) error {
+func (pt *imageTransformer) Transform(m resmap.ResMap) error {
 	if len(pt.images) == 0 {
 		return nil
 	}
-	for _, res := range resources {
-		err := pt.findAndReplaceImage(res.Map())
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-/*
- findAndReplaceImage replaces the image name and tags inside one object
- It searches the object for container session
- then loops though all images inside containers session,
- finds matched ones and update the image name and tag name
-*/
-func (pt *imageTransformer) findAndReplaceImage(obj map[string]interface{}) error {
-	found := false
-	for _, fs := range pt.fieldSpecs {
-		path := fs.Path
-		_, found = obj[path]
-		if found {
-			err := pt.updateContainers(obj, path)
+	for id := range m {
+		objMap := m[id].Map()
+		for _, path := range pt.fieldSpecs {
+			if !id.Gvk().IsSelected(&path.Gvk) {
+				continue
+			}
+			err := mutateField(objMap, path.PathSlice(), false, pt.updateContainers)
 			if err != nil {
 				return err
 			}
+
 		}
 	}
-	if !found {
-		return pt.findContainers(obj)
-	}
 	return nil
 }
 
-func (pt *imageTransformer) updateContainers(obj map[string]interface{}, path string) error {
-	containers, ok := obj[path].([]interface{})
+func (pt *imageTransformer) updateContainers(in interface{}) (interface{}, error) {
+	containers, ok := in.([]interface{})
 	if !ok {
-		return fmt.Errorf("containers path is not of type []interface{} but %T", obj[path])
+		return nil, fmt.Errorf("containers path is not of type []interface{} but %T", in)
 	}
 	for i := range containers {
 		container := containers[i].(map[string]interface{})
@@ -108,31 +91,7 @@ func (pt *imageTransformer) updateContainers(obj map[string]interface{}, path st
 			break
 		}
 	}
-	return nil
-}
-
-func (pt *imageTransformer) findContainers(obj map[string]interface{}) error {
-	for key := range obj {
-		switch typedV := obj[key].(type) {
-		case map[string]interface{}:
-			err := pt.findAndReplaceImage(typedV)
-			if err != nil {
-				return err
-			}
-		case []interface{}:
-			for i := range typedV {
-				item := typedV[i]
-				typedItem, ok := item.(map[string]interface{})
-				if ok {
-					err := pt.findAndReplaceImage(typedItem)
-					if err != nil {
-						return err
-					}
-				}
-			}
-		}
-	}
-	return nil
+	return containers, nil
 }
 
 func isImageMatched(s, t string) bool {