refactor: propagate context through admission handlers (#5392)

Signed-off-by: Charles-Edouard Brétéché <charles.edouard@nirmata.com>

cmd/cleanup-controller/handlers.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"time"
 
 	"github.com/go-logr/logr"
@@ -20,7 +21,7 @@ func NewHandlers(client dclient.Interface) CleanupPolicyHandlers {
 	}
 }
 
-func (h *cleanupPolicyHandlers) Validate(logger logr.Logger, request *admissionv1.AdmissionRequest, _ time.Time) *admissionv1.AdmissionResponse {
+func (h *cleanupPolicyHandlers) Validate(ctx context.Context, logger logr.Logger, request *admissionv1.AdmissionRequest, _ time.Time) *admissionv1.AdmissionResponse {
 	if request.SubResource != "" {
 		logger.V(4).Info("skip policy validation on status update")
 		return admissionutils.ResponseSuccess()

cmd/cleanup-controller/server.go

@@ -26,7 +26,7 @@ type Server interface {
 
 type CleanupPolicyHandlers interface {
 	// Validate performs the validation check on policy resources
-	Validate(logr.Logger, *admissionv1.AdmissionRequest, time.Time) *admissionv1.AdmissionResponse
+	Validate(context.Context, logr.Logger, *admissionv1.AdmissionRequest, time.Time) *admissionv1.AdmissionResponse
 }
 
 type server struct {

pkg/webhooks/handlers/admission.go

@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -13,7 +14,7 @@ import (
 	admissionv1 "k8s.io/api/admission/v1"
 )
 
-type AdmissionHandler func(logr.Logger, *admissionv1.AdmissionRequest, time.Time) *admissionv1.AdmissionResponse
+type AdmissionHandler func(context.Context, logr.Logger, *admissionv1.AdmissionRequest, time.Time) *admissionv1.AdmissionResponse
 
 func (h AdmissionHandler) WithAdmission(logger logr.Logger) http.HandlerFunc {
 	return withAdmission(logger, h)
@@ -21,7 +22,6 @@ func (h AdmissionHandler) WithAdmission(logger logr.Logger) http.HandlerFunc {
 
 func withAdmission(logger logr.Logger, inner AdmissionHandler) http.HandlerFunc {
 	return func(writer http.ResponseWriter, request *http.Request) {
-		ctx := request.Context()
 		startTime := time.Now()
 		if request.Body == nil {
 			logger.Info("empty body", "req", request.URL.String())
@@ -59,19 +59,9 @@ func withAdmission(logger logr.Logger, inner AdmissionHandler) http.HandlerFunc
 			Allowed: true,
 			UID:     admissionReview.Request.UID,
 		}
-		adminssionResponse := inner(logger, admissionReview.Request, startTime)
-		if adminssionResponse != nil {
-			admissionReview.Response = adminssionResponse
-		}
-		responseJSON, err := json.Marshal(admissionReview)
-		if err != nil {
-			http.Error(writer, fmt.Sprintf("Could not encode response: %v", err), http.StatusInternalServerError)
-			return
-		}
-
 		// start span from request context
-		_, span := tracing.StartSpan(
+		ctx, span := tracing.StartSpan(
-			ctx,
+			request.Context(),
 			"admission_webhook_operations",
 			string(admissionReview.Request.Operation),
 			attribute.String("kind", admissionReview.Request.Kind.Kind),
@@ -81,7 +71,15 @@ func withAdmission(logger logr.Logger, inner AdmissionHandler) http.HandlerFunc
 			attribute.String("uid", string(admissionReview.Request.UID)),
 		)
 		defer span.End()
-
+		adminssionResponse := inner(ctx, logger, admissionReview.Request, startTime)
+		if adminssionResponse != nil {
+			admissionReview.Response = adminssionResponse
+		}
+		responseJSON, err := json.Marshal(admissionReview)
+		if err != nil {
+			http.Error(writer, fmt.Sprintf("Could not encode response: %v", err), http.StatusInternalServerError)
+			return
+		}
 		writer.Header().Set("Content-Type", "application/json; charset=utf-8")
 		if _, err := writer.Write(responseJSON); err != nil {
 			http.Error(writer, fmt.Sprintf("could not write response: %v", err), http.StatusInternalServerError)

pkg/webhooks/handlers/dump.go

@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"context"
 	"strings"
 	"time"
 
@@ -101,8 +102,8 @@ func (h AdmissionHandler) WithDump(enabled bool) AdmissionHandler {
 }
 
 func withDump(inner AdmissionHandler) AdmissionHandler {
-	return func(logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
+	return func(ctx context.Context, logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
-		response := inner(logger, request, startTime)
+		response := inner(ctx, logger, request, startTime)
 		dumpPayload(logger, request, response)
 		return response
 	}

pkg/webhooks/handlers/filter.go

@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"context"
 	"time"
 
 	"github.com/go-logr/logr"
@@ -13,10 +14,10 @@ func (h AdmissionHandler) WithFilter(configuration config.Configuration) Admissi
 }
 
 func withFilter(c config.Configuration, inner AdmissionHandler) AdmissionHandler {
-	return func(logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
+	return func(ctx context.Context, logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
 		if c.ToFilter(request.Kind.Kind, request.Namespace, request.Name) {
 			return nil
 		}
-		return inner(logger, request, startTime)
+		return inner(ctx, logger, request, startTime)
 	}
 }

pkg/webhooks/handlers/metrics.go

@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"context"
 	"time"
 
 	"github.com/go-logr/logr"
@@ -15,9 +16,9 @@ func (h AdmissionHandler) WithMetrics(metricsConfig *metrics.MetricsConfig) Admi
 }
 
 func withMetrics(metricsConfig *metrics.MetricsConfig, inner AdmissionHandler) AdmissionHandler {
-	return func(logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
+	return func(ctx context.Context, logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
 		defer admissionReviewDuration.Process(metricsConfig, request, int64(time.Since(startTime)))
 		admissionRequests.Process(metricsConfig, request)
-		return inner(logger, request, startTime)
+		return inner(ctx, logger, request, startTime)
 	}
 }

pkg/webhooks/handlers/protect.go

@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"time"
@@ -22,7 +23,7 @@ func (h AdmissionHandler) WithProtection(enabled bool) AdmissionHandler {
 }
 
 func withProtection(inner AdmissionHandler) AdmissionHandler {
-	return func(logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
+	return func(ctx context.Context, logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
 		newResource, oldResource, err := utils.ExtractResources(nil, request)
 		if err != nil {
 			logger.Error(err, "Failed to extract resources")
@@ -37,6 +38,6 @@ func withProtection(inner AdmissionHandler) AdmissionHandler {
 				}
 			}
 		}
-		return inner(logger, request, startTime)
+		return inner(ctx, logger, request, startTime)
 	}
 }

pkg/webhooks/handlers/verify.go

@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"context"
 	"time"
 
 	"github.com/go-logr/logr"
@@ -11,7 +12,7 @@ import (
 )
 
 func Verify() AdmissionHandler {
-	return func(logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
+	return func(ctx context.Context, logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
 		if request.Name != "kyverno-health" || request.Namespace != config.KyvernoNamespace() {
 			return admissionutils.ResponseSuccess()
 		}

pkg/webhooks/policy/handlers.go

@@ -1,6 +1,7 @@
 package policy
 
 import (
+	"context"
 	"time"
 
 	"github.com/go-logr/logr"
@@ -24,7 +25,7 @@ func NewHandlers(client dclient.Interface, openApiManager openapi.Manager) webho
 	}
 }
 
-func (h *handlers) Validate(logger logr.Logger, request *admissionv1.AdmissionRequest, _ time.Time) *admissionv1.AdmissionResponse {
+func (h *handlers) Validate(ctx context.Context, logger logr.Logger, request *admissionv1.AdmissionRequest, _ time.Time) *admissionv1.AdmissionResponse {
 	if request.SubResource != "" {
 		logger.V(4).Info("skip policy validation on status update")
 		return admissionutils.ResponseSuccess()
@@ -42,6 +43,6 @@ func (h *handlers) Validate(logger logr.Logger, request *admissionv1.AdmissionRe
 	return admissionutils.Response(err, warnings...)
 }
 
-func (h *handlers) Mutate(logger logr.Logger, request *admissionv1.AdmissionRequest, _ time.Time) *admissionv1.AdmissionResponse {
+func (h *handlers) Mutate(ctx context.Context, logger logr.Logger, request *admissionv1.AdmissionRequest, _ time.Time) *admissionv1.AdmissionResponse {
 	return admissionutils.ResponseSuccess()
 }

pkg/webhooks/resource/handlers.go

@@ -1,6 +1,7 @@
 package resource
 
 import (
+	"context"
 	"errors"
 	"time"
 
@@ -93,7 +94,7 @@ func NewHandlers(
 	}
 }
 
-func (h *handlers) Validate(logger logr.Logger, request *admissionv1.AdmissionRequest, failurePolicy string, startTime time.Time) *admissionv1.AdmissionResponse {
+func (h *handlers) Validate(ctx context.Context, logger logr.Logger, request *admissionv1.AdmissionRequest, failurePolicy string, startTime time.Time) *admissionv1.AdmissionResponse {
 	if webhookutils.ExcludeKyvernoResources(request.Kind.Kind) {
 		return admissionutils.ResponseSuccess()
 	}
@@ -143,7 +144,7 @@ func (h *handlers) Validate(logger logr.Logger, request *admissionv1.AdmissionRe
 	return admissionutils.ResponseSuccess(warnings...)
 }
 
-func (h *handlers) Mutate(logger logr.Logger, request *admissionv1.AdmissionRequest, failurePolicy string, startTime time.Time) *admissionv1.AdmissionResponse {
+func (h *handlers) Mutate(ctx context.Context, logger logr.Logger, request *admissionv1.AdmissionRequest, failurePolicy string, startTime time.Time) *admissionv1.AdmissionResponse {
 	if webhookutils.ExcludeKyvernoResources(request.Kind.Kind) {
 		return admissionutils.ResponseSuccess()
 	}

pkg/webhooks/resource/handlers_test.go

@@ -282,18 +282,18 @@ func Test_AdmissionResponseValid(t *testing.T) {
 		},
 	}
 
-	response := handlers.Mutate(logger, request, "", time.Now())
+	response := handlers.Mutate(ctx, logger, request, "", time.Now())
 	assert.Assert(t, response != nil)
 	assert.Equal(t, response.Allowed, true)
 
-	response = handlers.Validate(logger, request, "", time.Now())
+	response = handlers.Validate(ctx, logger, request, "", time.Now())
 	assert.Equal(t, response.Allowed, true)
 	assert.Equal(t, len(response.Warnings), 0)
 
 	validPolicy.Spec.ValidationFailureAction = "Enforce"
 	policyCache.Set(key, &validPolicy)
 
-	response = handlers.Validate(logger, request, "", time.Now())
+	response = handlers.Validate(ctx, logger, request, "", time.Now())
 	assert.Equal(t, response.Allowed, false)
 	assert.Equal(t, len(response.Warnings), 0)
 
@@ -326,7 +326,7 @@ func Test_AdmissionResponseInvalid(t *testing.T) {
 	invalidPolicy.Spec.ValidationFailureAction = "Enforce"
 	policyCache.Set(keyInvalid, &invalidPolicy)
 
-	response := handlers.Validate(logger, request, "", time.Now())
+	response := handlers.Validate(ctx, logger, request, "", time.Now())
 	assert.Equal(t, response.Allowed, false)
 	assert.Equal(t, len(response.Warnings), 0)
 
@@ -334,7 +334,7 @@ func Test_AdmissionResponseInvalid(t *testing.T) {
 	invalidPolicy.Spec.FailurePolicy = &ignore
 	policyCache.Set(keyInvalid, &invalidPolicy)
 
-	response = handlers.Validate(logger, request, "", time.Now())
+	response = handlers.Validate(ctx, logger, request, "", time.Now())
 	assert.Equal(t, response.Allowed, true)
 	assert.Equal(t, len(response.Warnings), 1)
 }
@@ -367,7 +367,7 @@ func Test_ImageVerify(t *testing.T) {
 	policy.Spec.ValidationFailureAction = "Enforce"
 	policyCache.Set(key, &policy)
 
-	response := handlers.Mutate(logger, request, "", time.Now())
+	response := handlers.Mutate(ctx, logger, request, "", time.Now())
 	assert.Equal(t, response.Allowed, false)
 	assert.Equal(t, len(response.Warnings), 0)
 
@@ -375,7 +375,7 @@ func Test_ImageVerify(t *testing.T) {
 	policy.Spec.FailurePolicy = &ignore
 	policyCache.Set(key, &policy)
 
-	response = handlers.Mutate(logger, request, "", time.Now())
+	response = handlers.Mutate(ctx, logger, request, "", time.Now())
 	assert.Equal(t, response.Allowed, false)
 	assert.Equal(t, len(response.Warnings), 0)
 }
@@ -405,7 +405,7 @@ func Test_MutateAndVerify(t *testing.T) {
 		},
 	}
 
-	response := handlers.Mutate(logger, request, "", time.Now())
+	response := handlers.Mutate(ctx, logger, request, "", time.Now())
 	assert.Equal(t, response.Allowed, true)
 	assert.Equal(t, len(response.Warnings), 0)
 }

pkg/webhooks/server.go

@@ -39,16 +39,16 @@ type Server interface {
 
 type PolicyHandlers interface {
 	// Mutate performs the mutation of policy resources
-	Mutate(logr.Logger, *admissionv1.AdmissionRequest, time.Time) *admissionv1.AdmissionResponse
+	Mutate(context.Context, logr.Logger, *admissionv1.AdmissionRequest, time.Time) *admissionv1.AdmissionResponse
 	// Validate performs the validation check on policy resources
-	Validate(logr.Logger, *admissionv1.AdmissionRequest, time.Time) *admissionv1.AdmissionResponse
+	Validate(context.Context, logr.Logger, *admissionv1.AdmissionRequest, time.Time) *admissionv1.AdmissionResponse
 }
 
 type ResourceHandlers interface {
 	// Mutate performs the mutation of kube resources
-	Mutate(logr.Logger, *admissionv1.AdmissionRequest, string, time.Time) *admissionv1.AdmissionResponse
+	Mutate(context.Context, logr.Logger, *admissionv1.AdmissionRequest, string, time.Time) *admissionv1.AdmissionResponse
 	// Validate performs the validation check on kube resources
-	Validate(logr.Logger, *admissionv1.AdmissionRequest, string, time.Time) *admissionv1.AdmissionResponse
+	Validate(context.Context, logr.Logger, *admissionv1.AdmissionRequest, string, time.Time) *admissionv1.AdmissionResponse
 }
 
 type server struct {
@@ -199,14 +199,14 @@ func registerWebhookHandlers(
 	basePath string,
 	configuration config.Configuration,
 	metricsConfig *metrics.MetricsConfig,
-	handlerFunc func(logr.Logger, *admissionv1.AdmissionRequest, string, time.Time) *admissionv1.AdmissionResponse,
+	handlerFunc func(context.Context, logr.Logger, *admissionv1.AdmissionRequest, string, time.Time) *admissionv1.AdmissionResponse,
 	debugModeOpts DebugModeOptions,
 ) {
 	mux.HandlerFunc(
 		"POST",
 		basePath,
-		handlers.AdmissionHandler(func(logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
+		handlers.AdmissionHandler(func(ctx context.Context, logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
-			return handlerFunc(logger, request, "all", startTime)
+			return handlerFunc(ctx, logger, request, "all", startTime)
 		}).
 			WithFilter(configuration).
 			WithProtection(toggle.ProtectManagedResources.Enabled()).
@@ -217,8 +217,8 @@ func registerWebhookHandlers(
 	mux.HandlerFunc(
 		"POST",
 		basePath+"/fail",
-		handlers.AdmissionHandler(func(logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
+		handlers.AdmissionHandler(func(ctx context.Context, logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
-			return handlerFunc(logger, request, "fail", startTime)
+			return handlerFunc(ctx, logger, request, "fail", startTime)
 		}).
 			WithFilter(configuration).
 			WithProtection(toggle.ProtectManagedResources.Enabled()).
@@ -229,8 +229,8 @@ func registerWebhookHandlers(
 	mux.HandlerFunc(
 		"POST",
 		basePath+"/ignore",
-		handlers.AdmissionHandler(func(logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
+		handlers.AdmissionHandler(func(ctx context.Context, logger logr.Logger, request *admissionv1.AdmissionRequest, startTime time.Time) *admissionv1.AdmissionResponse {
-			return handlerFunc(logger, request, "ignore", startTime)
+			return handlerFunc(ctx, logger, request, "ignore", startTime)
 		}).
 			WithFilter(configuration).
 			WithProtection(toggle.ProtectManagedResources.Enabled()).