update JSON context with patched resource for chained mutate rules

Signed-off-by: Shuting Zhao <shutting06@gmail.com>

pkg/engine/context/context.go

@@ -143,6 +143,25 @@ func (ctx *Context) AddResource(dataRaw []byte) error {
 	return ctx.AddJSON(objRaw)
 }
 
+func (ctx *Context) AddResourceAsObject(data interface{}) error {
+	modifiedResource := struct {
+		Request interface{} `json:"request"`
+	}{
+		Request: struct {
+			Object interface{} `json:"object"`
+		}{
+			Object: data,
+		},
+	}
+
+	objRaw, err := json.Marshal(modifiedResource)
+	if err != nil {
+		ctx.log.Error(err, "failed to marshal the resource")
+		return err
+	}
+	return ctx.AddJSON(objRaw)
+}
+
 //AddUserInfo adds userInfo at path request.userInfo
 func (ctx *Context) AddUserInfo(userRequestInfo kyverno.RequestInfo) error {
 	modifiedResource := struct {

pkg/engine/context/imageutils.go

@@ -1,11 +1,13 @@
 package context
 
 import (
+	"fmt"
 	"strconv"
 	"strings"
 
 	"github.com/distribution/distribution/reference"
 	"github.com/go-logr/logr"
+	engineutils "github.com/kyverno/kyverno/pkg/engine/utils"
 	"github.com/pkg/errors"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 )
@@ -189,3 +191,31 @@ func addDefaultDomain(name string) string {
 
 	return name
 }
+
+func MutateResourceWithImageInfo(raw []byte, ctx *Context) error {
+	images := ctx.ImageInfo()
+	if images == nil {
+		return nil
+	}
+
+	buildJSONPatch := func(op, path, value string) []byte {
+		p := fmt.Sprintf(`{ "op": "%s", "path": "%s", "value":"%s" }`, op, path, value)
+		return []byte(p)
+	}
+
+	var patches [][]byte
+	for _, info := range images.Containers {
+		patches = append(patches, buildJSONPatch("replace", info.JSONPointer, info.String()))
+	}
+
+	for _, info := range images.InitContainers {
+		patches = append(patches, buildJSONPatch("replace", info.JSONPointer, info.String()))
+	}
+
+	patchedResource, err := engineutils.ApplyPatches(raw, patches)
+	if err != nil {
+		return err
+	}
+
+	return ctx.AddResource(patchedResource)
+}

pkg/engine/mutation.go

@@ -70,7 +70,18 @@ func Mutate(policyContext *PolicyContext) (resp *response.EngineResponse) {
 
 		logger.V(3).Info("matched mutate rule")
 
+		// Restore() is meant for restoring context loaded from external lookup (APIServer & ConfigMap)
+		// while we need to keep updated resource in the JSON context as rules can be chained
+		resource, err := policyContext.JSONContext.Query("request.object")
 		policyContext.JSONContext.Restore()
+		if err == nil && resource != nil {
+			if err := ctx.AddResourceAsObject(resource.(map[string]interface{})); err != nil {
+				logger.WithName("RestoreContext").Error(err, "unable to update resource object")
+			}
+		} else {
+			logger.WithName("RestoreContext").Error(err, "failed to quey resource object")
+		}
+
 		if err := LoadContext(logger, rule.Context, resCache, policyContext, rule.Name); err != nil {
 			if _, ok := err.(gojmespath.NotFoundError); ok {
 				logger.V(3).Info("failed to load context", "reason", err.Error())
@@ -93,7 +104,7 @@ func Mutate(policyContext *PolicyContext) (resp *response.EngineResponse) {
 			continue
 		}
 
-		if rule, err = variables.SubstituteAllInRule(logger, policyContext.JSONContext, rule); err != nil {
+		if rule, err = variables.SubstituteAllInRule(logger, ctx, rule); err != nil {
 			ruleResp := response.RuleResponse{
 				Name:    rule.Name,
 				Type:    utils.Validation.String(),
@@ -120,6 +131,7 @@ func Mutate(policyContext *PolicyContext) (resp *response.EngineResponse) {
 			logger.V(4).Info("mutate rule applied successfully", "ruleName", rule.Name)
 		}
 
+		ctx.AddResourceAsObject(patchedResource.Object)
 		resp.PolicyResponse.Rules = append(resp.PolicyResponse.Rules, ruleResponse)
 		incrementAppliedRuleCount(resp)
 	}

pkg/engine/mutation_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/kyverno/kyverno/pkg/engine/utils"
 	"github.com/kyverno/kyverno/pkg/kyverno/store"
 	"gotest.tools/assert"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 )
 
 func Test_VariableSubstitutionOverlay(t *testing.T) {
@@ -262,6 +263,42 @@ func Test_variableSubstitutionCLI(t *testing.T) {
 	t.Log(string(expectedPatch))
 	t.Log(string(er.PolicyResponse.Rules[0].Patches[0]))
 	if !reflect.DeepEqual(expectedPatch, er.PolicyResponse.Rules[0].Patches[0]) {
-		t.Error("patches dont match")
+		t.Error("patches don't match")
 	}
 }
+
+// https://github.com/kyverno/kyverno/issues/2022
+func Test_chained_rules(t *testing.T) {
+	policyRaw := []byte(`{"apiVersion":"kyverno.io/v1","kind":"ClusterPolicy","metadata":{"name":"replace-image-registry","annotations":{"policies.kyverno.io/minversion":"1.4.2"}},"spec":{"background":false,"rules":[{"name":"replace-image-registry","match":{"resources":{"kinds":["Pod"]}},"mutate":{"patchStrategicMerge":{"spec":{"containers":[{"(name)":"*","image":"{{regex_replace_all('^[^/]+','{{@}}','myregistry.corp.com')}}"}]}}}},{"name":"replace-image-registry-chained","match":{"resources":{"kinds":["Pod"]}},"mutate":{"patchStrategicMerge":{"spec":{"containers":[{"(name)":"*","image":"{{regex_replace_all('\\b(myregistry.corp.com)\\b','{{@}}','otherregistry.corp.com')}}"}]}}}}]}}`)
+	var policy kyverno.ClusterPolicy
+	err := json.Unmarshal(policyRaw, &policy)
+	assert.NilError(t, err)
+
+	resourceRaw := []byte(`{"apiVersion":"v1","kind":"Pod","metadata":{"name":"test"},"spec":{"containers":[{"name":"test","image":"foo/bash:5.0"}]}}`)
+	resource, err := utils.ConvertToUnstructured(resourceRaw)
+	assert.NilError(t, err)
+
+	ctx := context.NewContext()
+	err = ctx.AddResourceAsObject(resource.Object)
+	assert.NilError(t, err)
+
+	policyContext := &PolicyContext{
+		Policy:      policy,
+		JSONContext: ctx,
+		NewResource: *resource,
+	}
+
+	err = ctx.AddImageInfo(resource)
+	assert.NilError(t, err)
+
+	err = context.MutateResourceWithImageInfo(resourceRaw, ctx)
+	assert.NilError(t, err)
+
+	er := Mutate(policyContext)
+	containers, _, err := unstructured.NestedSlice(er.PatchedResource.Object, "spec", "containers")
+	assert.NilError(t, err)
+	assert.Equal(t, containers[0].(map[string]interface{})["image"], "otherregistry.corp.com/foo/bash:5.0")
+
+	assert.Equal(t, string(er.PolicyResponse.Rules[0].Patches[0]), `{"op":"replace","path":"/spec/containers/0/image","value":"myregistry.corp.com/foo/bash:5.0"}`)
+	assert.Equal(t, string(er.PolicyResponse.Rules[1].Patches[0]), `{"op":"replace","path":"/spec/containers/0/image","value":"otherregistry.corp.com/foo/bash:5.0"}`)
+}

pkg/webhooks/server.go

@@ -21,7 +21,6 @@ import (
 	"github.com/kyverno/kyverno/pkg/engine"
 	enginectx "github.com/kyverno/kyverno/pkg/engine/context"
 	"github.com/kyverno/kyverno/pkg/engine/response"
-	engineutils "github.com/kyverno/kyverno/pkg/engine/utils"
 	"github.com/kyverno/kyverno/pkg/event"
 	"github.com/kyverno/kyverno/pkg/generate"
 	"github.com/kyverno/kyverno/pkg/metrics"
@@ -373,7 +372,7 @@ func (ws *WebhookServer) buildPolicyContext(request *v1beta1.AdmissionRequest, a
 		return nil, errors.Wrap(err, "failed to add image information to the policy rule context")
 	}
 
-	if err := mutateResourceWithImageInfo(request.Object.Raw, ctx); err != nil {
+	if err := enginectx.MutateResourceWithImageInfo(request.Object.Raw, ctx); err != nil {
 		ws.log.Error(err, "failed to patch images info to resource, policies that mutate images may be impacted")
 	}
 
@@ -627,31 +626,3 @@ func newVariablesContext(request *v1beta1.AdmissionRequest, userRequestInfo *v1.
 
 	return ctx, nil
 }
-
-func mutateResourceWithImageInfo(raw []byte, ctx *enginectx.Context) error {
-	images := ctx.ImageInfo()
-	if images == nil {
-		return nil
-	}
-
-	var patches [][]byte
-	for _, info := range images.Containers {
-		patches = append(patches, buildJSONPatch("replace", info.JSONPointer, info.String()))
-	}
-
-	for _, info := range images.InitContainers {
-		patches = append(patches, buildJSONPatch("replace", info.JSONPointer, info.String()))
-	}
-
-	patchedResource, err := engineutils.ApplyPatches(raw, patches)
-	if err != nil {
-		return err
-	}
-
-	return ctx.AddResource(patchedResource)
-}
-
-func buildJSONPatch(op, path, value string) []byte {
-	p := fmt.Sprintf(`{ "op": "%s", "path": "%s", "value":"%s" }`, op, path, value)
-	return []byte(p)
-}