From 925c552c945401b14b45727153188280ad686dd2 Mon Sep 17 00:00:00 2001
From: Chris Hein <me@chrishein.com>
Date: Wed, 13 Jan 2021 16:31:35 -0800
Subject: [PATCH] adding WithContextFunc for webhook server to pass info to the
 handler

Signed-off-by: Chris Hein <me@chrishein.com>
---
 pkg/webhook/admission/http.go      |  6 +++++-
 pkg/webhook/admission/http_test.go | 28 ++++++++++++++++++++++++++++
 pkg/webhook/admission/webhook.go   |  5 +++++
 3 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/pkg/webhook/admission/http.go b/pkg/webhook/admission/http.go
index 6899011c9a..245e2d8ae1 100644
--- a/pkg/webhook/admission/http.go
+++ b/pkg/webhook/admission/http.go
@@ -45,6 +45,10 @@ var _ http.Handler = &Webhook{}
 func (wh *Webhook) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	var body []byte
 	var err error
+	ctx := r.Context()
+	if wh.WithContextFunc != nil {
+		ctx = wh.WithContextFunc(ctx, r)
+	}
 
 	var reviewResponse Response
 	if r.Body != nil {
@@ -93,7 +97,7 @@ func (wh *Webhook) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	wh.log.V(1).Info("received request", "UID", req.UID, "kind", req.Kind, "resource", req.Resource)
 
 	// TODO: add panic-recovery for Handle
-	reviewResponse = wh.Handle(r.Context(), req)
+	reviewResponse = wh.Handle(ctx, req)
 	wh.writeResponseTyped(w, reviewResponse, actualAdmRevGVK)
 }
 
diff --git a/pkg/webhook/admission/http_test.go b/pkg/webhook/admission/http_test.go
index 5afd9adacb..7dd2d5bcfc 100644
--- a/pkg/webhook/admission/http_test.go
+++ b/pkg/webhook/admission/http_test.go
@@ -162,6 +162,34 @@ var _ = Describe("Admission Webhooks", func() {
 			webhook.ServeHTTP(respRecorder, req.WithContext(ctx))
 			Expect(respRecorder.Body.String()).To(Equal(expected))
 		})
+
+		It("should mutate the Context from the HTTP request, if func supplied", func() {
+			req := &http.Request{
+				Header: http.Header{"Content-Type": []string{"application/json"}},
+				Body:   nopCloser{Reader: bytes.NewBufferString(`{"request":{}}`)},
+			}
+			type ctxkey int
+			const key ctxkey = 1
+			webhook := &Webhook{
+				Handler: &fakeHandler{
+					fn: func(ctx context.Context, req Request) Response {
+						return Allowed(ctx.Value(key).(string))
+					},
+				},
+				WithContextFunc: func(ctx context.Context, r *http.Request) context.Context {
+					return context.WithValue(ctx, key, r.Header["Content-Type"][0])
+				},
+				log: logf.RuntimeLog.WithName("webhook"),
+			}
+
+			expected := fmt.Sprintf(`{%s,"response":{"uid":"","allowed":true,"status":{"metadata":{},"reason":%q,"code":200}}}
+`, gvkJSONv1, "application/json")
+
+			ctx, cancel := context.WithCancel(context.Background())
+			cancel()
+			webhook.ServeHTTP(respRecorder, req.WithContext(ctx))
+			Expect(respRecorder.Body.String()).To(Equal(expected))
+		})
 	})
 })
 
diff --git a/pkg/webhook/admission/webhook.go b/pkg/webhook/admission/webhook.go
index 485f1758a4..4430c3132c 100644
--- a/pkg/webhook/admission/webhook.go
+++ b/pkg/webhook/admission/webhook.go
@@ -115,6 +115,11 @@ type Webhook struct {
 	// and potentially patches to apply to the handler.
 	Handler Handler
 
+	// WithContextFunc will allow you to take the http.Request.Context() and
+	// add any additional information such as passing the request path or
+	// headers thus allowing you to read them from within the handler
+	WithContextFunc func(context.Context, *http.Request) context.Context
+
 	// decoder is constructed on receiving a scheme and passed down to then handler
 	decoder *Decoder