From 483d7a5b03d4d5bd59249ba0f3fa91caf721146b Mon Sep 17 00:00:00 2001 From: Mendy Danzinger Date: Fri, 15 Jan 2021 12:47:11 -0500 Subject: [PATCH] feat: add support for bearer_token auth (#239) --- .gitignore | 1 + cmd/do-agent/config.go | 28 +++++++++++- go.mod | 1 + pkg/clients/roundtrippers/bearer_token.go | 25 +++++++++++ .../roundtrippers/bearer_token_file.go | 33 ++++++++++++++ .../roundtrippers/bearer_token_file_test.go | 44 +++++++++++++++++++ .../roundtrippers/bearer_token_test.go | 29 ++++++++++++ pkg/clients/roundtrippers/testdata/token | 1 + pkg/clients/roundtrippers/utils.go | 18 ++++++++ pkg/collector/scraper.go | 37 +++++++++++++--- 10 files changed, 210 insertions(+), 7 deletions(-) create mode 100644 pkg/clients/roundtrippers/bearer_token.go create mode 100644 pkg/clients/roundtrippers/bearer_token_file.go create mode 100644 pkg/clients/roundtrippers/bearer_token_file_test.go create mode 100644 pkg/clients/roundtrippers/bearer_token_test.go create mode 100644 pkg/clients/roundtrippers/testdata/token create mode 100644 pkg/clients/roundtrippers/utils.go diff --git a/.gitignore b/.gitignore index 4607e88b..9f0341bd 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ repos/ sonar-agent.key .id_rsa .vault-token +.idea diff --git a/cmd/do-agent/config.go b/cmd/do-agent/config.go index 20c25907..98a3ba2a 100644 --- a/cmd/do-agent/config.go +++ b/cmd/do-agent/config.go @@ -28,6 +28,8 @@ var ( targets map[string]string metadataURL *url.URL authURL *url.URL + bearerToken string + bearerTokenFile string sonarEndpoint string stdoutOnly bool debug bool @@ -94,6 +96,12 @@ func init() { kingpin.Flag("k8s-metrics-path", "enable DO Kubernetes metrics collection (this must be a DOKS metrics endpoint)"). StringVar(&config.kubernetes) + kingpin.Flag("bearer-token", "sets the `Authorization` header on every scrape request with the configured bearer token (mutually exclusive with `bearer-token-file`)"). + StringVar(&config.bearerToken) + + kingpin.Flag("bearer-token-file", "sets the `Authorization` header on every scrape request with the bearer token read from the configured file (mutually exclusive with `bearer-token`)"). + StringVar(&config.bearerTokenFile) + kingpin.Flag("no-collector.processes", "disable processes cpu/memory collection"). Default("true"). BoolVar(&config.noProcesses) @@ -144,6 +152,11 @@ func checkConfig() error { return errors.Wrapf(err, "url for target %q is not valid", name) } } + + if config.bearerTokenFile != "" && config.bearerToken != "" { + return errors.New("both mutually exclusive flags --bearer-token and --bearer-token-file set") + } + return nil } @@ -282,7 +295,20 @@ func initCollectors() []prometheus.Collector { // appendKubernetesCollectors appends a kubernetes metrics collector if it can be initialized successfully func appendKubernetesCollectors(cols []prometheus.Collector) []prometheus.Collector { - k, err := collector.NewScraper("dokubernetes", config.kubernetes, nil, k8sWhitelist, collector.WithTimeout(defaultTimeout), collector.WithLogLevel(log.LevelDebug)) + opts := []collector.Option{ + collector.WithTimeout(defaultTimeout), + collector.WithLogLevel(log.LevelDebug), + } + + if config.bearerToken != "" { + opts = append(opts, collector.WithBearerToken(config.bearerToken)) + } + + if config.bearerTokenFile != "" { + opts = append(opts, collector.WithBearerTokenFile(config.bearerTokenFile)) + } + + k, err := collector.NewScraper("dokubernetes", config.kubernetes, nil, k8sWhitelist, opts...) if err != nil { log.Error("Failed to initialize DO Kubernetes metrics: %+v", err) return cols diff --git a/go.mod b/go.mod index eaaf26a5..21700438 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module github.com/digitalocean/do-agent go 1.12 require ( + github.com/go-kit/kit v0.10.0 github.com/golang/snappy v0.0.0-20180518054509-2e65f85255db github.com/pkg/errors v0.9.1 github.com/prometheus/client_golang v1.6.0 diff --git a/pkg/clients/roundtrippers/bearer_token.go b/pkg/clients/roundtrippers/bearer_token.go new file mode 100644 index 00000000..956d56c2 --- /dev/null +++ b/pkg/clients/roundtrippers/bearer_token.go @@ -0,0 +1,25 @@ +package roundtrippers + +import ( + "fmt" + "net/http" +) + +type bearerTokenRoundTripper struct { + token string + rt http.RoundTripper +} + +// RoundTrip implements http.RoundTripper's interface +func (rt *bearerTokenRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + if len(req.Header.Get("Authorization")) == 0 { + req = cloneRequest(req) + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", rt.token)) + } + return rt.rt.RoundTrip(req) +} + +// NewBearerToken returns an http.RoundTripper that adds the bearer token to a request's header +func NewBearerToken(token string, rt http.RoundTripper) http.RoundTripper { + return &bearerTokenRoundTripper{token, rt} +} diff --git a/pkg/clients/roundtrippers/bearer_token_file.go b/pkg/clients/roundtrippers/bearer_token_file.go new file mode 100644 index 00000000..5fae0cee --- /dev/null +++ b/pkg/clients/roundtrippers/bearer_token_file.go @@ -0,0 +1,33 @@ +package roundtrippers + +import ( + "fmt" + "io/ioutil" + "net/http" + "strings" +) + +type bearerTokenFileRoundTripper struct { + tokenFile string + rt http.RoundTripper +} + +// RoundTrip implements http.RoundTripper's interface +func (rt *bearerTokenFileRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + t, err := ioutil.ReadFile(rt.tokenFile) + if err != nil { + return nil, fmt.Errorf("unable to read bearer token file %s: %s", rt.tokenFile, err) + } + + token := strings.TrimSpace(string(t)) + + req = cloneRequest(req) + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token)) + + return rt.rt.RoundTrip(req) +} + +// NewBearerTokenFile returns an http.RoundTripper that adds the bearer token from a file to a request's header +func NewBearerTokenFile(tokenFile string, rt http.RoundTripper) http.RoundTripper { + return &bearerTokenFileRoundTripper{tokenFile, rt} +} diff --git a/pkg/clients/roundtrippers/bearer_token_file_test.go b/pkg/clients/roundtrippers/bearer_token_file_test.go new file mode 100644 index 00000000..e252818a --- /dev/null +++ b/pkg/clients/roundtrippers/bearer_token_file_test.go @@ -0,0 +1,44 @@ +package roundtrippers + +import ( + "fmt" + "net/http" + "net/http/httptest" + "testing" +) + +const ( + tokenPath = "testdata/token" + invalidTokenPath = "testdata/missingToken" +) + +func Test_bearerTokenFileRoundTripper_RoundTrip_Happy_Path(t *testing.T) { + rt := NewBearerTokenFile(tokenPath, http.DefaultTransport) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + expectedHeader := fmt.Sprintf("Bearer %s", token) + if r.Header.Get("Authorization") != expectedHeader { + t.Errorf("Header.Authorization = %s, want %s", r.Header.Get("Authorization"), expectedHeader) + } + })) + defer ts.Close() + + _, err := rt.RoundTrip(httptest.NewRequest(http.MethodGet, ts.URL, nil)) + if err != nil { + t.Errorf("Unexpected error: %v", err) + } +} + +func Test_bearerTokenFileRoundTripper_RoundTrip_Missing_File(t *testing.T) { + rt := NewBearerTokenFile(invalidTokenPath, http.DefaultTransport) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + return + })) + defer ts.Close() + + _, err := rt.RoundTrip(httptest.NewRequest(http.MethodGet, ts.URL, nil)) + if err == nil { + t.Errorf("Expected error, got none") + } +} diff --git a/pkg/clients/roundtrippers/bearer_token_test.go b/pkg/clients/roundtrippers/bearer_token_test.go new file mode 100644 index 00000000..d1a8a225 --- /dev/null +++ b/pkg/clients/roundtrippers/bearer_token_test.go @@ -0,0 +1,29 @@ +package roundtrippers + +import ( + "fmt" + "net/http" + "net/http/httptest" + "testing" +) + +const ( + token = "test-token-value" +) + +func TestBearerTokenRoundTripper_RoundTrip_Happy_Path(t *testing.T) { + rt := NewBearerToken(token, http.DefaultTransport) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + expectedHeader := fmt.Sprintf("Bearer %s", token) + if r.Header.Get("Authorization") != expectedHeader { + t.Errorf("Header.Authorization = %s, want %s", r.Header.Get("Authorization"), expectedHeader) + } + })) + defer ts.Close() + + _, err := rt.RoundTrip(httptest.NewRequest(http.MethodGet, ts.URL, nil)) + if err != nil { + t.Errorf("Unexpected error: %v", err) + } +} diff --git a/pkg/clients/roundtrippers/testdata/token b/pkg/clients/roundtrippers/testdata/token new file mode 100644 index 00000000..cc78decc --- /dev/null +++ b/pkg/clients/roundtrippers/testdata/token @@ -0,0 +1 @@ +test-token-value diff --git a/pkg/clients/roundtrippers/utils.go b/pkg/clients/roundtrippers/utils.go new file mode 100644 index 00000000..ba2f5b36 --- /dev/null +++ b/pkg/clients/roundtrippers/utils.go @@ -0,0 +1,18 @@ +package roundtrippers + +import "net/http" + +func cloneRequest(req *http.Request) *http.Request { + r := new(http.Request) + + // shallow clone + *r = *req + + // deep copy headers + r.Header = make(http.Header) + for k, v := range req.Header { + r.Header[k] = v + } + + return r +} diff --git a/pkg/collector/scraper.go b/pkg/collector/scraper.go index 60f222e1..07ae9a5b 100644 --- a/pkg/collector/scraper.go +++ b/pkg/collector/scraper.go @@ -10,25 +10,41 @@ import ( "strings" "time" + "github.com/digitalocean/do-agent/internal/log" + "github.com/digitalocean/do-agent/pkg/clients" + "github.com/digitalocean/do-agent/pkg/clients/roundtrippers" "github.com/pkg/errors" "github.com/prometheus/client_golang/prometheus" dto "github.com/prometheus/client_model/go" "github.com/prometheus/common/expfmt" - - "github.com/digitalocean/do-agent/internal/log" - "github.com/digitalocean/do-agent/pkg/clients" ) var defaultScrapeTimeout = 5 * time.Second type scraperOpts struct { - timeout time.Duration - logLevel log.Level + timeout time.Duration + logLevel log.Level + bearerToken string + bearerTokenFile string } // Option is used to configure optional scraper options. type Option func(o *scraperOpts) +// WithBearerToken configures a scraper to use a bearer token +func WithBearerToken(token string) Option { + return func(o *scraperOpts) { + o.bearerToken = token + } +} + +// WithBearerTokenFile configures a scraper to use a bearer token read from a file +func WithBearerTokenFile(tokenFile string) Option { + return func(o *scraperOpts) { + o.bearerTokenFile = tokenFile + } +} + // WithTimeout configures a scraper with a timeout for scraping metrics. func WithTimeout(d time.Duration) Option { return func(o *scraperOpts) { @@ -54,6 +70,15 @@ func NewScraper(name, metricsEndpoint string, extraMetricLabels []*dto.LabelPair opt(defOpts) } + // setup http client, add auth roundtrippers + client := clients.NewHTTP(defOpts.timeout) + if defOpts.bearerTokenFile != "" { + client.Transport = roundtrippers.NewBearerTokenFile(defOpts.bearerTokenFile, client.Transport) + } + if defOpts.bearerToken != "" { + client.Transport = roundtrippers.NewBearerToken(defOpts.bearerToken, client.Transport) + } + metricsEndpoint = strings.TrimRight(metricsEndpoint, "/") req, err := http.NewRequest("GET", metricsEndpoint, nil) if err != nil { @@ -71,7 +96,7 @@ func NewScraper(name, metricsEndpoint string, extraMetricLabels []*dto.LabelPair whitelist: whitelist, timeout: defOpts.timeout, logLevel: defOpts.logLevel, - client: clients.NewHTTP(defOpts.timeout), + client: client, scrapeDurationDesc: prometheus.NewDesc( prometheus.BuildFQName(name, "scrape", "collector_duration_seconds"), fmt.Sprintf("%s: Duration of a collector scrape.", name),