k8s: add a locator for image objects (#3525)

This is moving towards #3427,
but also is a good model of how to write new kinds of image locators

internal/k8s/image_test.go

@@ -271,13 +271,15 @@ func TestEntityHasImage(t *testing.T) {
 		t.Fatal(err)
 	}
 	assert.False(t, match, "deployment yaml should not match image %s", img.String())
+}
 
-	entities, err = ParseYAMLFromString(testyaml.CRDYAML)
+func TestCRDExtract(t *testing.T) {
+	entities, err := ParseYAMLFromString(testyaml.CRDYAML)
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	img = container.MustParseTaggedSelector("docker.io/bitnami/minideb:latest")
+	img := container.MustParseTaggedSelector("docker.io/bitnami/minideb:latest")
 	e := entities[0]
 	selector, err := NewPartialMatchObjectSelector("", "", "projects.example.martin-helmich.de", "")
 	require.NoError(t, err)
@@ -287,18 +289,20 @@ func TestEntityHasImage(t *testing.T) {
 		"{.spec.validation.openAPIV3Schema.properties.spec.properties.image}")
 	require.NoError(t, err)
 
-	match, err = e.HasImage(img, []ImageLocator{jp}, false)
+	match, err := e.HasImage(img, []ImageLocator{jp}, false)
 	require.NoError(t, err)
 
 	assert.True(t, match, "CRD yaml should match image %s", img.String())
+}
 
