Skip to content

Commit

Permalink
Merge pull request #409 from prometheus/mem/proxy_header
Browse files Browse the repository at this point in the history
Add support for proxy connect headers
  • Loading branch information
roidelapluie authored Dec 14, 2022
2 parents 18281a2 + 4a0d730 commit 296ec92
Show file tree
Hide file tree
Showing 8 changed files with 330 additions and 1 deletion.
24 changes: 24 additions & 0 deletions config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ package config

import (
"encoding/json"
"net/http"
"path/filepath"
)

Expand Down Expand Up @@ -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 {
Expand Down
192 changes: 192 additions & 0 deletions config/config_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,13 @@
package config

import (
"bytes"
"encoding/json"
"net/http"
"reflect"
"testing"

"gopkg.in/yaml.v2"
)

func TestJSONMarshalSecret(t *testing.T) {
Expand Down Expand Up @@ -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)
}
})
}
}
12 changes: 11 additions & 1 deletion config/http_config.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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 {
Expand Down Expand Up @@ -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
}

Expand Down Expand Up @@ -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,
Expand Down
67 changes: 67 additions & 0 deletions config/http_config_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
}
11 changes: 11 additions & 0 deletions config/testdata/http.conf.proxy-headers.bad.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
{
"proxy_connect_header": {
"single": [
"value_0"
],
"multi": [
"value_1",
"value_2"
]
}
}
6 changes: 6 additions & 0 deletions config/testdata/http.conf.proxy-headers.bad.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
proxy_connect_header:
single:
- value_0
multi:
- value_1
- value_2
12 changes: 12 additions & 0 deletions config/testdata/http.conf.proxy-headers.good.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
{
"proxy_url": "http://remote.host",
"proxy_connect_header": {
"single": [
"value_0"
],
"multi": [
"value_1",
"value_2"
]
}
}
Loading

0 comments on commit 296ec92

Please sign in to comment.