Json context optimizations

go.mod

@@ -28,7 +28,7 @@ require (
 	github.com/jmoiron/jsonq v0.0.0-20150511023944-e874b168d07e
 	github.com/julienschmidt/httprouter v1.3.0
 	github.com/kataras/tablewriter v0.0.0-20180708051242-e063d29b7c23
-	github.com/kyverno/go-jmespath v0.4.1-0.20230705123211-d067dc3d6613
+	github.com/kyverno/go-jmespath v0.4.1-0.20230906134905-62fa64b71f91
 	github.com/lensesio/tableprinter v0.0.0-20201125135848-89e81fc956e7
 	github.com/notaryproject/notation-core-go v1.0.0
 	github.com/notaryproject/notation-go v1.0.0

go.sum

@@ -1022,8 +1022,8 @@ 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.20230705123211-d067dc3d6613 h1:M0uOLuCAZydi/vZy7uvNhwaIge0HFMdfqQYOKw7kgnQ=
-github.com/kyverno/go-jmespath v0.4.1-0.20230705123211-d067dc3d6613/go.mod h1:yzDHaKovQy16rjN4kFnjF+IdNoN4p1ndw+va6+B8zUU=
+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=

pkg/engine/context/context.go

@@ -2,30 +2,38 @@
 
 import (
 	"encoding/csv"
-	"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
@@ -100,50 +108,69 @@
 
 	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 {
@@ -200,12 +227,21 @@
 
 // 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
@@ -225,33 +261,14 @@
 		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
 }
@@ -267,8 +284,8 @@
 	data = map[string]interface{}{
 		"element":          data,
 		nestedElement:      data,
-		"elementIndex":     index,
-		nestedElementIndex: index,
+		"elementIndex":     int64(index),
+		nestedElementIndex: int64(index),
 	}
 	return addToContext(ctx, data)
 }
@@ -295,9 +312,33 @@
 		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) {
@@ -321,13 +362,23 @@
 // 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)
@@ -344,20 +395,19 @@
 	}
 }
 
-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

@@ -95,8 +95,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(),
 	}
 }
@@ -290,7 +290,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++ {
@@ -299,7 +299,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)

pkg/engine/context/evaluate.go

@@ -1,7 +1,6 @@
 package context
 
 import (
-	"encoding/json"
 	"fmt"
 	"strings"
 
@@ -25,13 +24,7 @@ func (ctx *context) Query(query string) (interface{}, error) {
 		return nil, fmt.Errorf("incorrect query %s: %v", query, err)
 	}
 	// search
-	ctx.mutex.RLock()
-	defer ctx.mutex.RUnlock()
-	var data interface{}
-	if err := json.Unmarshal(ctx.jsonRaw, &data); err != nil {
-		return nil, fmt.Errorf("failed to unmarshal context: %w", err)
-	}
-	result, err := queryPath.Search(data)
+	result, err := queryPath.Search(ctx.jsonRaw)
 	if err != nil {
 		return nil, fmt.Errorf("JMESPath query failed: %w", err)
 	}

pkg/engine/context/evaluate_test.go

@@ -3,6 +3,7 @@ package context
 import (
 	"testing"
 
+	kyvernov1 "github.com/kyverno/kyverno/api/kyverno/v1"
 	"github.com/stretchr/testify/assert"
 	admissionv1 "k8s.io/api/admission/v1"
 )
@@ -65,3 +66,23 @@ func createTestContext(obj, oldObj string) Interface {
 	ctx.AddRequest(request)
 	return ctx
 }
+
+func TestQueryOperation(t *testing.T) {
+	ctx := createTestContext(`{"a": {"b": 1, "c": 2}, "d": 3}`, `{"a": {"b": 2, "c": 2}, "d": 4}`)
+	assert.Equal(t, ctx.QueryOperation(), "UPDATE")
+	request := admissionv1.AdmissionRequest{
+		Operation: admissionv1.Delete,
+	}
+
+	err := ctx.AddRequest(request)
+	assert.Nil(t, err)
+	assert.Equal(t, ctx.QueryOperation(), "DELETE")
+
+	err = ctx.AddOperation(string(kyvernov1.Connect))
+	assert.Nil(t, err)
+	assert.Equal(t, ctx.QueryOperation(), "CONNECT")
+
+	err = ctx.AddRequest(admissionv1.AdmissionRequest{})
+	assert.Nil(t, err)
+	assert.Equal(t, ctx.QueryOperation(), "")
+}