Skip to content

Commit

Permalink
Merge pull request #416 from roidelapluie/headers
Browse files Browse the repository at this point in the history
Add HTTP headers support to common HTTP client.
  • Loading branch information
roidelapluie authored Apr 29, 2024
2 parents 5c108f8 + aeb616a commit 36303a8
Show file tree
Hide file tree
Showing 13 changed files with 292 additions and 26 deletions.
4 changes: 2 additions & 2 deletions config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -60,9 +60,9 @@ func (s Secret) MarshalJSON() ([]byte, error) {
return json.Marshal(secretToken)
}

type Header map[string][]Secret
type ProxyHeader map[string][]Secret

func (h *Header) HTTPHeader() http.Header {
func (h *ProxyHeader) HTTPHeader() http.Header {
if h == nil || *h == nil {
return nil
}
Expand Down
44 changes: 22 additions & 22 deletions config/config_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -83,11 +83,11 @@ func TestJSONMarshalSecret(t *testing.T) {

func TestHeaderHTTPHeader(t *testing.T) {
testcases := map[string]struct {
header Header
header ProxyHeader
expected http.Header
}{
"basic": {
header: Header{
header: ProxyHeader{
"single": []Secret{"v1"},
"multi": []Secret{"v1", "v2"},
"empty": []Secret{},
Expand Down Expand Up @@ -119,32 +119,32 @@ func TestHeaderHTTPHeader(t *testing.T) {
func TestHeaderYamlUnmarshal(t *testing.T) {
testcases := map[string]struct {
input string
expected Header
expected ProxyHeader
}{
"void": {
input: ``,
},
"simple": {
input: "single:\n- a\n",
expected: Header{"single": []Secret{"a"}},
expected: ProxyHeader{"single": []Secret{"a"}},
},
"multi": {
input: "multi:\n- a\n- b\n",
expected: Header{"multi": []Secret{"a", "b"}},
expected: ProxyHeader{"multi": []Secret{"a", "b"}},
},
"empty": {
input: "{}",
expected: Header{},
expected: ProxyHeader{},
},
"empty value": {
input: "empty:\n",
expected: Header{"empty": nil},
expected: ProxyHeader{"empty": nil},
},
}

for name, tc := range testcases {
t.Run(name, func(t *testing.T) {
var actual Header
var actual ProxyHeader
err := yaml.Unmarshal([]byte(tc.input), &actual)
if err != nil {
t.Fatalf("error unmarshaling %s: %s", tc.input, err)
Expand All @@ -158,23 +158,23 @@ func TestHeaderYamlUnmarshal(t *testing.T) {

func TestHeaderYamlMarshal(t *testing.T) {
testcases := map[string]struct {
input Header
input ProxyHeader
expected []byte
}{
"void": {
input: nil,
expected: []byte("{}\n"),
},
"simple": {
input: Header{"single": []Secret{"a"}},
input: ProxyHeader{"single": []Secret{"a"}},
expected: []byte("single:\n- <secret>\n"),
},
"multi": {
input: Header{"multi": []Secret{"a", "b"}},
input: ProxyHeader{"multi": []Secret{"a", "b"}},
expected: []byte("multi:\n- <secret>\n- <secret>\n"),
},
"empty": {
input: Header{"empty": nil},
input: ProxyHeader{"empty": nil},
expected: []byte("empty: []\n"),
},
}
Expand All @@ -195,32 +195,32 @@ func TestHeaderYamlMarshal(t *testing.T) {
func TestHeaderJsonUnmarshal(t *testing.T) {
testcases := map[string]struct {
input string
expected Header
expected ProxyHeader
}{
"void": {
input: `null`,
},
"simple": {
input: `{"single": ["a"]}`,
expected: Header{"single": []Secret{"a"}},
expected: ProxyHeader{"single": []Secret{"a"}},
},
"multi": {
input: `{"multi": ["a", "b"]}`,
expected: Header{"multi": []Secret{"a", "b"}},
expected: ProxyHeader{"multi": []Secret{"a", "b"}},
},
"empty": {
input: `{}`,
expected: Header{},
expected: ProxyHeader{},
},
"empty value": {
input: `{"empty":null}`,
expected: Header{"empty": nil},
expected: ProxyHeader{"empty": nil},
},
}

for name, tc := range testcases {
t.Run(name, func(t *testing.T) {
var actual Header
var actual ProxyHeader
err := json.Unmarshal([]byte(tc.input), &actual)
if err != nil {
t.Fatalf("error unmarshaling %s: %s", tc.input, err)
Expand All @@ -234,23 +234,23 @@ func TestHeaderJsonUnmarshal(t *testing.T) {

func TestHeaderJsonMarshal(t *testing.T) {
testcases := map[string]struct {
input Header
input ProxyHeader
expected []byte
}{
"void": {
input: nil,
expected: []byte("null"),
},
"simple": {
input: Header{"single": []Secret{"a"}},
input: ProxyHeader{"single": []Secret{"a"}},
expected: []byte("{\"single\":[\"\\u003csecret\\u003e\"]}"),
},
"multi": {
input: Header{"multi": []Secret{"a", "b"}},
input: ProxyHeader{"multi": []Secret{"a", "b"}},
expected: []byte("{\"multi\":[\"\\u003csecret\\u003e\",\"\\u003csecret\\u003e\"]}"),
},
"empty": {
input: Header{"empty": nil},
input: ProxyHeader{"empty": nil},
expected: []byte(`{"empty":null}`),
},
}
Expand Down
134 changes: 134 additions & 0 deletions config/headers.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,134 @@
// Copyright 2024 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"
"os"
"strings"
)

// reservedHeaders that change the connection, are set by Prometheus, or can
// be changed otherwise.
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]Header `yaml:",inline" json:",inline"`
dir string
}

// Headers represents the configuration for HTTP headers.
type Header struct {
Values []string `yaml:"values,omitempty" json:"values,omitempty"`
Secrets []Secret `yaml:"secrets,omitempty" json:"secrets,omitempty"`
Files []string `yaml:"files,omitempty" json:"files,omitempty"`
}

// SetDirectory records the directory to make headers file relative to the
// configuration file.
func (h *Headers) SetDirectory(dir string) {
if h == nil {
return
}
h.dir = dir
}

// Validate validates the Headers config.
func (h *Headers) Validate() error {
for n, header := range h.Headers {
if _, ok := reservedHeaders[http.CanonicalHeaderKey(n)]; ok {
return fmt.Errorf("setting header %q is not allowed", http.CanonicalHeaderKey(n))
}
for _, v := range header.Files {
f := JoinDir(h.dir, v)
_, err := os.ReadFile(f)
if err != nil {
return fmt.Errorf("unable to read header %q from file %s: %w", http.CanonicalHeaderKey(n), f, err)
}
}
}
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) == 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 n, h := range rt.config.Headers {
for _, v := range h.Values {
req.Header.Add(n, v)
}
for _, v := range h.Secrets {
req.Header.Add(n, string(v))
}
for _, v := range h.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.Add(n, 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()
}
}
31 changes: 31 additions & 0 deletions config/headers_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
// Copyright 2024 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 (
"net/http"
"testing"
)

func TestReservedHeaders(t *testing.T) {
for k := range reservedHeaders {
l := http.CanonicalHeaderKey(k)
if k != l {
t.Errorf("reservedHeaders keys should be lowercase: got %q, expected %q", k, http.CanonicalHeaderKey(k))
}
}
}
15 changes: 14 additions & 1 deletion config/http_config.go
Original file line number Diff line number Diff line change
Expand Up @@ -311,6 +311,9 @@ type HTTPClientConfig struct {
EnableHTTP2 bool `yaml:"enable_http2" json:"enable_http2"`
// Proxy configuration.
ProxyConfig `yaml:",inline"`
// 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.
Expand All @@ -322,6 +325,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)
}

Expand Down Expand Up @@ -388,6 +392,11 @@ func (c *HTTPClientConfig) Validate() error {
if err := c.ProxyConfig.Validate(); err != nil {
return err
}
if c.HTTPHeaders != nil {
if err := c.HTTPHeaders.Validate(); err != nil {
return err
}
}
return nil
}

Expand Down Expand Up @@ -572,6 +581,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)
}
Expand Down Expand Up @@ -1236,7 +1249,7 @@ type ProxyConfig struct {
// 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"`
ProxyConnectHeader ProxyHeader `yaml:"proxy_connect_header,omitempty" json:"proxy_connect_header,omitempty"`

proxyFunc func(*http.Request) (*url.URL, error)
}
Expand Down
Loading

0 comments on commit 36303a8

Please sign in to comment.