Add support for DD_API_KEY_SECRET_ARN environment variable

README.md

@@ -83,11 +83,28 @@ A more complete example can be found in the `ddlambda_example_test.go` file.
 
 ### DD_FLUSH_TO_LOG
 
-Set to `true` (recommended) to send custom metrics asynchronously (with no added latency to your Lambda function executions) through CloudWatch Logs with the help of [Datadog Forwarder](https://github.com/DataDog/datadog-serverless-functions/tree/master/aws/logs_monitoring). Defaults to `false`. If set to `false`, you also need to set `DD_API_KEY` and `DD_SITE`.
+Set to `true` (recommended) to send custom metrics asynchronously (with no added latency to your Lambda function executions) through CloudWatch Logs with the help of [Datadog Forwarder](https://github.com/DataDog/datadog-serverless-functions/tree/master/aws/logs_monitoring). Defaults to `false`. If set to `false`, you also need to set `DD_SITE` and one of the
+following `DD_API_KEY`, `DD_KMS_API_KEY` or `DD_API_KEY_SECRET_ARN`.
 
 ### DD_API_KEY
 
-If `DD_FLUSH_TO_LOG` is set to `false` (not recommended), the Datadog API Key must be defined.
+If `DD_FLUSH_TO_LOG` is set to `false` (not recommended), the Datadog API key must be
+set to `DD_API_KEY`, or either `DD_KMS_API_KEY` or `DD_API_KEY_SECRET_ARN` below must
+be defined instead.
+
+### DD_KMS_API_KEY
+
+If `DD_FLUSH_TO_LOG` is set to `false` (not recommended) and `DD_API_KEY` is not set,
+the Datadog API key encrypted using the AWS Key Management Service (KMS) must be set to
+`DD_KMS_API_KEY`, or the following `DD_API_KEY_SECRET_ARN` must be defined instead.
+The encryption key used to encrypt the API key must be a symmetric KMS key.
+
+### DD_API_KEY_SECRET_ARN
+
+If `DD_FLUSH_TO_LOG` is set to `false` (not recommended) and neither `DD_API_KEY` nor
+`DD_KMS_API_KEY` is set, the ARN of an AWS Secrets Manager secret where the Datadog API
+key is stored must be set to `DD_API_KEY_SECRET_ARN`. The secret value must be just the
+API key string itself (no double quotes), not a JSON object.
 
 ### DD_SITE
 

ddlambda.go

@@ -33,6 +33,9 @@ type (
 		APIKey string
 		// KMSAPIKey is your Datadog API key, encrypted using the AWS KMS service. This is used for sending metrics.
 		KMSAPIKey string
+		// APIKeySecretARN is the ARN of an AWS Secrets Manager secret where your Datadog API key is stored.
+		// This is used for sending metrics.
+		APIKeySecretARN string
 		// ShouldRetryOnFailure is used to turn on retry logic when sending metrics via the API. This can negatively effect the performance of your lambda,
 		// and should only be turned on if you can't afford to lose metrics data under poor network conditions.
 		ShouldRetryOnFailure bool
@@ -79,6 +82,8 @@ const (
 	DatadogAPIKeyEnvVar = "DD_API_KEY"
 	// DatadogKMSAPIKeyEnvVar is the environment variable that will be sent to KMS for decryption, then used as an API key.
 	DatadogKMSAPIKeyEnvVar = "DD_KMS_API_KEY"
+	// DatadogAPIKeySecretARNEnvVar is the environment variable that will be used to retreive the API key from AWS Secrets Manager.
+	DatadogAPIKeySecretARNEnvVar = "DD_API_KEY_SECRET_ARN"
 	// DatadogSiteEnvVar is the environment variable that will be used as the API host.
 	DatadogSiteEnvVar = "DD_SITE"
 	// LogLevelEnvVar is the environment variable that will be used to set the log level.
@@ -234,6 +239,7 @@ func (cfg *Config) toMetricsConfig() metrics.Config {
 		mc.ShouldRetryOnFailure = cfg.ShouldRetryOnFailure
 		mc.APIKey = cfg.APIKey
 		mc.KMSAPIKey = cfg.KMSAPIKey
+		mc.APIKeySecretARN = cfg.APIKeySecretARN
 		mc.Site = cfg.Site
 		mc.ShouldUseLogForwarder = cfg.ShouldUseLogForwarder
 		mc.HttpClientTimeout = cfg.HttpClientTimeout
@@ -258,13 +264,17 @@ func (cfg *Config) toMetricsConfig() metrics.Config {
 
 	if mc.APIKey == "" {
 		mc.APIKey = os.Getenv(DatadogAPIKeyEnvVar)
-
 	}
 	if mc.KMSAPIKey == "" {
 		mc.KMSAPIKey = os.Getenv(DatadogKMSAPIKeyEnvVar)
 	}
-	if mc.APIKey == "" && mc.KMSAPIKey == "" && !mc.ShouldUseLogForwarder {
-		logger.Error(fmt.Errorf("couldn't read DD_API_KEY or DD_KMS_API_KEY from environment"))
+	if mc.APIKeySecretARN == "" {
+		mc.APIKeySecretARN = os.Getenv(DatadogAPIKeySecretARNEnvVar)
+	}
+	if mc.APIKey == "" && mc.KMSAPIKey == "" && mc.APIKeySecretARN == "" && !mc.ShouldUseLogForwarder {
+		logger.Error(fmt.Errorf(
+			"couldn't read DD_API_KEY, DD_KMS_API_KEY or DD_API_KEY_SECRET_ARN from environment",
+		))
 	}
 
 	enhancedMetrics := os.Getenv("DD_ENHANCED_METRICS")

internal/metrics/api.go

@@ -28,11 +28,11 @@ type (
 
 	// APIClient send metrics to Datadog, via the Datadog API
 	APIClient struct {
-		apiKey            string
-		apiKeyDecryptChan <-chan string
-		baseAPIURL        string
-		httpClient        *http.Client
-		context           context.Context
+		apiKey             string
+		apiKeyLazyLoadChan <-chan string
+		baseAPIURL         string
+		httpClient         *http.Client
+		context            context.Context
 	}
 
 	// APIClientOptions contains instantiation options from creating an APIClient.
@@ -41,6 +41,8 @@ type (
 		apiKey            string
 		kmsAPIKey         string
 		decrypter         Decrypter
+		apiKeySecretARN   string
+		secretFetcher     SecretFetcher
 		httpClientTimeout time.Duration
 	}
 
@@ -60,20 +62,58 @@ func MakeAPIClient(ctx context.Context, options APIClientOptions) *APIClient {
 		httpClient: httpClient,
 		context:    ctx,
 	}
-	if len(options.apiKey) == 0 && len(options.kmsAPIKey) != 0 {
-		client.apiKeyDecryptChan = client.decryptAPIKey(options.decrypter, options.kmsAPIKey)
+	if options.apiKey != "" {
+		return client
+	}
+
+	if options.kmsAPIKey != "" {
+		client.apiKeyLazyLoadChan = decryptAPIKey(options.decrypter, options.kmsAPIKey)
+	} else if options.apiKeySecretARN != "" {
+		client.apiKeyLazyLoadChan = fetchAPIKey(options.secretFetcher, options.apiKeySecretARN)
 	}
 
 	return client
 }
 
+func decryptAPIKey(decrypter Decrypter, kmsAPIKey string) <-chan string {
+	ch := make(chan string)
+
+	go func() {
+		result, err := decrypter.Decrypt(kmsAPIKey)
+		if err != nil {
+			logger.Error(fmt.Errorf("Couldn't decrypt api kms key: %s", err))
+		}
+
+		ch <- result
+		close(ch)
+	}()
+
+	return ch
+}
+
+func fetchAPIKey(fetcher SecretFetcher, apiKeySecretARN string) <-chan string {
+	ch := make(chan string)
+
+	go func() {
+		result, err := fetcher.FetchSecret(apiKeySecretARN)
+		if err != nil {
+			logger.Error(fmt.Errorf("Couldn't retrieve api key secret: %s", err))
+		}
+
+		ch <- result
+		close(ch)
+	}()
+
+	return ch
+}
+
 // SendMetrics posts a batch metrics payload to the Datadog API
 func (cl *APIClient) SendMetrics(metrics []APIMetric) error {
 
 	// If the api key was provided as a kms key, wait for it to finish decrypting
-	if cl.apiKeyDecryptChan != nil {
-		cl.apiKey = <-cl.apiKeyDecryptChan
-		cl.apiKeyDecryptChan = nil
+	if cl.apiKeyLazyLoadChan != nil {
+		cl.apiKey = <-cl.apiKeyLazyLoadChan
+		cl.apiKeyLazyLoadChan = nil
 	}
 
 	content, err := marshalAPIMetricsModel(metrics)
@@ -118,21 +158,6 @@ func (cl *APIClient) SendMetrics(metrics []APIMetric) error {
 	return err
 }
 
-func (cl *APIClient) decryptAPIKey(decrypter Decrypter, kmsAPIKey string) <-chan string {
-
-	ch := make(chan string)
-
-	go func() {
-		result, err := decrypter.Decrypt(kmsAPIKey)
-		if err != nil {
-			logger.Error(fmt.Errorf("Couldn't decrypt api kms key %s", err))
-		}
-		ch <- result
-		close(ch)
-	}()
-	return ch
-}
-
 func (cl *APIClient) addAPICredentials(req *http.Request) {
 	query := req.URL.Query()
 	query.Add(apiKeyParam, cl.apiKey)

internal/metrics/api_test.go

@@ -22,6 +22,7 @@ const (
 	mockAPIKey          = "12345"
 	mockEncryptedAPIKey = "mockEncrypted"
 	mockDecryptedAPIKey = "mockDecrypted"
+	mockAPIKeySecretARN = "arn:aws:secretsmanager:us-east-1:123456789012:secret:apiKey"
 )
 
 type (
@@ -35,6 +36,15 @@ func (md *mockDecrypter) Decrypt(cipherText string) (string, error) {
 	return md.returnValue, md.returnError
 }
 
+type mockSecretFetcher struct {
+	returnValue string
+	returnError error
+}
+
+func (m *mockSecretFetcher) FetchSecret(secretID string) (string, error) {
+	return m.returnValue, m.returnError
+}
+
 func TestAddAPICredentials(t *testing.T) {
 	cl := MakeAPIClient(context.Background(), APIClientOptions{baseAPIURL: "", apiKey: mockAPIKey})
 	req, _ := http.NewRequest("GET", "http://some-api.com/endpoint", nil)
@@ -140,33 +150,57 @@ func TestSendMetricsCantReachServer(t *testing.T) {
 	assert.False(t, called)
 }
 
-func TestDecryptsUsingKMSKey(t *testing.T) {
-	called := false
-	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		called = true
-		assert.Equal(t, "/distribution_points?api_key=mockDecrypted", r.URL.String())
-	}))
-	defer server.Close()
-
-	am := []APIMetric{
-		{
-			Name:       "metric-1",
-			Host:       nil,
-			Tags:       []string{"a", "b", "c"},
-			MetricType: DistributionType,
-			Points: []interface{}{
-				[]interface{}{float64(1), []interface{}{float64(2)}},
-				[]interface{}{float64(3), []interface{}{float64(4)}},
-				[]interface{}{float64(5), []interface{}{float64(6)}},
+func TestLazyLoadAPIKey(t *testing.T) {
+	tests := map[string]struct {
+		clientOptions  APIClientOptions
+		expectedAPIKey string
+	}{
+		"decrypt using KMS key": {
+			clientOptions: APIClientOptions{
+				kmsAPIKey: mockEncryptedAPIKey,
+				decrypter: &mockDecrypter{returnValue: mockDecryptedAPIKey},
 			},
+			expectedAPIKey: mockDecryptedAPIKey,
+		},
+		"fetch from secret ARN": {
+			clientOptions: APIClientOptions{
+				apiKeySecretARN: mockAPIKeySecretARN,
+				secretFetcher:   &mockSecretFetcher{returnValue: mockAPIKey},
+			},
+			expectedAPIKey: mockAPIKey,
 		},
 	}
-	md := mockDecrypter{}
-	md.returnValue = mockDecryptedAPIKey
-
-	cl := MakeAPIClient(context.Background(), APIClientOptions{baseAPIURL: server.URL, apiKey: "", kmsAPIKey: mockEncryptedAPIKey, decrypter: &md})
-	err := cl.SendMetrics(am)
 
-	assert.NoError(t, err)
-	assert.True(t, called)
+	for name, tt := range tests {
+		tt := tt // https://go.dev/doc/faq#closures_and_goroutines
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			called := false
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				called = true
+				assert.Equal(t, "/distribution_points?api_key="+tt.expectedAPIKey, r.URL.String())
+			}))
+			defer server.Close()
+
+			am := []APIMetric{
+				{
+					Name:       "metric-1",
+					Host:       nil,
+					Tags:       []string{"a", "b", "c"},
+					MetricType: DistributionType,
+					Points: []interface{}{
+						[]interface{}{float64(1), []interface{}{float64(2)}},
+						[]interface{}{float64(3), []interface{}{float64(4)}},
+						[]interface{}{float64(5), []interface{}{float64(6)}},
+					},
+				},
+			}
+			tt.clientOptions.baseAPIURL = server.URL
+			err := MakeAPIClient(context.Background(), tt.clientOptions).SendMetrics(am)
+
+			assert.NoError(t, err)
+			assert.True(t, called)
+		})
+	}
 }

internal/metrics/kms_decrypter.go

@@ -13,7 +13,7 @@ import (
 	"os"
 
 	"github.com/DataDog/datadog-lambda-go/internal/logger"
-	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/aws/client"
 	"github.com/aws/aws-sdk-go/service/kms"
 	"github.com/aws/aws-sdk-go/service/kms/kmsiface"
 )
@@ -35,10 +35,10 @@ const functionNameEnvVar = "AWS_LAMBDA_FUNCTION_NAME"
 // encryptionContextKey is the key added to the encryption context by the Lambda console UI
 const encryptionContextKey = "LambdaFunctionName"
 
-// MakeKMSDecrypter creates a new decrypter which uses the AWS KMS service to decrypt variables
-func MakeKMSDecrypter() Decrypter {
+// MakeKMSDecrypter creates a new decrypter which uses the AWS KMS service to decrypt variables.
+func MakeKMSDecrypter(p client.ConfigProvider) Decrypter {
 	return &kmsDecrypter{
-		kmsClient: kms.New(session.New(nil)),
+		kmsClient: kms.New(p),
 	}
 }
 

internal/metrics/listener.go

@@ -18,6 +18,7 @@ import (
 	"time"
 
 	"github.com/aws/aws-lambda-go/lambdacontext"
+	"github.com/aws/aws-sdk-go/aws/session"
 
 	"github.com/DataDog/datadog-go/statsd"
 	"github.com/DataDog/datadog-lambda-go/internal/extension"
@@ -40,6 +41,7 @@ type (
 	Config struct {
 		APIKey                      string
 		KMSAPIKey                   string
+		APIKeySecretARN             string
 		Site                        string
 		ShouldRetryOnFailure        bool
 		ShouldUseLogForwarder       bool
@@ -61,12 +63,14 @@ type (
 
 // MakeListener initializes a new metrics lambda listener
 func MakeListener(config Config, extensionManager *extension.ExtensionManager) Listener {
-
+	sess := session.Must(session.NewSession())
 	apiClient := MakeAPIClient(context.Background(), APIClientOptions{
 		baseAPIURL:        config.Site,
 		apiKey:            config.APIKey,
-		decrypter:         MakeKMSDecrypter(),
+		decrypter:         MakeKMSDecrypter(sess),
 		kmsAPIKey:         config.KMSAPIKey,
+		secretFetcher:     MakeSecretsManagerSecretFetcher(sess),
+		apiKeySecretARN:   config.APIKeySecretARN,
 		httpClientTimeout: config.HttpClientTimeout,
 	})
 	if config.HttpClientTimeout <= 0 {