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

ddlambda.go

@@ -101,13 +101,15 @@ const (
 // 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 +291,30 @@ 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 constext (invocation details span tags) is avaialble on AppSec traces.
+func setupAppSec() {
+	const ServerlessAppSecEnabledEnvVar = "DD_SERVERLESS_APPSEC_ENABLED"
+	const AwsLambdaRuntimeApiEnvVar = "AWS_LAMBDA_RUNTIME_API"
+	const DatadogAgentUrl = "127.0.0.1:9000"
+
+	enabled := false
+	if env := os.Getenv(ServerlessAppSecEnabledEnvVar); env != "" {
+		if on, err := strconv.ParseBool(env); err == nil {
+			enabled = on
+		}
+	}
+
+	if !enabled {
+		return
+	}
+
+	_ = os.Setenv(AwsLambdaRuntimeApiEnvVar, DatadogAgentUrl)
+
+	if val := os.Getenv(UniversalInstrumentation); val == "" {
+		_ = os.Setenv(UniversalInstrumentation, "1")
+	}
+}