From 1f72d5d61330dd6b4f1241b96ad44aab5131f7fe Mon Sep 17 00:00:00 2001
From: Di Jin <[email protected]>
Date: Thu, 1 Sep 2022 15:25:26 -0700
Subject: [PATCH] Add an option for aggregator

---
 cmd/kube-apiserver/app/aggregator.go          |  9 ++-
 cmd/kube-apiserver/app/options/options.go     |  9 ++-
 .../app/options/options_test.go               | 17 ++--
 .../pkg/util/proxy/upgradeaware.go            | 27 +++++++
 .../pkg/util/proxy/upgradeaware_test.go       | 77 +++++++++++++++++++
 .../pkg/apiserver/apiserver.go                |  7 ++
 .../pkg/apiserver/handler_proxy.go            |  6 ++
 7 files changed, 138 insertions(+), 14 deletions(-)

diff --git a/cmd/kube-apiserver/app/aggregator.go b/cmd/kube-apiserver/app/aggregator.go
index 4289ce48..1ac17065 100644
--- a/cmd/kube-apiserver/app/aggregator.go
+++ b/cmd/kube-apiserver/app/aggregator.go
@@ -109,10 +109,11 @@ func createAggregatorConfig(
 			SharedInformerFactory: externalInformers,
 		},
 		ExtraConfig: aggregatorapiserver.ExtraConfig{
-			ProxyClientCertFile: commandOptions.ProxyClientCertFile,
-			ProxyClientKeyFile:  commandOptions.ProxyClientKeyFile,
-			ServiceResolver:     serviceResolver,
-			ProxyTransport:      proxyTransport,
+			ProxyClientCertFile:       commandOptions.ProxyClientCertFile,
+			ProxyClientKeyFile:        commandOptions.ProxyClientKeyFile,
+			ServiceResolver:           serviceResolver,
+			ProxyTransport:            proxyTransport,
+			RejectForwardingRedirects: commandOptions.AggregatorRejectForwardingRedirects,
 		},
 	}
 
diff --git a/cmd/kube-apiserver/app/options/options.go b/cmd/kube-apiserver/app/options/options.go
index 68a714bf..653179fd 100644
--- a/cmd/kube-apiserver/app/options/options.go
+++ b/cmd/kube-apiserver/app/options/options.go
@@ -79,7 +79,8 @@ type ServerRunOptions struct {
 	ProxyClientCertFile string
 	ProxyClientKeyFile  string
 
-	EnableAggregatorRouting bool
+	EnableAggregatorRouting             bool
+	AggregatorRejectForwardingRedirects bool
 
 	MasterCount            int
 	EndpointReconcilerType string
@@ -134,7 +135,8 @@ func NewServerRunOptions() *ServerRunOptions {
 			},
 			HTTPTimeout: time.Duration(5) * time.Second,
 		},
-		ServiceNodePortRange: kubeoptions.DefaultServiceNodePortRange,
+		ServiceNodePortRange:                kubeoptions.DefaultServiceNodePortRange,
+		AggregatorRejectForwardingRedirects: true,
 	}
 
 	// Overwrite the default for storage data format.
