optimize JSON context processing using in memory maps

Signed-off-by: Jim Bugwadia <[email protected]>

go.mod

@@ -25,6 +25,7 @@ require (
 	github.com/in-toto/in-toto-golang v0.8.0
 	github.com/jmespath/go-jmespath v0.4.0
 	github.com/jmoiron/jsonq v0.0.0-20150511023944-e874b168d07e
+	github.com/json-iterator/go v1.1.12
 	github.com/julienschmidt/httprouter v1.3.0
 	github.com/kataras/tablewriter v0.0.0-20180708051242-e063d29b7c23
 	github.com/lensesio/tableprinter v0.0.0-20201125135848-89e81fc956e7
@@ -42,7 +43,7 @@ require (
 	github.com/sigstore/k8s-manifest-sigstore v0.4.4
 	github.com/sigstore/sigstore v1.6.3
 	github.com/spf13/cobra v1.7.0
-	github.com/stretchr/testify v1.8.2
+	github.com/stretchr/testify v1.8.4
 	github.com/zach-klippenstein/goregen v0.0.0-20160303162051-795b5e3961ea
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.40.0
 	go.opentelemetry.io/otel v1.14.0
@@ -220,7 +221,6 @@ require (
 	github.com/jellydator/ttlcache/v3 v3.0.1 // indirect
 	github.com/jinzhu/copier v0.3.5 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
-	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/kevinburke/ssh_config v1.2.0 // indirect
 	github.com/klauspost/compress v1.16.3 // indirect
 	github.com/leodido/go-urn v1.2.3 // indirect
@@ -329,6 +329,6 @@ require (
 )
 
 replace (
-	github.com/jmespath/go-jmespath => github.com/kyverno/go-jmespath v0.4.1-0.20230204162932-3ee946b9433d
+	github.com/jmespath/go-jmespath => github.com/kyverno/go-jmespath v0.4.1-0.20230906134905-62fa64b71f91
 	github.com/sigstore/cosign => github.com/nirmata/cosign v1.13.2-0.20230726092108-615d4da057d8
 )

go.sum

@@ -849,8 +849,6 @@ github.com/jingyugao/rowserrcheck v0.0.0-20210315055705-d907ca737bb1/go.mod h1:T
 github.com/jinzhu/copier v0.3.5 h1:GlvfUwHk62RokgqVNvYsku0TATCF7bAHVwEXoBh3iJg=
 github.com/jinzhu/copier v0.3.5/go.mod h1:DfbEm0FYsaqBcKcFuvmOZb218JkPGtvSHsKg8S8hyyg=
 github.com/jirfag/go-printf-func-name v0.0.0-20200119135958-7558a9eaa5af/go.mod h1:HEWGJkRDzjJY2sqdDwxccsGicWEf9BQOZsq2tV+xzM0=
-github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8=
-github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U=
 github.com/jmhodges/clock v1.2.0 h1:eq4kys+NI0PLngzaHEe7AmPT90XMGIEySD1JfV1PDIs=
 github.com/jmoiron/jsonq v0.0.0-20150511023944-e874b168d07e h1:ZZCvgaRDZg1gC9/1xrsgaJzQUCQgniKtw0xjWywWAOE=
 github.com/jmoiron/jsonq v0.0.0-20150511023944-e874b168d07e/go.mod h1:+rHyWac2R9oAZwFe1wGY2HBzFJJy++RHBg1cU23NkD8=
@@ -914,8 +912,10 @@ github.com/kunwardeep/paralleltest v1.0.2/go.mod h1:ZPqNm1fVHPllh5LPVujzbVz1JN2G
 github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
 github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
 github.com/kyoh86/exportloopref v0.1.8/go.mod h1:1tUcJeiioIs7VWe5gcOObrux3lb66+sBqGZrRkMwPgg=
-github.com/kyverno/go-jmespath v0.4.1-0.20230204162932-3ee946b9433d h1:g63VNwOo6yYRY1n3mgF2ou4cjnwyonsIKqnbBM9pTRA=
-github.com/kyverno/go-jmespath v0.4.1-0.20230204162932-3ee946b9433d/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=
+github.com/kyverno/go-jmespath v0.4.1-0.20230906134905-62fa64b71f91 h1:n63aowZk61f65e6OJxuql8BS/hCv8LZxz/eCO4+2NfM=
+github.com/kyverno/go-jmespath v0.4.1-0.20230906134905-62fa64b71f91/go.mod h1:yzDHaKovQy16rjN4kFnjF+IdNoN4p1ndw+va6+B8zUU=
+github.com/kyverno/go-jmespath/internal/testify v1.5.2-0.20230630133209-945021c749d9 h1:lL311dF3a2aeNibJj8v+uhFU3XkvRHZmCtAdSPOrQYY=
+github.com/kyverno/go-jmespath/internal/testify v1.5.2-0.20230630133209-945021c749d9/go.mod h1:XRxUGHIiCy1WYma1CdfdO1WOhIe8dLPTENaZr5D1ex4=
 github.com/ldez/gomoddirectives v0.2.1/go.mod h1:sGicqkRgBOg//JfpXwkB9Hj0X5RyJ7mlACM5B9f6Me4=
 github.com/ldez/tagliatelle v0.2.0/go.mod h1:8s6WJQwEYHbKZDsp/LjArytKOG8qaMrKQQ3mFukHs88=
 github.com/lensesio/tableprinter v0.0.0-20201125135848-89e81fc956e7 h1:k/1ku0yehLCPqERCHkIHMDqDg1R02AcCScRuHbamU3s=
@@ -1344,8 +1344,9 @@ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/stretchr/testify v1.7.2/go.mod h1:R6va5+xMeoiuVRoj+gSkQ7d3FALtqAAGI1FQKckRals=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
-github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8=
 github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
+github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 github.com/subosito/gotenv v1.2.0/go.mod h1:N0PQaV/YGNqwC0u51sEeR/aUtSLEXKX9iv69rRypqCw=
 github.com/subosito/gotenv v1.4.2 h1:X1TuBLAMDFbaTAChgCBLu3DU3UPyELpnF2jjJ2cz/S8=
 github.com/subosito/gotenv v1.4.2/go.mod h1:ayKnFf/c6rvx/2iiLrJUk1e6plDbT3edrFNGqEflhK0=

pkg/engine/context/context.go

@@ -1,30 +1,38 @@
 package context
 
 import (
-	"encoding/json"
 	"fmt"
+	"regexp"
 	"strings"
-	"sync"
 
-	jsonpatch "github.com/evanphx/json-patch/v5"
+	jsoniter "github.com/json-iterator/go"
 	kyvernov1 "github.com/kyverno/kyverno/api/kyverno/v1"
 	kyvernov1beta1 "github.com/kyverno/kyverno/api/kyverno/v1beta1"
 	"github.com/kyverno/kyverno/pkg/config"
 	"github.com/kyverno/kyverno/pkg/engine/jmespath"
+	"github.com/kyverno/kyverno/pkg/engine/jsonutils"
 	"github.com/kyverno/kyverno/pkg/logging"
 	apiutils "github.com/kyverno/kyverno/pkg/utils/api"
 	admissionv1 "k8s.io/api/admission/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/apimachinery/pkg/runtime"
 )
 
-var logger = logging.WithName("context")
+var (
+	logger       = logging.WithName("context")
+	json         = jsoniter.ConfigCompatibleWithStandardLibrary
+	ReservedKeys = regexp.MustCompile(`request|serviceAccountName|serviceAccountNamespace|element|elementIndex|@|images|image|([a-z_0-9]+\()[^{}]`)
+)
 
 // EvalInterface is used to query and inspect context data
 // TODO: move to contextapi to prevent circular dependencies
 type EvalInterface interface {
 	// Query accepts a JMESPath expression and returns matching data
 	Query(query string) (interface{}, error)
 
+	// Operation returns the admission operation i.e. "request.operation"
+	QueryOperation() string
+
 	// HasChanged accepts a JMESPath expression and compares matching data in the
 	// request.object and request.oldObject context fields. If the data has changed
 	// it return `true`. If the data has not changed it returns false. If either
@@ -99,50 +107,69 @@ type Interface interface {
 
 	EvalInterface
 
-	// AddJSON  merges the json with context
-	addJSON(dataRaw []byte) error
+	// AddJSON  merges the json map with context
+	addJSON(dataMap map[string]interface{}) error
 }
 
 // Context stores the data resources as JSON
 type context struct {
 	jp                 jmespath.Interface
-	mutex              sync.RWMutex
-	jsonRaw            []byte
-	jsonRawCheckpoints [][]byte
+	jsonRaw            map[string]interface{}
+	jsonRawCheckpoints []map[string]interface{}
 	images             map[string]map[string]apiutils.ImageInfo
+	operation          kyvernov1.AdmissionOperation
 	deferred           DeferredLoaders
 }
 
 // NewContext returns a new context
 func NewContext(jp jmespath.Interface) Interface {
-	return NewContextFromRaw(jp, []byte(`{}`))
+	return NewContextFromRaw(jp, map[string]interface{}{})
 }
 
 // NewContextFromRaw returns a new context initialized with raw data
-func NewContextFromRaw(jp jmespath.Interface, raw []byte) Interface {
+func NewContextFromRaw(jp jmespath.Interface, raw map[string]interface{}) Interface {
 	return &context{
 		jp:                 jp,
 		jsonRaw:            raw,
-		jsonRawCheckpoints: make([][]byte, 0),
+		jsonRawCheckpoints: make([]map[string]interface{}, 0),
 		deferred:           NewDeferredLoaders(),
 	}
 }
 
 // addJSON merges json data
-func (ctx *context) addJSON(dataRaw []byte) error {
-	ctx.mutex.Lock()
-	defer ctx.mutex.Unlock()
-	json, err := jsonpatch.MergeMergePatches(ctx.jsonRaw, dataRaw)
-	if err != nil {
-		return fmt.Errorf("failed to merge JSON data: %w", err)
-	}
-	ctx.jsonRaw = json
+func (ctx *context) addJSON(dataMap map[string]interface{}) error {
+	mergeMaps(dataMap, ctx.jsonRaw)
 	return nil
 }
 
+func (ctx *context) QueryOperation() string {
+	if ctx.operation != "" {
+		return string(ctx.operation)
+	}
+
+	if requestMap, val := ctx.jsonRaw["request"].(map[string]interface{}); val {
+		if op, val := requestMap["operation"].(string); val {
+			return op
+		}
+	}
+
+	return ""
+}
+
 // AddRequest adds an admission request to context
 func (ctx *context) AddRequest(request admissionv1.AdmissionRequest) error {
-	return addToContext(ctx, request, "request")
+	// an AdmissionRequest needs to be marshaled / unmarshaled as
+	// JSON to properly convert types of runtime.RawExtension
+	mapObj, err := jsonutils.DocumentToUntyped(request)
+	if err != nil {
+		return err
+	}
+	if err := addToContext(ctx, mapObj, "request"); err != nil {
+		return err
+	}
+
+	ctx.operation = kyvernov1.AdmissionOperation(request.Operation)
+	return nil
 }
 
 func (ctx *context) AddVariable(key string, value interface{}) error {
@@ -193,12 +220,21 @@ func (ctx *context) SetTargetResource(data map[string]interface{}) error {
 
 // AddOperation data at path: request.operation
 func (ctx *context) AddOperation(data string) error {
-	return addToContext(ctx, data, "request", "operation")
+	if err := addToContext(ctx, data, "request", "operation"); err != nil {
+		return err
+	}
+
+	ctx.operation = kyvernov1.AdmissionOperation(data)
+	return nil
 }
 
 // AddUserInfo adds userInfo at path request.userInfo
 func (ctx *context) AddUserInfo(userRequestInfo kyvernov1beta1.RequestInfo) error {
-	return addToContext(ctx, userRequestInfo, "request")
+	if data, err := toUnstructured(&userRequestInfo); err == nil {
+		return addToContext(ctx, data, "request")
+	} else {
+		return err
+	}
 }
 
 // AddServiceAccount removes prefix 'system:serviceaccount:' and namespace, then loads only SA name and SA namespace
@@ -218,33 +254,14 @@ func (ctx *context) AddServiceAccount(userName string) error {
 		saName = groups[1]
 		saNamespace = groups[0]
 	}
-	saNameObj := struct {
-		SA string `json:"serviceAccountName"`
-	}{
-		SA: saName,
-	}
-	saNameRaw, err := json.Marshal(saNameObj)
-	if err != nil {
-		logger.Error(err, "failed to marshal the SA")
-		return err
+	data := map[string]interface{}{
+		"serviceAccountName":      saName,
+		"serviceAccountNamespace": saNamespace,
 	}
-	if err := ctx.addJSON(saNameRaw); err != nil {
+	if err := ctx.addJSON(data); err != nil {
 		return err
 	}
 
-	saNsObj := struct {
-		SA string `json:"serviceAccountNamespace"`
-	}{
-		SA: saNamespace,
-	}
-	saNsRaw, err := json.Marshal(saNsObj)
-	if err != nil {
-		logger.Error(err, "failed to marshal the SA namespace")
-		return err
-	}
-	if err := ctx.addJSON(saNsRaw); err != nil {
-		return err
-	}
 	logger.V(4).Info("Adding service account", "service account name", saName, "service account namespace", saNamespace)
 	return nil
 }
@@ -260,8 +277,8 @@ func (ctx *context) AddElement(data interface{}, index, nesting int) error {
 	data = map[string]interface{}{
 		"element":          data,
 		nestedElement:      data,
-		"elementIndex":     index,
-		nestedElementIndex: index,
+		"elementIndex":     int64(index),
+		nestedElementIndex: int64(index),
 	}
 	return addToContext(ctx, data)
 }
@@ -288,9 +305,33 @@ func (ctx *context) AddImageInfos(resource *unstructured.Unstructured, cfg confi
 		return nil
 	}
 	ctx.images = images
+	utm, err := convertImagesToUntyped(images)
+	if err != nil {
+		return err
+	}
+
+	logging.V(4).Info("updated image info", "images", utm)
+	return addToContext(ctx, utm, "images")
+}
 
-	logging.V(4).Info("updated image info", "images", images)
-	return addToContext(ctx, images, "images")
+func convertImagesToUntyped(images map[string]map[string]apiutils.ImageInfo) (map[string]interface{}, error) {
+	results := map[string]interface{}{}
+	for containerType, v := range images {
+		imgMap := map[string]interface{}{}
+		for containerName, imageInfo := range v {
+			img, err := toUnstructured(&imageInfo.ImageInfo)
+			if err != nil {
+				return nil, err
+			}
+
+			img["jsonPointer"] = imageInfo.Pointer
+			imgMap[containerName] = img
+		}
+
+		results[containerType] = imgMap
+	}
+
+	return results, nil
 }
 
 func (ctx *context) GenerateCustomImageInfo(resource *unstructured.Unstructured, imageExtractorConfigs kyvernov1.ImageExtractorConfigs, cfg config.Configuration) (map[string]map[string]apiutils.ImageInfo, error) {
@@ -314,13 +355,23 @@ func (ctx *context) ImageInfo() map[string]map[string]apiutils.ImageInfo {
 // Checkpoint creates a copy of the current internal state and
 // pushes it into a stack of stored states.
 func (ctx *context) Checkpoint() {
-	ctx.mutex.Lock()
-	defer ctx.mutex.Unlock()
-	jsonRawCheckpoint := make([]byte, len(ctx.jsonRaw))
-	copy(jsonRawCheckpoint, ctx.jsonRaw)
+	jsonRawCheckpoint := ctx.copyContext(ctx.jsonRaw)
 	ctx.jsonRawCheckpoints = append(ctx.jsonRawCheckpoints, jsonRawCheckpoint)
 }
 
+func (ctx *context) copyContext(in map[string]interface{}) map[string]interface{} {
+	out := make(map[string]interface{}, len(in))
+	for k, v := range in {
+		if ReservedKeys.MatchString(k) {
+			out[k] = v
+		} else {
+			out[k] = runtime.DeepCopyJSONValue(v)
+		}
+	}
+
+	return out
+}
+
 // Restore sets the internal state to the last checkpoint, and removes the checkpoint.
 func (ctx *context) Restore() {
 	ctx.reset(true)
@@ -337,20 +388,19 @@ func (ctx *context) reset(restore bool) {
 	}
 }
 
-func (ctx *context) resetCheckpoint(removeCheckpoint bool) bool {
-	ctx.mutex.Lock()
-	defer ctx.mutex.Unlock()
-
+func (ctx *context) resetCheckpoint(restore bool) bool {
 	if len(ctx.jsonRawCheckpoints) == 0 {
 		return false
 	}
 
 	n := len(ctx.jsonRawCheckpoints) - 1
 	jsonRawCheckpoint := ctx.jsonRawCheckpoints[n]
-	ctx.jsonRaw = make([]byte, len(jsonRawCheckpoint))
-	copy(ctx.jsonRaw, jsonRawCheckpoint)
-	if removeCheckpoint {
+
+	if restore {
 		ctx.jsonRawCheckpoints = ctx.jsonRawCheckpoints[:n]
+		ctx.jsonRaw = jsonRawCheckpoint
+	} else {
+		ctx.jsonRaw = ctx.copyContext(jsonRawCheckpoint)
 	}
 
 	return true

pkg/engine/context/deferred_test.go

@@ -94,8 +94,8 @@ func TestDeferredLoaderMismatch(t *testing.T) {
 func newContext() *context {
 	return &context{
 		jp:                 jp,
-		jsonRaw:            []byte(`{}`),
-		jsonRawCheckpoints: make([][]byte, 0),
+		jsonRaw:            make(map[string]interface{}),
+		jsonRawCheckpoints: make([]map[string]interface{}, 0),
 		deferred:           NewDeferredLoaders(),
 	}
 }
@@ -289,7 +289,7 @@ func TestDeferredCheckpointRestore(t *testing.T) {
 
 func TestDeferredForloop(t *testing.T) {
 	ctx := newContext()
-	addDeferred(ctx, "value", -1)
+	addDeferred(ctx, "value", float64(-1))
 
 	ctx.Checkpoint()
 	for i := 0; i < 5; i++ {
@@ -298,7 +298,7 @@ func TestDeferredForloop(t *testing.T) {
 		assert.Equal(t, float64(i-1), val)
 
 		ctx.Reset()
-		mock, _ := addDeferred(ctx, "value", i)
+		mock, _ := addDeferred(ctx, "value", float64(i))
 		val, err = ctx.Query("value")
 		assert.NilError(t, err)
 		assert.Equal(t, float64(i), val)