[Serverless] Fix KMS key detection (#20382)

cmd/serverless/main.go

@@ -10,7 +10,6 @@ import (
 	"os"
 	"os/signal"
 	"strconv"
-	"strings"
 	"sync"
 	"syscall"
 	"time"
@@ -20,9 +19,11 @@ import (
 	configUtils "github.com/DataDog/datadog-agent/pkg/config/utils"
 	pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/trace"
 	"github.com/DataDog/datadog-agent/pkg/serverless"
+	"github.com/DataDog/datadog-agent/pkg/serverless/apikey"
 	"github.com/DataDog/datadog-agent/pkg/serverless/appsec"
 	"github.com/DataDog/datadog-agent/pkg/serverless/appsec/httpsec"
 	"github.com/DataDog/datadog-agent/pkg/serverless/daemon"
+	"github.com/DataDog/datadog-agent/pkg/serverless/debug"
 	"github.com/DataDog/datadog-agent/pkg/serverless/flush"
 	"github.com/DataDog/datadog-agent/pkg/serverless/invocationlifecycle"
 	serverlessLogs "github.com/DataDog/datadog-agent/pkg/serverless/logs"
@@ -38,12 +39,9 @@ import (
 )
 
 var (
-	kmsAPIKeyEnvVar            = "DD_KMS_API_KEY"
-	secretsManagerAPIKeyEnvVar = "DD_API_KEY_SECRET_ARN"
-	apiKeyEnvVar               = "DD_API_KEY"
-	logLevelEnvVar             = "DD_LOG_LEVEL"
-	flushStrategyEnvVar        = "DD_SERVERLESS_FLUSH_STRATEGY"
-	logsLogsTypeSubscribed     = "DD_LOGS_CONFIG_LAMBDA_LOGS_TYPE"
+	logLevelEnvVar         = "DD_LOG_LEVEL"
+	flushStrategyEnvVar    = "DD_SERVERLESS_FLUSH_STRATEGY"
+	logsLogsTypeSubscribed = "DD_LOGS_CONFIG_LAMBDA_LOGS_TYPE"
 
 	// AWS Lambda is writing the Lambda function files in /var/task, we want the
 	// configuration file to be at the root of this directory.
@@ -53,8 +51,6 @@ var (
 const (
 	loggerName config.LoggerName = "DD_EXTENSION"
 
-	runtimeAPIEnvVar = "AWS_LAMBDA_RUNTIME_API"
-
 	extensionRegistrationRoute   = "/2020-01-01/extension/register"
 	extensionRegistrationTimeout = 5 * time.Second
 
@@ -116,9 +112,9 @@ func runAgent(stopCh chan struct{}) (serverlessDaemon *daemon.Daemon, err error)
 		}
 	}
 
-	outputDatadogEnvVariablesForDebugging()
+	debug.OutputDatadogEnvVariablesForDebugging()
 
-	if !hasApiKey() {
+	if !apikey.HasAPIKey() {
 		log.Errorf("Can't start the Datadog extension as no API Key has been detected, or API Key could not be decrypted. Data will not be sent to Datadog.")
 		// we still need to register the extension but let's return after (no-op)
 		id, _, registrationError := registration.RegisterExtension(extensionRegistrationRoute, extensionRegistrationTimeout)
@@ -165,28 +161,13 @@ func runAgent(stopCh chan struct{}) (serverlessDaemon *daemon.Daemon, err error)
 	// KMS > Secrets Manager > Plaintext API key
 	// If one is set but failing, the next will be tried
 
-	// some useful warnings first
-
-	var apikeySetIn = []string{}
-	if os.Getenv(kmsAPIKeyEnvVar) != "" {
-		apikeySetIn = append(apikeySetIn, "KMS")
-	}
-	if os.Getenv(secretsManagerAPIKeyEnvVar) != "" {
-		apikeySetIn = append(apikeySetIn, "SSM")
-	}
-	if os.Getenv(apiKeyEnvVar) != "" {
-		apikeySetIn = append(apikeySetIn, "environment variable")
-	}
-
-	if len(apikeySetIn) > 1 {
-		log.Warn("An API Key has been set in multiple places:", strings.Join(apikeySetIn, ", "))
-	}
+	apikey.CheckForSingleAPIKey()
 
 	config.LoadProxyFromEnv(config.Datadog)
 
 	// Set secrets from the environment that are suffixed with
 	// KMS_ENCRYPTED or SECRET_ARN
-	setSecretsFromEnv(os.Environ())
+	apikey.SetSecretsFromEnv(os.Environ())
 
 	// adaptive flush configuration
 	if v, exists := os.LookupEnv(flushStrategyEnvVar); exists {

cmd/serverless/api_key.go → pkg/serverless/apikey/api_key.go

@@ -3,14 +3,15 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/).
 // Copyright 2016-present Datadog, Inc.
 
-package main
+package apikey
 
 import (
 	"encoding/base64"
 	"fmt"
 	"net/http"
 	"os"
 	"regexp"
+	"strings"
 
 	"github.com/DataDog/datadog-agent/pkg/config"
 	datadogHttp "github.com/DataDog/datadog-agent/pkg/util/http"
@@ -30,6 +31,12 @@ const encryptionContextKey = "LambdaFunctionName"
 // functionNameEnvVar is the environment variable that stores the function name.
 const functionNameEnvVar = "AWS_LAMBDA_FUNCTION_NAME"
 
+// one of those env variable must be set
+const apiKeyEnvVar = "DD_API_KEY"
+const apiKeySecretManagerEnvVar = "DD_API_KEY_SECRET_ARN"
+const apiKeyKmsEnvVar = "DD_KMS_API_KEY"
+const apiKeyKmsEncryptedEnvVar = "DD_API_KEY_KMS_ENCRYPTED"
+
 // kmsKeySuffix is the suffix of all environment variables which should be decrypted by KMS
 const kmsKeySuffix = "_KMS_ENCRYPTED"
 
@@ -129,11 +136,11 @@ func readAPIKeyFromSecretsManager(arn string) (string, error) {
 		return secretString, nil
 	} else if output.SecretBinary != nil {
 		decodedBinarySecretBytes := make([]byte, base64.StdEncoding.DecodedLen(len(output.SecretBinary)))
-		len, err := base64.StdEncoding.Decode(decodedBinarySecretBytes, output.SecretBinary)
+		secretLen, err := base64.StdEncoding.Decode(decodedBinarySecretBytes, output.SecretBinary)
 		if err != nil {
 			return "", fmt.Errorf("Can't base64 decode Secrets Manager secret: %s", err)
 		}
-		return string(decodedBinarySecretBytes[:len]), nil
+		return string(decodedBinarySecretBytes[:secretLen]), nil
 	}
 	// should not happen but let's handle this gracefully
 	log.Warn("Secrets Manager returned something but there seems to be no data available")
@@ -156,8 +163,32 @@ func extractRegionFromSecretsManagerArn(secretsManagerArn string) (string, error
 	return arnObject.Region, nil
 }
 
-func hasApiKey() bool {
+// HasAPIKey returns true if an API key has been set in any of the supported ways.
+func HasAPIKey() bool {
 	return config.Datadog.IsSet("api_key") ||
-		len(os.Getenv(kmsAPIKeyEnvVar)) > 0 ||
-		len(os.Getenv(secretsManagerAPIKeyEnvVar)) > 0
+		len(os.Getenv(apiKeyKmsEncryptedEnvVar)) > 0 ||
+		len(os.Getenv(apiKeyKmsEnvVar)) > 0 ||
+		len(os.Getenv(apiKeySecretManagerEnvVar)) > 0 ||
+		len(os.Getenv(apiKeyEnvVar)) > 0
+}
+
+// CheckForSingleAPIKey checks if an API key has been set in multiple places and logs a warning if so.
+func CheckForSingleAPIKey() {
+	var apikeySetIn = []string{}
+	if len(os.Getenv(apiKeyKmsEncryptedEnvVar)) > 0 {
+		apikeySetIn = append(apikeySetIn, "KMS_ENCRYPTED")
+	}
+	if len(os.Getenv(apiKeyKmsEnvVar)) > 0 {
+		apikeySetIn = append(apikeySetIn, "KMS")
+	}
+	if len(os.Getenv(apiKeySecretManagerEnvVar)) > 0 {
+		apikeySetIn = append(apikeySetIn, "SSM")
+	}
+	if len(os.Getenv(apiKeyEnvVar)) > 0 {
+		apikeySetIn = append(apikeySetIn, "environment variable")
+	}
+
+	if len(apikeySetIn) > 1 {
+		log.Warn("An API Key has been set in multiple places:", strings.Join(apikeySetIn, ", "))
+	}
 }

cmd/serverless/api_key_test.go → pkg/serverless/apikey/api_key_test.go

@@ -3,7 +3,8 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/).
 // Copyright 2016-present Datadog, Inc.
 
-package main
+// Package apikey holds logic around api keys in the Lambda Extension
+package apikey
 
 import (
 	"bytes"
@@ -105,22 +106,33 @@ func TestExtractRegionFromMalformedPrefixSecretsManagerArnPrefix(t *testing.T) {
 
 func TestDDApiKey(t *testing.T) {
 	t.Setenv("DD_API_KEY", "abc")
-	assert.True(t, hasApiKey())
+	assert.True(t, HasAPIKey())
 }
 
 func TestHasDDApiKeySecretArn(t *testing.T) {
 	t.Setenv("DD_API_KEY_SECRET_ARN", "abc")
-	assert.True(t, hasApiKey())
+	assert.True(t, HasAPIKey())
 }
 
-func TestHasDDKmsApiKey(t *testing.T) {
+func TestHasDDKmsApiKeyEncrypted(t *testing.T) {
+	t.Setenv("DD_API_KEY_KMS_ENCRYPTED", "abc")
+	assert.True(t, HasAPIKey())
+}
+
+func TestHasDDKmspiKey(t *testing.T) {
 	t.Setenv("DD_KMS_API_KEY", "abc")
-	assert.True(t, hasApiKey())
+	assert.True(t, HasAPIKey())
+}
+
+func TestHasAPIKey(t *testing.T) {
+	t.Setenv("DD_API_KEY", "abc")
+	assert.True(t, HasAPIKey())
 }
 
 func TestHasNoKeys(t *testing.T) {
-	t.Setenv("DD_KMS_API_KEY", "")
+	t.Setenv("DD_API_KEY_KMS_ENCRYPTED", "")
 	t.Setenv("DD_API_KEY_SECRET_ARN", "")
+	t.Setenv("DD_KMS_API_KEY", "")
 	t.Setenv("DD_API_KEY", "")
-	assert.False(t, hasApiKey())
+	assert.False(t, HasAPIKey())
 }

cmd/serverless/env.go → pkg/serverless/apikey/env.go

@@ -3,7 +3,7 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/).
 // Copyright 2022-present Datadog, Inc.
 
-package main
+package apikey
 
 import (
 	"os"
@@ -35,6 +35,14 @@ func getSecretEnvVars(envVars []string, kmsFunc decryptFunc, smFunc decryptFunc)
 			}
 			decryptedEnvVars[strings.TrimSuffix(envKey, kmsKeySuffix)] = secretVal
 		}
+		if envKey == apiKeyKmsEnvVar {
+			secretVal, err := kmsFunc(envVal)
+			if err != nil {
+				log.Debugf("Couldn't read API key from KMS: %v", err)
+				continue
+			}
+			decryptedEnvVars[apiKeyEnvVar] = secretVal
+		}
 		if strings.HasSuffix(envKey, secretArnSuffix) {
 			log.Debugf("Retrieving %v from secrets manager", envVar)
 			secretVal, err := smFunc(envVal)
@@ -48,15 +56,15 @@ func getSecretEnvVars(envVars []string, kmsFunc decryptFunc, smFunc decryptFunc)
 	return decryptedEnvVars
 }
 
-// The agent is going to get any environment variables ending with _KMS_ENCRYPTED and _SECRET_ARN,
+// SetSecretsFromEnv - The agent is going to get any environment variables ending with _KMS_ENCRYPTED and _SECRET_ARN,
 // get the contents of the environment variable, and query SM/KMS to retrieve the value. This allows us
 // to read arbitrarily encrypted environment variables and use the decrypted version in the extension.
 // Right now, this feature is used for dual shipping, since customers need to set DD_LOGS_CONFIGURATION
 // and a few other variables, which include an API key. The user can set DD_LOGS_CONFIGURATION_SECRET_ARN
 // or DD_LOGS_CONFIGURATION_KMS_ENCRYPTED, which will get converted in the extension to a plaintext
 // DD_LOGS_CONFIGURATION, and will have dual shipping enabled without exposing
 // their API key in plaintext through environment variables.
-func setSecretsFromEnv(envVars []string) {
+func SetSecretsFromEnv(envVars []string) {
 	for envKey, envVal := range getSecretEnvVars(envVars, readAPIKeyFromKMS, readAPIKeyFromSecretsManager) {
 		os.Setenv(envKey, envVal)
 	}

cmd/serverless/env_test.go → pkg/serverless/apikey/env_test.go

@@ -3,7 +3,7 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/).
 // Copyright 2022-present Datadog, Inc.
 
-package main
+package apikey
 
 import (
 	"testing"
@@ -23,12 +23,14 @@ func TestGetSecretEnvVars(t *testing.T) {
 		"TEST_SM_SECRET_ARN=123",
 		"",
 		"MALFORMED=ENV=VAR",
+		"DD_KMS_API_KEY=123",
 	}
 
 	decryptedEnvVars := getSecretEnvVars(testEnvVars, getFunc, getFunc)
 
 	assert.Equal(t, map[string]string{
-		"TEST_KMS": "DECRYPTED_VAL",
-		"TEST_SM":  "DECRYPTED_VAL",
+		"TEST_KMS":   "DECRYPTED_VAL",
+		"TEST_SM":    "DECRYPTED_VAL",
+		"DD_API_KEY": "DECRYPTED_VAL",
 	}, decryptedEnvVars)
 }

cmd/serverless/debug.go → pkg/serverless/debug/debug.go

@@ -3,7 +3,8 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/).
 // Copyright 2016-present Datadog, Inc.
 
-package main
+// Package debug holds logic around debug information in the Lambda Extension
+package debug
 
 import (
 	"fmt"
@@ -17,7 +18,8 @@ import (
 
 const ddPrefix = "DD_"
 
-func outputDatadogEnvVariablesForDebugging() {
+// OutputDatadogEnvVariablesForDebugging outputs the Datadog environment variables for debugging purposes
+func OutputDatadogEnvVariablesForDebugging() {
 	log.Debug(buildDebugString())
 }
 

cmd/serverless/debug_test.go → pkg/serverless/debug/debug_test.go

@@ -3,7 +3,7 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/).
 // Copyright 2016-present Datadog, Inc.
 
-package main
+package debug
 
 import (
 	"testing"