Repository: prometheus/common
Pull request #416 (2022-11-28 10:20:09)

Add HTTP headers support to common HTTP client.
This is named `http_headers` so it does not clash with blackbox exporter's headers and Prometheus remote client's headers, which are simple maps.

Signed-off-by: Julien Pivotto <roidelapluie@o11y.eu>

roidelapluie: Mailing list discussion: https://groups.google.com/g/prometheus-developers/c/p9HTlPzwUEo
zetaab: @roidelapluie @LeviHarrison what is status in this PR? At least I am interested of this feature
Dronec: Me too! Adding headers to requests is a relatively trivial operation and multiple APIs support that
LeviHarrison: Just some small revisions

@roidelapluie friendly ping

Patch:
diff --git a/config/config.go b/config/config.go
index 3da4854..7588da5 100644
--- a/config/config.go
+++ b/config/config.go
@@ -60,9 +60,9 @@ func (s Secret) MarshalJSON() ([]byte, error) {
 	return json.Marshal(secretToken)
 }
 
-type Header map[string][]Secret
+type ProxyHeader map[string][]Secret
 
-func (h *Header) HTTPHeader() http.Header {
+func (h *ProxyHeader) HTTPHeader() http.Header {
 	if h == nil || *h == nil {
 		return nil
 	}
diff --git a/config/headers.go b/config/headers.go
new file mode 100644
index 0000000..7c3149e
--- /dev/null
+++ b/config/headers.go
@@ -0,0 +1,134 @@
+// Copyright 2024 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// This package no longer handles safe yaml parsing. In order to
+// ensure correct yaml unmarshalling, use "yaml.UnmarshalStrict()".
+
+package config
+
+import (
+	"fmt"
+	"net/http"
+	"os"
+	"strings"
+)
+
+// reservedHeaders that change the connection, are set by Prometheus, or can
+// be changed otherwise.
+var reservedHeaders = map[string]struct{}{
+	"Authorization":                       {},
+	"Host":                                {},
+	"Content-Encoding":                    {},
+	"Content-Length":                      {},
+	"Content-Type":                        {},
+	"User-Agent":                          {},
+	"Connection":                          {},
+	"Keep-Alive":                          {},
+	"Proxy-Authenticate":                  {},
+	"Proxy-Authorization":                 {},
+	"Www-Authenticate":                    {},
+	"Accept-Encoding":                     {},
+	"X-Prometheus-Remote-Write-Version":   {},
+	"X-Prometheus-Remote-Read-Version":    {},
+	"X-Prometheus-Scrape-Timeout-Seconds": {},
+
+	// Added by SigV4.
+	"X-Amz-Date":           {},
+	"X-Amz-Security-Token": {},
+	"X-Amz-Content-Sha256": {},
+}
+
+// Headers represents the configuration for HTTP headers.
+type Headers struct {
+	Headers map[string]Header `yaml:",inline" json:",inline"`
+	dir     string
+}
+
+// Headers represents the configuration for HTTP headers.
+type Header struct {
+	Values  []string `yaml:"values,omitempty" json:"values,omitempty"`
+	Secrets []Secret `yaml:"secrets,omitempty" json:"secrets,omitempty"`
+	Files   []string `yaml:"files,omitempty" json:"files,omitempty"`
+}
+
+// SetDirectory records the directory to make headers file relative to the
+// configuration file.
+func (h *Headers) SetDirectory(dir string) {
+	if h == nil {
+		return
+	}
+	h.dir = dir
+}
+
+// Validate validates the Headers config.
+func (h *Headers) Validate() error {
+	for n, header := range h.Headers {
+		if _, ok := reservedHeaders[http.CanonicalHeaderKey(n)]; ok {
+			return fmt.Errorf("setting header %q is not allowed", http.CanonicalHeaderKey(n))
+		}
+		for _, v := range header.Files {
+			f := JoinDir(h.dir, v)
+			_, err := os.ReadFile(f)
+			if err != nil {
+				return fmt.Errorf("unable to read header %q from file %s: %w", http.CanonicalHeaderKey(n), f, err)
+			}
+		}
+	}
+	return nil
+}
+
+// NewHeadersRoundTripper returns a RoundTripper that sets HTTP headers on
+// requests as configured.
+func NewHeadersRoundTripper(config *Headers, next http.RoundTripper) http.RoundTripper {
+	if len(config.Headers) == 0 {
+		return next
+	}
+	return &headersRoundTripper{
+		config: config,
+		next:   next,
+	}
+}
+
+type headersRoundTripper struct {
+	next   http.RoundTripper
+	config *Headers
+}
+
+// RoundTrip implements http.RoundTripper.
+func (rt *headersRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	req = cloneRequest(req)
+	for n, h := range rt.config.Headers {
+		for _, v := range h.Values {
+			req.Header.Add(n, v)
+		}
+		for _, v := range h.Secrets {
+			req.Header.Add(n, string(v))
+		}
+		for _, v := range h.Files {
+			f := JoinDir(rt.config.dir, v)
+			b, err := os.ReadFile(f)
+			if err != nil {
+				return nil, fmt.Errorf("unable to read headers file %s: %w", f, err)
+			}
+			req.Header.Add(n, strings.TrimSpace(string(b)))
+		}
+	}
+	return rt.next.RoundTrip(req)
+}
+
+// CloseIdleConnections implements closeIdler.
+func (rt *headersRoundTripper) CloseIdleConnections() {
+	if ci, ok := rt.next.(closeIdler); ok {
+		ci.CloseIdleConnections()
+	}
+}
diff --git a/config/http_config.go b/config/http_config.go
index 2044181..3241b1a 100644
--- a/config/http_config.go
+++ b/config/http_config.go
@@ -311,6 +311,9 @@ type HTTPClientConfig struct {
 	EnableHTTP2 bool `yaml:"enable_http2" json:"enable_http2"`
 	// Proxy configuration.
 	ProxyConfig `yaml:",inline"`
+	// HTTPHeaders specify headers to inject in the requests. Those headers
+	// could be marshalled back to the users.
+	HTTPHeaders *Headers `yaml:"http_headers" json:"http_headers"`
 }
 
 // SetDirectory joins any relative file paths with dir.
@@ -322,6 +325,7 @@ func (c *HTTPClientConfig) SetDirectory(dir string) {
 	c.BasicAuth.SetDirectory(dir)
 	c.Authorization.SetDirectory(dir)
 	c.OAuth2.SetDirectory(dir)
+	c.HTTPHeaders.SetDirectory(dir)
 	c.BearerTokenFile = JoinDir(dir, c.BearerTokenFile)
 }
 
@@ -388,6 +392,11 @@ func (c *HTTPClientConfig) Validate() error {
 	if err := c.ProxyConfig.Validate(); err != nil {
 		return err
 	}
+	if c.HTTPHeaders != nil {
+		if err := c.HTTPHeaders.Validate(); err != nil {
+			return err
+		}
+	}
 	return nil
 }
 
@@ -572,6 +581,10 @@ func NewRoundTripperFromConfig(cfg HTTPClientConfig, name string, optFuncs ...HT
 			rt = NewOAuth2RoundTripper(cfg.OAuth2, rt, &opts)
 		}
 
+		if cfg.HTTPHeaders != nil {
+			rt = NewHeadersRoundTripper(cfg.HTTPHeaders, rt)
+		}
+
 		if opts.userAgent != "" {
 			rt = NewUserAgentRoundTripper(opts.userAgent, rt)
 		}
@@ -1236,7 +1249,7 @@ type ProxyConfig struct {
 	// proxies during CONNECT requests. Assume that at least _some_ of
 	// these headers are going to contain secrets and use Secret as the
 	// value type instead of string.
-	ProxyConnectHeader Header `yaml:"proxy_connect_header,omitempty" json:"proxy_connect_header,omitempty"`
+	ProxyConnectHeader ProxyHeader `yaml:"proxy_connect_header,omitempty" json:"proxy_connect_header,omitempty"`
 
 	proxyFunc func(*http.Request) (*url.URL, error)
 }


Test patch:
diff --git a/config/config_test.go b/config/config_test.go
index 51d2340..1ab58d7 100644
--- a/config/config_test.go
+++ b/config/config_test.go
@@ -83,11 +83,11 @@ func TestJSONMarshalSecret(t *testing.T) {
 
 func TestHeaderHTTPHeader(t *testing.T) {
 	testcases := map[string]struct {
-		header   Header
+		header   ProxyHeader
 		expected http.Header
 	}{
 		"basic": {
-			header: Header{
+			header: ProxyHeader{
 				"single": []Secret{"v1"},
 				"multi":  []Secret{"v1", "v2"},
 				"empty":  []Secret{},
@@ -119,32 +119,32 @@ func TestHeaderHTTPHeader(t *testing.T) {
 func TestHeaderYamlUnmarshal(t *testing.T) {
 	testcases := map[string]struct {
 		input    string
-		expected Header
+		expected ProxyHeader
 	}{
 		"void": {
 			input: ``,
 		},
 		"simple": {
 			input:    "single:\n- a\n",
-			expected: Header{"single": []Secret{"a"}},
+			expected: ProxyHeader{"single": []Secret{"a"}},
 		},
 		"multi": {
 			input:    "multi:\n- a\n- b\n",
-			expected: Header{"multi": []Secret{"a", "b"}},
+			expected: ProxyHeader{"multi": []Secret{"a", "b"}},
 		},
 		"empty": {
 			input:    "{}",
-			expected: Header{},
+			expected: ProxyHeader{},
 		},
 		"empty value": {
 			input:    "empty:\n",
-			expected: Header{"empty": nil},
+			expected: ProxyHeader{"empty": nil},
 		},
 	}
 
 	for name, tc := range testcases {
 		t.Run(name, func(t *testing.T) {
-			var actual Header
+			var actual ProxyHeader
 			err := yaml.Unmarshal([]byte(tc.input), &actual)
 			if err != nil {
 				t.Fatalf("error unmarshaling %s: %s", tc.input, err)
@@ -158,7 +158,7 @@ func TestHeaderYamlUnmarshal(t *testing.T) {
 
 func TestHeaderYamlMarshal(t *testing.T) {
 	testcases := map[string]struct {
-		input    Header
+		input    ProxyHeader
 		expected []byte
 	}{
 		"void": {
@@ -166,15 +166,15 @@ func TestHeaderYamlMarshal(t *testing.T) {
 			expected: []byte("{}\n"),
 		},
 		"simple": {
-			input:    Header{"single": []Secret{"a"}},
+			input:    ProxyHeader{"single": []Secret{"a"}},
 			expected: []byte("single:\n- <secret>\n"),
 		},
 		"multi": {
-			input:    Header{"multi": []Secret{"a", "b"}},
+			input:    ProxyHeader{"multi": []Secret{"a", "b"}},
 			expected: []byte("multi:\n- <secret>\n- <secret>\n"),
 		},
 		"empty": {
-			input:    Header{"empty": nil},
+			input:    ProxyHeader{"empty": nil},
 			expected: []byte("empty: []\n"),
 		},
 	}
@@ -195,32 +195,32 @@ func TestHeaderYamlMarshal(t *testing.T) {
 func TestHeaderJsonUnmarshal(t *testing.T) {
 	testcases := map[string]struct {
 		input    string
-		expected Header
+		expected ProxyHeader
 	}{
 		"void": {
 			input: `null`,
 		},
 		"simple": {
 			input:    `{"single": ["a"]}`,
-			expected: Header{"single": []Secret{"a"}},
+			expected: ProxyHeader{"single": []Secret{"a"}},
 		},
 		"multi": {
 			input:    `{"multi": ["a", "b"]}`,
-			expected: Header{"multi": []Secret{"a", "b"}},
+			expected: ProxyHeader{"multi": []Secret{"a", "b"}},
 		},
 		"empty": {
 			input:    `{}`,
-			expected: Header{},
+			expected: ProxyHeader{},
 		},
 		"empty value": {
 			input:    `{"empty":null}`,
-			expected: Header{"empty": nil},
+			expected: ProxyHeader{"empty": nil},
 		},
 	}
 
 	for name, tc := range testcases {
 		t.Run(name, func(t *testing.T) {
-			var actual Header
+			var actual ProxyHeader
 			err := json.Unmarshal([]byte(tc.input), &actual)
 			if err != nil {
 				t.Fatalf("error unmarshaling %s: %s", tc.input, err)
@@ -234,7 +234,7 @@ func TestHeaderJsonUnmarshal(t *testing.T) {
 
 func TestHeaderJsonMarshal(t *testing.T) {
 	testcases := map[string]struct {
-		input    Header
+		input    ProxyHeader
 		expected []byte
 	}{
 		"void": {
@@ -242,15 +242,15 @@ func TestHeaderJsonMarshal(t *testing.T) {
 			expected: []byte("null"),
 		},
 		"simple": {
-			input:    Header{"single": []Secret{"a"}},
+			input:    ProxyHeader{"single": []Secret{"a"}},
 			expected: []byte("{\"single\":[\"\\u003csecret\\u003e\"]}"),
 		},
 		"multi": {
-			input:    Header{"multi": []Secret{"a", "b"}},
+			input:    ProxyHeader{"multi": []Secret{"a", "b"}},
 			expected: []byte("{\"multi\":[\"\\u003csecret\\u003e\",\"\\u003csecret\\u003e\"]}"),
 		},
 		"empty": {
-			input:    Header{"empty": nil},
+			input:    ProxyHeader{"empty": nil},
 			expected: []byte(`{"empty":null}`),
 		},
 	}
diff --git a/config/headers_test.go b/config/headers_test.go
new file mode 100644
index 0000000..39c6f9f
--- /dev/null
+++ b/config/headers_test.go
@@ -0,0 +1,31 @@
+// Copyright 2024 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// This package no longer handles safe yaml parsing. In order to
+// ensure correct yaml unmarshalling, use "yaml.UnmarshalStrict()".
+
+package config
+
+import (
+	"net/http"
+	"testing"
+)
+
+func TestReservedHeaders(t *testing.T) {
+	for k := range reservedHeaders {
+		l := http.CanonicalHeaderKey(k)
+		if k != l {
+			t.Errorf("reservedHeaders keys should be lowercase: got %q, expected %q", k, http.CanonicalHeaderKey(k))
+		}
+	}
+}
diff --git a/config/http_config_test.go b/config/http_config_test.go
index cd13a18..01dc7d6 100644
--- a/config/http_config_test.go
+++ b/config/http_config_test.go
@@ -131,6 +131,10 @@ var invalidHTTPClientConfigs = []struct {
 		httpClientConfigFile: "testdata/http.conf.no-proxy-without-proxy-url.bad.yaml",
 		errMsg:               "if no_proxy is configured, proxy_url must also be configured",
 	},
+	{
+		httpClientConfigFile: "testdata/http.conf.headers-reserved.bad.yaml",
+		errMsg:               `setting header "User-Agent" is not allowed`,
+	},
 }
 
 func newTestServer(handler func(w http.ResponseWriter, r *http.Request)) (*httptest.Server, error) {
@@ -1439,7 +1443,7 @@ func TestInvalidHTTPConfigs(t *testing.T) {
 	for _, ee := range invalidHTTPClientConfigs {
 		_, _, err := LoadHTTPConfigFile(ee.httpClientConfigFile)
 		if err == nil {
-			t.Error("Expected error with config but got none")
+			t.Errorf("Expected error with config %q but got none", ee.httpClientConfigFile)
 			continue
 		}
 		if !strings.Contains(err.Error(), ee.errMsg) {
@@ -2230,3 +2234,63 @@ func readFile(t *testing.T, filename string) string {
 
 	return string(content)
 }
+
+func TestHeaders(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		for k, v := range map[string]string{
+			"One":   "value1",
+			"Two":   "value2",
+			"Three": "value3",
+		} {
+			if r.Header.Get(k) != v {
+				t.Errorf("expected %q, got %q", v, r.Header.Get(k))
+			}
+		}
+		w.WriteHeader(http.StatusNoContent)
+	}))
+	t.Cleanup(ts.Close)
+
+	cfg, _, err := LoadHTTPConfigFile("testdata/http.conf.headers.good.yaml")
+	if err != nil {
+		t.Fatalf("Error loading HTTP client config: %v", err)
+	}
+	client, err := NewClientFromConfig(*cfg, "test")
+	if err != nil {
+		t.Fatalf("Error creating HTTP Client: %v", err)
+	}
+
+	_, err = client.Get(ts.URL)
+	if err != nil {
+		t.Fatalf("can't fetch URL: %v", err)
+	}
+}
+
+func TestMultipleHeaders(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		for k, v := range map[string][]string{
+			"One":   {"value1a", "value1b", "value1c"},
+			"Two":   {"value2a", "value2b", "value2c"},
+			"Three": {"value3a", "value3b", "value3c"},
+		} {
+			if !reflect.DeepEqual(r.Header.Values(k), v) {
+				t.Errorf("expected %v, got %v", v, r.Header.Values(k))
+			}
+		}
+		w.WriteHeader(http.StatusNoContent)
+	}))
+	t.Cleanup(ts.Close)
+
+	cfg, _, err := LoadHTTPConfigFile("testdata/http.conf.headers-multiple.good.yaml")
+	if err != nil {
+		t.Fatalf("Error loading HTTP client config: %v", err)
+	}
+	client, err := NewClientFromConfig(*cfg, "test")
+	if err != nil {
+		t.Fatalf("Error creating HTTP Client: %v", err)
+	}
+
+	_, err = client.Get(ts.URL)
+	if err != nil {
+		t.Fatalf("can't fetch URL: %v", err)
+	}
+}
diff --git a/config/testdata/headers-file b/config/testdata/headers-file
new file mode 100644
index 0000000..44cd5c8
--- /dev/null
+++ b/config/testdata/headers-file
@@ -0,0 +1,1 @@
+value3
diff --git a/config/testdata/headers-file-a b/config/testdata/headers-file-a
new file mode 100644
index 0000000..0d2a9cd
--- /dev/null
+++ b/config/testdata/headers-file-a
@@ -0,0 +1,3 @@
+value3a
+
+
diff --git a/config/testdata/headers-file-b b/config/testdata/headers-file-b
new file mode 100644
index 0000000..c1240b4
--- /dev/null
+++ b/config/testdata/headers-file-b
@@ -0,0 +1,1 @@
+value3b 
diff --git a/config/testdata/headers-file-c b/config/testdata/headers-file-c
new file mode 100644
index 0000000..264b69d
--- /dev/null
+++ b/config/testdata/headers-file-c
@@ -0,0 +1,1 @@
+value3c
diff --git a/config/testdata/http.conf.headers-multiple.good.yaml b/config/testdata/http.conf.headers-multiple.good.yaml
new file mode 100644
index 0000000..f9dcdc0
--- /dev/null
+++ b/config/testdata/http.conf.headers-multiple.good.yaml
@@ -0,0 +1,8 @@
+http_headers:
+  one:
+    values: [value1a, value1b, value1c]
+  two:
+    values: [value2a]
+    secrets: [value2b, value2c]
+  three:
+    files: [testdata/headers-file-a, testdata/headers-file-b, testdata/headers-file-c]
diff --git a/config/testdata/http.conf.headers-reserved.bad.yaml b/config/testdata/http.conf.headers-reserved.bad.yaml
new file mode 100644
index 0000000..4e48856
--- /dev/null
+++ b/config/testdata/http.conf.headers-reserved.bad.yaml
@@ -0,0 +1,3 @@
+http_headers:
+  user-Agent:
+    values: [bar]
diff --git a/config/testdata/http.conf.headers.good.yaml b/config/testdata/http.conf.headers.good.yaml
new file mode 100644
index 0000000..f8a0236
--- /dev/null
+++ b/config/testdata/http.conf.headers.good.yaml
@@ -0,0 +1,7 @@
+http_headers:
+  one:
+    values: [value1]
+  two:
+    secrets: [value2]
+  three:
+    files: [testdata/headers-file]