apiserver: introduce abstraction for wrapping ResponseWriter

staging/src/k8s.io/apiserver/pkg/endpoints/responsewriter/fake.go

@@ -0,0 +1,54 @@
+/*
+Copyright 2021 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package responsewriter
+
+import (
+	"bufio"
+	"net"
+	"net/http"
+)
+
+var _ http.ResponseWriter = &FakeResponseWriter{}
+
+// FakeResponseWriter implements http.ResponseWriter,
+// it is used for testing purpose only
+type FakeResponseWriter struct{}
+
+func (fw *FakeResponseWriter) Header() http.Header          { return http.Header{} }
+func (fw *FakeResponseWriter) WriteHeader(code int)         {}
+func (fw *FakeResponseWriter) Write(bs []byte) (int, error) { return len(bs), nil }
+
+// For HTTP2 an http.ResponseWriter object implements
+// http.Flusher and http.CloseNotifier.
+// It is used for testing purpose only
+type FakeResponseWriterFlusherCloseNotifier struct {
+	*FakeResponseWriter
+}
+
+func (fw *FakeResponseWriterFlusherCloseNotifier) Flush()                   {}
+func (fw *FakeResponseWriterFlusherCloseNotifier) CloseNotify() <-chan bool { return nil }
+
+// For HTTP/1.x an http.ResponseWriter object implements
+// http.Flusher, http.CloseNotifier and http.Hijacker.
+// It is used for testing purpose only
+type FakeResponseWriterFlusherCloseNotifierHijacker struct {
+	*FakeResponseWriterFlusherCloseNotifier
+}
+
+func (fw *FakeResponseWriterFlusherCloseNotifierHijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return nil, nil, nil
+}

staging/src/k8s.io/apiserver/pkg/endpoints/responsewriter/wrapper.go

@@ -0,0 +1,180 @@
+/*
+Copyright 2021 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package responsewriter
+
+import (
+	"bufio"
+	"net"
+	"net/http"
+)
+
+// UserProvidedDecorator represensts a user (client that uses this package)
+// provided decorator that wraps an inner http.ResponseWriter object.
+// The user-provided decorator object must return the inner (decorated)
+// http.ResponseWriter object via the Unwrap function.
+type UserProvidedDecorator interface {
+	http.ResponseWriter
+
+	// Unwrap returns the inner http.ResponseWriter object associated
+	// with the user-provided decorator.
+	Unwrap() http.ResponseWriter
+}
+
+// WrapForHTTP1Or2 accepts a user-provided decorator of an "inner" http.responseWriter
+// object and potentially wraps the user-provided decorator with a new http.ResponseWriter
+// object that implements http.CloseNotifier, http.Flusher, and/or http.Hijacker by
+// delegating to the user-provided decorator (if it implements the relevant method) or
+// the inner http.ResponseWriter (otherwise), so that the returned http.ResponseWriter
+// object implements the same subset of those interfaces as the inner http.ResponseWriter.
+//
+// This function handles the following three casses.
+// - The inner ResponseWriter implements `http.CloseNotifier`, `http.Flusher`,
+//   and `http.Hijacker` (an HTTP/1.1 sever provides such a ResponseWriter).
+// - The inner ResponseWriter implements `http.CloseNotifier` and `http.Flusher`
+//   but not `http.Hijacker` (an HTTP/2 server provides such a ResponseWriter).
+// - All the other cases collapse to this one, in which the given ResponseWriter is returned.
+//
+// There are three applicable terms:
+//  - "outer": this is the ResponseWriter object returned by the WrapForHTTP1Or2 function.
+//  - "user-provided decorator" or "middle": this is the user-provided decorator
+//     that decorates an inner ResponseWriter object. A user-provided decorator
+//     implements the UserProvidedDecorator interface. A user-provided decorator
+//     may or may not implement http.CloseNotifier, http.Flusher or http.Hijacker.
+//  - "inner": the ResponseWriter that the user-provided decorator extends.
+func WrapForHTTP1Or2(decorator UserProvidedDecorator) http.ResponseWriter {
+	// from go net/http documentation:
+	// The default HTTP/1.x and HTTP/2 ResponseWriter implementations support Flusher
+	// Handlers should always test for this ability at runtime.
+	//
+	// The Hijacker interface is implemented by ResponseWriters that allow an HTTP handler
+	// to take over the connection.
+	// The default ResponseWriter for HTTP/1.x connections supports Hijacker, but HTTP/2 connections
+	// intentionally do not. ResponseWriter wrappers may also not support Hijacker.
+	// Handlers should always test for this ability at runtime
+	//
+	// The CloseNotifier interface is implemented by ResponseWriters which allow detecting
+	// when the underlying connection has gone away.
+	// Deprecated: the CloseNotifier interface predates Go's context package.
+	// New code should use Request.Context instead.
+	inner := decorator.Unwrap()
+	if innerNotifierFlusher, ok := inner.(CloseNotifierFlusher); ok {
+		// for HTTP/2 request, the default ResponseWriter object (http2responseWriter)
+		// implements Flusher and CloseNotifier.
+		outerHTTP2 := outerWithCloseNotifyAndFlush{
+			UserProvidedDecorator:     decorator,
+			InnerCloseNotifierFlusher: innerNotifierFlusher,
+		}
+
+		if innerHijacker, hijackable := inner.(http.Hijacker); hijackable {
+			// for HTTP/1.x request the default implementation of ResponseWriter
+			// also implement CloseNotifier, Flusher and Hijacker
+			return &outerWithCloseNotifyFlushAndHijack{
+				outerWithCloseNotifyAndFlush: outerHTTP2,
+				InnerHijacker:                innerHijacker,
+			}
+		}
+
+		return outerHTTP2
+	}
+
+	// we should never be here for either http/1.x or http2 request
+	return decorator
+}
+
+// CloseNotifierFlusher is a combination of http.CloseNotifier and http.Flusher
+// This applies to both http/1.x and http2 requests.
+type CloseNotifierFlusher interface {
+	http.CloseNotifier
+	http.Flusher
+}
+
+// GetOriginal goes through the chain of wrapped http.ResponseWriter objects
+// and returns the original http.ResponseWriter object provided to the first
+// request handler in the filter chain.
+func GetOriginal(w http.ResponseWriter) http.ResponseWriter {
+	decorator, ok := w.(UserProvidedDecorator)
+	if !ok {
+		return w
+	}
+
+	inner := decorator.Unwrap()
+	if inner == w {
+		// infinite cycle here, we should never be here though.
+		panic("http.ResponseWriter decorator chain has a cycle")
+	}
+
+	return GetOriginal(inner)
+}
+
+//lint:ignore SA1019 backward compatibility
+var _ http.CloseNotifier = outerWithCloseNotifyAndFlush{}
+var _ http.Flusher = outerWithCloseNotifyAndFlush{}
+var _ http.ResponseWriter = outerWithCloseNotifyAndFlush{}
+var _ UserProvidedDecorator = outerWithCloseNotifyAndFlush{}
+
+// outerWithCloseNotifyAndFlush is the outer object that extends the
+// user provied decorator with http.CloseNotifier and http.Flusher only.
+type outerWithCloseNotifyAndFlush struct {
+	// UserProvidedDecorator is the user-provided object, it decorates
+	// an inner ResponseWriter object.
+	UserProvidedDecorator
+
+	// http.CloseNotifier and http.Flusher for the inner object
+	InnerCloseNotifierFlusher CloseNotifierFlusher
+}
+
+func (wr outerWithCloseNotifyAndFlush) CloseNotify() <-chan bool {
+	if notifier, ok := wr.UserProvidedDecorator.(http.CloseNotifier); ok {
+		return notifier.CloseNotify()
+	}
+
+	return wr.InnerCloseNotifierFlusher.CloseNotify()
+}
+
+func (wr outerWithCloseNotifyAndFlush) Flush() {
+	if flusher, ok := wr.UserProvidedDecorator.(http.Flusher); ok {
+		flusher.Flush()
+		return
+	}
+
+	wr.InnerCloseNotifierFlusher.Flush()
+}
+
+//lint:file-ignore SA1019 Keep supporting deprecated http.CloseNotifier
+var _ http.CloseNotifier = outerWithCloseNotifyFlushAndHijack{}
+var _ http.Flusher = outerWithCloseNotifyFlushAndHijack{}
+var _ http.Hijacker = outerWithCloseNotifyFlushAndHijack{}
+var _ http.ResponseWriter = outerWithCloseNotifyFlushAndHijack{}
+var _ UserProvidedDecorator = outerWithCloseNotifyFlushAndHijack{}
+
+// outerWithCloseNotifyFlushAndHijack is the outer object that extends the
+// user-provided decorator with http.CloseNotifier, http.Flusher and http.Hijacker.
+// This applies to http/1.x requests only.
+type outerWithCloseNotifyFlushAndHijack struct {
+	outerWithCloseNotifyAndFlush
+
+	// http.Hijacker for the inner object
+	InnerHijacker http.Hijacker
+}
+
+func (wr outerWithCloseNotifyFlushAndHijack) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if hijacker, ok := wr.UserProvidedDecorator.(http.Hijacker); ok {
+		return hijacker.Hijack()
+	}
+
+	return wr.InnerHijacker.Hijack()
+}

staging/src/k8s.io/apiserver/pkg/endpoints/responsewriter/wrapper_test.go

@@ -0,0 +1,298 @@
+/*
+Copyright 2021 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package responsewriter
+
+import (
+	"bufio"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+	"time"
+)
+
+func TestWithHTTP1(t *testing.T) {
+	var originalWant http.ResponseWriter
+	counterGot := &counter{}
+	chain := func(h http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if originalWant == nil {
+				originalWant = w
+			}
+
+			assertCloseNotifierFlusherHijacker(t, true, w)
+
+			decorator := &fakeResponseWriterDecorator{
+				ResponseWriter: w,
+				counter:        counterGot,
+			}
+			wrapped := WrapForHTTP1Or2(decorator)
+
+			assertCloseNotifierFlusherHijacker(t, true, wrapped)
+
+			originalGot := GetOriginal(wrapped)
+			if originalWant != originalGot {
+				t.Errorf("Expected GetOriginal to return the original ResponseWriter object")
+				return
+			}
+
+			h.ServeHTTP(wrapped, r)
+		})
+	}
+
+	// wrap the original http.ResponseWriter multiple times
+	handler := chain(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// at this point, the original ResponseWriter object has been wrapped three times
+		// so each decorator is expected to tick the count by one for each method.
+		defer counterGot.assert(t, &counter{FlushInvoked: 3, CloseNotifyInvoked: 3, HijackInvoked: 3})
+
+		//lint:ignore SA1019 backward compatibility
+		w.(http.CloseNotifier).CloseNotify()
+		w.(http.Flusher).Flush()
+
+		conn, _, err := w.(http.Hijacker).Hijack()
+		if err != nil {
+			t.Errorf("Expected Hijack to succeed, but got error: %v", err)
+			return
+		}
+		conn.Close()
+	}))
+	handler = chain(handler)
+	handler = chain(handler)
+
+	server := newServer(t, handler, false)
+	defer server.Close()
+
+	sendRequest(t, server)
+}
+
+func TestWithHTTP2(t *testing.T) {
+	var originalWant http.ResponseWriter
+	counterGot := &counter{}
+	chain := func(h http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if originalWant == nil {
+				originalWant = w
+			}
+
+			assertCloseNotifierFlusherHijacker(t, false, w)
+
+			decorator := &fakeResponseWriterDecorator{
+				ResponseWriter: w,
+				counter:        counterGot,
+			}
+			wrapped := WrapForHTTP1Or2(decorator)
+
+			assertCloseNotifierFlusherHijacker(t, false, wrapped)
+
+			originalGot := GetOriginal(wrapped)
+			if originalWant != originalGot {
+				t.Errorf("Expected GetOriginal to return the original ResponseWriter object")
+				return
+			}
+
+			h.ServeHTTP(wrapped, r)
+		})
+	}
+
+	// wrap the original http.ResponseWriter multiple times
+	handler := chain(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// at this point, the original ResponseWriter object has been wrapped three times
+		// so each decorator is expected to tick the count by one for each method.
+		defer counterGot.assert(t, &counter{FlushInvoked: 3, CloseNotifyInvoked: 3, HijackInvoked: 0})
+
+		//lint:ignore SA1019 backward compatibility
+		w.(http.CloseNotifier).CloseNotify()
+		w.(http.Flusher).Flush()
+
+	}))
+	handler = chain(handler)
+	handler = chain(handler)
+
+	server := newServer(t, handler, true)
+	defer server.Close()
+
+	sendRequest(t, server)
+}
+
+func TestGetOriginal(t *testing.T) {
+	tests := []struct {
+		name          string
+		wrap          func() (http.ResponseWriter, http.ResponseWriter)
+		panicExpected bool
+	}{
+		{
+			name: "not wrapped",
+			wrap: func() (http.ResponseWriter, http.ResponseWriter) {
+				original := &FakeResponseWriter{}
+				return original, original
+			},
+		},
+		{
+			name: "wrapped once",
+			wrap: func() (http.ResponseWriter, http.ResponseWriter) {
+				original := &FakeResponseWriter{}
+				return original, &fakeResponseWriterDecorator{
+					ResponseWriter: original,
+				}
+			},
+		},
+		{
+			name: "wrapped multiple times",
+			wrap: func() (http.ResponseWriter, http.ResponseWriter) {
+				original := &FakeResponseWriter{}
+				return original, &fakeResponseWriterDecorator{
+					ResponseWriter: &fakeResponseWriterDecorator{
+						ResponseWriter: &fakeResponseWriterDecorator{
+							ResponseWriter: original,
+						},
+					},
+				}
+			},
+		},
+		{
+			name: "wraps itself",
+			wrap: func() (http.ResponseWriter, http.ResponseWriter) {
+				faulty := &fakeResponseWriterDecorator{}
+				faulty.ResponseWriter = faulty
+				return faulty, &fakeResponseWriterDecorator{
+					ResponseWriter: faulty,
+				}
+			},
+			panicExpected: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			originalExpected, wrapped := test.wrap()
+
+			func() {
+				defer func() {
+					err := recover()
+					switch {
+					case err != nil:
+						if !test.panicExpected {
+							t.Errorf("Expected no panic, but got: %v", err)
+						}
+					default:
+						if test.panicExpected {
+							t.Errorf("Expected a panic")
+						}
+					}
+				}()
+
+				originalGot := GetOriginal(wrapped)
+				if originalExpected != originalGot {
+					t.Errorf("Expected to get tehe original http.ResponseWriter object")
+				}
+			}()
+		})
+	}
+}
+
+func newServer(t *testing.T, h http.Handler, http2 bool) *httptest.Server {
+	server := httptest.NewUnstartedServer(h)
+	if http2 {
+		server.EnableHTTP2 = true
+		server.StartTLS()
+	} else {
+		server.Start()
+	}
+	_, err := url.Parse(server.URL)
+	if err != nil {
+		t.Fatalf("Expected the server to have a valid URL, but got: %s", server.URL)
+	}
+	return server
+}
+
+func sendRequest(t *testing.T, server *httptest.Server) {
+	req, err := http.NewRequest("GET", server.URL, nil)
+	if err != nil {
+		t.Fatalf("error creating request: %v", err)
+	}
+
+	client := server.Client()
+	client.Timeout = 30 * time.Second
+	_, err = client.Do(req)
+	if err != nil {
+		t.Fatalf("Unexpected non-nil err from client.Do: %v", err)
+	}
+}
+
+func assertCloseNotifierFlusherHijacker(t *testing.T, hijackableExpected bool, w http.ResponseWriter) {
+	// the http.ResponseWriter object for both http/1.x and http2
+	// implement http.Flusher and http.CloseNotifier
+	if _, ok := w.(http.Flusher); !ok {
+		t.Errorf("Expected the http.ResponseWriter object to implement http.Flusher")
+	}
+
+	//lint:ignore SA1019 backward compatibility
+	if _, ok := w.(http.CloseNotifier); !ok {
+		t.Errorf("Expected the http.ResponseWriter object to implement http.CloseNotifier")
+	}
+
+	// http/1.x implements http.Hijacker, not http2
+	if _, ok := w.(http.Hijacker); ok != hijackableExpected {
+		t.Errorf("Unexpected http.Hijacker implementation, expected: %t, but got: %t", hijackableExpected, ok)
+	}
+}
+
+type counter struct {
+	FlushInvoked       int
+	HijackInvoked      int
+	CloseNotifyInvoked int
+}
+
+func (c *counter) assert(t *testing.T, expected *counter) {
+	if expected.FlushInvoked != c.FlushInvoked {
+		t.Errorf("Expected Flush() count to match, wanted: %d, but got: %d", expected.FlushInvoked, c.FlushInvoked)
+	}
+	if expected.CloseNotifyInvoked != c.CloseNotifyInvoked {
+		t.Errorf("Expected CloseNotify() count to match, wanted: %d, but got: %d", expected.CloseNotifyInvoked, c.CloseNotifyInvoked)
+	}
+	if expected.HijackInvoked != c.HijackInvoked {
+		t.Errorf("Expected Hijack() count to match, wanted: %d, but got: %d", expected.HijackInvoked, c.HijackInvoked)
+	}
+}
+
+type fakeResponseWriterDecorator struct {
+	http.ResponseWriter
+	counter *counter
+}
+
+func (fw *fakeResponseWriterDecorator) Unwrap() http.ResponseWriter { return fw.ResponseWriter }
+func (fw *fakeResponseWriterDecorator) Flush() {
+	if fw.counter != nil {
+		fw.counter.FlushInvoked++
+	}
+	fw.ResponseWriter.(http.Flusher).Flush()
+}
+func (fw *fakeResponseWriterDecorator) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if fw.counter != nil {
+		fw.counter.HijackInvoked++
+	}
+	return fw.ResponseWriter.(http.Hijacker).Hijack()
+}
+func (fw *fakeResponseWriterDecorator) CloseNotify() <-chan bool {
+	if fw.counter != nil {
+		fw.counter.CloseNotifyInvoked++
+	}
+	//lint:ignore SA1019 backward compatibility
+	return fw.ResponseWriter.(http.CloseNotifier).CloseNotify()
+}

vendor/modules.txt

@@ -1503,6 +1503,7 @@ k8s.io/apiserver/pkg/endpoints/handlers/responsewriters
 k8s.io/apiserver/pkg/endpoints/metrics
 k8s.io/apiserver/pkg/endpoints/openapi
 k8s.io/apiserver/pkg/endpoints/request
+k8s.io/apiserver/pkg/endpoints/responsewriter
 k8s.io/apiserver/pkg/endpoints/warning
 k8s.io/apiserver/pkg/features
 k8s.io/apiserver/pkg/quota/v1