Merge pull request #409 from prometheus/mem/proxy_header

Add support for proxy connect headers

config/config.go

@@ -18,6 +18,7 @@ package config
 
 import (
 	"encoding/json"
+	"net/http"
 	"path/filepath"
 )
 
@@ -48,6 +49,29 @@ func (s Secret) MarshalJSON() ([]byte, error) {
 	return json.Marshal(secretToken)
 }
 
+type Header map[string][]Secret
+
+func (h *Header) HTTPHeader() http.Header {
+	if h == nil || *h == nil {
+		return nil
+	}
+
+	header := make(http.Header)
+
+	for name, values := range *h {
+		var s []string
+		if values != nil {
+			s = make([]string, 0, len(values))
+			for _, value := range values {
+				s = append(s, string(value))
+			}
+		}
+		header[name] = s
+	}
+
+	return header
+}
+
 // DirectorySetter is a config type that contains file paths that may
 // be relative to the file containing the config.
 type DirectorySetter interface {

config/config_test.go

@@ -14,8 +14,13 @@
 package config
 
 import (
+	"bytes"
 	"encoding/json"
+	"net/http"
+	"reflect"
 	"testing"
+
+	"gopkg.in/yaml.v2"
 )
 
 func TestJSONMarshalSecret(t *testing.T) {
@@ -51,3 +56,190 @@ func TestJSONMarshalSecret(t *testing.T) {
 		})
 	}
 }
+
+func TestHeaderHTTPHeader(t *testing.T) {
+	testcases := map[string]struct {
+		header   Header
+		expected http.Header
+	}{
+		"basic": {
+			header: Header{
+				"single": []Secret{"v1"},
+				"multi":  []Secret{"v1", "v2"},
+				"empty":  []Secret{},
+				"nil":    nil,
+			},
+			expected: http.Header{
+				"single": []string{"v1"},
+				"multi":  []string{"v1", "v2"},
+				"empty":  []string{},
+				"nil":    nil,
+			},
+		},
+		"nil": {
+			header:   nil,
+			expected: nil,
+		},
+	}
+
+	for name, tc := range testcases {
+		t.Run(name, func(t *testing.T) {
+			actual := tc.header.HTTPHeader()
+			if !reflect.DeepEqual(actual, tc.expected) {
+				t.Fatalf("expecting: %#v, actual: %#v", tc.expected, actual)
+			}
+		})
+	}
+}
+
+func TestHeaderYamlUnmarshal(t *testing.T) {
+	testcases := map[string]struct {
+		input    string
+		expected Header
+	}{
+		"void": {
+			input: ``,
+		},
+		"simple": {
+			input:    "single:\n- a\n",
+			expected: Header{"single": []Secret{"a"}},
+		},
+		"multi": {
+			input:    "multi:\n- a\n- b\n",
+			expected: Header{"multi": []Secret{"a", "b"}},
+		},
+		"empty": {
+			input:    "{}",
+			expected: Header{},
+		},
+		"empty value": {
+			input:    "empty:\n",
+			expected: Header{"empty": nil},
+		},
+	}
+
+	for name, tc := range testcases {
+		t.Run(name, func(t *testing.T) {
+			var actual Header
+			err := yaml.Unmarshal([]byte(tc.input), &actual)
+			if err != nil {
+				t.Fatalf("error unmarshaling %s: %s", tc.input, err)
+			}
+			if !reflect.DeepEqual(actual, tc.expected) {
+				t.Fatalf("expecting: %#v, actual: %#v", tc.expected, actual)
+			}
+		})
+	}
+}
+
+func TestHeaderYamlMarshal(t *testing.T) {
+	testcases := map[string]struct {
+		input    Header
+		expected []byte
+	}{
+		"void": {
+			input:    nil,
+			expected: []byte("{}\n"),
+		},
+		"simple": {
+			input:    Header{"single": []Secret{"a"}},
+			expected: []byte("single:\n- <secret>\n"),
+		},
+		"multi": {
+			input:    Header{"multi": []Secret{"a", "b"}},
+			expected: []byte("multi:\n- <secret>\n- <secret>\n"),
+		},
+		"empty": {
+			input:    Header{"empty": nil},
+			expected: []byte("empty: []\n"),
+		},
+	}
+
+	for name, tc := range testcases {
+		t.Run(name, func(t *testing.T) {
+			actual, err := yaml.Marshal(tc.input)
+			if err != nil {
+				t.Fatalf("error unmarshaling %#v: %s", tc.input, err)
+			}
+			if !bytes.Equal(actual, tc.expected) {
+				t.Fatalf("expecting: %q, actual: %q", tc.expected, actual)
+			}
+		})
+	}
+}
+
+func TestHeaderJsonUnmarshal(t *testing.T) {
+	testcases := map[string]struct {
+		input    string
+		expected Header
+	}{
+		"void": {
+			input: `null`,
+		},
+		"simple": {
+			input:    `{"single": ["a"]}`,
+			expected: Header{"single": []Secret{"a"}},
+		},
+		"multi": {
+			input:    `{"multi": ["a", "b"]}`,
+			expected: Header{"multi": []Secret{"a", "b"}},
+		},
+		"empty": {
+			input:    `{}`,
+			expected: Header{},
+		},
+		"empty value": {
+			input:    `{"empty":null}`,
+			expected: Header{"empty": nil},
+		},
+	}
+
+	for name, tc := range testcases {
+		t.Run(name, func(t *testing.T) {
+			var actual Header
+			err := json.Unmarshal([]byte(tc.input), &actual)
+			if err != nil {
+				t.Fatalf("error unmarshaling %s: %s", tc.input, err)
+			}
+			if !reflect.DeepEqual(actual, tc.expected) {
+				t.Fatalf("expecting: %#v, actual: %#v", tc.expected, actual)
+			}
+		})
+	}
+}
+
+func TestHeaderJsonMarshal(t *testing.T) {
+	testcases := map[string]struct {
+		input    Header
+		expected []byte
+	}{
+		"void": {
+			input:    nil,
+			expected: []byte("null"),
+		},
+		"simple": {
+			input:    Header{"single": []Secret{"a"}},
+			expected: []byte("{\"single\":[\"\\u003csecret\\u003e\"]}"),
+		},
+		"multi": {
+			input:    Header{"multi": []Secret{"a", "b"}},
+			expected: []byte("{\"multi\":[\"\\u003csecret\\u003e\",\"\\u003csecret\\u003e\"]}"),
+		},
+		"empty": {
+			input:    Header{"empty": nil},
+			expected: []byte(`{"empty":null}`),
+		},
+	}
+
+	for name, tc := range testcases {
+		t.Run(name, func(t *testing.T) {
+			actual, err := json.Marshal(tc.input)
+			if err != nil {
+				t.Fatalf("error marshaling %#v: %s", tc.input, err)
+			}
+			if !bytes.Equal(actual, tc.expected) {
+				t.Fatalf("expecting: %q, actual: %q", tc.expected, actual)
+			}
+		})
+	}
+}

