Make token providers modular

backend/stdout.go

@@ -0,0 +1,16 @@
+package backend
+
+import "github.com/buildkite/buildkite-agent-metrics/collector"
+
+// Stdout is a "metrics collection backend" that simply prints metrics to standard output.
+// Useful for testing and local debugging
+type Stdout struct{}
+
+func NewStdoutBackend() *Stdout {
+	return &Stdout{}
+}
+
+func (s *Stdout) Collect(r *collector.Result) error {
+	r.Dump()
+	return nil
+}

go.mod

@@ -10,8 +10,10 @@ require (
 	github.com/golang/mock v1.2.0
 	github.com/golang/protobuf v1.2.0
 	github.com/jmespath/go-jmespath v0.0.0-20180206201540-c2b33e8439af // indirect
+	github.com/kr/pretty v0.3.0 // indirect
 	github.com/newrelic/go-agent v3.0.0+incompatible
 	github.com/prometheus/client_golang v0.9.3-0.20190127221311-3c4408c8b829
 	github.com/prometheus/client_model v0.0.0-20190115171406-56726106282f
+	github.com/stretchr/testify v1.7.1
 	google.golang.org/genproto v0.0.0-20190401181712-f467c93bbac2
 )

go.sum

@@ -25,6 +25,7 @@ github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDk
 github.com/coreos/go-systemd v0.0.0-20181012123002-c6f51f82210d/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
 github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU=
 github.com/cpuguy83/go-md2man/v2 v2.0.0/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU=
+github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -77,7 +78,14 @@ github.com/julienschmidt/httprouter v1.2.0/go.mod h1:SYymIcj16QtmaHHD7aYtjjsJG7V
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
 github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc=
+github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
+github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0=
+github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk=
+github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/pty v1.1.3/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
+github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
+github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
+github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/matttproud/golang_protobuf_extensions v1.0.1 h1:4hp9jkHxhMHkqkrB3Ix0jegS5sx/RkqARlsWZ6pIwiU=
 github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
 github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
@@ -108,14 +116,17 @@ github.com/prometheus/procfs v0.0.0-20181005140218-185b4288413d/go.mod h1:c3At6R
 github.com/prometheus/procfs v0.0.0-20190117184657-bf6a532e95b1 h1:/K3IL0Z1quvmJ7X0A1AwNEK7CRkVK3YwfOU/QAL4WGg=
 github.com/prometheus/procfs v0.0.0-20190117184657-bf6a532e95b1/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk=
 github.com/rcrowley/go-metrics v0.0.0-20181016184325-3113b8401b8a/go.mod h1:bCqnVzQkZxMG4s8nGwiZ5l3QUCyqpo9Y+/ZMZ9VjZe4=
+github.com/rogpeppe/go-internal v1.6.1 h1:/FiVV8dS/e+YqF2JvO3yXRFbBLTIuSDkuC7aBOAvL+k=
+github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc=
 github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/shurcooL/sanitized_anchor_name v1.0.0/go.mod h1:1NzhyTcUVG4SuEtjjoZeVRXNmyL/1OwPU0+IJeTBvfc=
 github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
-github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0=
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY=
+github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/tarm/serial v0.0.0-20180830185346-98f6abe2eb07/go.mod h1:kDXzergiv9cbyO7IOYJZWg1U88JhDg3PB6klq9Hg2pA=
 github.com/urfave/cli/v2 v2.2.0/go.mod h1:SE9GqnLQmjVa0iPEY0f1w3ygNIYcIJ0OKPMoW2caLfQ=
 go.opencensus.io v0.18.0/go.mod h1:vKdFvxhtzZ9onBp9VKHK8z/sRpBMnKAsufL7wlDrCOA=
@@ -199,7 +210,9 @@ google.golang.org/grpc v1.19.0 h1:cfg4PD8YEdSFnm7qLV4++93WcmhH2nIUhMjhdCvl3j8=
 google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c=
 gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
 gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys=
 gopkg.in/inf.v0 v0.9.1/go.mod h1:cWUDdTG/fYaXco+Dcufb5Vnc6Gp2YChqWtbxRZE0mXw=
 gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw=

lambda/main.go

