From 2dcc79a20352722a4d72669ac3aeaafabb4d4613 Mon Sep 17 00:00:00 2001
From: Anushka Mittal <anushka@nirmata.com>
Date: Tue, 27 Aug 2024 17:06:34 +0530
Subject: [PATCH] NDEV-19864: deferred loading of images

Signed-off-by: Anushka Mittal <anushka@nirmata.com>
---
 pkg/engine/context/context.go      | 44 +++++++++++++++++++++-
 pkg/engine/context/context_test.go | 59 +++++++++++++++++++++++++++++-
 2 files changed, 100 insertions(+), 3 deletions(-)

diff --git a/pkg/engine/context/context.go b/pkg/engine/context/context.go
index 155be884532a..8ee1741dcb3c 100644
--- a/pkg/engine/context/context.go
+++ b/pkg/engine/context/context.go
@@ -304,11 +304,45 @@ 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
+	}
+	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
+}
+
+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 ctx.addImageInfos(images)
+
+	return l.eCtx.addImageInfos(images)
 }
 
 func (ctx *context) addImageInfos(images map[string]map[string]apiutils.ImageInfo) error {
@@ -361,6 +395,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
 }
 
diff --git a/pkg/engine/context/context_test.go b/pkg/engine/context/context_test.go
index d24f4980760c..5920efbab9fa 100644
--- a/pkg/engine/context/context_test.go
+++ b/pkg/engine/context/context_test.go
@@ -8,9 +8,13 @@ import (
 	"github.com/kyverno/kyverno/pkg/config"
 	"github.com/kyverno/kyverno/pkg/engine/jmespath"
 	authenticationv1 "k8s.io/api/authentication/v1"
+	kubeutils "github.com/kyverno/kyverno/pkg/utils/kube"
 )
 
-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
@@ -123,3 +127,56 @@ func Test_addResourceAndUserContext(t *testing.T) {
 		t.Error("expected result does not match")
 	}
 }
+
+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")
+}