feat: Provide label value by http header (#118)

* feat: Provide label value by http header

Signed-off-by: Jan-Otto Kröpke <[email protected]>

Signed-off-by: Simon Pasquier <[email protected]>

* Define name of query-param and header-name by CLI parameter

Signed-off-by: Jan-Otto Kröpke <[email protected]>

Signed-off-by: Simon Pasquier <[email protected]>

* Add LabelExtracter interface

The interface abstracts how the label value is extracted from the
incoming request. Currently 3 implementations are provided:
1. Get the value from the HTTP form/query parameters.
2. Get the value from the HTTP headers.
3. Get a static value.

Signed-off-by: Simon Pasquier <[email protected]>

* Apply suggestions from code review

Co-authored-by: Simon Pasquier <[email protected]>
Signed-off-by: Jan-Otto Kröpke <[email protected]>

* remove extra spaces

Signed-off-by: Simon Pasquier <[email protected]>

Signed-off-by: Simon Pasquier <[email protected]>
Signed-off-by: Jan-Otto Kröpke <[email protected]>
Co-authored-by: Simon Pasquier <[email protected]>

Author: jkroepke

README.md

@@ -61,36 +61,61 @@ When started with the `-enable-label-apis` flag, the application can also proxy
 * `/api/v1/labels` for GET and POST methods (Prometheus/Thanos)
 * `/api/v1/label/<name>/values` for GET method (Prometheus/Thanos)
 
-Particularly, you can run `prom-label-proxy` with label `tenant` and point to example, demo Prometheus server e.g:
+You can run `prom-label-proxy` to enforce the value of the `tenant` label
+provided in the client's request via the `tenant` HTTP query/form parameter:
 
 ```
 prom-label-proxy \
+   -query-param tenant \
    -label tenant \
    -upstream http://demo.do.prometheus.io:9090 \
    -insecure-listen-address 127.0.0.1:8080
 ```
 
-Accessing demo Prometheus APIs on `127.0.0.1:8080` will now expect `tenant` query parameter to be set in the URL:
+Accessing the demo Prometheus APIs on `http://127.0.0.1:8080` will now expect
+that the client's request provides the `tenant` label value using the `tenant`
+HTTP query parameter:
 
 ```bash
 ➜  ~ curl http://127.0.0.1:8080/api/v1/query\?query="up"
-The "tenant" query parameter must be provided.
+{"error":"The \"tenant\" query parameter must be provided.","errorType":"prom-label-proxy","status":"error"}
 ➜  ~ curl http://127.0.0.1:8080/api/v1/query\?query="up"\&tenant\="something"
 {"status":"success","data":{"resultType":"vector","result":[]}}%
 ```
 
-You can also provide a static value for a label. For example, running `prom-label-proxy` with
+It also works with POST requests:
+
+```bash
+➜  ~ curl http://127.0.0.1:8080/api/v1/query" -d "tenant=foo"
+{"status":"success","data":{"resultType":"vector","result":[]}}%
+```
+
+Alternatively, `prom-label-proxy` can use a custom HTTP header instead HTTP parameters:
+
 ```
 prom-label-proxy \
+   -header-name X-Tenant \
    -label tenant \
-   -value prometheus \
    -upstream http://demo.do.prometheus.io:9090 \
    -insecure-listen-address 127.0.0.1:8080
 ```
-will enforce `tenant=prometheus` in all requests.
 
+```bash
+➜  ~ curl -H 'X-Tenant=something' http://127.0.0.1:8080/api/v1/query\?query="up"
+{"status":"success","data":{"resultType":"vector","result":[]}}%
+```
+
+A last option is to provide a static value for the label:
+
+```
+prom-label-proxy \
+   -label tenant \
+   -label-value prometheus \
+   -upstream http://demo.do.prometheus.io:9090 \
+   -insecure-listen-address 127.0.0.1:8080
+```
 
-In this mode, sending the label value as a query parameter will result in the request getting rejected as a 400 Bad Request.
+Now prom-label-proxy enforces the `tenant="prometheus"` label in all requests.
 
 Once again for clarity: **this project only enforces a particular label in the respective calls to Prometheus, it in itself does not authenticate or
 authorize the requesting entity in any way, this has to be built around this project.**

injectproxy/alerts_test.go

@@ -65,7 +65,8 @@ func TestGetAlerts(t *testing.T) {
 		t.Run(strings.Join(tc.filters, "&"), func(t *testing.T) {
 			m := newMockUpstream(checkQueryHandler("", tc.queryParam, tc.expQueryValues...))
 			defer m.Close()
-			r, err := NewRoutes(m.url, proxyLabel)
+
+			r, err := NewRoutes(m.url, proxyLabel, HTTPFormEnforcer{ParameterName: proxyLabel})
 			if err != nil {
 				t.Fatalf("unexpected error: %v", err)
 			}

injectproxy/routes.go

@@ -37,18 +37,17 @@ const (
 )
 
 type routes struct {
-	upstream   *url.URL
-	handler    http.Handler
-	label      string
-	labelValue string
+	upstream *url.URL
+	handler  http.Handler
+	label    string
+	el       ExtractLabeler
 
 	mux            http.Handler
 	modifiers      map[string]func(*http.Response) error
 	errorOnReplace bool
 }
 
 type options struct {
-	labelValue       string
 	enableLabelAPIs  bool
 	passthroughPaths []string
 	errorOnReplace   bool
@@ -96,14 +95,6 @@ func WithErrorOnReplace() Option {
 	})
 }
 
-// WithLabelValue enforces a specific value for the multi-tenancy label.
-// If not specified, the value has to be provided as a URL parameter.
-func WithLabelValue(value string) Option {
-	return optionFunc(func(o *options) {
-		o.labelValue = value
-	})
-}
-
 // mux abstracts away the behavior we expect from the http.ServeMux type in this package.
 type mux interface {
 	http.Handler
@@ -170,11 +161,110 @@ func (i *instrumentedMux) Handle(pattern string, handler http.Handler) {
 	i.mux.Handle(pattern, i.i.NewHandler(prometheus.Labels{"handler": pattern}, handler))
 }
 
-func NewRoutes(upstream *url.URL, label string, opts ...Option) (*routes, error) {
+// ExtractLabeler is an HTTP handler that extract the label value to be
+// enforced from the HTTP request.  If a valid label value is found, it should
+// store it in the request's context.  Otherwise it should return an error in
+// the HTTP response (usually 400 or 500).
+type ExtractLabeler interface {
+	ExtractLabel(next http.HandlerFunc) http.Handler
+}
+
+// HTTPFormEnforcer enforces a label value extracted from the HTTP form and query parameters.
+type HTTPFormEnforcer struct {
+	ParameterName string
+}
+
+// ExtractLabel implements the ExtractLabeler interface.
+func (hff HTTPFormEnforcer) ExtractLabel(next http.HandlerFunc) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		labelValue, err := hff.getLabelValue(r)
+		if err != nil {
+			prometheusAPIError(w, humanFriendlyErrorMessage(err), http.StatusBadRequest)
+			return
+		}
+
+		// Remove the proxy label from the query parameters.
+		q := r.URL.Query()
+		q.Del(hff.ParameterName)
+		r.URL.RawQuery = q.Encode()
+
+		// Remove the param from the PostForm.
+		if r.Method == http.MethodPost {
+			if err := r.ParseForm(); err != nil {
+				prometheusAPIError(w, fmt.Sprintf("Failed to parse the PostForm: %v", err), http.StatusInternalServerError)
+				return
+			}
+			if r.PostForm.Get(hff.ParameterName) != "" {
+				r.PostForm.Del(hff.ParameterName)
+				newBody := r.PostForm.Encode()
+				// We are replacing request body, close previous one (r.FormValue ensures it is read fully and not nil).
+				_ = r.Body.Close()
+				r.Body = io.NopCloser(strings.NewReader(newBody))
+				r.ContentLength = int64(len(newBody))
+			}
+		}
+
+		next.ServeHTTP(w, r.WithContext(WithLabelValue(r.Context(), labelValue)))
+	})
+}
+
+func (hff HTTPFormEnforcer) getLabelValue(r *http.Request) (string, error) {
+	formValue := r.FormValue(hff.ParameterName)
+	if formValue == "" {
+		return "", fmt.Errorf("the %q query parameter must be provided", hff.ParameterName)
+	}
+
+	return formValue, nil
+}
+
+// HTTPHeaderEnforcer enforces a label value extracted from the HTTP headers.
+type HTTPHeaderEnforcer struct {
+	Name string
+}
+
+// ExtractLabel implements the ExtractLabeler interface.
+func (hhe HTTPHeaderEnforcer) ExtractLabel(next http.HandlerFunc) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		labelValue, err := hhe.getLabelValue(r)
+		if err != nil {
+			prometheusAPIError(w, humanFriendlyErrorMessage(err), http.StatusBadRequest)
+			return
+		}
+
+		next.ServeHTTP(w, r.WithContext(WithLabelValue(r.Context(), labelValue)))
+	})
+}
+
+func (hhe HTTPHeaderEnforcer) getLabelValue(r *http.Request) (string, error) {
+	headerValues := r.Header[hhe.Name]
+
+	if len(headerValues) == 0 {
+		return "", fmt.Errorf("missing HTTP header %q", hhe.Name)
+	}
+
+	if len(headerValues) > 1 {
+		return "", fmt.Errorf("multiple values for the http header %q", hhe.Name)
+	}
+
+	return headerValues[0], nil
+}
+
+// StaticLabelEnforcer enforces a static label value.
+type StaticLabelEnforcer string
+
+// ExtractLabel implements the ExtractLabeler interface.
+func (sle StaticLabelEnforcer) ExtractLabel(next http.HandlerFunc) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		next(w, r.WithContext(WithLabelValue(r.Context(), string(sle))))
+	})
+}
+
+func NewRoutes(upstream *url.URL, label string, extractLabeler ExtractLabeler, opts ...Option) (*routes, error) {
 	opt := options{}
 	for _, o := range opts {
 		o.apply(&opt)
 	}
+
 	if opt.registerer == nil {
 		opt.registerer = prometheus.NewRegistry()
 	}
@@ -185,35 +275,35 @@ func NewRoutes(upstream *url.URL, label string, opts ...Option) (*routes, error)
 		upstream:       upstream,
 		handler:        proxy,
 		label:          label,
-		labelValue:     opt.labelValue,
+		el:             extractLabeler,
 		errorOnReplace: opt.errorOnReplace,
 	}
 	mux := newStrictMux(newInstrumentedMux(http.NewServeMux(), opt.registerer))
 
 	errs := merrors.New(
-		mux.Handle("/federate", r.enforceLabel(enforceMethods(r.matcher, "GET"))),
-		mux.Handle("/api/v1/query", r.enforceLabel(enforceMethods(r.query, "GET", "POST"))),
-		mux.Handle("/api/v1/query_range", r.enforceLabel(enforceMethods(r.query, "GET", "POST"))),
-		mux.Handle("/api/v1/alerts", r.enforceLabel(enforceMethods(r.passthrough, "GET"))),
-		mux.Handle("/api/v1/rules", r.enforceLabel(enforceMethods(r.passthrough, "GET"))),
-		mux.Handle("/api/v1/series", r.enforceLabel(enforceMethods(r.matcher, "GET", "POST"))),
-		mux.Handle("/api/v1/query_exemplars", r.enforceLabel(enforceMethods(r.query, "GET", "POST"))),
+		mux.Handle("/federate", r.el.ExtractLabel(enforceMethods(r.matcher, "GET"))),
+		mux.Handle("/api/v1/query", r.el.ExtractLabel(enforceMethods(r.query, "GET", "POST"))),
+		mux.Handle("/api/v1/query_range", r.el.ExtractLabel(enforceMethods(r.query, "GET", "POST"))),
+		mux.Handle("/api/v1/alerts", r.el.ExtractLabel(enforceMethods(r.passthrough, "GET"))),
+		mux.Handle("/api/v1/rules", r.el.ExtractLabel(enforceMethods(r.passthrough, "GET"))),
+		mux.Handle("/api/v1/series", r.el.ExtractLabel(enforceMethods(r.matcher, "GET", "POST"))),
+		mux.Handle("/api/v1/query_exemplars", r.el.ExtractLabel(enforceMethods(r.query, "GET", "POST"))),
 	)
 
 	if opt.enableLabelAPIs {
 		errs.Add(
-			mux.Handle("/api/v1/labels", r.enforceLabel(enforceMethods(r.matcher, "GET", "POST"))),
+			mux.Handle("/api/v1/labels", r.el.ExtractLabel(enforceMethods(r.matcher, "GET", "POST"))),
 			// Full path is /api/v1/label/<label_name>/values but http mux does not support patterns.
 			// This is fine though as we don't care about name for matcher injector.
-			mux.Handle("/api/v1/label/", r.enforceLabel(enforceMethods(r.matcher, "GET"))),
+			mux.Handle("/api/v1/label/", r.el.ExtractLabel(enforceMethods(r.matcher, "GET"))),
 		)
 	}
 
 	errs.Add(
-		mux.Handle("/api/v2/silences", r.enforceLabel(enforceMethods(r.silences, "GET", "POST"))),
-		mux.Handle("/api/v2/silence/", r.enforceLabel(enforceMethods(r.deleteSilence, "DELETE"))),
-		mux.Handle("/api/v2/alerts/groups", r.enforceLabel(enforceMethods(r.enforceFilterParameter, "GET"))),
-		mux.Handle("/api/v2/alerts", r.enforceLabel(enforceMethods(r.alerts, "GET"))),
+		mux.Handle("/api/v2/silences", r.el.ExtractLabel(enforceMethods(r.silences, "GET", "POST"))),
+		mux.Handle("/api/v2/silence/", r.el.ExtractLabel(enforceMethods(r.deleteSilence, "DELETE"))),
+		mux.Handle("/api/v2/alerts/groups", r.el.ExtractLabel(enforceMethods(r.enforceFilterParameter, "GET"))),
+		mux.Handle("/api/v2/alerts", r.el.ExtractLabel(enforceMethods(r.alerts, "GET"))),
 	)
 
 	errs.Add(
@@ -256,62 +346,6 @@ func NewRoutes(upstream *url.URL, label string, opts ...Option) (*routes, error)
 	return r, nil
 }
 
