From 7cab0ad2d2b2688575c1d6c8b5ecee2bfa5a39ff Mon Sep 17 00:00:00 2001
From: Abu Kashem <akashem@redhat.com>
Date: Thu, 26 Jan 2023 08:56:10 -0500
Subject: [PATCH] apiserver: warning should not panic when request times out

---
 .../src/k8s.io/apiserver/pkg/server/config.go |  5 +-
 .../pkg/server/genericapiserver_test.go       | 91 +++++++++++++++++++
 2 files changed, 95 insertions(+), 1 deletion(-)

diff --git a/staging/src/k8s.io/apiserver/pkg/server/config.go b/staging/src/k8s.io/apiserver/pkg/server/config.go
index ea0b82a350a..d42baab63bc 100644
--- a/staging/src/k8s.io/apiserver/pkg/server/config.go
+++ b/staging/src/k8s.io/apiserver/pkg/server/config.go
@@ -1011,6 +1011,10 @@ func DefaultBuildHandlerChain(apiHandler http.Handler, c *Config) http.Handler {
 
 	handler = genericfilters.WithCORS(handler, c.CorsAllowedOriginList, nil, nil, nil, "true")
 
+	// WithWarningRecorder must be wrapped by the timeout handler
+	// to make the addition of warning headers threadsafe
+	handler = genericapifilters.WithWarningRecorder(handler)
+
 	// WithTimeoutForNonLongRunningRequests will call the rest of the request handling in a go-routine with the
 	// context with deadline. The go-routine can keep running, while the timeout logic will return a timeout to the client.
 	handler = genericfilters.WithTimeoutForNonLongRunningRequests(handler, c.LongRunningFunc)
@@ -1024,7 +1028,6 @@ func DefaultBuildHandlerChain(apiHandler http.Handler, c *Config) http.Handler {
 	if c.SecureServing != nil && !c.SecureServing.DisableHTTP2 && c.GoawayChance > 0 {
 		handler = genericfilters.WithProbabilisticGoaway(handler, c.GoawayChance)
 	}
-	handler = genericapifilters.WithWarningRecorder(handler)
 	handler = genericapifilters.WithCacheControl(handler)
 	handler = genericfilters.WithHSTS(handler, c.HSTSDirectives)
 	if c.ShutdownSendRetryAfter {
diff --git a/staging/src/k8s.io/apiserver/pkg/server/genericapiserver_test.go b/staging/src/k8s.io/apiserver/pkg/server/genericapiserver_test.go
index 73709558b26..091349182e7 100644
--- a/staging/src/k8s.io/apiserver/pkg/server/genericapiserver_test.go
+++ b/staging/src/k8s.io/apiserver/pkg/server/genericapiserver_test.go
@@ -48,6 +48,7 @@ import (
 	openapinamer "k8s.io/apiserver/pkg/endpoints/openapi"
 	"k8s.io/apiserver/pkg/registry/rest"
 	genericfilters "k8s.io/apiserver/pkg/server/filters"
+	"k8s.io/apiserver/pkg/warning"
 	"k8s.io/client-go/informers"
 	"k8s.io/client-go/kubernetes/fake"
 	restclient "k8s.io/client-go/rest"
@@ -682,3 +683,93 @@ func TestGracefulShutdown(t *testing.T) {
 		t.Errorf("Timed out waiting for response.")
 	}
 }
+
+func TestWarningWithRequestTimeout(t *testing.T) {
+	type result struct {
+		err        interface{}
+		stackTrace string
+	}
+	clientDoneCh, resultCh := make(chan struct{}), make(chan result, 1)
+	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// this will catch recoverable panic like 'Header called after Handler finished'.
+		// go runtime crashes the program if it detects a program-ending
+		// panic like 'concurrent map iteration and map write', so this
+		// panic can not be caught.
+		defer func() {
+			result := result{}
+			result.err = recover()
+			if result.err != nil {
+				// Same as stdlib http server code. Manually allocate stack
+				// trace buffer size to prevent excessively large logs
+				const size = 64 << 10
+				buf := make([]byte, size)
+				buf = buf[:goruntime.Stack(buf, false)]
+				result.stackTrace = string(buf)
+			}
+			resultCh <- result
+		}()
+
+		// add warnings while we're waiting for the request to timeout to catch read/write races
+	loop:
+		for {
+			select {
+			case <-r.Context().Done():
+				break loop
+			default:
+				warning.AddWarning(r.Context(), "a", "1")
+			}
+		}
+		// the request has just timed out, write to catch read/write races
+		warning.AddWarning(r.Context(), "agent", "text")
+
+		// give time for the timeout response to be written, then try to
+		// write a response header to catch the "Header after Handler finished" panic
+		<-clientDoneCh
+
+		warning.AddWarning(r.Context(), "agent", "text")
+	})
+	handler := newGenericAPIServerHandlerChain(t, "/test", testHandler)
+
+	server := httptest.NewUnstartedServer(handler)
+	server.EnableHTTP2 = true
+	server.StartTLS()
+	defer server.Close()
+
+	request, err := http.NewRequest("GET", server.URL+"/test?timeout=100ms", nil)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	client := server.Client()
+	response, err := client.Do(request)
+	close(clientDoneCh)
+	if err != nil {
+		t.Errorf("expected server to return an HTTP response: %v", err)
+	}
+	if want := http.StatusGatewayTimeout; response == nil || response.StatusCode != want {
+		t.Errorf("expected server to return %d, but got: %v", want, response)
+	}
+
+	var resultGot result
+	select {
+	case resultGot = <-resultCh:
+	case <-time.After(wait.ForeverTestTimeout):
+		t.Errorf("the handler never returned a result")
+	}
+	if resultGot.err != nil {
+		t.Errorf("Expected no panic, but got: %v", resultGot.err)
+		t.Errorf("Stack Trace: %s", resultGot.stackTrace)
+	}
+}
+
+// builds a handler chain with the given user handler as used by GenericAPIServer.
+func newGenericAPIServerHandlerChain(t *testing.T, path string, handler http.Handler) http.Handler {
+	config, _ := setUp(t)
+	s, err := config.Complete(nil).New("test", NewEmptyDelegate())
+	if err != nil {
+		t.Fatalf("Error in bringing up the server: %v", err)
+	}
+
+	s.Handler.NonGoRestfulMux.Handle(path, handler)
+	return s.Handler
+}