From 9c66e4664d4041a3d5b987f0f7238ff849430bb8 Mon Sep 17 00:00:00 2001
From: Alexander Zielenski <zielenski@google.com>
Date: Mon, 26 Feb 2024 16:56:19 -0800
Subject: [PATCH] refactor compilatoin for real types

---
 .../plugin/policy/mutating/compilation.go     | 41 +++++++++----------
 .../policy/mutating/compilation_test.go       | 11 +++--
 2 files changed, 25 insertions(+), 27 deletions(-)

diff --git a/staging/src/k8s.io/apiserver/pkg/admission/plugin/policy/mutating/compilation.go b/staging/src/k8s.io/apiserver/pkg/admission/plugin/policy/mutating/compilation.go
index aa5992b192f7a..d05365a85d3fa 100644
--- a/staging/src/k8s.io/apiserver/pkg/admission/plugin/policy/mutating/compilation.go
+++ b/staging/src/k8s.io/apiserver/pkg/admission/plugin/policy/mutating/compilation.go
@@ -41,13 +41,17 @@ import (
 // any error is stored and delayed until invocation.
 func compilePolicy(policy *Policy) PolicyEvaluator {
 	// removal not yet supported
-	e := &evaluator{policy: policy}
-	e.compiledEvaluator, e.err = e.compile(plugincel.OptionalVariableDeclarations{HasParams: e.hasParams()})
-	return e.Invoke
+	var res []MutationEvaluationFunc
+	for _, m := range policy.Spec.Mutations {
+		e := &evaluator{mutation: &m}
+		e.compiledEvaluator, e.err = e.compile(plugincel.OptionalVariableDeclarations{HasParams: policy.Spec.ParamKind != nil})
+		res = append(res, e.Invoke)
+	}
+	return res
 }
 
 type evaluator struct {
-	policy *Policy
+	mutation *PolicyMutation
 
 	compiledEvaluator *compiledEvaluator
 	// err holds the error during the creation of compiledEvaluator
@@ -58,10 +62,6 @@ type compiledEvaluator struct {
 	programs []cel.Program
 }
 
-func (e *evaluator) hasParams() bool {
-	return e.policy.Spec.ParamKind != nil
-}
-
 func (e *evaluator) compile(vars plugincel.OptionalVariableDeclarations) (*compiledEvaluator, error) {
 	envSet, err := createEnvSet(vars)
 	if err != nil {
@@ -72,20 +72,19 @@ func (e *evaluator) compile(vars plugincel.OptionalVariableDeclarations) (*compi
 		return nil, err
 	}
 	var programs []cel.Program
-	for _, m := range e.policy.Spec.Mutations {
-		if m.PatchType != "Apply" {
-			return nil, fmt.Errorf("unsupported mutation type %q", m.PatchType)
-		}
-		ast, issues := env.Compile(m.Expression)
-		if issues != nil {
-			return nil, fmt.Errorf("cannot compile CEL expression: %v", issues)
-		}
-		program, err := env.Program(ast)
-		if err != nil {
-			return nil, fmt.Errorf("cannot initiate program: %w", err)
-		}
-		programs = append(programs, program)
+	// if m.PatchType != "Apply" {
+	// 	return nil, fmt.Errorf("unsupported mutation type %q", m.PatchType)
+	// }
+	ast, issues := env.Compile(e.mutation.Expression)
+	if issues != nil {
+		return nil, fmt.Errorf("cannot compile CEL expression: %v", issues)
 	}
+	program, err := env.Program(ast)
+	if err != nil {
+		return nil, fmt.Errorf("cannot initiate program: %w", err)
+	}
+	programs = append(programs, program)
+
 	return &compiledEvaluator{
 		programs: programs,
 	}, nil
diff --git a/staging/src/k8s.io/apiserver/pkg/admission/plugin/policy/mutating/compilation_test.go b/staging/src/k8s.io/apiserver/pkg/admission/plugin/policy/mutating/compilation_test.go
index 1ae7f115d8312..997d950b9c015 100644
--- a/staging/src/k8s.io/apiserver/pkg/admission/plugin/policy/mutating/compilation_test.go
+++ b/staging/src/k8s.io/apiserver/pkg/admission/plugin/policy/mutating/compilation_test.go
@@ -21,6 +21,7 @@ import (
 	"reflect"
 	"testing"
 
+	"k8s.io/api/admissionregistration/v1alpha1"
 	appsv1 "k8s.io/api/apps/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	plugincel "k8s.io/apiserver/pkg/admission/plugin/cel"
@@ -42,9 +43,8 @@ func TestCompilation(t *testing.T) {
 		{
 			name: "refer to object",
 			policy: &Policy{
-				Spec: MutatingAdmissionPolicySpec{Mutations: []Mutation{
+				Spec: v1alpha1.MutatingAdmissionPolicySpec{Mutations: []v1alpha1.Mutation{
 					{
-						PatchType: "Apply",
 						Expression: `Object{
 							spec: Object.spec{
 								replicas: object.spec.replicas % 2 == 0?object.spec.replicas + 1:object.spec.replicas
@@ -63,9 +63,8 @@ func TestCompilation(t *testing.T) {
 		{
 			name: "refer to oldObject",
 			policy: &Policy{
-				Spec: MutatingAdmissionPolicySpec{Mutations: []Mutation{
+				Spec: v1alpha1.MutatingAdmissionPolicySpec{Mutations: []v1alpha1.Mutation{
 					{
-						PatchType: "Apply",
 						Expression: `Object{
 							spec: Object.spec{
 								replicas: oldObject.spec.replicas % 2 == 0?oldObject.spec.replicas + 1:oldObject.spec.replicas
@@ -84,8 +83,8 @@ func TestCompilation(t *testing.T) {
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			e := &evaluator{policy: tc.policy}
-			c, err := e.compile(plugincel.OptionalVariableDeclarations{HasParams: e.hasParams()})
+			e := &evaluator{mutation: &tc.policy.Spec.Mutations[0]}
+			c, err := e.compile(plugincel.OptionalVariableDeclarations{HasParams: tc.policy.Spec.ParamKind != nil})
 			if err != nil {
 				if tc.expectedErr == "" {
 					t.Fatalf("unexpected error: %v", err)