Add support for proxy connect headers

Some proxy configurations require additional headers to be able to use
them (e.g. authorization token specific to the proxy).

Fixes: #402
Co-authored-by: Julien Pivotto <roidelapluie@o11y.eu>
Signed-off-by: Marcelo E. Magallon <marcelo.magallon@grafana.com>

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 {
+		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,110 @@ 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 TestHeaderUnmarshal(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:    "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 TestHeaderMarshal(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)
+			}
+		})
+	}
+}

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,37 @@ func TestNewClientFromConfig(t *testing.T) {
 	}
 }
 
+func TestProxyConfiguration(t *testing.T) {
+	testcases := map[string]struct {
+		testdata string
+		isValid  bool
+	}{
+		"good": {
+			testdata: "testdata/http.conf.proxy-headers.good.yml",
+			isValid:  true,
+		},
+		"bad": {
+			testdata: "testdata/http.conf.proxy-headers.bad.yml",
+			isValid:  false,
+		},
+	}
+
+	for name, tc := range testcases {
+		t.Run(name, func(t *testing.T) {
+			_, _, err := LoadHTTPConfigFile(tc.testdata)
+			if tc.isValid {
+				if err != nil {
+					t.Fatalf("Error validating %s: %s", tc.testdata, err)
+				}
+			} else {
+				if err == nil {
+					t.Fatalf("Expecting error validating %s but got %s", tc.testdata, err)
+				}
+			}
+		})
+	}
+}
+
 func TestNewClientFromInvalidConfig(t *testing.T) {
 	var newClientInvalidConfig = []struct {
 		clientConfig HTTPClientConfig

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.yml

@@ -0,0 +1,7 @@
+proxy_url: "http://remote.host"
+proxy_connect_header:
+  single:
+    - value_0
+  multi:
+    - value_1
+    - value_2