From 9599179a6f42a05c1820af98f872eaba6950d082 Mon Sep 17 00:00:00 2001
From: "Marcelo E. Magallon" <marcelo.magallon@grafana.com>
Date: Thu, 10 Nov 2022 19:33:52 -0600
Subject: [PATCH] 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
Signed-off-by: Marcelo E. Magallon <marcelo.magallon@grafana.com>
---
 config/http_config.go      | 10 +++++++++-
 config/http_config_test.go | 31 +++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/config/http_config.go b/config/http_config.go
index 69c0b671..cc1591fc 100644
--- a/config/http_config.go
+++ b/config/http_config.go
@@ -264,6 +264,9 @@ 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.
+	ProxyConnectHeader http.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.
@@ -289,7 +292,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 {
@@ -347,6 +351,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
 }
 
@@ -475,6 +482,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,
 			MaxIdleConns:          20000,
 			MaxIdleConnsPerHost:   1000, // see https://github.com/golang/go/issues/13801
 			DisableKeepAlives:     !opts.keepAlivesEnabled,
diff --git a/config/http_config_test.go b/config/http_config_test.go
index 036fd2dc..efaaf411 100644
--- a/config/http_config_test.go
+++ b/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