config/http_config.go

@@ -289,6 +289,11 @@ type HTTPClientConfig struct {
 	BearerTokenFile string `yaml:"bearer_token_file,omitempty" json:"bearer_token_file,omitempty"`
 	// HTTP proxy server to use to connect to the targets.
 	ProxyURL URL `yaml:"proxy_url,omitempty" json:"proxy_url,omitempty"`
+	// ProxyConnectHeader optionally specifies headers to send to
+	// 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"`
 	// TLSConfig to use to connect to the targets.
 	TLSConfig TLSConfig `yaml:"tls_config,omitempty" json:"tls_config,omitempty"`
 	// FollowRedirects specifies whether the client should follow HTTP 3xx redirects.
@@ -314,7 +319,8 @@ func (c *HTTPClientConfig) SetDirectory(dir string) {
 }
 
 // Validate validates the HTTPClientConfig to check only one of BearerToken,
-// BasicAuth and BearerTokenFile is configured.
+// BasicAuth and BearerTokenFile is configured. It also validates that ProxyURL
+// is set if ProxyConnectHeader is set.
 func (c *HTTPClientConfig) Validate() error {
 	// Backwards compatibility with the bearer_token field.
 	if len(c.BearerToken) > 0 && len(c.BearerTokenFile) > 0 {
@@ -372,6 +378,9 @@ func (c *HTTPClientConfig) Validate() error {
 			return fmt.Errorf("at most one of oauth2 client_secret & client_secret_file must be configured")
 		}
 	}
+	if len(c.ProxyConnectHeader) > 0 && (c.ProxyURL.URL == nil || c.ProxyURL.String() == "") {
+		return fmt.Errorf("if proxy_connect_header is configured proxy_url must also be configured")
+	}
 	return nil
 }
 
@@ -500,6 +509,7 @@ func NewRoundTripperFromConfig(cfg HTTPClientConfig, name string, optFuncs ...HT
 		// It is applied on request. So we leave out any timings here.
 		var rt http.RoundTripper = &http.Transport{
 			Proxy:                 http.ProxyURL(cfg.ProxyURL.URL),
+			ProxyConnectHeader:    cfg.ProxyConnectHeader.HTTPHeader(),
 			MaxIdleConns:          20000,
 			MaxIdleConnsPerHost:   1000, // see https://github.com/golang/go/issues/13801
 			DisableKeepAlives:     !opts.keepAlivesEnabled,

config/http_config_test.go

@@ -447,6 +447,50 @@ func TestNewClientFromConfig(t *testing.T) {
 	}
 }
 
+func TestProxyConfiguration(t *testing.T) {
+	testcases := map[string]struct {
+		testFn  string
+		loader  func(string) (*HTTPClientConfig, []byte, error)
+		isValid bool
+	}{
+		"good yaml": {
+			testFn:  "testdata/http.conf.proxy-headers.good.yml",
+			loader:  LoadHTTPConfigFile,
+			isValid: true,
+		},
+		"bad yaml": {
+			testFn:  "testdata/http.conf.proxy-headers.bad.yml",
+			loader:  LoadHTTPConfigFile,
+			isValid: false,
+		},
+		"good json": {
+			testFn:  "testdata/http.conf.proxy-headers.good.json",
+			loader:  loadHTTPConfigJSONFile,
+			isValid: true,
+		},
+		"bad json": {
+			testFn:  "testdata/http.conf.proxy-headers.bad.json",
+			loader:  loadHTTPConfigJSONFile,
+			isValid: false,
+		},
+	}
+
+	for name, tc := range testcases {
+		t.Run(name, func(t *testing.T) {
+			_, _, err := tc.loader(tc.testFn)
+			if tc.isValid {
+				if err != nil {
+					t.Fatalf("Error validating %s: %s", tc.testFn, err)
+				}
+			} else {
+				if err == nil {
+					t.Fatalf("Expecting error validating %s but got %s", tc.testFn, err)
+				}
+			}
+		})
+	}
+}
+
 func TestNewClientFromInvalidConfig(t *testing.T) {
 	var newClientInvalidConfig = []struct {
 		clientConfig HTTPClientConfig
@@ -1622,3 +1666,26 @@ func TestModifyTLSCertificates(t *testing.T) {
 		})
 	}
 }
+
+// loadHTTPConfigJSON parses the JSON input s into a HTTPClientConfig.
+func loadHTTPConfigJSON(buf []byte) (*HTTPClientConfig, error) {
+	cfg := &HTTPClientConfig{}
+	err := json.Unmarshal(buf, cfg)
+	if err != nil {
+		return nil, err
+	}
+	return cfg, nil
+}
+
+// loadHTTPConfigJSONFile parses the given JSON file into a HTTPClientConfig.
+func loadHTTPConfigJSONFile(filename string) (*HTTPClientConfig, []byte, error) {
+	content, err := os.ReadFile(filename)
+	if err != nil {
+		return nil, nil, err
+	}
+	cfg, err := loadHTTPConfigJSON(content)
+	if err != nil {
+		return nil, nil, err
+	}
+	return cfg, content, nil
+}

config/testdata/http.conf.proxy-headers.bad.json

@@ -0,0 +1,11 @@
+{
+    "proxy_connect_header": {
+	"single": [
+	    "value_0"
+	],
+	"multi": [
+	    "value_1",
+	    "value_2"
+	]
+    }
+}

config/testdata/http.conf.proxy-headers.bad.yml

@@ -0,0 +1,6 @@
+proxy_connect_header:
+  single:
+    - value_0
+  multi:
+    - value_1
+    - value_2

config/testdata/http.conf.proxy-headers.good.json

@@ -0,0 +1,12 @@
+{
+    "proxy_url": "http://remote.host",
+    "proxy_connect_header": {
+	"single": [
+	    "value_0"
+	],
+	"multi": [
+	    "value_1",
+	    "value_2"
+	]
+    }
+}