From 822e4f52797505245d5e74a251e686333f0c6165 Mon Sep 17 00:00:00 2001
From: Traian Schiau <traian_schiau@epam.com>
Date: Thu, 24 Oct 2024 11:17:21 +0300
Subject: [PATCH] Add `DefaulterPreserveUnknownFields`

An option stops the defaulter from pruning the fields
that are not recognized in the local scheme.
---
 pkg/webhook/admission/defaulter_custom.go     | 40 +++++++++++++----
 .../admission/defaulter_custom_test.go        | 43 ++++++++++++++++---
 2 files changed, 69 insertions(+), 14 deletions(-)

diff --git a/pkg/webhook/admission/defaulter_custom.go b/pkg/webhook/admission/defaulter_custom.go
index e3839176c7..6050dd4d13 100644
--- a/pkg/webhook/admission/defaulter_custom.go
+++ b/pkg/webhook/admission/defaulter_custom.go
@@ -36,17 +36,33 @@ type CustomDefaulter interface {
 	Default(ctx context.Context, obj runtime.Object) error
 }
 
+type defaulterOptions struct {
+	preserveUnknownFields bool
+}
+
+type defaulterOption func(*defaulterOptions)
+
+// DefaulterPreserveUnknownFields stops the defaulter from pruning the fields that are not recognized in the local scheme.
+func DefaulterPreserveUnknownFields(o *defaulterOptions) {
+	o.preserveUnknownFields = true
+}
+
 // WithCustomDefaulter creates a new Webhook for a CustomDefaulter interface.
-func WithCustomDefaulter(scheme *runtime.Scheme, obj runtime.Object, defaulter CustomDefaulter) *Webhook {
+func WithCustomDefaulter(scheme *runtime.Scheme, obj runtime.Object, defaulter CustomDefaulter, opts ...defaulterOption) *Webhook {
+	options := &defaulterOptions{}
+	for _, o := range opts {
+		o(options)
+	}
 	return &Webhook{
-		Handler: &defaulterForType{object: obj, defaulter: defaulter, decoder: NewDecoder(scheme)},
+		Handler: &defaulterForType{object: obj, defaulter: defaulter, decoder: NewDecoder(scheme), preserveUnknownFields: options.preserveUnknownFields},
 	}
 }
 
 type defaulterForType struct {
-	defaulter CustomDefaulter
-	object    runtime.Object
-	decoder   Decoder
+	defaulter             CustomDefaulter
+	object                runtime.Object
+	decoder               Decoder
+	preserveUnknownFields bool
 }
 
 // Handle handles admission requests.
@@ -79,8 +95,11 @@ func (h *defaulterForType) Handle(ctx context.Context, req Request) Response {
 		return Errored(http.StatusBadRequest, err)
 	}
 
-	// Keep a copy of the object
-	originalObj := obj.DeepCopyObject()
+	// Keep a copy of the object if needed
+	var originalObj runtime.Object
+	if h.preserveUnknownFields {
+		originalObj = obj.DeepCopyObject()
+	}
 
 	// Default the object
 	if err := h.defaulter.Default(ctx, obj); err != nil {
@@ -96,9 +115,12 @@ func (h *defaulterForType) Handle(ctx context.Context, req Request) Response {
 	if err != nil {
 		return Errored(http.StatusInternalServerError, err)
 	}
-	handlerResponse := PatchResponseFromRaw(req.Object.Raw, marshalled)
 
-	return h.dropSchemeRemovals(handlerResponse, originalObj, req.Object.Raw)
+	handlerResponse := PatchResponseFromRaw(req.Object.Raw, marshalled)
+	if h.preserveUnknownFields {
+		handlerResponse = h.dropSchemeRemovals(handlerResponse, originalObj, req.Object.Raw)
+	}
+	return handlerResponse
 }
 
 func (h *defaulterForType) dropSchemeRemovals(r Response, original runtime.Object, raw []byte) Response {
diff --git a/pkg/webhook/admission/defaulter_custom_test.go b/pkg/webhook/admission/defaulter_custom_test.go
index b454411968..7dda675fd7 100644
--- a/pkg/webhook/admission/defaulter_custom_test.go
+++ b/pkg/webhook/admission/defaulter_custom_test.go
@@ -29,7 +29,7 @@ import (
 
 var _ = Describe("Defaulter Handler", func() {
 
-	It("should not lose unknown fields", func() {
+	It("should not preserve unknown fields by default", func() {
 		obj := &TestDefaulter{}
 		handler := WithCustomDefaulter(admissionScheme, obj, &TestCustomDefaulter{})
 
@@ -42,17 +42,50 @@ var _ = Describe("Defaulter Handler", func() {
 			},
 		})
 		Expect(resp.Allowed).Should(BeTrue())
-		Expect(resp.Patches).To(Equal([]jsonpatch.JsonPatchOperation{
-			{
+		Expect(resp.Patches).To(HaveLen(3))
+		Expect(resp.Patches).To(ContainElements(
+			jsonpatch.JsonPatchOperation{
 				Operation: "add",
 				Path:      "/replica",
 				Value:     2.0,
 			},
-			{
+			jsonpatch.JsonPatchOperation{
+				Operation: "remove",
+				Path:      "/newField",
+			},
+			jsonpatch.JsonPatchOperation{
+				Operation: "remove",
+				Path:      "/totalReplicas",
+			},
+		))
+		Expect(resp.Result.Code).Should(Equal(int32(http.StatusOK)))
+	})
+
+	It("should preserve unknown fields when DefaulterPreserveUnknownFields is passed", func() {
+		obj := &TestDefaulter{}
+		handler := WithCustomDefaulter(admissionScheme, obj, &TestCustomDefaulter{}, DefaulterPreserveUnknownFields)
+
+		resp := handler.Handle(context.TODO(), Request{
+			AdmissionRequest: admissionv1.AdmissionRequest{
+				Operation: admissionv1.Create,
+				Object: runtime.RawExtension{
+					Raw: []byte(`{"newField":"foo", "totalReplicas":5}`),
+				},
+			},
+		})
+		Expect(resp.Allowed).Should(BeTrue())
+		Expect(resp.Patches).To(HaveLen(2))
+		Expect(resp.Patches).To(ContainElements(
+			jsonpatch.JsonPatchOperation{
+				Operation: "add",
+				Path:      "/replica",
+				Value:     2.0,
+			},
+			jsonpatch.JsonPatchOperation{
 				Operation: "remove",
 				Path:      "/totalReplicas",
 			},
-		}))
+		))
 		Expect(resp.Result.Code).Should(Equal(int32(http.StatusOK)))
 	})