Add support for propagating headers to downstream servergroups

Fixes #322

cmd/promxy/main.go

@@ -7,6 +7,7 @@ import (
 	"path"
 	"regexp"
 
+	"github.com/jacksontj/promxy/pkg/middleware"
 	"github.com/jacksontj/promxy/pkg/server"
 
 	"go.uber.org/atomic"
@@ -85,7 +86,8 @@ type cliOpts struct {
 	WebCORSOriginRegex string        `long:"web.cors.origin" description:"Regex for CORS origin. It is fully anchored." default:".*"`
 	WebReadTimeout     time.Duration `long:"web.read-timeout" description:"Maximum duration before timing out read of the request, and closing idle connections." default:"5m"`
 
-	MetricsPath string `long:"metrics-path" description:"URL path for the prometheus metrics endpoint." default:"/metrics"`
+	MetricsPath  string   `long:"metrics-path" description:"URL path for the prometheus metrics endpoint." default:"/metrics"`
+	ProxyHeaders []string `long:"proxy-headers" env:"PROXY_HEADERS" description:"a list of headers to proxy to downstream servergroups."`
 
 	ExternalURL     string `long:"web.external-url" description:"The URL under which Prometheus is externally reachable (for example, if Prometheus is served via a reverse proxy). Used for generating relative and absolute links back to Prometheus itself. If the URL has a path portion, it will be used to prefix all HTTP endpoints served by Prometheus. If omitted, relevant URL components will be derived automatically."`
 	EnableLifecycle bool   `long:"web.enable-lifecycle" description:"Enable shutdown and reload via HTTP request."`
@@ -433,7 +435,7 @@ func main() {
 		logrus.Fatalf("Invalid AccessLogDestination: %s", opts.AccessLogDestination)
 	}
 
-	srv, err := server.CreateAndStart(opts.BindAddr, opts.LogFormat, opts.WebReadTimeout, accessLogOut, r, opts.WebConfigFile)
+	srv, err := server.CreateAndStart(opts.BindAddr, opts.LogFormat, opts.WebReadTimeout, accessLogOut, middleware.NewProxyHeaders(r, opts.ProxyHeaders), opts.WebConfigFile)
 	if err != nil {
 		logrus.Fatalf("Error creating server: %v", err)
 	}

pkg/middleware/headers.go

@@ -0,0 +1,42 @@
+package middleware
+
+import (
+	"context"
+	"net/http"
+)
+
+type contextKey int
+
+const headerKey contextKey = 0
+
+func NewProxyHeaders(h http.Handler, headers []string) *ProxyHeaders {
+	return &ProxyHeaders{
+		h:       h,
+		headers: headers,
+	}
+}
+
+type ProxyHeaders struct {
+	h       http.Handler
+	headers []string
+}
+
+func (p *ProxyHeaders) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	hdrs := make(map[string]string, len(p.headers))
+	for _, header := range p.headers {
+		if v := r.Header.Get(header); v != "" {
+			hdrs[header] = v
+		}
+	}
+
+	p.h.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), headerKey, hdrs)))
+}
+
+func GetHeaders(ctx context.Context) map[string]string {
+	v := ctx.Value(headerKey)
+	if v == nil {
+		return nil
+	}
+
+	return v.(map[string]string)
+}

pkg/server/api.go

@@ -7,15 +7,14 @@ import (
 	"net/http"
 	"time"
 
-	"github.com/julienschmidt/httprouter"
 	"github.com/prometheus/exporter-toolkit/web"
 	"github.com/sirupsen/logrus"
 	"gopkg.in/yaml.v2"
 
 	"github.com/jacksontj/promxy/pkg/logging"
 )
 
-func CreateAndStart(bindAddr string, logFormat string, webReadTimeout time.Duration, accessLogOut io.Writer, router *httprouter.Router, tlsConfigFile string) (*http.Server, error) {
+func CreateAndStart(bindAddr string, logFormat string, webReadTimeout time.Duration, accessLogOut io.Writer, router http.Handler, tlsConfigFile string) (*http.Server, error) {
 	handler := createHandler(accessLogOut, router, logFormat)
 
 	srv := &http.Server{
@@ -31,7 +30,7 @@ func CreateAndStart(bindAddr string, logFormat string, webReadTimeout time.Durat
 	return createAndStartHTTPS(srv, tlsConfigFile)
 }
 
-func createHandler(accessLogOut io.Writer, router *httprouter.Router, logFormat string) http.Handler {
+func createHandler(accessLogOut io.Writer, router http.Handler, logFormat string) http.Handler {
 	var handler http.Handler
 	if accessLogOut == nil {
 		handler = router

pkg/servergroup/servergroup.go

@@ -23,6 +23,7 @@ import (
 	"github.com/prometheus/prometheus/storage/remote"
 	"github.com/sirupsen/logrus"
 
+	"github.com/jacksontj/promxy/pkg/middleware"
 	"github.com/jacksontj/promxy/pkg/promclient"
 	//	sd_config "github.com/prometheus/prometheus/discovery/config"
 )
@@ -95,6 +96,14 @@ func (s *ServerGroup) Cancel() {
 	s.ctxCancel()
 }
 
+// RoundTrip allows us to intercept and mutate downstream HTTP requests at the transport level
+func (s *ServerGroup) RoundTrip(r *http.Request) (*http.Response, error) {
+	for k, v := range middleware.GetHeaders(r.Context()) {
+		r.Header.Set(k, v)
+	}
+	return s.client.Transport.RoundTrip(r)
+}
+
 // Sync updates the targets from our discovery manager
 func (s *ServerGroup) Sync() {
 	syncCh := s.targetManager.SyncCh()
@@ -148,7 +157,7 @@ SYNC_LOOP:
 
 					targets = append(targets, u.Host)
 
-					client, err := api.NewClient(api.Config{Address: u.String(), RoundTripper: s.client.Transport})
+					client, err := api.NewClient(api.Config{Address: u.String(), RoundTripper: s})
 					if err != nil {
 						panic(err) // TODO: shouldn't be possible? If this happens I guess we log and skip?
 					}