Merge 95f979a into 1ca51d5

Author: RomainMuller

ddlambda.go

@@ -96,18 +96,30 @@ 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"
 )
 
 // 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 +301,39 @@ 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 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))
+		}
+	}
+}