@@ -280,6 +282,9 @@ func (s *ServerRunOptions) Flags() (fss cliflag.NamedFlagSets) {
 	fs.BoolVar(&s.EnableAggregatorRouting, "enable-aggregator-routing", s.EnableAggregatorRouting,
 		"Turns on aggregator routing requests to endpoints IP rather than cluster IP.")
 
+	fs.BoolVar(&s.AggregatorRejectForwardingRedirects, "aggregator-reject-forwarding-redirect", s.AggregatorRejectForwardingRedirects,
+		"Aggregator reject forwarding redirect response back to client.")
+
 	fs.StringVar(&s.ServiceAccountSigningKeyFile, "service-account-signing-key-file", s.ServiceAccountSigningKeyFile, ""+
 		"Path to the file that contains the current private key of the service account token issuer. The issuer will sign issued ID tokens with this private key.")
 
diff --git a/cmd/kube-apiserver/app/options/options_test.go b/cmd/kube-apiserver/app/options/options_test.go
index 8ad927bb..1f188590 100644
--- a/cmd/kube-apiserver/app/options/options_test.go
+++ b/cmd/kube-apiserver/app/options/options_test.go
@@ -299,14 +299,15 @@ func TestAddFlags(t *testing.T) {
 		EgressSelector: &apiserveroptions.EgressSelectorOptions{
 			ConfigFile: "/var/run/kubernetes/egress-selector/connectivity.yaml",
 		},
-		EnableLogsHandler:                 false,
-		EnableAggregatorRouting:           true,
-		ProxyClientKeyFile:                "/var/run/kubernetes/proxy.key",
-		ProxyClientCertFile:               "/var/run/kubernetes/proxy.crt",
-		Metrics:                           &metrics.Options{},
-		Logs:                              logs.NewOptions(),
-		IdentityLeaseDurationSeconds:      3600,
-		IdentityLeaseRenewIntervalSeconds: 10,
+		EnableLogsHandler:                   false,
+		EnableAggregatorRouting:             true,
+		ProxyClientKeyFile:                  "/var/run/kubernetes/proxy.key",
+		ProxyClientCertFile:                 "/var/run/kubernetes/proxy.crt",
+		Metrics:                             &metrics.Options{},
+		Logs:                                logs.NewOptions(),
+		IdentityLeaseDurationSeconds:        3600,
+		IdentityLeaseRenewIntervalSeconds:   10,
+		AggregatorRejectForwardingRedirects: true,
 	}
 
 	if !reflect.DeepEqual(expected, s) {
diff --git a/staging/src/k8s.io/apimachinery/pkg/util/proxy/upgradeaware.go b/staging/src/k8s.io/apimachinery/pkg/util/proxy/upgradeaware.go
index 43f528d9..4ed108e9 100644
--- a/staging/src/k8s.io/apimachinery/pkg/util/proxy/upgradeaware.go
+++ b/staging/src/k8s.io/apimachinery/pkg/util/proxy/upgradeaware.go
@@ -80,6 +80,8 @@ type UpgradeAwareHandler struct {
 	MaxBytesPerSec int64
 	// Responder is passed errors that occur while setting up proxying.
 	Responder ErrorResponder
+	// Reject to forward redirect response
+	RejectForwardingRedirects bool
 }
 
 const defaultFlushInterval = 200 * time.Millisecond
@@ -232,6 +234,31 @@ func (h *UpgradeAwareHandler) ServeHTTP(w http.ResponseWriter, req *http.Request
 	proxy.Transport = h.Transport
 	proxy.FlushInterval = h.FlushInterval
 	proxy.ErrorLog = log.New(noSuppressPanicError{}, "", log.LstdFlags)
+	if h.RejectForwardingRedirects {
+		oldModifyResponse := proxy.ModifyResponse
+		proxy.ModifyResponse = func(response *http.Response) error {
+			code := response.StatusCode
+			if code >= 300 && code <= 399 {
+				// close the original response
+				response.Body.Close()
+				msg := "the backend attempted to redirect this request, which is not permitted"
+				// replace the response
+				*response = http.Response{
+					StatusCode:    http.StatusBadGateway,
+					Status:        fmt.Sprintf("%d %s", response.StatusCode, http.StatusText(response.StatusCode)),
+					Body:          io.NopCloser(strings.NewReader(msg)),
+					ContentLength: int64(len(msg)),
+				}
+			} else {
+				if oldModifyResponse != nil {
+					if err := oldModifyResponse(response); err != nil {
+						return err
+					}
+				}
+			}
+			return nil
+		}
+	}
 	if h.Responder != nil {
 		// if an optional error interceptor/responder was provided wire it
 		// the custom responder might be used for providing a unified error reporting
diff --git a/staging/src/k8s.io/apimachinery/pkg/util/proxy/upgradeaware_test.go b/staging/src/k8s.io/apimachinery/pkg/util/proxy/upgradeaware_test.go
index 3f2d41f9..33051ffa 100644
--- a/staging/src/k8s.io/apimachinery/pkg/util/proxy/upgradeaware_test.go
+++ b/staging/src/k8s.io/apimachinery/pkg/util/proxy/upgradeaware_test.go
@@ -663,6 +663,83 @@ func TestProxyUpgradeErrorResponse(t *testing.T) {
 	}
 }
 
+func TestRejectForwardingRedirectsOption(t *testing.T) {
+	originalBody := []byte(`some data`)
+	testCases := []struct {
+		name                      string
+		rejectForwardingRedirects bool
+		serverStatusCode          int
+		expectStatusCode          int
+		expectBody                []byte
+	}{
+		{
+			name:                      "reject redirection enabled in proxy, backend server sending 200 response",
+			rejectForwardingRedirects: true,
+			serverStatusCode:          200,
+			expectStatusCode:          200,
+			expectBody:                originalBody,
+		},
+		{
+			name:                      "reject redirection enabled in proxy, backend server sending 301 response",
+			rejectForwardingRedirects: true,
+			serverStatusCode:          301,
+			expectStatusCode:          502,
+			expectBody:                []byte(`the backend attempted to redirect this request, which is not permitted`),
+		},
+		{
+			name:                      "reject redirection disabled in proxy, backend server sending 200 response",
+			rejectForwardingRedirects: false,
+			serverStatusCode:          200,
+			expectStatusCode:          200,
+			expectBody:                originalBody,
+		},
+		{
+			name:                      "reject redirection disabled in proxy, backend server sending 301 response",
+			rejectForwardingRedirects: false,
+			serverStatusCode:          301,
+			expectStatusCode:          301,
+			expectBody:                originalBody,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Set up a backend server
+			backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(tc.serverStatusCode)
+				w.Write(originalBody)
+			}))
+			defer backendServer.Close()
+			backendServerURL, _ := url.Parse(backendServer.URL)
+
+			// Set up a proxy pointing to the backend
+			proxyHandler := NewUpgradeAwareHandler(backendServerURL, nil, false, false, &fakeResponder{t: t})
+			proxyHandler.RejectForwardingRedirects = tc.rejectForwardingRedirects
+			proxy := httptest.NewServer(proxyHandler)
+			defer proxy.Close()
+			proxyURL, _ := url.Parse(proxy.URL)
+
+			conn, err := net.Dial("tcp", proxyURL.Host)
+			require.NoError(t, err)
+			bufferedReader := bufio.NewReader(conn)
+
+			req, _ := http.NewRequest("GET", proxyURL.String(), nil)
+			require.NoError(t, req.Write(conn))
+			// Verify we get the correct response and message body content
+			resp, err := http.ReadResponse(bufferedReader, nil)
+			require.NoError(t, err)
+			assert.Equal(t, tc.expectStatusCode, resp.StatusCode)
+			data, err := ioutil.ReadAll(resp.Body)
+			require.NoError(t, err)
+			assert.Equal(t, tc.expectBody, data)
+			assert.Equal(t, int64(len(tc.expectBody)), resp.ContentLength)
+			resp.Body.Close()
+
+			// clean up
+			conn.Close()
+		})
+	}
+}
+
 func TestDefaultProxyTransport(t *testing.T) {
 	tests := []struct {
 		name,
diff --git a/staging/src/k8s.io/kube-aggregator/pkg/apiserver/apiserver.go b/staging/src/k8s.io/kube-aggregator/pkg/apiserver/apiserver.go
index d0ab3186..50a85820 100644
--- a/staging/src/k8s.io/kube-aggregator/pkg/apiserver/apiserver.go
+++ b/staging/src/k8s.io/kube-aggregator/pkg/apiserver/apiserver.go
@@ -80,6 +80,8 @@ type ExtraConfig struct {
 
 	// Mechanism by which the Aggregator will resolve services. Required.
 	ServiceResolver ServiceResolver
+
+	RejectForwardingRedirects bool
 }
 
 // Config represents the configuration needed to create an APIAggregator.
@@ -143,6 +145,9 @@ type APIAggregator struct {
 	// egressSelector selects the proper egress dialer to communicate with the custom apiserver
 	// overwrites proxyTransport dialer if not nil
 	egressSelector *egressselector.EgressSelector
+
+	// rejectForwardingRedirects is whether to allow to forward redirect response
+	rejectForwardingRedirects bool
 }
 
 // Complete fills in any fields not set that are required to have valid data. It's mutating the receiver.
@@ -194,6 +199,7 @@ func (c completedConfig) NewWithDelegate(delegationTarget genericapiserver.Deleg
 		openAPIConfig:              openAPIConfig,
 		egressSelector:             c.GenericConfig.EgressSelector,
 		proxyCurrentCertKeyContent: func() (bytes []byte, bytes2 []byte) { return nil, nil },
+		rejectForwardingRedirects:  c.ExtraConfig.RejectForwardingRedirects,
 	}
 
 	apiGroupInfo := apiservicerest.NewRESTStorage(c.GenericConfig.MergedResourceConfig, c.GenericConfig.RESTOptionsGetter)
@@ -384,6 +390,7 @@ func (s *APIAggregator) AddAPIService(apiService *v1.APIService) error {
 		proxyTransport:             s.proxyTransport,
 		serviceResolver:            s.serviceResolver,
 		egressSelector:             s.egressSelector,
+		rejectForwardingRedirects:  s.rejectForwardingRedirects,
 	}
 	proxyHandler.updateAPIService(apiService)
 	if s.openAPIAggregationController != nil {
diff --git a/staging/src/k8s.io/kube-aggregator/pkg/apiserver/handler_proxy.go b/staging/src/k8s.io/kube-aggregator/pkg/apiserver/handler_proxy.go
index e09f0d1a..a4257420 100644
--- a/staging/src/k8s.io/kube-aggregator/pkg/apiserver/handler_proxy.go
+++ b/staging/src/k8s.io/kube-aggregator/pkg/apiserver/handler_proxy.go
@@ -68,6 +68,9 @@ type proxyHandler struct {
 	// egressSelector selects the proper egress dialer to communicate with the custom apiserver
 	// overwrites proxyTransport dialer if not nil
 	egressSelector *egressselector.EgressSelector
+
+	// reject to forward redirect response
+	rejectForwardingRedirects bool
 }
 
 type proxyHandlingInfo struct {
@@ -176,6 +179,9 @@ func (r *proxyHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	}
 
 	handler := proxy.NewUpgradeAwareHandler(location, proxyRoundTripper, true, upgrade, &responder{w: w})
+	if r.rejectForwardingRedirects {
+		handler.RejectForwardingRedirects = true
+	}
 	handler.ServeHTTP(w, newReq)
 }
 
-- 
2.25.1