From 2163b321e087b51168f0f71e65d0b5e5344c5986 Mon Sep 17 00:00:00 2001
From: Graeme Christie <gchristie@bunnings.com.au>
Date: Fri, 14 Jul 2023 22:49:22 +0800
Subject: [PATCH] Refactored to use options pattern for new command line args

---
 injectproxy/routes.go | 43 +++++++++++++++++++++++++++++--------------
 main.go               | 10 +++++++++-
 2 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/injectproxy/routes.go b/injectproxy/routes.go
index 76c404c4..d6c69f77 100644
--- a/injectproxy/routes.go
+++ b/injectproxy/routes.go
@@ -18,7 +18,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
-	"log"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
@@ -50,10 +49,12 @@ type routes struct {
 }
 
 type options struct {
-	enableLabelAPIs  bool
-	passthroughPaths []string
-	errorOnReplace   bool
-	registerer       prometheus.Registerer
+	enableLabelAPIs   bool
+	passthroughPaths  []string
+	errorOnReplace    bool
+	registerer        prometheus.Registerer
+	extraHttpHeaders  map[string]string
+	rewriteHostHeader string
 }
 
 type Option interface {
@@ -97,6 +98,24 @@ func WithErrorOnReplace() Option {
 	})
 }
 
+func WithExtraHttpHeaders(headers []string) Option {
+	return optionFunc(func(o *options) {
+		o.extraHttpHeaders = make(map[string]string)
+		for _, headerArg := range headers {
+			header, val, found := strings.Cut(headerArg, ":")
+			if found {
+				o.extraHttpHeaders[strings.TrimSpace(header)] = strings.TrimSpace(val)
+			}
+		}
+	})
+}
+
+func WithRewriteHostHeader(host string) Option {
+	return optionFunc(func(o *options) {
+		o.rewriteHostHeader = host
+	})
+}
+
 // mux abstracts away the behavior we expect from the http.ServeMux type in this package.
 type mux interface {
 	http.Handler
@@ -262,7 +281,8 @@ func (sle StaticLabelEnforcer) ExtractLabel(next http.HandlerFunc) http.Handler
 	})
 }
 
-func NewRoutes(upstream *url.URL, label string, extractLabeler ExtractLabeler, extraHttpHeaders []string, rewriteHostHeader string, opts ...Option) (*routes, error) {
+func NewRoutes(upstream *url.URL, label string, extractLabeler ExtractLabeler, opts ...Option) (*routes, error) {
+	//extraHttpHeaders []string, rewriteHostHeader string
 	opt := options{}
 	for _, o := range opts {
 		o.apply(&opt)
@@ -275,17 +295,12 @@ func NewRoutes(upstream *url.URL, label string, extractLabeler ExtractLabeler, e
 	proxy := &httputil.ReverseProxy{
 		Rewrite: func(r *httputil.ProxyRequest) {
 			r.SetURL(upstream)
-			if len(strings.TrimSpace(rewriteHostHeader)) == 0 {
+			if len(opt.rewriteHostHeader) == 0 {
 				r.Out.Host = r.In.Host
 			} else {
-				r.Out.Host = strings.TrimSpace(rewriteHostHeader)
+				r.Out.Host = opt.rewriteHostHeader
 			}
-			for _, headerArg := range extraHttpHeaders {
-				header, val, found := strings.Cut(headerArg, ":")
-				if !found {
-					log.Printf("Header %s specified but ':' delimited not found", headerArg)
-					continue
-				}
+			for header, val := range opt.extraHttpHeaders {
 				r.Out.Header[strings.TrimSpace(header)] = []string{strings.TrimSpace(val)}
 			}
 		},
diff --git a/main.go b/main.go
index ad4533ce..c4534c5b 100644
--- a/main.go
+++ b/main.go
@@ -137,6 +137,14 @@ func main() {
 		opts = append(opts, injectproxy.WithErrorOnReplace())
 	}
 
+	if len(extraHttpHeaders) > 0 {
+		opts = append(opts, injectproxy.WithExtraHttpHeaders(extraHttpHeaders))
+	}
+
+	if len(rewriteHostHeader) > 0 {
+		opts = append(opts, injectproxy.WithRewriteHostHeader(rewriteHostHeader))
+	}
+
 	var extractLabeler injectproxy.ExtractLabeler
 	switch {
 	case len(labelValues) > 0:
@@ -151,7 +159,7 @@ func main() {
 
 	{
 		// Run the insecure HTTP server.
-		routes, err := injectproxy.NewRoutes(upstreamURL, label, extractLabeler, extraHttpHeaders, rewriteHostHeader, opts...)
+		routes, err := injectproxy.NewRoutes(upstreamURL, label, extractLabeler, opts...)
 		if err != nil {
 			log.Fatalf("Failed to create injectproxy Routes: %v", err)
 		}