From b2515fa9eb1bfb4aafce06b8b5a052ad80c7cd02 Mon Sep 17 00:00:00 2001
From: shuting <shutting06@gmail.com>
Date: Tue, 20 Jul 2021 21:20:37 -0700
Subject: [PATCH] Add default image registry to patched resource (#2166)

---
 cmd/kyverno/main.go    |  3 ++-
 pkg/webhooks/server.go | 36 ++++++++++++++++++++++++++++++++++--
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/cmd/kyverno/main.go b/cmd/kyverno/main.go
index b61936effd..6fead7455b 100755
--- a/cmd/kyverno/main.go
+++ b/cmd/kyverno/main.go
@@ -4,13 +4,14 @@ import (
 	"context"
 	"flag"
 	"fmt"
-	"github.com/kyverno/kyverno/pkg/cosign"
 	"net/http"
 	_ "net/http/pprof"
 	"os"
 	"strings"
 	"time"
 
+	"github.com/kyverno/kyverno/pkg/cosign"
+
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	kubeinformers "k8s.io/client-go/informers"
 	"k8s.io/klog/v2"
diff --git a/pkg/webhooks/server.go b/pkg/webhooks/server.go
index 02166662d2..df1806a878 100644
--- a/pkg/webhooks/server.go
+++ b/pkg/webhooks/server.go
@@ -9,8 +9,6 @@ import (
 	"net/http"
 	"time"
 
-	"github.com/kyverno/kyverno/pkg/engine"
-
 	"github.com/go-logr/logr"
 	"github.com/julienschmidt/httprouter"
 	v1 "github.com/kyverno/kyverno/pkg/api/kyverno/v1"
@@ -20,8 +18,10 @@ import (
 	"github.com/kyverno/kyverno/pkg/common"
 	"github.com/kyverno/kyverno/pkg/config"
 	client "github.com/kyverno/kyverno/pkg/dclient"
+	"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,6 +373,10 @@ 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 {
+		ws.log.Error(err, "failed to patch images info to resource, policies that mutate images may be impacted")
+	}
+
 	policyContext := &engine.PolicyContext{
 		NewResource:         resource,
 		AdmissionInfo:       userRequestInfo,
@@ -623,3 +627,31 @@ 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.JSONPath, info.String()))
+	}
+
+	for _, info := range images.InitContainers {
+		patches = append(patches, buildJSONPatch("replace", info.JSONPath, 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)
+}