-func (r *routes) enforceLabel(h http.HandlerFunc) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
-		lvalue, err := r.getLabelValue(req)
-		if err != nil {
-			prometheusAPIError(w, humanFriendlyErrorMessage(err), http.StatusBadRequest)
-			return
-		}
-
-		req = req.WithContext(withLabelValue(req.Context(), lvalue))
-
-		// Remove the proxy label from the query parameters.
-		q := req.URL.Query()
-		if q.Get(r.label) != "" {
-			q.Del(r.label)
-		}
-		req.URL.RawQuery = q.Encode()
-		// Remove the proxy label from the PostForm.
-		if req.Method == http.MethodPost {
-			if err := req.ParseForm(); err != nil {
-				prometheusAPIError(w, fmt.Sprintf("Failed to parse the PostForm: %v", err), http.StatusInternalServerError)
-				return
-			}
-			if req.PostForm.Get(r.label) != "" {
-				req.PostForm.Del(r.label)
-				newBody := req.PostForm.Encode()
-				// We are replacing request body, close previous one (req.FormValue ensures it is read fully and not nil).
-				_ = req.Body.Close()
-				req.Body = io.NopCloser(strings.NewReader(newBody))
-				req.ContentLength = int64(len(newBody))
-			}
-		}
-
-		h.ServeHTTP(w, req)
-	})
-}
-
-// getLabelValue returns the statically set label value, or the label value
-// sent through a URL parameter.
-// It returns an error when either the value is found in both places, or is not found at all.
-func (r *routes) getLabelValue(req *http.Request) (string, error) {
-	formValue := req.FormValue(r.label)
-	if r.labelValue != "" && formValue != "" {
-		return "", fmt.Errorf("a static value for the %s label has already been specified", r.label)
-	}
-
-	if r.labelValue == "" && formValue == "" {
-		return "", fmt.Errorf("the %q query parameter must be provided", r.label)
-	}
-
-	if r.labelValue != "" {
-		return r.labelValue, nil
-	}
-
-	return formValue, nil
-}
-
 func (r *routes) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	r.mux.ServeHTTP(w, req)
 }