-	entities, err = ParseYAMLFromString(testyaml.SanchoImageInEnvYAML)
+func TestEnvExtract(t *testing.T) {
+	entities, err := ParseYAMLFromString(testyaml.SanchoImageInEnvYAML)
 	if err != nil {
 		t.Fatal(err)
 	}
-	img = container.MustParseSelector("gcr.io/some-project-162817/sancho")
-	e = entities[0]
-	match, err = e.HasImage(img, nil, false)
+	img := container.MustParseSelector("gcr.io/some-project-162817/sancho")
+	e := entities[0]
+	match, err := e.HasImage(img, nil, false)
 	if err != nil {
 		t.Fatal(err)
 	}

internal/k8s/json_path.go

@@ -45,29 +45,36 @@ func (jp JSONPath) FindStrings(obj interface{}) ([]string, error) {
 //
 // Returns an error if the object at the specified path isn't a string.
 func (jp JSONPath) VisitStrings(obj interface{}, visit func(val jsonpath.Value, str string) error) error {
+	return jp.Visit(obj, func(match jsonpath.Value) error {
+		val := match.Interface()
+		str, ok := val.(string)
+		if !ok {
+			return fmt.Errorf("May only match strings (json_path=%q)\nGot Type: %T\nGot Value: %s",
+				jp.path, val, val)
+		}
+
+		return visit(match, str)
+	})
+}
+
+// Visit all the values from the given object on this path.
+func (jp JSONPath) Visit(obj interface{}, visit func(val jsonpath.Value) error) error {
 	// JSONPath is stateful and not thread-safe, so we need to parse a new one
 	// each time
 	matcher := jsonpath.New("jp")
 	err := matcher.Parse(jp.path)
 	if err != nil {
-		return fmt.Errorf("Matching strings (json_path=%q): %v", jp.path, err)
+		return fmt.Errorf("Matching (json_path=%q): %v", jp.path, err)
 	}
 
 	matches, err := matcher.FindResults(obj)
 	if err != nil {
-		return fmt.Errorf("Matching strings (json_path=%q): %v", jp.path, err)
+		return fmt.Errorf("Matching (json_path=%q): %v", jp.path, err)
 	}
 
 	for _, matchSet := range matches {
 		for _, match := range matchSet {
-			val := match.Interface()
-			str, ok := val.(string)
-			if !ok {
-				return fmt.Errorf("May only match strings (json_path=%q)\nGot Type: %T\nGot Value: %s",
-					jp.path, val, val)
-			}
-
-			err := visit(match, str)
+			err := visit(match)
 			if err != nil {
 				return err
 			}

internal/k8s/locator.go

@@ -1,6 +1,7 @@
 package k8s
 
 import (
+	"fmt"
 	"reflect"
 
 	"github.com/docker/distribution/reference"
@@ -68,16 +69,8 @@ func (l *JSONPathImageLocator) EqualsImageLocator(other interface{}) bool {
 		return false
 	}
 
-	if l.path.path != otherL.path.path {
-		return false
-	}
-
-	o1 := l.selector
-	o2 := otherL.selector
-	return o1.name == o2.name &&
-		o1.namespace == o2.namespace &&
-		o1.kind == o2.kind &&
-		o1.apiVersion == o2.apiVersion
+	return l.path.path == otherL.path.path &&
+		l.selector.EqualsSelector(otherL.selector)
 }
 
 func (l *JSONPathImageLocator) MatchesType(e K8sEntity) bool {
@@ -138,3 +131,125 @@ func (l *JSONPathImageLocator) Inject(e K8sEntity, selector container.RefSelecto
 }
 
 var _ ImageLocator = &JSONPathImageLocator{}
+
+type JSONPathImageObjectLocator struct {
+	selector  ObjectSelector
+	path      JSONPath
+	repoField string
+	tagField  string
+}
+
+func MustJSONPathImageObjectLocator(selector ObjectSelector, path, repoField, tagField string) *JSONPathImageObjectLocator {
+	locator, err := NewJSONPathImageObjectLocator(selector, path, repoField, tagField)
+	if err != nil {
+		panic(err)
+	}
+	return locator
+}
+
+func NewJSONPathImageObjectLocator(selector ObjectSelector, path, repoField, tagField string) (*JSONPathImageObjectLocator, error) {
+	p, err := NewJSONPath(path)
+	if err != nil {
+		return nil, err
+	}
+	return &JSONPathImageObjectLocator{
+		selector:  selector,
+		path:      p,
+		repoField: repoField,
+		tagField:  tagField,
+	}, nil
+}
+
+func (l *JSONPathImageObjectLocator) EqualsImageLocator(other interface{}) bool {
+	otherL, ok := other.(*JSONPathImageObjectLocator)
+	if !ok {
+		return false
+	}
+	return l.path.path == otherL.path.path &&
+		l.repoField == otherL.repoField &&
+		l.tagField == otherL.tagField &&
+		l.selector.EqualsSelector(otherL.selector)
+}
+
+func (l *JSONPathImageObjectLocator) MatchesType(e K8sEntity) bool {
+	return l.selector.Matches(e)
+}
+
+func (l *JSONPathImageObjectLocator) unpack(e K8sEntity) interface{} {
+	if u, ok := e.Obj.(runtime.Unstructured); ok {
+		return u.UnstructuredContent()
+	}
+	return e.Obj
+}
+
+func (l *JSONPathImageObjectLocator) extractImageFromMap(val jsonpath.Value) (reference.Named, error) {
+	m, ok := val.Interface().(map[string]interface{})
+	if !ok {
+		return nil, fmt.Errorf("May only match maps (json_path=%q)\nGot Type: %s\nGot Value: %s",
+			l.path.path, val.Type(), val)
+	}
+
+	repoField, ok := m[l.repoField].(string)
+	imageString := ""
+	if ok {
+		imageString = repoField
+	}
+
+	tagField, ok := m[l.tagField].(string)
+	if ok && tagField != "" {
+		imageString = fmt.Sprintf("%s:%s", repoField, tagField)
+	}
+
+	return container.ParseNamed(imageString)
+}
+
+func (l *JSONPathImageObjectLocator) Extract(e K8sEntity) ([]reference.Named, error) {
+	if !l.selector.Matches(e) {
+		return nil, nil
+	}
+
+	result := make([]reference.Named, 0)
+	err := l.path.Visit(l.unpack(e), func(val jsonpath.Value) error {
+		ref, err := l.extractImageFromMap(val)
+		if err != nil {
+			return err
+		}
+		result = append(result, ref)
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return result, nil
+}
+
+func (l *JSONPathImageObjectLocator) Inject(e K8sEntity, selector container.RefSelector, injectRef reference.Named) (K8sEntity, bool, error) {
+	if !l.selector.Matches(e) {
+		return e, false, nil
+	}
+
+	tagged, isTagged := injectRef.(reference.Tagged)
+
+	modified := false
+	err := l.path.Visit(l.unpack(e), func(val jsonpath.Value) error {
+		ref, err := l.extractImageFromMap(val)
+		if err != nil {
+			return err
+		}
+		if selector.Matches(ref) {
+			m := val.Interface().(map[string]interface{})
+			m[l.repoField] = reference.FamiliarName(injectRef)
+			if isTagged {
+				m[l.tagField] = tagged.Tag()
+			}
+			modified = true
+		}
+		return nil
+	})
+	if err != nil {
+		return e, false, err
+	}
+	return e, modified, nil
+}
+
+var _ ImageLocator = &JSONPathImageObjectLocator{}

internal/k8s/locator_test.go

@@ -0,0 +1,34 @@
+package k8s
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/tilt-dev/tilt/internal/container"
+	"github.com/tilt-dev/tilt/internal/k8s/testyaml"
+)
+
+func TestCRDImageObjectInjection(t *testing.T) {
+	entities, err := ParseYAMLFromString(testyaml.CRDImageObjectYAML)
+	require.NoError(t, err)
+
+	e := entities[0]
+	selector := MustKindSelector("UselessMachine")
+	locator := MustJSONPathImageObjectLocator(selector, "{.spec.imageObject}", "repo", "tag")
+	images, err := locator.Extract(e)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(images))
+	assert.Equal(t, "docker.io/library/frontend", images[0].String())
+
+	e, modified, err := locator.Inject(e, container.MustParseSelector("frontend"),
+		container.MustParseNamed("frontend:tilt-123"))
+	require.NoError(t, err)
+	assert.True(t, modified)
+
+	images, err = locator.Extract(e)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(images))
+	assert.Equal(t, "docker.io/library/frontend:tilt-123", images[0].String())
+}

internal/k8s/object_selector.go

@@ -120,6 +120,13 @@ func NewPartialMatchObjectSelector(apiVersion string, kind string, name string,
 	return ret, nil
 }
 
+func (o1 ObjectSelector) EqualsSelector(o2 ObjectSelector) bool {
+	return o1.name == o2.name &&
+		o1.namespace == o2.namespace &&
+		o1.kind == o2.kind &&
+		o1.apiVersion == o2.apiVersion
+}
+
 func (k ObjectSelector) Matches(e K8sEntity) bool {
 	gvk := e.GVK()
 	return k.apiVersion.MatchString(gvk.GroupVersion().String()) &&

internal/k8s/testyaml/testyaml.go

@@ -1424,6 +1424,15 @@ spec:
 
 const CRDImage = "docker.io/bitnami/minideb:latest"
 
+const CRDImageObjectYAML = `apiVersion: tilt.dev/v1alpha1
+kind: UselessMachine
+metadata:
+  name: um
+spec:
+  imageObject:
+    repo: frontend
+`
+
 const MyNamespaceYAML = `apiVersion: v1
 kind: Namespace
 metadata: