Make LossLessDefaulter selective to only prevent dropping fields outside of schema

pkg/controller/jobframework/webhook/defaulter.go

@@ -18,9 +18,13 @@ package webhook
 
 import (
 	"context"
+	"encoding/json"
+	"net/http"
+	"slices"
 
-	jsonpatch "gomodules.xyz/jsonpatch/v2"
+	"gomodules.xyz/jsonpatch/v2"
 	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/util/sets"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
 )
 
@@ -29,11 +33,15 @@ import (
 func WithLosslessDefaulter(scheme *runtime.Scheme, obj runtime.Object, defaulter admission.CustomDefaulter) admission.Handler {
 	return &losslessDefaulter{
 		Handler: admission.WithCustomDefaulter(scheme, obj, defaulter).Handler,
+		decoder: admission.NewDecoder(scheme),
+		object:  obj,
 	}
 }
 
 type losslessDefaulter struct {
 	admission.Handler
+	decoder admission.Decoder
+	object  runtime.Object
 }
 
 // Handle handles admission requests, **dropping** remove operations from patches produced by controller-runtime.
@@ -43,18 +51,27 @@ type losslessDefaulter struct {
 // released CRDs.
 // Dropping the "remove" operations is safe because Kueue's job mutators never remove fields.
 func (h *losslessDefaulter) Handle(ctx context.Context, req admission.Request) admission.Response {
+	const opRemove = "remove"
 	response := h.Handler.Handle(ctx, req)
-	if response.Allowed {
-		var patches []jsonpatch.Operation
-		for _, p := range response.Patches {
-			if p.Operation != "remove" {
-				patches = append(patches, p)
-			}
+	if response.Allowed && slices.ContainsFunc(response.Patches, func(p jsonpatch.JsonPatchOperation) bool { return p.Operation == opRemove }) {
+		// get the schema caused patch
+		obj := h.object.DeepCopyObject()
+		if err := h.decoder.Decode(req, obj); err != nil {
+			return admission.Errored(http.StatusBadRequest, err)
+		}
+
+		marshalled, err := json.Marshal(obj)
+		if err != nil {
+			return admission.Errored(http.StatusInternalServerError, err)
 		}
-		if len(patches) == 0 {
-			response.PatchType = nil
+		schemePatch := admission.PatchResponseFromRaw(req.Object.Raw, marshalled)
+		if len(schemePatch.Patches) > 0 {
+			removeOperation := sets.New(slices.DeleteFunc(schemePatch.Patches, func(p jsonpatch.JsonPatchOperation) bool { return p.Operation != opRemove })...)
+			response.Patches = slices.DeleteFunc(response.Patches, func(p jsonpatch.JsonPatchOperation) bool { return removeOperation.Has(p) })
+			if len(response.Patches) == 0 {
+				response.PatchType = nil
+			}
 		}
-		response.Patches = patches
 	}
 	return response
 }

pkg/controller/jobframework/webhook/defaulter_test.go

@@ -21,11 +21,13 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
-	jsonpatch "gomodules.xyz/jsonpatch/v2"
+	"github.com/google/go-cmp/cmp/cmpopts"
+	"gomodules.xyz/jsonpatch/v2"
 	admissionv1 "k8s.io/api/admission/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/utils/ptr"
 	"sigs.k8s.io/controller-runtime/pkg/scheme"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
 )
@@ -36,7 +38,28 @@ var (
 )
 
 type TestResource struct {
-	Foo string `json:"foo,omitempty"`
+	metav1.ObjectMeta `json:",inline"`
+
+	Foo string   `json:"foo,omitempty"`
+	Bar string   `json:"bar,omitempty"`
+	Baz []string `json:"baz,omitempty"`
+
+	Conditions  []metav1.Condition `json:"conditions,omitempty"`
+	Items       []TestResourceItem `json:"items,omitempty"`
+	SubResource TestSubResource    `json:"subresource"`
+}
+
+type TestResourceItem struct {
+	SubItems []TestResourceSubItem `json:"subItems,omitempty"`
+}
+
+type TestResourceSubItem struct {
+	Baz []string `json:"baz,omitempty"`
+}
+
+type TestSubResource struct {
+	Foo string  `json:"foo"`
+	Bar *string `json:"bar"`
 }
 
 func (d *TestResource) GetObjectKind() schema.ObjectKind { return d }
@@ -57,43 +80,148 @@ type TestCustomDefaulter struct{}
 func (*TestCustomDefaulter) Default(ctx context.Context, obj runtime.Object) error {
 	d := obj.(*TestResource)
 	if d.Foo == "" {
-		d.Foo = "bar"
+		d.Foo = "foo"
+	}
+	if d.Bar == "bar" {
+		d.Bar = ""
+	}
+	if len(d.Baz) > 0 {
+		d.Baz = nil
 	}
-	return nil
-}
 
-func TestLossLessDefaulter(t *testing.T) {
-	sch := runtime.NewScheme()
-	builder := scheme.Builder{GroupVersion: testResourceGVK.GroupVersion()}
-	builder.Register(&TestResource{})
-	if err := builder.AddToScheme(sch); err != nil {
-		t.Fatalf("Couldn't add types to scheme: %v", err)
+	if d.Labels != nil {
+		delete(d.Labels, "example.com/foo")
 	}
 
-	handler := WithLosslessDefaulter(sch, &TestResource{}, &TestCustomDefaulter{})
+	if len(d.Finalizers) > 0 {
+		finalizers := make([]string, 0, len(d.Finalizers))
+		for _, val := range d.Finalizers {
+			if val != "foo" {
+				finalizers = append(finalizers, val)
+			}
+		}
+		d.Finalizers = finalizers
+	}
 
-	req := admission.Request{
-		AdmissionRequest: admissionv1.AdmissionRequest{
-			Kind: metav1.GroupVersionKind(testResourceGVK),
-			Object: runtime.RawExtension{
-				// This raw object has a field not defined in the go type.
-				// controller-runtime CustomDefaulter would have added a remove operation for it.
-				Raw: []byte(`{"baz": "qux"}`),
-			},
-		},
+	if len(d.Conditions) > 0 {
+		conditions := make([]metav1.Condition, 0, len(d.Conditions))
+		for _, cond := range d.Conditions {
+			if cond.Type == "foo" {
+				cond.ObservedGeneration = 0
+				conditions = append(conditions, cond)
+			}
+		}
+		d.Conditions = conditions
+	}
+
+	if len(d.Items) > 0 && len(d.Items[0].SubItems) > 0 {
+		baz := make([]string, 0, len(d.Items[0].SubItems[0].Baz))
+		for _, val := range d.Items[0].SubItems[0].Baz {
+			if val != "foo" {
+				baz = append(baz, val)
+			}
+		}
+		d.Items[0].SubItems[0].Baz = baz
+	}
+
+	if d.SubResource.Foo == "foo" {
+		d.SubResource.Foo = ""
 	}
-	resp := handler.Handle(context.Background(), req)
-	if !resp.Allowed {
-		t.Errorf("Response not allowed")
+	if ptr.Deref(d.SubResource.Bar, "") == "bar" {
+		d.SubResource.Bar = nil
 	}
-	wantPatches := []jsonpatch.Operation{
-		{
-			Operation: "add",
-			Path:      "/foo",
-			Value:     "bar",
+
+	return nil
+}
+
+func TestLossLessDefaulter(t *testing.T) {
+	testCases := map[string]struct {
+		request     admission.Request
+		wantAllowed bool
+		wantResult  *metav1.Status
+		wantPatches []jsonpatch.Operation
+	}{
+		"valid request with unknown fields": {
+			request: admission.Request{
+				AdmissionRequest: admissionv1.AdmissionRequest{
+					Kind: metav1.GroupVersionKind(testResourceGVK),
+					Object: runtime.RawExtension{
+						// This raw object has fields not defined in the go type.
+						// controller-runtime CustomDefaulter would have added remove operations for it.
+						Raw: []byte(`{
+	"unknown1": "unknown",
+	"unknown2": ["unknown"],
+	"unknown/unknown": "unknown",
+	"bar": "bar", 
+	"baz": ["foo"],
+	"finalizers": ["foo","bar"],
+	"labels": {"example.com/foo": "foo", "example.com/bar": "bar", "unknown": "unknown"},
+	"subresource": {"unknown1": "unknown", "unknown2": ["unknown"], "foo": "foo", "bar": "bar"},
+	"conditions": [
+		{"type": "foo", "message": "foo", "reason": "", "status": "", "lastTransitionTime": null, "observedGeneration": 1, "unknown": "unknown"}, 
+		{"type": "bar", "message": "bar", "reason": "", "status": "", "lastTransitionTime": null, "observedGeneration": 1, "unknown": "unknown"}
+	],
+	"items": [{"subItems": [{ "unknown1": "unknown", "baz": ["foo"] }] }]	
+}`),
+					},
+				},
+			},
+			wantAllowed: true,
+			wantPatches: []jsonpatch.Operation{
+				{Operation: "add", Path: "/creationTimestamp"},
+				{Operation: "add", Path: "/foo", Value: "foo"},
+				{Operation: "remove", Path: "/bar"},
+				{Operation: "remove", Path: "/baz"},
+				{Operation: "replace", Path: "/finalizers/0", Value: "bar"},
+				{Operation: "remove", Path: "/finalizers/1"},
+				{Operation: "remove", Path: "/labels/example.com~1foo"},
+				{Operation: "replace", Path: "/subresource/foo", Value: ""},
+				{Operation: "replace", Path: "/subresource/bar"},
+				{Operation: "remove", Path: "/conditions/0/observedGeneration"},
+				{Operation: "remove", Path: "/conditions/1"},
+				{Operation: "remove", Path: "/items/0/subItems/0/baz"},
+			},
+		},
+		"invalid request": {
+			request: admission.Request{
+				AdmissionRequest: admissionv1.AdmissionRequest{
+					Kind: metav1.GroupVersionKind(testResourceGVK),
+					Object: runtime.RawExtension{
+						Raw: []byte(`{"foo": 1}`),
+					},
+				},
+			},
+			wantResult: &metav1.Status{
+				Message: "json: cannot unmarshal number into Go struct field TestResource.foo of type string",
+				Code:    400,
+			},
 		},
 	}
-	if diff := cmp.Diff(wantPatches, resp.Patches); diff != "" {
-		t.Errorf("Unexpected patches (-want, +got): %s", diff)
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			sch := runtime.NewScheme()
+			builder := scheme.Builder{GroupVersion: testResourceGVK.GroupVersion()}
+			builder.Register(&TestResource{})
+			if err := builder.AddToScheme(sch); err != nil {
+				t.Fatalf("Couldn't add types to scheme: %v", err)
+			}
+
+			handler := WithLosslessDefaulter(sch, &TestResource{}, &TestCustomDefaulter{})
+			resp := handler.Handle(context.Background(), tc.request)
+
+			if diff := cmp.Diff(tc.wantAllowed, resp.Allowed); diff != "" {
+				t.Errorf("Unexpected allowed option (-want, +got): %s", diff)
+			}
+
+			if diff := cmp.Diff(tc.wantResult, resp.Result); diff != "" {
+				t.Errorf("Unexpected result (-want, +got): %s", diff)
+			}
+
+			if diff := cmp.Diff(tc.wantPatches, resp.Patches, cmpopts.SortSlices(func(a, b jsonpatch.Operation) bool {
+				return a.Path < b.Path
+			})); diff != "" {
+				t.Errorf("Unexpected patches (-want, +got): %s", diff)
+			}
+		})
 	}
 }