@@ -341,7 +375,10 @@ type ctxKey int
 
 const keyLabel ctxKey = iota
 
-func mustLabelValue(ctx context.Context) string {
+// MustLabelValue returns a label (previously stored using WithLabelValue())
+// from the given context.
+// It will panic if no label is found or the value is empty.
+func MustLabelValue(ctx context.Context) string {
 	label, ok := ctx.Value(keyLabel).(string)
 	if !ok {
 		panic(fmt.Sprintf("can't find the %q value in the context", keyLabel))
@@ -352,7 +389,8 @@ func mustLabelValue(ctx context.Context) string {
 	return label
 }
 
-func withLabelValue(ctx context.Context, label string) context.Context {
+// WithLabelValue stores a label in the given context.
+func WithLabelValue(ctx context.Context, label string) context.Context {
 	return context.WithValue(ctx, keyLabel, label)
 }
 
@@ -365,7 +403,7 @@ func (r *routes) query(w http.ResponseWriter, req *http.Request) {
 		[]*labels.Matcher{{
 			Name:  r.label,
 			Type:  labels.MatchEqual,
-			Value: mustLabelValue(req.Context()),
+			Value: MustLabelValue(req.Context()),
 		}}...)
 
 	// The `query` can come in the URL query string and/or the POST body.
@@ -451,7 +489,7 @@ func (r *routes) matcher(w http.ResponseWriter, req *http.Request) {
 	matcher := &labels.Matcher{
 		Name:  r.label,
 		Type:  labels.MatchEqual,
-		Value: mustLabelValue(req.Context()),
+		Value: MustLabelValue(req.Context()),
 	}
 	q := req.URL.Query()