app/vmauth: add src_headers option at url_map, which allows routing incoming requests to different backends depending on request headers

2024-11-23 12:31:07 +01:00 · 2024-03-06 21:56:32 +02:00 · 2024-03-06 21:56:32 +02:00 · 76ef84fcae
commit 76ef84fcae
parent d8688c9e82
7 changed files with 116 additions and 90 deletions
--- a/app/vmauth/auth_config.go
+++ b/app/vmauth/auth_config.go
@ -136,7 +136,7 @@ func (h *Header) MarshalYAML() (interface{}, error) {

 // URLMap is a mapping from source paths to target urls.
 type URLMap struct {
-	// SrcPaths is the list of regular expressions, which must match the request path.
+	// SrcPaths is an optional list of regular expressions, which must match the request path.
 	SrcPaths []*Regex `yaml:"src_paths,omitempty"`

 	// SrcHosts is an optional list of regular expressions, which must match the request hostname.
@ -145,6 +145,9 @@ type URLMap struct {
 	// SrcQueryArgs is an optional list of query args, which must match request URL query args.
 	SrcQueryArgs []QueryArg `yaml:"src_query_args,omitempty"`

+	// SrcHeaders is an optional list of headers, which must match request headers.
+	SrcHeaders []Header `yaml:"src_headers,omitempty"`
+
 	// UrlPrefix contains backend url prefixes for the proxied request url.
 	URLPrefix *URLPrefix `yaml:"url_prefix,omitempty"`

@ -170,8 +173,8 @@ type Regex struct {

 // QueryArg represents HTTP query arg
 type QueryArg struct {
-	Name  string `yaml:"name"`
-	Value string `yaml:"value,omitempty"`
+	Name  string
+	Value string

 	sOriginal string
 }
@ -711,8 +714,8 @@ func (ui *UserInfo) initURLs() error {
 		}
 	}
 	for _, e := range ui.URLMaps {
-		if len(e.SrcPaths) == 0 && len(e.SrcHosts) == 0 && len(e.SrcQueryArgs) == 0 {
-			return fmt.Errorf("missing `src_paths`, `src_hosts` and `src_query_args` in `url_map`")
+		if len(e.SrcPaths) == 0 && len(e.SrcHosts) == 0 && len(e.SrcQueryArgs) == 0 && len(e.SrcHeaders) == 0 {
+			return fmt.Errorf("missing `src_paths`, `src_hosts`, `src_query_args` and `src_headers` in `url_map`")
 		}
 		if e.URLPrefix == nil {
 			return fmt.Errorf("missing `url_prefix` in `url_map`")
--- a/app/vmauth/auth_config_test.go
+++ b/app/vmauth/auth_config_test.go
@ -219,6 +219,15 @@ users:
    url_prefix: http://foobar
 `)

+	// Invalid src_headers
+	f(`
+users:
+- username: a
+  url_map:
+  - src_headers: abc
+    url_prefix: http://foobar
+`)
+
 	// Invalid headers in url_map (missing ':')
 	f(`
 users:
@ -332,6 +341,47 @@ users:
 	})

 	// non-empty URLMap
+	sharedUserInfo := &UserInfo{
+		BearerToken: "foo",
+		URLMaps: []URLMap{
+			{
+				SrcPaths:  getRegexs([]string{"/api/v1/query", "/api/v1/query_range", "/api/v1/label/[^./]+/.+"}),
+				URLPrefix: mustParseURL("http://vmselect/select/0/prometheus"),
+			},
+			{
+				SrcHosts: getRegexs([]string{"foo\\.bar", "baz:1234"}),
+				SrcPaths: getRegexs([]string{"/api/v1/write"}),
+				SrcQueryArgs: []QueryArg{
+					{
+						Name:  "foo",
+						Value: "bar",
+					},
+				},
+				SrcHeaders: []Header{
+					{
+						Name:  "TenantID",
+						Value: "345",
+					},
+				},
+				URLPrefix: mustParseURLs([]string{
+					"http://vminsert1/insert/0/prometheus",
+					"http://vminsert2/insert/0/prometheus",
+				}),
+				HeadersConf: HeadersConf{
+					RequestHeaders: []Header{
+						{
+							Name:  "foo",
+							Value: "bar",
+						},
+						{
+							Name:  "xxx",
+							Value: "y",
+						},
+					},
+				},
+			},
+		},
+	}
 	f(`
 users:
 - bearer_token: foo
@ -340,83 +390,15 @@ users:
    url_prefix: http://vmselect/select/0/prometheus
  - src_paths: ["/api/v1/write"]
    src_hosts: ["foo\\.bar", "baz:1234"]
-    src_query_args:
-    - 'foo=bar'
+    src_query_args: ['foo=bar']
+    src_headers: ['TenantID: 345']
    url_prefix: ["http://vminsert1/insert/0/prometheus","http://vminsert2/insert/0/prometheus"]
    headers:
    - "foo: bar"
    - "xxx: y"
 `, map[string]*UserInfo{
-		getHTTPAuthBearerToken("foo"): {
-			BearerToken: "foo",
-			URLMaps: []URLMap{
-				{
-					SrcPaths:  getRegexs([]string{"/api/v1/query", "/api/v1/query_range", "/api/v1/label/[^./]+/.+"}),
-					URLPrefix: mustParseURL("http://vmselect/select/0/prometheus"),
-				},
-				{
-					SrcHosts: getRegexs([]string{"foo\\.bar", "baz:1234"}),
-					SrcPaths: getRegexs([]string{"/api/v1/write"}),
-					SrcQueryArgs: []QueryArg{
-						{
-							Name:  "foo",
-							Value: "bar",
-						},
-					},
-					URLPrefix: mustParseURLs([]string{
-						"http://vminsert1/insert/0/prometheus",
-						"http://vminsert2/insert/0/prometheus",
-					}),
-					HeadersConf: HeadersConf{
-						RequestHeaders: []Header{
-							{
-								Name:  "foo",
-								Value: "bar",
-							},
-							{
-								Name:  "xxx",
-								Value: "y",
-							},
-						},
-					},
-				},
-			},
-		},
-		getHTTPAuthBasicToken("foo", ""): {
-			BearerToken: "foo",
-			URLMaps: []URLMap{
-				{
-					SrcPaths:  getRegexs([]string{"/api/v1/query", "/api/v1/query_range", "/api/v1/label/[^./]+/.+"}),
-					URLPrefix: mustParseURL("http://vmselect/select/0/prometheus"),
-				},
-				{
-					SrcHosts: getRegexs([]string{"foo\\.bar", "baz:1234"}),
-					SrcPaths: getRegexs([]string{"/api/v1/write"}),
-					SrcQueryArgs: []QueryArg{
-						{
-							Name:  "foo",
-							Value: "bar",
-						},
-					},
-					URLPrefix: mustParseURLs([]string{
-						"http://vminsert1/insert/0/prometheus",
-						"http://vminsert2/insert/0/prometheus",
-					}),
-					HeadersConf: HeadersConf{
-						RequestHeaders: []Header{
-							{
-								Name:  "foo",
-								Value: "bar",
-							},
-							{
-								Name:  "xxx",
-								Value: "y",
-							},
-						},
-					},
-				},
-			},
-		},
+		getHTTPAuthBearerToken("foo"):    sharedUserInfo,
+		getHTTPAuthBasicToken("foo", ""): sharedUserInfo,
 	})

 	// Multiple users with the same name - this should work, since these users have different passwords
--- a/app/vmauth/main.go
+++ b/app/vmauth/main.go
@ -184,7 +184,7 @@ func processUserRequest(w http.ResponseWriter, r *http.Request, ui *UserInfo) {

 func processRequest(w http.ResponseWriter, r *http.Request, ui *UserInfo) {
 	u := normalizeURL(r.URL)
-	up, hc := ui.getURLPrefixAndHeaders(u)
+	up, hc := ui.getURLPrefixAndHeaders(u, r.Header)
 	isDefault := false
 	if up == nil {
 		if ui.DefaultURL == nil {
--- a/app/vmauth/target_url.go
+++ b/app/vmauth/target_url.go
@ -1,6 +1,7 @@
 package main

 import (
+	"net/http"
 	"net/url"
 	"path"
 	"slices"
@ -50,11 +51,22 @@ func dropPrefixParts(path string, parts int) string {
 	return path
 }

-func (ui *UserInfo) getURLPrefixAndHeaders(u *url.URL) (*URLPrefix, HeadersConf) {
+func (ui *UserInfo) getURLPrefixAndHeaders(u *url.URL, h http.Header) (*URLPrefix, HeadersConf) {
 	for _, e := range ui.URLMaps {
-		if matchAnyRegex(e.SrcHosts, u.Host) && matchAnyRegex(e.SrcPaths, u.Path) && matchAnyQueryArg(e.SrcQueryArgs, u.Query()) {
-			return e.URLPrefix, e.HeadersConf
+		if !matchAnyRegex(e.SrcHosts, u.Host) {
+			continue
 		}
+		if !matchAnyRegex(e.SrcPaths, u.Path) {
+			continue
+		}
+		if !matchAnyQueryArg(e.SrcQueryArgs, u.Query()) {
+			continue
+		}
+		if !matchAnyHeader(e.SrcHeaders, h) {
+			continue
+		}
+
+		return e.URLPrefix, e.HeadersConf
 	}
 	if ui.URLPrefix != nil {
 		return ui.URLPrefix, ui.HeadersConf
@ -86,6 +98,18 @@ func matchAnyQueryArg(qas []QueryArg, args url.Values) bool {
 	return false
 }

+func matchAnyHeader(headers []Header, h http.Header) bool {
+	if len(headers) == 0 {
+		return true
+	}
+	for _, header := range headers {
+		if slices.Contains(h.Values(header.Name), header.Value) {
+			return true
+		}
+	}
+	return false
+}
+
 func normalizeURL(uOrig *url.URL) *url.URL {
 	u := *uOrig
 	// Prevent from attacks with using `..` in r.URL.Path
--- a/app/vmauth/target_url_test.go
+++ b/app/vmauth/target_url_test.go
@ -89,7 +89,7 @@ func TestCreateTargetURLSuccess(t *testing.T) {
 			t.Fatalf("cannot parse %q: %s", requestURI, err)
 		}
 		u = normalizeURL(u)
-		up, hc := ui.getURLPrefixAndHeaders(u)
+		up, hc := ui.getURLPrefixAndHeaders(u, nil)
 		if up == nil {
 			t.Fatalf("cannot determie backend: %s", err)
 		}
@ -249,7 +249,7 @@ func TestCreateTargetURLFailure(t *testing.T) {
 			t.Fatalf("cannot parse %q: %s", requestURI, err)
 		}
 		u = normalizeURL(u)
-		up, hc := ui.getURLPrefixAndHeaders(u)
+		up, hc := ui.getURLPrefixAndHeaders(u, nil)
 		if up != nil {
 			t.Fatalf("unexpected non-empty up=%#v", up)
 		}
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@ -33,6 +33,7 @@ See also [LTS releases](https://docs.victoriametrics.com/lts-releases/).
 * SECURITY: upgrade Go builder from Go1.21.7 to Go1.22.1. See [the list of issues addressed in Go1.22.1](https://github.com/golang/go/issues?q=milestone%3AGo1.22.1+label%3ACherryPickApproved).

 * FEATURE: [vmauth](https://docs.victoriametrics.com/vmauth/): allow routing incoming requests bassed on HTTP [query args](https://en.wikipedia.org/wiki/Query_string) via `src_query_args` option at `url_map`. See [these docs](https://docs.victoriametrics.com/vmauth/#generic-http-proxy-for-different-backends) and [this feature request](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/5878).
+* FEATURE: [vmauth](https://docs.victoriametrics.com/vmauth/): allow routing incoming requests bassed on HTTP request headers via `src_headers` option at `url_map`. See [these docs](https://docs.victoriametrics.com/vmauth/#generic-http-proxy-for-different-backends).
 * FEATURE: [stream aggregation](https://docs.victoriametrics.com/stream-aggregation/): reduce memory usage by up to 5x when aggregating over big number of unique [time series](https://docs.victoriametrics.com/keyconcepts/#time-series). The memory usage reduction is most visible when [stream deduplication](https://docs.victoriametrics.com/stream-aggregation/#deduplication) is enabled. The downside is increased CPU usage by up to 30%.
 * FEATURE: [stream aggregation](https://docs.victoriametrics.com/stream-aggregation/): allow using `-streamAggr.dedupInterval` and `-remoteWrite.streamAggr.dedupInterval` command-line flags without the need to specify `-streamAggr.config` and `-remoteWrite.streamAggr.config`. See [these docs](https://docs.victoriametrics.com/stream-aggregation/#deduplication).
 * FEATURE: [stream aggregation](https://docs.victoriametrics.com/stream-aggregation/): add `-streamAggr.dropInputLabels` command-line flag, which can be used for dropping the listed labels from input samples before applying stream [de-duplication](https://docs.victoriametrics.com/stream-aggregation/#deduplication) and aggregation. This is faster and easier to use alternative to [input_relabel_configs](https://docs.victoriametrics.com/stream-aggregation/#relabeling). See [these docs](https://docs.victoriametrics.com/stream-aggregation/#dropping-unneeded-labels).
--- a/docs/vmauth.md
+++ b/docs/vmauth.md
@ -75,7 +75,7 @@ unauthorized_user:

 ### Generic HTTP proxy for different backends

-`vmauth` can proxy requests to different backends depending on the requested host, path and [query args](https://en.wikipedia.org/wiki/Query_string).
+`vmauth` can proxy requests to different backends depending on the requested host, path, [query args](https://en.wikipedia.org/wiki/Query_string) and any HTTP request header.
 For example, the following [`-auth.config`](#auth-config) instructs `vmauth` to make the following:

 - Requests starting with `/app1/` are proxied to `http://app1-backend/`, while the `/app1/` path prefix is dropped according to [`drop_src_path_prefix_parts`](#dropping-request-path-prefix).
@ -124,18 +124,34 @@ while routing requests with `db=bar` query arg to `http://app2-backend`:
 ```yaml
 unauthorized_user:
  url_map:
-  - src_query_args:
-    - "db=foo"
+  - src_query_args: ["db=foo"]
    url_prefix: "http://app1-backend/"
-  - src_query_args:
-    - "db=bar"
+  - src_query_args: ["db=bar"]
    url_prefix: "http://app2-backend/"
 ```

 If `src_query_args` contains multiple entries, then it is enough to match only a single entry in order to route the request to the given `url_prefix`.

-If `src_hosts` and/or `src_paths` are specified together with `src_query_args`, then the request is routed to the given `url_prefix` if its host, path and query args
-match the given lists simultaneously.
+If `src_query_args` are specified together with `src_hosts`, `src_paths` or `src_headers`, then the request is routed to the given `url_prefix`
+if its query args, host, path and headers match the given lists simultaneously.
+
+An optional `src_headers` can be used for routing requests based on HTTP request headers additionally to hostname, path and [HTTP query args](https://en.wikipedia.org/wiki/Query_string).
+For example, the following config routes requests to `http://app1-backend` if `TenantID` request header equals to `42`, while routing requests to `http://app2-backend`
+if `TenantID` request header equals to `123:456`:
+
+```yaml
+unauthorized_user:
+  url_map:
+  - src_headers: ["TenantID: 42"]
+    url_prefix: "http://app1-backend/"
+  - src_headers: ["TenantID: 123:456"]
+    url_prefix: "http://app2-backend/"
+```
+
+If `src_headers` contains multiple entries, then it is enough to match only a single entry in order to route the request to the given `url_prefix`.
+
+If `src_headers` are specified together with `src_hosts`, `src_paths` or `src_query_args`, then the request is routed to the given `url_prefix`
+if its headers, host, path and query args match the given lists simultaneously.

 ### Generic HTTP load balancer