feat: automate AppSec enablement setup (e.g: AWS_LAMBDA_RUNTIME_API)

awslambdanorpc.go

@@ -0,0 +1,14 @@
+//go:build lambda.norpc
+// +build lambda.norpc
+
+/*
+ * Unless explicitly stated otherwise all files in this repository are licensed
+ * under the Apache License Version 2.0.
+ *
+ * This product includes software developed at Datadog (https://www.datadoghq.com/).
+ * Copyright 2021 Datadog, Inc.
+ */
+
+package ddlambda
+
+const awsLambdaRpcSupport = false

awslambdawithrpc.go

@@ -0,0 +1,14 @@
+//go:build !lambda.norpc
+// +build !lambda.norpc
+
+/*
+* Unless explicitly stated otherwise all files in this repository are licensed
+* under the Apache License Version 2.0.
+*
+* This product includes software developed at Datadog (https://www.datadoghq.com/).
+* Copyright 2021 Datadog, Inc.
+ */
+
+package ddlambda
+
+const awsLambdaRpcSupport = true

ddlambda.go

@@ -96,18 +96,35 @@ const (
 	DefaultSite = "datadoghq.com"
 	// DefaultEnhancedMetrics enables enhanced metrics by default.
 	DefaultEnhancedMetrics = true
+
+	// serverlessAppSecEnabledEnvVar is the environment variable used to activate Serverless ASM through the use of an
+	// AWS Lambda runtime API proxy.
+	serverlessAppSecEnabledEnvVar = "DD_SERVERLESS_APPSEC_ENABLED"
+	// awsLambdaRuntimeApiEnvVar is the environment variable used to redirect AWS Lambda runtime API calls to the proxy.
+	awsLambdaRuntimeApiEnvVar = "AWS_LAMBDA_RUNTIME_API"
+	// datadogAgentUrl is the URL of the agent and proxy started by the Datadog lambda extension.
+	datadogAgentUrl = "127.0.0.1:9000"
+	// ddExtensionFilePath is the path on disk of the datadog lambda extension.
+	ddExtensionFilePath = "/opt/extensions/datadog-agent"
+
+	// awsLambdaServerPortEnvVar is the environment variable set by the go1.x Lambda Runtime to indicate which port the
+	// RCP server should listen on. This is used as a sign that a warning should be printed if customers want to enable
+	// ASM support, but did not enable the lambda.norpc build taf.
+	awsLambdaServerPortEnvVar = "_LAMBDA_SERVER_PORT"
 )
 
 // WrapLambdaHandlerInterface is used to instrument your lambda functions.
 // It returns a modified handler that can be passed directly to the lambda.StartHandler function from aws-lambda-go.
 func WrapLambdaHandlerInterface(handler lambda.Handler, cfg *Config) lambda.Handler {
+	setupAppSec()
 	listeners := initializeListeners(cfg)
 	return wrapper.WrapHandlerInterfaceWithListeners(handler, listeners...)
 }
 
 // WrapFunction is used to instrument your lambda functions.
 // It returns a modified handler that can be passed directly to the lambda.Start function from aws-lambda-go.
 func WrapFunction(handler interface{}, cfg *Config) interface{} {
+	setupAppSec()
 	listeners := initializeListeners(cfg)
 	return wrapper.WrapHandlerWithListeners(handler, listeners...)
 }
@@ -289,3 +306,45 @@ func (cfg *Config) toMetricsConfig(isExtensionRunning bool) metrics.Config {
 
 	return mc
 }
+
+// setupAppSec checks if DD_SERVERLESS_APPSEC_ENABLED is set (to true) and when that
+// is the case, redirects `AWS_LAMBDA_RUNTIME_API` to the agent extension, and turns
+// on universal instrumentation unless it was already configured by the customer, so
+// that the HTTP context (invocation details span tags) is available on AppSec traces.
+func setupAppSec() {
+	enabled := false
+	if env := os.Getenv(serverlessAppSecEnabledEnvVar); env != "" {
+		if on, err := strconv.ParseBool(env); err == nil {
+			enabled = on
+		}
+	}
+
+	if !enabled {
+		return
+	}
+
+	if _, err := os.Stat(ddExtensionFilePath); os.IsNotExist(err) {
+		logger.Debug(fmt.Sprintf("%s is enabled, but the Datadog extension was not found at %s", serverlessAppSecEnabledEnvVar, ddExtensionFilePath))
+		return
+	}
+
+	if awsLambdaRpcSupport {
+		if port := os.Getenv(awsLambdaServerPortEnvVar); port != "" {
+			logger.Warn(fmt.Sprintf("%s activation with the go1.x AWS Lambda runtime requires setting the `lambda.norpc` go build tag", serverlessAppSecEnabledEnvVar))
+		}
+	}
+
+	if err := os.Setenv(awsLambdaRuntimeApiEnvVar, datadogAgentUrl); err != nil {
+		logger.Debug(fmt.Sprintf("failed to set %s=%s: %v", awsLambdaRuntimeApiEnvVar, datadogAgentUrl, err))
+	} else {
+		logger.Debug(fmt.Sprintf("successfully set %s=%s", awsLambdaRuntimeApiEnvVar, datadogAgentUrl))
+	}
+
+	if val := os.Getenv(UniversalInstrumentation); val == "" {
+		if err := os.Setenv(UniversalInstrumentation, "1"); err != nil {
+			logger.Debug(fmt.Sprintf("failed to set %s=%d: %v", UniversalInstrumentation, 1, err))
+		} else {
+			logger.Debug(fmt.Sprintf("successfully set %s=%d", UniversalInstrumentation, 1))
+		}
+	}
+}

internal/logger/log.go

@@ -14,12 +14,12 @@ type LogLevel int
 const (
 	// LevelDebug logs all information
 	LevelDebug LogLevel = iota
-	// LevelError only logs errors
-	LevelError LogLevel = iota
+	// LevelWarn only logs warnings and errors
+	LevelWarn LogLevel = iota
 )
 
 var (
-	logLevel           = LevelError
+	logLevel           = LevelWarn
 	output   io.Writer = os.Stdout
 )
 
@@ -36,12 +36,6 @@ func SetOutput(w io.Writer) {
 
 // Error logs a structured error message to stdout
 func Error(err error) {
-
-	type logStructure struct {
-		Status  string `json:"status"`
-		Message string `json:"message"`
-	}
-
 	finalMessage := logStructure{
 		Status:  "error",
 		Message: fmt.Sprintf("datadog: %s", err.Error()),
@@ -56,10 +50,6 @@ func Debug(message string) {
 	if logLevel > LevelDebug {
 		return
 	}
-	type logStructure struct {
-		Status  string `json:"status"`
-		Message string `json:"message"`
-	}
 	finalMessage := logStructure{
 		Status:  "debug",
 		Message: fmt.Sprintf("datadog: %s", message),
@@ -70,7 +60,27 @@ func Debug(message string) {
 	log.Println(string(result))
 }
 
+// Warn logs a structured log message to stdout
+func Warn(message string) {
+	if logLevel > LevelWarn {
+		return
+	}
+	finalMessage := logStructure{
+		Status:  "warning",
+		Message: fmt.Sprintf("datadog: %s", message),
+	}
+
+	result, _ := json.Marshal(finalMessage)
+
+	log.Println(string(result))
+}
+
 // Raw prints a raw message to the logs.
 func Raw(message string) {
 	fmt.Fprintln(output, message)
 }
+
+type logStructure struct {
+	Status  string `json:"status"`
+	Message string `json:"message"`
+}