Refactored to use options pattern for new command line args

prometheus-community · Jul 14, 2023 · 2163b32 · 2163b32
1 parent 88275aa
commit 2163b32
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 15 deletions.
diff --git 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
@@ -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)
 		}