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>

config/headers.go

@@ -0,0 +1,139 @@
+// Copyright 2022 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"
+	"net/textproto"
+	"os"
+	"strings"
+)
+
+// reservedHeaders that change the connection, are set by Prometheus, or car be set
+// otherwise can't be changed.
+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]string `yaml:"headers,omitempty" json:"headers,omitempty"`
+	SecretHeaders map[string]Secret `yaml:"secret_headers,omitempty" json:"secret_headers,omitempty"`
+	Files         map[string]string `yaml:"files,omitempty" json:"files,omitempty"`
+	dir           string
+}
+
+// SetDirectory records the directory to make headers file relative to the
+// configuration file.
+func (h *Headers) SetDirectory(dir string) {
+	h.dir = dir
+}
+
+// Validate validates the Headers config.
+func (h *Headers) Validate() error {
+	uniqueHeaders := make(map[string]struct{}, len(h.Headers))
+	for k := range h.Headers {
+		uniqueHeaders[strings.ToLower(k)] = struct{}{}
+	}
+	for k := range h.SecretHeaders {
+		if _, ok := uniqueHeaders[strings.ToLower(k)]; ok {
+			return fmt.Errorf("header %q is defined in multiple sections", textproto.CanonicalMIMEHeaderKey(k))
+		}
+		uniqueHeaders[strings.ToLower(k)] = struct{}{}
+	}
+	for k, v := range h.Files {
+		if _, ok := uniqueHeaders[strings.ToLower(k)]; ok {
+			return fmt.Errorf("header %q is defined in multiple sections", textproto.CanonicalMIMEHeaderKey(k))
+		}
+		uniqueHeaders[strings.ToLower(k)] = struct{}{}
+		f := JoinDir(h.dir, v)
+		_, err := os.ReadFile(f)
+		if err != nil {
+			return fmt.Errorf("unable to read header %q from file %s: %w", textproto.CanonicalMIMEHeaderKey(k), f, err)
+		}
+	}
+	for k := range uniqueHeaders {
+		if _, ok := reservedHeaders[strings.ToLower(k)]; ok {
+			return fmt.Errorf("setting header %q is not allowed", textproto.CanonicalMIMEHeaderKey(k))
+		}
+	}
+	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)+len(config.SecretHeaders)+len(config.Files) == 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 k, v := range rt.config.Headers {
+		req.Header.Set(textproto.CanonicalMIMEHeaderKey(k), v)
+	}
+	for k, v := range rt.config.SecretHeaders {
+		req.Header.Set(textproto.CanonicalMIMEHeaderKey(k), string(v))
+	}
+	for k, v := range rt.config.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.Set(textproto.CanonicalMIMEHeaderKey(k), 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()
+	}
+}

config/headers_test.go

@@ -0,0 +1,31 @@
+// Copyright 2022 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 (
+	"strings"
+	"testing"
+)
+
+func TestReservedHeaders(t *testing.T) {
+	for k := range reservedHeaders {
+		l := strings.ToLower(k)
+		if k != l {
+			t.Errorf("reservedHeaders keys should be lowercase: got %q, expected %q", k, strings.ToLower(k))
+		}
+	}
+}

config/http_config.go

@@ -274,6 +274,9 @@ type HTTPClientConfig struct {
 	// The omitempty flag is not set, because it would be hidden from the
 	// marshalled configuration when set to false.
 	EnableHTTP2 bool `yaml:"enable_http2" json:"enable_http2"`
+	// 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.
@@ -285,6 +288,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)
 }
 
@@ -347,6 +351,11 @@ func (c *HTTPClientConfig) Validate() error {
 			return fmt.Errorf("at most one of oauth2 client_secret & client_secret_file must be configured")
 		}
 	}
+	if c.HTTPHeaders != nil {
+		if err := c.HTTPHeaders.Validate(); err != nil {
+			return err
+		}
+	}
 	return nil
 }
 
@@ -522,6 +531,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)
 		}

config/http_config_test.go

@@ -119,6 +119,34 @@ var invalidHTTPClientConfigs = []struct {
 		httpClientConfigFile: "testdata/http.conf.oauth2-no-token-url.bad.yaml",
 		errMsg:               "oauth2 token_url must be configured",
 	},
+	{
+		httpClientConfigFile: "testdata/http.conf.headers-duplicate-1.bad.yaml",
+		errMsg:               `header "Foo" is defined in multiple sections`,
+	},
+	{
+		httpClientConfigFile: "testdata/http.conf.headers-duplicate-2.bad.yaml",
+		errMsg:               `header "Foo" is defined in multiple sections`,
+	},
+	{
+		httpClientConfigFile: "testdata/http.conf.headers-duplicate-3.bad.yaml",
+		errMsg:               `header "Foo" is defined in multiple sections`,
+	},
+	{
+		httpClientConfigFile: "testdata/http.conf.headers-duplicate-4.bad.yaml",
+		errMsg:               `header "Foo" is defined in multiple sections`,
+	},
+	{
+		httpClientConfigFile: "testdata/http.conf.headers-reserved-1.bad.yaml",
+		errMsg:               `setting header "User-Agent" is not allowed`,
+	},
+	{
+		httpClientConfigFile: "testdata/http.conf.headers-reserved-2.bad.yaml",
+		errMsg:               `setting header "User-Agent" is not allowed`,
+	},
+	{
+		httpClientConfigFile: "testdata/http.conf.headers-reserved-3.bad.yaml",
+		errMsg:               `setting header "User-Agent" is not allowed`,
+	},
 }
 
 func newTestServer(handler func(w http.ResponseWriter, r *http.Request)) (*httptest.Server, error) {
@@ -1645,3 +1673,33 @@ func TestModifyTLSCertificates(t *testing.T) {
 		})
 	}
 }
+
+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)
+	}
+}

config/testdata/headers-file

@@ -0,0 +1 @@
+value3

config/testdata/http.conf.headers-duplicate-1.bad.yaml

@@ -0,0 +1,5 @@
+http_headers:
+  headers:
+    Foo: bar
+  secret_headers:
+    foO: bar

config/testdata/http.conf.headers-duplicate-2.bad.yaml

@@ -0,0 +1,5 @@
+http_headers:
+  files:
+    Foo: bar
+  secret_headers:
+    foO: bar

config/testdata/http.conf.headers-duplicate-3.bad.yaml

@@ -0,0 +1,5 @@
+http_headers:
+  headers:
+    Foo: bar
+  files:
+    foO: bar

config/testdata/http.conf.headers-duplicate-4.bad.yaml

@@ -0,0 +1,7 @@
+http_headers:
+  headers:
+    Foo: bar
+  secret_headers:
+    foO: bar
+  files:
+    foO: bar

config/testdata/http.conf.headers-reserved-1.bad.yaml

@@ -0,0 +1,3 @@
+http_headers:
+  headers:
+    user-Agent: bar

config/testdata/http.conf.headers-reserved-2.bad.yaml

@@ -0,0 +1,3 @@
+http_headers:
+  secret_headers:
+    user-Agent: bar

config/testdata/http.conf.headers-reserved-3.bad.yaml

@@ -0,0 +1,3 @@
+http_headers:
+  files:
+    user-Agent: testdata/headers-file

config/testdata/http.conf.headers.good.yaml

@@ -0,0 +1,7 @@
+http_headers:
+  headers:
+    one: value1
+  secret_headers:
+    two: value2
+  files:
+    three: testdata/headers-file