feat: add support for bearer_token auth (#239)

.gitignore

@@ -3,3 +3,4 @@ repos/
 sonar-agent.key
 .id_rsa
 .vault-token
+.idea

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

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

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}
+}

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}
+}

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")
+	}
+}

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)
+	}
+}

pkg/clients/roundtrippers/testdata/token

@@ -0,0 +1 @@
+test-token-value

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
+}

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),