@@ -11,7 +11,6 @@ import (
 	"time"
 
 	"github.com/aws/aws-lambda-go/lambda"
-	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/secretsmanager"
 	"github.com/aws/aws-sdk-go/service/ssm"
@@ -46,8 +45,6 @@ func main() {
 
 func Handler(ctx context.Context, evt json.RawMessage) (string, error) {
 	var b backend.Backend
-	var provider token.Provider
-	var bkToken string
 	var err error
 
 	awsRegion := os.Getenv("AWS_REGION")
@@ -69,12 +66,14 @@ func Handler(ctx context.Context, evt json.RawMessage) (string, error) {
 		return "", nil
 	}
 
-	provider, err = initTokenProvider(awsRegion)
-	if err != nil {
-		return "", err
-	}
+	sess := session.Must(session.NewSession())
+	tokenProviderChain := token.NewProviderChain(
+		token.NewEnvVar(),
+		token.NewSecretsManager(secretsmanager.New(sess)),
+		token.NewSSM(ssm.New(sess)),
+	)
 
-	bkToken, err = provider.Get()
+	bkToken, err := tokenProviderChain.Resolve()
 	if err != nil {
 		return "", err
 	}
@@ -112,6 +111,8 @@ func Handler(ctx context.Context, evt json.RawMessage) (string, error) {
 			fmt.Printf("Error starting New Relic client: %v\n", err)
 			os.Exit(1)
 		}
+	case "stdout":
+		b = backend.NewStdoutBackend()
 	default:
 		dimensions, err := backend.ParseCloudWatchDimensions(clwDimensions)
 		if err != nil {
@@ -147,59 +148,3 @@ func Handler(ctx context.Context, evt json.RawMessage) (string, error) {
 
 	return "", nil
 }
-
-func initTokenProvider(awsRegion string) (token.Provider, error) {
-	mutuallyExclusiveEnvVars := []string{
-		BKAgentTokenEnvVar,
-		BKAgentTokenSSMKeyEnvVar,
-		BKAgentTokenSecretsManagerSecretIDEnvVar,
-	}
-
-	if err := checkMutuallyExclusiveEnvVars(mutuallyExclusiveEnvVars...); err != nil {
-		return nil, err
-	}
-
-	if bkToken := os.Getenv(BKAgentTokenEnvVar); bkToken != "" {
-		return token.NewInMemory(bkToken)
-	}
-
-	if ssmKey := os.Getenv(BKAgentTokenSSMKeyEnvVar); ssmKey != "" {
-		sess, err := session.NewSession(&aws.Config{Region: aws.String(awsRegion)})
-		if err != nil {
-			return nil, err
-		}
-		client := ssm.New(sess)
-		return token.NewSSM(client, ssmKey)
-	}
-
-	if secretsManagerSecretID := os.Getenv(BKAgentTokenSecretsManagerSecretIDEnvVar); secretsManagerSecretID != "" {
-		jsonKey := os.Getenv(BKAgentTokenSecretsManagerJSONKeyEnvVar)
-		sess, err := session.NewSession(&aws.Config{Region: aws.String(awsRegion)})
-		if err != nil {
-			return nil, err
-		}
-		client := secretsmanager.New(sess)
-		if jsonKey == "" {
-			return token.NewSecretsManager(client, secretsManagerSecretID)
-		} else {
-			return token.NewSecretsManager(client, secretsManagerSecretID, token.WithSecretsManagerJSONSecret(jsonKey))
-		}
-	}
-
-	return nil, fmt.Errorf("failed to initialize Buildkite token provider: one of the [%s] environment variables "+
-		"must be provided", strings.Join(mutuallyExclusiveEnvVars, ","))
-}
-
-func checkMutuallyExclusiveEnvVars(varNames ...string) error {
-	foundVars := make([]string, 0)
-	for _, varName := range varNames {
-		value := os.Getenv(varName)
-		if value != "" {
-			foundVars = append(foundVars, value)
-		}
-	}
-	if len(foundVars) > 1 {
-		return fmt.Errorf("the environment variables [%s] are mutually exclusive", strings.Join(foundVars, ","))
-	}
-	return nil
-}

main.go

@@ -9,23 +9,27 @@ import (
 	"strings"
 	"time"
 
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/secretsmanager"
+	"github.com/aws/aws-sdk-go/service/ssm"
 	"github.com/buildkite/buildkite-agent-metrics/backend"
 	"github.com/buildkite/buildkite-agent-metrics/collector"
+	"github.com/buildkite/buildkite-agent-metrics/token"
 	"github.com/buildkite/buildkite-agent-metrics/version"
 )
 
 var bk backend.Backend
 
 func main() {
 	var (
-		token       = flag.String("token", "", "A Buildkite Agent Registration Token")
-		interval    = flag.Duration("interval", 0, "Update metrics every interval, rather than once")
-		showVersion = flag.Bool("version", false, "Show the version")
-		quiet       = flag.Bool("quiet", false, "Only print errors")
-		debug       = flag.Bool("debug", false, "Show debug output")
-		debugHttp   = flag.Bool("debug-http", false, "Show full http traces")
-		dryRun      = flag.Bool("dry-run", false, "Whether to only print metrics")
-		endpoint    = flag.String("endpoint", "https://agent.buildkite.com/v3", "A custom Buildkite Agent API endpoint")
+		tokenFromFlag = flag.String("token", "", "A Buildkite Agent Registration Token")
+		interval      = flag.Duration("interval", 0, "Update metrics every interval, rather than once")
+		showVersion   = flag.Bool("version", false, "Show the version")
+		quiet         = flag.Bool("quiet", false, "Only print errors")
+		debug         = flag.Bool("debug", false, "Show debug output")
+		debugHttp     = flag.Bool("debug-http", false, "Show full http traces")
+		dryRun        = flag.Bool("dry-run", false, "Whether to only print metrics")
+		endpoint      = flag.String("endpoint", "https://agent.buildkite.com/v3", "A custom Buildkite Agent API endpoint")
 
 		// backend config
 		backendOpt     = flag.String("backend", "cloudwatch", "Specify the backend to use: cloudwatch, statsd, prometheus, stackdriver")
@@ -51,16 +55,23 @@ func main() {
 		os.Exit(0)
 	}
 
-	if *token == "" {
-		if bkToken := os.Getenv("BUILDKITE_AGENT_TOKEN"); bkToken != "" {
-			*token = bkToken
-		} else {
-			fmt.Println("Must provide a token")
-			os.Exit(1)
-		}
+	chain := token.NewProviderChain(
+		token.NewStaticCredentialsProvider(tokenFromFlag),
+		token.NewEnvVar(),
+	)
+
+	sess, err := session.NewSession()
+	if err == nil {
+		chain.Append(token.NewSSM(ssm.New(sess)))
+		chain.Append(token.NewSecretsManager(secretsmanager.New(sess)))
+	}
+
+	token, err := chain.Resolve()
+	if err != nil {
+		fmt.Printf("Error: %v\n", err)
+		os.Exit(1)
 	}
 
-	var err error
 	switch strings.ToLower(*backendOpt) {
 	case "cloudwatch":
 		region := *clwRegion
@@ -95,6 +106,8 @@ func main() {
 			fmt.Printf("Error starting New Relic client: %v\n", err)
 			os.Exit(1)
 		}
+	case "stdout":
+		bk = backend.NewStdoutBackend()
 	default:
 		fmt.Println("Must provide a supported backend: cloudwatch, statsd, prometheus, stackdriver, newrelic")
 		os.Exit(1)
@@ -112,7 +125,7 @@ func main() {
 	c := collector.Collector{
 		UserAgent: userAgent,
 		Endpoint:  *endpoint,
-		Token:     *token,
+		Token:     token,
 		Queues:    []string(queues),
 		Quiet:     *quiet,
 		Debug:     *debug,

token/env_var.go

@@ -0,0 +1,29 @@
+package token
+
+import (
+	"os"
+)
+
+type envVarProvider struct {
+	// The token value to provide on each Get call.
+	Token string
+}
+
+const TokenEnvVarKey = "BUILDKITE_AGENT_TOKEN"
+
+// NewEnvVar constructs a Buildkite API token provider backed by the environment variable BUILDKITE_AGENT_TOKEN
+func NewEnvVar() Provider {
+	return &envVarProvider{Token: os.Getenv(TokenEnvVarKey)}
+}
+
+func (p envVarProvider) Get() (string, error) {
+	if p.Token == "" {
+		return "", errNotUsable
+	}
+
+	return p.Token, nil
+}
+
+func (p envVarProvider) String() string {
+	return "EnvVarProvider"
+}

token/env_var_test.go

@@ -0,0 +1,29 @@
+package token
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+const (
+	testToken = "some-token"
+)
+
+func TestEnvVarProvider_Get(t *testing.T) {
+	t.Setenv(TokenEnvVarKey, testToken)
+
+	provider := NewEnvVar()
+	token, err := provider.Get()
+	assert.NoError(t, err)
+
+	assert.Equal(t, testToken, token)
+}
+
+func TestEnvVarProvider_WithNoEnvVar_IsUnusable(t *testing.T) {
+	provider := NewEnvVar()
+	_, err := provider.Get()
+
+	assert.Error(t, err)
+	assert.ErrorIs(t, err, errNotUsable)
+}