From c6b3d482cca6dd7566ab5e92a6f59cafc8db2f92 Mon Sep 17 00:00:00 2001
From: William Petit <wpetit@cadoles.com>
Date: Sat, 1 Jul 2023 13:41:11 -0600
Subject: [PATCH] feat: allow modification of proxy via custom factory function

---
 options.go | 20 ++++++++++++++++++++
 proxy.go   |  8 +++++++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/options.go b/options.go
index 3dd168f..9c63d2f 100644
--- a/options.go
+++ b/options.go
@@ -1,9 +1,16 @@
 package proxy
 
+import (
+	"context"
+	"net/http/httputil"
+	"net/url"
+)
+
 type Options struct {
 	Middlewares          []Middleware
 	RequestTransformers  []RequestTransformer
 	ResponseTransformers []ResponseTransformer
+	ReverseProxyFactory  ReverseProxyFactory
 }
 
 func defaultOptions() *Options {
@@ -11,6 +18,7 @@ func defaultOptions() *Options {
 		Middlewares:          make([]Middleware, 0),
 		RequestTransformers:  make([]RequestTransformer, 0),
 		ResponseTransformers: make([]ResponseTransformer, 0),
+		ReverseProxyFactory:  DefaultReverseProxyFactory,
 	}
 }
 
@@ -33,3 +41,15 @@ func WithResponseTransformers(transformers ...ResponseTransformer) OptionFunc {
 		o.ResponseTransformers = transformers
 	}
 }
+
+type ReverseProxyFactory func(ctx context.Context, target *url.URL) *httputil.ReverseProxy
+
+func WithReverseProxyFactory(fn ReverseProxyFactory) OptionFunc {
+	return func(o *Options) {
+		o.ReverseProxyFactory = fn
+	}
+}
+
+func DefaultReverseProxyFactory(ctx context.Context, target *url.URL) *httputil.ReverseProxy {
+	return httputil.NewSingleHostReverseProxy(target)
+}
diff --git a/proxy.go b/proxy.go
index 04b518a..9875aa3 100644
--- a/proxy.go
+++ b/proxy.go
@@ -16,6 +16,7 @@ type Proxy struct {
 	handler              http.Handler
 	responseTransformers []ResponseTransformer
 	requestTransformers  []RequestTransformer
+	reverseProxyFactory  ReverseProxyFactory
 }
 
 // ServeHTTP implements http.Handler
@@ -27,6 +28,7 @@ func (p *Proxy) proxyRequest(w http.ResponseWriter, r *http.Request) {
 	var reverser *httputil.ReverseProxy
 
 	key := fmt.Sprintf("%s://%s", r.URL.Scheme, r.URL.Host)
+	ctx := r.Context()
 
 	createAndStore := func() {
 		target := &url.URL{
@@ -38,7 +40,10 @@ func (p *Proxy) proxyRequest(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 
-		reverser = httputil.NewSingleHostReverseProxy(target)
+		reverser = p.reverseProxyFactory(ctx, target)
+		if reverser == nil {
+			return
+		}
 
 		originalDirector := reverser.Director
 
@@ -97,6 +102,7 @@ func New(funcs ...OptionFunc) *Proxy {
 	proxy.handler = createMiddlewareChain(handler, opts.Middlewares)
 	proxy.requestTransformers = opts.RequestTransformers
 	proxy.responseTransformers = opts.ResponseTransformers
+	proxy.reverseProxyFactory = opts.ReverseProxyFactory
 
 	return proxy
 }