Make queue-proxy's reverse proxy handler extendable (#16097)

* Make queue-proxy's reverse proxy handler extendable

This builds on #13133 to make it possible to adjust settings on the
*httputil.ReverseProxy that queue-proxy uses, or to replace it entirely
with any http.Handler, using the out-of-tree extension pattern.

* Ensure configured Transport is used

If an integrator customizes the Transport in an Options function, we want to
apply the provided Transport to our default httputil.ReverseProxy
regardless of which Option sets the Transport and regardless of whether
or not other Option functions replace our default http.Handler.

Add some tests for this since the interactions beteween
Options.Transport and Options.ProxyHandler can be subtle.

Author: mbaynton

pkg/queue/sharedmain/handlers.go

@@ -18,7 +18,6 @@ package sharedmain
 
 import (
 	"context"
-	"net"
 	"net/http"
 	"time"
 
@@ -28,34 +27,24 @@ import (
 	"go.uber.org/zap"
 
 	netheader "knative.dev/networking/pkg/http/header"
-	netproxy "knative.dev/networking/pkg/http/proxy"
 	netstats "knative.dev/networking/pkg/http/stats"
 	pkghandler "knative.dev/pkg/network/handlers"
-	"knative.dev/serving/pkg/activator"
-	pkghttp "knative.dev/serving/pkg/http"
 	"knative.dev/serving/pkg/http/handler"
 	"knative.dev/serving/pkg/queue"
 	"knative.dev/serving/pkg/queue/health"
 )
 
 func mainHandler(
 	env config,
-	transport http.RoundTripper,
+	d Defaults,
 	prober func() bool,
 	stats *netstats.RequestStats,
 	logger *zap.SugaredLogger,
 	mp metric.MeterProvider,
 	tp trace.TracerProvider,
 ) (http.Handler, *pkghandler.Drainer) {
-	target := net.JoinHostPort("127.0.0.1", env.UserPort)
 	tracer := tp.Tracer("knative.dev/serving/pkg/queue")
 
-	httpProxy := pkghttp.NewHeaderPruningReverseProxy(target, pkghttp.NoHostOverride, activator.RevisionHeaders, false /* use HTTP */)
-	httpProxy.Transport = transport
-	httpProxy.ErrorHandler = pkghandler.Error(logger)
-	httpProxy.BufferPool = netproxy.NewBufferPool()
-	httpProxy.FlushInterval = netproxy.FlushInterval
-
 	breaker := buildBreaker(logger, env)
 
 	timeout := time.Duration(env.RevisionTimeoutSeconds) * time.Second
@@ -69,7 +58,7 @@ func mainHandler(
 	}
 	// Create queue handler chain.
 	// Note: innermost handlers are specified first, ie. the last handler in the chain will be executed first.
-	var composedHandler http.Handler = httpProxy
+	composedHandler := d.ProxyHandler
 
 	composedHandler = requestAppMetricsHandler(logger, composedHandler, breaker, mp)
 	composedHandler = queue.ProxyHandler(tracer, breaker, stats, composedHandler)

pkg/queue/sharedmain/main.go

@@ -22,11 +22,17 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"net"
 	"net/http"
+	"net/http/httputil"
 	"os"
 	"strconv"
 	"time"
 
+	netproxy "knative.dev/networking/pkg/http/proxy"
+	pkghandler "knative.dev/pkg/network/handlers"
+	"knative.dev/serving/pkg/activator"
+
 	"github.com/kelseyhightower/envconfig"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
 	"go.opentelemetry.io/otel/metric"
@@ -145,6 +151,12 @@ type Defaults struct {
 	// If Transport is wrapped, the new RoundTripper should replace the value of Transport.
 	// The new Transport will then be used by other Options (called next) and by QP.
 	Transport http.RoundTripper
+
+	// ProxyHandler provides Options with the QP's main ReverseProxy
+	// The default value is of type *httputil.ReverseProxy
+	// An Option may type assert to *httutil.ReverseProxy and adjust attributes of the ReverseProxy,
+	// or replace the handler entirely, typically to insert additional handler(s) into the chain.
+	ProxyHandler http.Handler
 }
 
 type Option func(*Defaults)
@@ -204,11 +216,9 @@ func Main(opts ...Option) error {
 	}()
 
 	d.Transport = buildTransport(env, tp, mp)
+	proxyHandler := buildProxyHandler(logger, env, d.Transport)
 
-	// allow extensions to read d and return modified context and transport
-	for _, opts := range opts {
-		opts(&d)
-	}
+	applyOptions(&d, proxyHandler, opts...)
 
 	protoStatReporter := queue.NewProtobufStatsReporter(env.ServingPod, reportingPeriod)
 
@@ -232,7 +242,7 @@ func Main(opts ...Option) error {
 	// Enable TLS when certificate is mounted.
 	tlsEnabled := exists(logger, certPath) && exists(logger, keyPath)
 
-	mainHandler, drainer := mainHandler(env, d.Transport, probe, stats, logger, mp, tp)
+	mainHandler, drainer := mainHandler(env, d, probe, stats, logger, mp, tp)
 	adminHandler := adminHandler(d.Ctx, logger, drainer)
 
 	// Enable TLS server when activator server certs are mounted.
@@ -325,6 +335,17 @@ func Main(opts ...Option) error {
 	return nil
 }
 
+func applyOptions(d *Defaults, proxyHandler *httputil.ReverseProxy, opts ...Option) {
+	d.ProxyHandler = proxyHandler
+
+	// allow extensions to read d and return modified context, transport, and proxy handler.
+	for _, opts := range opts {
+		opts(d)
+	}
+
+	proxyHandler.Transport = d.Transport
+}
+
 func exists(logger *zap.SugaredLogger, filename string) bool {
 	_, err := os.Stat(filename)
 	if err != nil && !os.IsNotExist(err) {
@@ -359,6 +380,17 @@ func buildTransport(env config, tp trace.TracerProvider, mp metric.MeterProvider
 	)
 }
 
+func buildProxyHandler(logger *zap.SugaredLogger, env config, transport http.RoundTripper) *httputil.ReverseProxy {
+	target := net.JoinHostPort("127.0.0.1", env.UserPort)
+	httpProxy := pkghttp.NewHeaderPruningReverseProxy(target, pkghttp.NoHostOverride, activator.RevisionHeaders, false /* use HTTP */)
+	httpProxy.Transport = transport
+	httpProxy.ErrorHandler = pkghandler.Error(logger)
+	httpProxy.BufferPool = netproxy.NewBufferPool()
+	httpProxy.FlushInterval = netproxy.FlushInterval
+
+	return httpProxy
+}
+
 func buildBreaker(logger *zap.SugaredLogger, env config) *queue.Breaker {
 	if env.ContainerConcurrency < 1 {
 		return nil

pkg/queue/sharedmain/main_test.go

@@ -221,3 +221,61 @@ func getFieldValue(cfg *config, fieldName string) reflect.Value {
 	f := reflect.Indirect(rVal).FieldByName(fieldName)
 	return f
 }
+
+func TestApplyOptions(t *testing.T) {
+	t.Run("option sets Transport on proxyHandler", func(t *testing.T) {
+		d := &Defaults{}
+		proxyHandler := &httputil.ReverseProxy{}
+		customTransport := &http.Transport{}
+
+		opt := func(d *Defaults) {
+			d.Transport = customTransport
+		}
+
+		applyOptions(d, proxyHandler, opt)
+
+		if proxyHandler.Transport != customTransport {
+			t.Errorf("proxyHandler.Transport = %v, want %v", proxyHandler.Transport, customTransport)
+		}
+	})
+
+	t.Run("option receives non-nil ProxyHandler", func(t *testing.T) {
+		d := &Defaults{}
+		proxyHandler := &httputil.ReverseProxy{}
+		var receivedHandler http.Handler
+
+		opt := func(d *Defaults) {
+			receivedHandler = d.ProxyHandler
+		}
+
+		applyOptions(d, proxyHandler, opt)
+
+		if receivedHandler == nil {
+			t.Error("option function received nil ProxyHandler, want non-nil")
+		}
+		if receivedHandler != proxyHandler {
+			t.Errorf("option function received %v, want %v", receivedHandler, proxyHandler)
+		}
+	})
+
+	t.Run("multiple options with handler and transport changes", func(t *testing.T) {
+		d := &Defaults{}
+		proxyHandler := &httputil.ReverseProxy{}
+		customTransport := &http.Transport{}
+		wrappedHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
+
+		opt1 := func(d *Defaults) {
+			d.ProxyHandler = wrappedHandler
+		}
+
+		opt2 := func(d *Defaults) {
+			d.Transport = customTransport
+		}
+
+		applyOptions(d, proxyHandler, opt1, opt2)
+
+		if proxyHandler.Transport != customTransport {
+			t.Errorf("proxyHandler.Transport = %v, want %v", proxyHandler.Transport, customTransport)
+		}
+	})
+}