From b18eed34274dbd6fda0497afe43bb3b7e500edd0 Mon Sep 17 00:00:00 2001
From: Alexander Marshalov <_@marshalov.org>
Date: Thu, 31 Aug 2023 14:26:51 +0200
Subject: [PATCH] vmauth: added ability to set and remove response headers
 (#4825) (#4914)

* added ability to set and clear response headers (#4825)

Signed-off-by: Alexander Marshalov <_@marshalov.org>

* added ability to set and clear response headers (#4825)

Signed-off-by: Alexander Marshalov <_@marshalov.org>

* fix review comment

Signed-off-by: Alexander Marshalov <_@marshalov.org>

---------

Signed-off-by: Alexander Marshalov <_@marshalov.org>
---
 app/vmauth/README.md           | 10 ++++-
 app/vmauth/auth_config.go      | 30 ++++++++------
 app/vmauth/auth_config_test.go | 72 +++++++++++++++++++---------------
 app/vmauth/example_config.yml  |  5 +++
 app/vmauth/main.go             | 23 +++++++----
 app/vmauth/target_url.go       |  8 ++--
 app/vmauth/target_url_test.go  | 19 +++++----
 docs/CHANGELOG.md              |  1 +
 docs/vmauth.md                 | 10 ++++-
 9 files changed, 113 insertions(+), 65 deletions(-)

diff --git a/app/vmauth/README.md b/app/vmauth/README.md
index dea54e0e83..ea14aeff6f 100644
--- a/app/vmauth/README.md
+++ b/app/vmauth/README.md
@@ -117,11 +117,16 @@ users:
 
   # Requests with the 'Authorization: Bearer YYY' header are proxied to http://localhost:8428 ,
   # The `X-Scope-OrgID: foobar` http header is added to every proxied request.
+  # The `X-Server-Hostname` http header is removed from the proxied response.
   # For example, http://vmauth:8427/api/v1/query is proxied to http://localhost:8428/api/v1/query
 - bearer_token: "YYY"
   url_prefix: "http://localhost:8428"
+  # extra headers to add to the request or remove from the request (if header value is empty)
   headers:
-  - "X-Scope-OrgID: foobar"
+    - "X-Scope-OrgID: foobar"
+  # extra headers to add to the response or remove from the response (if header value is empty)
+  response_headers:
+    - "X-Server-Hostname:" # empty value means the header will be removed from the response
 
   # All the requests to http://vmauth:8427 with the given Basic Auth (username:password)
   # are proxied to http://localhost:8428 .
@@ -175,6 +180,7 @@ users:
   #
   # - Requests to http://vmauth:8427/api/v1/write are proxied to http://vminsert:8480/insert/42/prometheus/api/v1/write .
   #   The "X-Scope-OrgID: abc" http header is added to these requests.
+  #   The "X-Server-Hostname" http header is removed from the proxied response.
   #
   # Request which do not match `src_paths` from the `url_map` are proxied to the urls from `default_url`
   # in a round-robin manner. The original request path is passed in `request_path` query arg.
@@ -194,6 +200,8 @@ users:
     url_prefix: "http://vminsert:8480/insert/42/prometheus"
     headers:
     - "X-Scope-OrgID: abc"
+    response_headers:
+    - "X-Server-Hostname:" # empty value means the header will be removed from the response
     ip_filters:
       deny_list: [127.0.0.1]
   default_url:
diff --git a/app/vmauth/auth_config.go b/app/vmauth/auth_config.go
index 6fab91dc27..af1b0573f6 100644
--- a/app/vmauth/auth_config.go
+++ b/app/vmauth/auth_config.go
@@ -37,15 +37,15 @@ type AuthConfig struct {
 
 // UserInfo is user information read from authConfigPath
 type UserInfo struct {
-	Name                  string     `yaml:"name,omitempty"`
-	BearerToken           string     `yaml:"bearer_token,omitempty"`
-	Username              string     `yaml:"username,omitempty"`
-	Password              string     `yaml:"password,omitempty"`
-	URLPrefix             *URLPrefix `yaml:"url_prefix,omitempty"`
-	URLMaps               []URLMap   `yaml:"url_map,omitempty"`
-	Headers               []Header   `yaml:"headers,omitempty"`
-	MaxConcurrentRequests int        `yaml:"max_concurrent_requests,omitempty"`
-	DefaultURL            *URLPrefix `yaml:"default_url,omitempty"`
+	Name                  string      `yaml:"name,omitempty"`
+	BearerToken           string      `yaml:"bearer_token,omitempty"`
+	Username              string      `yaml:"username,omitempty"`
+	Password              string      `yaml:"password,omitempty"`
+	URLPrefix             *URLPrefix  `yaml:"url_prefix,omitempty"`
+	URLMaps               []URLMap    `yaml:"url_map,omitempty"`
+	HeadersConf           HeadersConf `yaml:",inline"`
+	MaxConcurrentRequests int         `yaml:"max_concurrent_requests,omitempty"`
+	DefaultURL            *URLPrefix  `yaml:"default_url,omitempty"`
 
 	concurrencyLimitCh      chan struct{}
 	concurrencyLimitReached *metrics.Counter
@@ -54,6 +54,12 @@ type UserInfo struct {
 	requestsDuration *metrics.Summary
 }
 
+// HeadersConf represents config for request and response headers.
+type HeadersConf struct {
+	RequestHeaders  []Header `yaml:"headers,omitempty"`
+	ResponseHeaders []Header `yaml:"response_headers,omitempty"`
+}
+
 func (ui *UserInfo) beginConcurrencyLimit() error {
 	select {
 	case ui.concurrencyLimitCh <- struct{}{}:
@@ -105,9 +111,9 @@ func (h *Header) MarshalYAML() (interface{}, error) {
 
 // URLMap is a mapping from source paths to target urls.
 type URLMap struct {
-	SrcPaths  []*SrcPath `yaml:"src_paths,omitempty"`
-	URLPrefix *URLPrefix `yaml:"url_prefix,omitempty"`
-	Headers   []Header   `yaml:"headers,omitempty"`
+	SrcPaths    []*SrcPath  `yaml:"src_paths,omitempty"`
+	URLPrefix   *URLPrefix  `yaml:"url_prefix,omitempty"`
+	HeadersConf HeadersConf `yaml:",inline"`
 }
 
 // SrcPath represents an src path
diff --git a/app/vmauth/auth_config_test.go b/app/vmauth/auth_config_test.go
index e794231aad..dc465bcfcf 100644
--- a/app/vmauth/auth_config_test.go
+++ b/app/vmauth/auth_config_test.go
@@ -299,14 +299,16 @@ users:
 						"http://vminsert1/insert/0/prometheus",
 						"http://vminsert2/insert/0/prometheus",
 					}),
-					Headers: []Header{
-						{
-							Name:  "foo",
-							Value: "bar",
-						},
-						{
-							Name:  "xxx",
-							Value: "y",
+					HeadersConf: HeadersConf{
+						RequestHeaders: []Header{
+							{
+								Name:  "foo",
+								Value: "bar",
+							},
+							{
+								Name:  "xxx",
+								Value: "y",
+							},
 						},
 					},
 				},
@@ -325,14 +327,16 @@ users:
 						"http://vminsert1/insert/0/prometheus",
 						"http://vminsert2/insert/0/prometheus",
 					}),
-					Headers: []Header{
-						{
-							Name:  "foo",
-							Value: "bar",
-						},
-						{
-							Name:  "xxx",
-							Value: "y",
+					HeadersConf: HeadersConf{
+						RequestHeaders: []Header{
+							{
+								Name:  "foo",
+								Value: "bar",
+							},
+							{
+								Name:  "xxx",
+								Value: "y",
+							},
 						},
 					},
 				},
@@ -389,14 +393,16 @@ users:
 						"http://vminsert1/insert/0/prometheus",
 						"http://vminsert2/insert/0/prometheus",
 					}),
-					Headers: []Header{
-						{
-							Name:  "foo",
-							Value: "bar",
-						},
-						{
-							Name:  "xxx",
-							Value: "y",
+					HeadersConf: HeadersConf{
+						RequestHeaders: []Header{
+							{
+								Name:  "foo",
+								Value: "bar",
+							},
+							{
+								Name:  "xxx",
+								Value: "y",
+							},
 						},
 					},
 				},
@@ -419,14 +425,16 @@ users:
 						"http://vminsert1/insert/0/prometheus",
 						"http://vminsert2/insert/0/prometheus",
 					}),
-					Headers: []Header{
-						{
-							Name:  "foo",
-							Value: "bar",
-						},
-						{
-							Name:  "xxx",
-							Value: "y",
+					HeadersConf: HeadersConf{
+						RequestHeaders: []Header{
+							{
+								Name:  "foo",
+								Value: "bar",
+							},
+							{
+								Name:  "xxx",
+								Value: "y",
+							},
 						},
 					},
 				},
diff --git a/app/vmauth/example_config.yml b/app/vmauth/example_config.yml
index e7721be0ea..f9e149626f 100644
--- a/app/vmauth/example_config.yml
+++ b/app/vmauth/example_config.yml
@@ -12,11 +12,16 @@ users:
 
   # Requests with the 'Authorization: Bearer YYY' header are proxied to http://localhost:8428 ,
   # The `X-Scope-OrgID: foobar` http header is added to every proxied request.
+  # The `X-Server-Hostname:` http header is removed from the proxied response.
   # For example, http://vmauth:8427/api/v1/query is proxied to http://localhost:8428/api/v1/query
 - bearer_token: "YYY"
   url_prefix: "http://localhost:8428"
+  # extra headers to add to the request or remove from the request (if header value is empty)
   headers:
   - "X-Scope-OrgID: foobar"
+  # extra headers to add to the response or remove from the response (if header value is empty)
+  response_headers:
+  - "X-Server-Hostname:" # empty value means the header will be removed from the response
 
   # All the requests to http://vmauth:8427 with the given Basic Auth (username:password)
   # are proxied to http://localhost:8428 .
diff --git a/app/vmauth/main.go b/app/vmauth/main.go
index 925bfb7c60..1f7d685a04 100644
--- a/app/vmauth/main.go
+++ b/app/vmauth/main.go
@@ -151,7 +151,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, headers := ui.getURLPrefixAndHeaders(u)
+	up, headersConf := ui.getURLPrefixAndHeaders(u)
 	isDefault := false
 	if up == nil {
 		missingRouteRequests.Inc()
@@ -159,7 +159,7 @@ func processRequest(w http.ResponseWriter, r *http.Request, ui *UserInfo) {
 			httpserver.Errorf(w, r, "missing route for %q", u.String())
 			return
 		}
-		up, headers = ui.DefaultURL, ui.Headers
+		up, headersConf = ui.DefaultURL, ui.HeadersConf
 		isDefault = true
 	}
 	r.Body = &readTrackingBody{
@@ -178,7 +178,7 @@ func processRequest(w http.ResponseWriter, r *http.Request, ui *UserInfo) {
 		} else { // Update path for regular routes.
 			targetURL = mergeURLs(targetURL, u)
 		}
-		ok := tryProcessingRequest(w, r, targetURL, headers)
+		ok := tryProcessingRequest(w, r, targetURL, headersConf)
 		bu.put()
 		if ok {
 			return
@@ -192,13 +192,11 @@ func processRequest(w http.ResponseWriter, r *http.Request, ui *UserInfo) {
 	httpserver.Errorf(w, r, "%s", err)
 }
 
-func tryProcessingRequest(w http.ResponseWriter, r *http.Request, targetURL *url.URL, headers []Header) bool {
+func tryProcessingRequest(w http.ResponseWriter, r *http.Request, targetURL *url.URL, headersConf HeadersConf) bool {
 	// This code has been copied from net/http/httputil/reverseproxy.go
 	req := sanitizeRequestHeaders(r)
 	req.URL = targetURL
-	for _, h := range headers {
-		req.Header.Set(h.Name, h.Value)
-	}
+	updateHeadersByConfig(req.Header, headersConf.RequestHeaders)
 	transportOnce.Do(transportInit)
 	res, err := transport.RoundTrip(req)
 	if err != nil {
@@ -222,6 +220,7 @@ func tryProcessingRequest(w http.ResponseWriter, r *http.Request, targetURL *url
 	}
 	removeHopHeaders(res.Header)
 	copyHeader(w.Header(), res.Header)
+	updateHeadersByConfig(w.Header(), headersConf.ResponseHeaders)
 	w.WriteHeader(res.StatusCode)
 
 	copyBuf := copyBufPool.Get()
@@ -247,6 +246,16 @@ func copyHeader(dst, src http.Header) {
 	}
 }
 
+func updateHeadersByConfig(headers http.Header, config []Header) {
+	for _, h := range config {
+		if h.Value == "" {
+			headers.Del(h.Name)
+		} else {
+			headers.Set(h.Name, h.Value)
+		}
+	}
+}
+
 func sanitizeRequestHeaders(r *http.Request) *http.Request {
 	// This code has been copied from net/http/httputil/reverseproxy.go
 	req := r.Clone(r.Context())
diff --git a/app/vmauth/target_url.go b/app/vmauth/target_url.go
index 7e4afd24aa..365484d952 100644
--- a/app/vmauth/target_url.go
+++ b/app/vmauth/target_url.go
@@ -29,18 +29,18 @@ func mergeURLs(uiURL, requestURI *url.URL) *url.URL {
 	return &targetURL
 }
 
-func (ui *UserInfo) getURLPrefixAndHeaders(u *url.URL) (*URLPrefix, []Header) {
+func (ui *UserInfo) getURLPrefixAndHeaders(u *url.URL) (*URLPrefix, HeadersConf) {
 	for _, e := range ui.URLMaps {
 		for _, sp := range e.SrcPaths {
 			if sp.match(u.Path) {
-				return e.URLPrefix, e.Headers
+				return e.URLPrefix, e.HeadersConf
 			}
 		}
 	}
 	if ui.URLPrefix != nil {
-		return ui.URLPrefix, ui.Headers
+		return ui.URLPrefix, ui.HeadersConf
 	}
-	return nil, nil
+	return nil, HeadersConf{}
 }
 
 func normalizeURL(uOrig *url.URL) *url.URL {
diff --git a/app/vmauth/target_url_test.go b/app/vmauth/target_url_test.go
index 482138e17b..41c4ea1af8 100644
--- a/app/vmauth/target_url_test.go
+++ b/app/vmauth/target_url_test.go
@@ -24,7 +24,7 @@ func TestCreateTargetURLSuccess(t *testing.T) {
 		if target.String() != expectedTarget {
 			t.Fatalf("unexpected target; got %q; want %q", target, expectedTarget)
 		}
-		headersStr := fmt.Sprintf("%q", headers)
+		headersStr := fmt.Sprintf("%q", headers.RequestHeaders)
 		if headersStr != expectedHeaders {
 			t.Fatalf("unexpected headers; got %s; want %s", headersStr, expectedHeaders)
 		}
@@ -35,10 +35,10 @@ func TestCreateTargetURLSuccess(t *testing.T) {
 	}, "", "http://foo.bar/.", "[]")
 	f(&UserInfo{
 		URLPrefix: mustParseURL("http://foo.bar"),
-		Headers: []Header{{
+		HeadersConf: HeadersConf{RequestHeaders: []Header{{
 			Name:  "bb",
 			Value: "aaa",
-		}},
+		}}},
 	}, "/", "http://foo.bar", `[{"bb" "aaa"}]`)
 	f(&UserInfo{
 		URLPrefix: mustParseURL("http://foo.bar/federate"),
@@ -62,7 +62,7 @@ func TestCreateTargetURLSuccess(t *testing.T) {
 			{
 				SrcPaths:  getSrcPaths([]string{"/api/v1/query"}),
 				URLPrefix: mustParseURL("http://vmselect/0/prometheus"),
-				Headers: []Header{
+				HeadersConf: HeadersConf{RequestHeaders: []Header{
 					{
 						Name:  "xx",
 						Value: "aa",
@@ -71,7 +71,7 @@ func TestCreateTargetURLSuccess(t *testing.T) {
 						Name:  "yy",
 						Value: "asdf",
 					},
-				},
+				}},
 			},
 			{
 				SrcPaths:  getSrcPaths([]string{"/api/v1/write"}),
@@ -79,10 +79,10 @@ func TestCreateTargetURLSuccess(t *testing.T) {
 			},
 		},
 		URLPrefix: mustParseURL("http://default-server"),
-		Headers: []Header{{
+		HeadersConf: HeadersConf{RequestHeaders: []Header{{
 			Name:  "bb",
 			Value: "aaa",
-		}},
+		}}},
 	}
 	f(ui, "/api/v1/query?query=up", "http://vmselect/0/prometheus/api/v1/query?query=up", `[{"xx" "aa"} {"yy" "asdf"}]`)
 	f(ui, "/api/v1/write", "http://vminsert/0/prometheus/api/v1/write", "[]")
@@ -128,7 +128,10 @@ func TestCreateTargetURLFailure(t *testing.T) {
 		if up != nil {
 			t.Fatalf("unexpected non-empty up=%#v", up)
 		}
-		if headers != nil {
+		if headers.RequestHeaders != nil {
+			t.Fatalf("unexpected non-empty headers=%q", headers)
+		}
+		if headers.ResponseHeaders != nil {
 			t.Fatalf("unexpected non-empty headers=%q", headers)
 		}
 	}
diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index 304ed980de..df5e0983ed 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -36,6 +36,7 @@ The following `tip` changes can be tested by building VictoriaMetrics components
 * FEATURE: [vmui](https://docs.victoriametrics.com/#vmui): improve accessibility score to 100 according to [Google's Lighthouse](https://developer.chrome.com/docs/lighthouse/accessibility/) tests.
 * FEATURE: [vmui](https://docs.victoriametrics.com/#vmui): organize `min`, `max`, `median` values on the chart legend and tooltips for better visibility.
 * FEATURE: dashboards: provide copies of Grafana dashboards alternated with VictoriaMetrics datasource at [dashboards/vm](https://github.com/VictoriaMetrics/VictoriaMetrics/tree/master/dashboards/vm).
+* FEATURE: [vmauth](https://docs.victoriametrics.com/vmauth.html): added ability to set and clear response headers. See [this issue](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/4825) and [these docs](https://docs.victoriametrics.com/vmauth.html#auth-config) for details.
 
 * BUGFIX: [build](https://docs.victoriametrics.com/): fix Docker builds for old Docker releases. See [this issue](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/4907).
 * BUGFIX: [vmagent](https://docs.victoriametrics.com/vmagent.html): consistently set `User-Agent` header to `vm_promscrape` during scraping with enabled or disabled [stream parsing mode](https://docs.victoriametrics.com/vmagent.html#stream-parsing-mode). See [this issue](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/4884).
diff --git a/docs/vmauth.md b/docs/vmauth.md
index 4dd06f501b..e8f4cb243d 100644
--- a/docs/vmauth.md
+++ b/docs/vmauth.md
@@ -128,11 +128,16 @@ users:
 
   # Requests with the 'Authorization: Bearer YYY' header are proxied to http://localhost:8428 ,
   # The `X-Scope-OrgID: foobar` http header is added to every proxied request.
+  # The `X-Server-Hostname` http header is removed from the proxied response.
   # For example, http://vmauth:8427/api/v1/query is proxied to http://localhost:8428/api/v1/query
 - bearer_token: "YYY"
   url_prefix: "http://localhost:8428"
+  # extra headers to add to the request or remove from the request (if header value is empty)
   headers:
-  - "X-Scope-OrgID: foobar"
+    - "X-Scope-OrgID: foobar"
+  # extra headers to add to the response or remove from the response (if header value is empty)
+  response_headers:
+    - "X-Server-Hostname:" # empty value means the header will be removed from the response
 
   # All the requests to http://vmauth:8427 with the given Basic Auth (username:password)
   # are proxied to http://localhost:8428 .
@@ -186,6 +191,7 @@ users:
   #
   # - Requests to http://vmauth:8427/api/v1/write are proxied to http://vminsert:8480/insert/42/prometheus/api/v1/write .
   #   The "X-Scope-OrgID: abc" http header is added to these requests.
+  #   The "X-Server-Hostname" http header is removed from the proxied response.
   #
   # Request which do not match `src_paths` from the `url_map` are proxied to the urls from `default_url`
   # in a round-robin manner. The original request path is passed in `request_path` query arg.
@@ -205,6 +211,8 @@ users:
     url_prefix: "http://vminsert:8480/insert/42/prometheus"
     headers:
     - "X-Scope-OrgID: abc"
+    response_headers:
+    - "X-Server-Hostname:" # empty value means the header will be removed from the response
     ip_filters:
       deny_list: [127.0.0.1]
   default_url: