Deferred loading for image info in policy context (#10787)

* fix: add deferred loading for image info in policy context

Signed-off-by: anushkamittal20 <anumittal4641@gmail.com>

* chore: move logic to context.go

Signed-off-by: anushkamittal20 <anumittal4641@gmail.com>

* chore: add unit test

Signed-off-by: anushkamittal20 <anumittal4641@gmail.com>

* fix: cleanup code

Signed-off-by: anushkamittal20 <anumittal4641@gmail.com>

* fix: unit test

Signed-off-by: anushkamittal20 <anumittal4641@gmail.com>

* fix: unit test

Signed-off-by: anushkamittal20 <anumittal4641@gmail.com>

---------

Signed-off-by: anushkamittal20 <anumittal4641@gmail.com>
Co-authored-by: anushkamittal20 <anumittal4641@gmail.com>
Co-authored-by: Vishal Choudhary <vishal.choudhary@nirmata.com>
Co-authored-by: Jim Bugwadia <jim@nirmata.com>

pkg/engine/context/context.go

@@ -1,6 +1,7 @@
 package context
 
 import (
+	cont "context"
 	"encoding/csv"
 	"fmt"
 	"regexp"
@@ -13,6 +14,7 @@ import (
 	"github.com/kyverno/kyverno/pkg/engine/jmespath"
 	"github.com/kyverno/kyverno/pkg/engine/jsonutils"
 	"github.com/kyverno/kyverno/pkg/logging"
+	"github.com/kyverno/kyverno/pkg/toggle"
 	apiutils "github.com/kyverno/kyverno/pkg/utils/api"
 	admissionv1 "k8s.io/api/admission/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
@@ -304,12 +306,25 @@ func (ctx *context) AddImageInfo(info apiutils.ImageInfo, cfg config.Configurati
 }
 
 func (ctx *context) AddImageInfos(resource *unstructured.Unstructured, cfg config.Configuration) error {
-	images, err := apiutils.ExtractImagesFromResource(*resource, nil, cfg)
+	imageInfoLoader := &ImageInfoLoader{
+		resource: resource,
+		eCtx:     ctx,
+		cfg:      cfg,
+	}
+	dl, err := NewDeferredLoader("images", imageInfoLoader, logger)
 	if err != nil {
 		return err
 	}
-
-	return ctx.addImageInfos(images)
+	if toggle.FromContext(cont.Background()).EnableDeferredLoading() {
+		if err := ctx.AddDeferredLoader(dl); err != nil {
+			return err
+		}
+	} else {
+		if err := imageInfoLoader.LoadData(); err != nil {
+			return err
+		}
+	}
+	return nil
 }
 
 func (ctx *context) addImageInfos(images map[string]map[string]apiutils.ImageInfo) error {
@@ -321,11 +336,30 @@ func (ctx *context) addImageInfos(images map[string]map[string]apiutils.ImageInf
 	if err != nil {
 		return err
 	}
-
 	logging.V(4).Info("updated image info", "images", utm)
 	return addToContext(ctx, utm, false, "images")
 }
 
+type ImageInfoLoader struct {
+	resource  *unstructured.Unstructured
+	hasLoaded bool
+	eCtx      *context
+	cfg       config.Configuration
+}
+
+func (l *ImageInfoLoader) HasLoaded() bool {
+	return l.hasLoaded
+}
+
+func (l *ImageInfoLoader) LoadData() error {
+	images, err := apiutils.ExtractImagesFromResource(*l.resource, nil, l.cfg)
+	if err != nil {
+		return err
+	}
+
+	return l.eCtx.addImageInfos(images)
+}
+
 func convertImagesToUnstructured(images map[string]map[string]apiutils.ImageInfo) (map[string]interface{}, error) {
 	results := map[string]interface{}{}
 	for containerType, v := range images {
@@ -363,6 +397,12 @@ func (ctx *context) GenerateCustomImageInfo(resource *unstructured.Unstructured,
 }
 
 func (ctx *context) ImageInfo() map[string]map[string]apiutils.ImageInfo {
+	// force load of image info from deferred loader
+	if len(ctx.images) == 0 {
+		if err := ctx.loadDeferred("images"); err != nil {
+			return map[string]map[string]apiutils.ImageInfo{}
+		}
+	}
 	return ctx.images
 }
 

pkg/engine/context/context_test.go

@@ -7,11 +7,15 @@ import (
 	urkyverno "github.com/kyverno/kyverno/api/kyverno/v2"
 	"github.com/kyverno/kyverno/pkg/config"
 	"github.com/kyverno/kyverno/pkg/engine/jmespath"
+	kubeutils "github.com/kyverno/kyverno/pkg/utils/kube"
 	"github.com/stretchr/testify/assert"
 	authenticationv1 "k8s.io/api/authentication/v1"
 )
 
-var jp = jmespath.New(config.NewDefaultConfiguration(false))
+var (
+	jp  = jmespath.New(config.NewDefaultConfiguration(false))
+	cfg = config.NewDefaultConfiguration(false)
+)
 
 func Test_addResourceAndUserContext(t *testing.T) {
 	var err error
@@ -227,3 +231,56 @@ func TestAddVariable(t *testing.T) {
 		})
 	}
 }
+
+func Test_ImageInfoLoader(t *testing.T) {
+	resource1, err := kubeutils.BytesToUnstructured([]byte(`{
+		"apiVersion": "v1",
+		"kind": "Pod",
+		"metadata": {
+		  "name": "test-pod",
+		  "namespace": "default"
+		},
+		"spec": {
+		  "containers": [{
+			"name": "test_container",
+			"image": "nginx:latest"
+		  }]
+		}
+	}`))
+	assert.Nil(t, err)
+	newctx := newContext()
+	err = newctx.AddImageInfos(resource1, cfg)
+	assert.Nil(t, err)
+	// images not loaded
+	assert.Nil(t, newctx.images)
+	// images loaded on Query
+	name, err := newctx.Query("images.containers.test_container.name")
+	assert.Nil(t, err)
+	assert.Equal(t, name, "nginx")
+}
+
+func Test_ImageInfoLoader_OnDirectCall(t *testing.T) {
+	resource1, err := kubeutils.BytesToUnstructured([]byte(`{
+		"apiVersion": "v1",
+		"kind": "Pod",
+		"metadata": {
+		  "name": "test-pod",
+		  "namespace": "default"
+		},
+		"spec": {
+		  "containers": [{
+			"name": "test_container",
+			"image": "nginx:latest"
+		  }]
+		}
+	}`))
+	assert.Nil(t, err)
+	newctx := newContext()
+	err = newctx.AddImageInfos(resource1, cfg)
+	assert.Nil(t, err)
+	// images not loaded
+	assert.Nil(t, newctx.images)
+	// images loaded on explicit call to ImageInfo
+	imageinfos := newctx.ImageInfo()
+	assert.Equal(t, imageinfos["containers"]["test_container"].Name, "nginx")
+}

pkg/engine/policycontext/policy_context.go

@@ -207,9 +207,11 @@ func NewPolicyContext(
 	if err := enginectx.AddNamespace(resource.GetNamespace()); err != nil {
 		return nil, err
 	}
+
 	if err := enginectx.AddImageInfos(&resource, configuration); err != nil {
 		return nil, err
 	}
+
 	if admissionInfo != nil {
 		if err := enginectx.AddUserInfo(*admissionInfo); err != nil {
 			return nil, err