diff --git a/ddlambda.go b/ddlambda.go index 4867bbdc..da69d3a8 100644 --- a/ddlambda.go +++ b/ddlambda.go @@ -14,7 +14,6 @@ import ( "fmt" "net/http" "os" - "os/exec" "strconv" "strings" "time" @@ -102,16 +101,16 @@ 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) - applyLambdaExecWrapperConfiguration() 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) - applyLambdaExecWrapperConfiguration() return wrapper.WrapHandlerWithListeners(handler, listeners...) } @@ -293,63 +292,29 @@ func (cfg *Config) toMetricsConfig(isExtensionRunning bool) metrics.Config { return mc } -// applyLambdaExecWrapperConfiguration applies environment variables set by the wrapper configured in the -// `AWS_LAMBDA_EXEC_WRAPPER` environment variable if present. This is done here because the AWS Lambda runtimes used by -// go applications (`go1.x` and `provided.al2`) do not honor this setting, while other runtimes do. This assumes the -// wrapper script does nothing other than setting environment variables, and will only honor environment variables with -// names starting with `DD_` or `AWS_LAMBDA_`. In particular, `AWS_LAMBDA_RUNTIME_API` is used to re-route the Lambda -// control flow API through a proxy, which is used by ASM. -// -// See: https://docs.aws.amazon.com/lambda/latest/dg/runtimes-modify.html#runtime-wrapper. -// See: https://github.com/DataDog/datadog-lambda-extension/blob/main/scripts/datadog_wrapper. -func applyLambdaExecWrapperConfiguration() { - const AwsExecutionEnvEnvVar = "AWS_EXECUTION_ENV" - const AwsLambdaExecWrapperEnvVar = "AWS_LAMBDA_EXEC_WRAPPER" - - // Only perform this operation if the Lambda runtime doesn't do it by itself, meaning we only do something for - // provided runtimes (in which `AWS_EXECUTION_ENV` is not set) and `go1.x`. - if env := os.Getenv(AwsExecutionEnvEnvVar); env != "" && env != "AWS_Lambda_go1.x" { - logger.Debug(fmt.Sprintf("Skipping applyLambdaExecWrapperConfiguration, runtime is %s", env)) - return +// 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 ServerlessAppSSecEnabledEnvVar = "DD_SERVERLESS_APPSEC_ENABLED" + const AwsLambdaRuntimeApiEnvVar = "AWS_LAMBDA_RUNTIME_API" + const DatadogAgentUrl = "127.0.0.1:9000" + + enabled := false + if env := os.Getenv(ServerlessAppSSecEnabledEnvVar); env != "" { + if on, err := strconv.ParseBool(env); err == nil { + enabled = on + } } - script := os.Getenv(AwsLambdaExecWrapperEnvVar) - if script == "" { - // Nothing to do + if !enabled { return } - cmd := exec.Command("env", "-u", AwsLambdaExecWrapperEnvVar, script, "sh", "-c", "env") - logger.Debug(fmt.Sprintf("[%s] Command: %s", AwsLambdaExecWrapperEnvVar, cmd.String())) + _ = os.Setenv(AwsLambdaRuntimeApiEnvVar, DatadogAgentUrl) - if stdout, err := cmd.Output(); err != nil { - logger.Debug(fmt.Sprintf("[%s] Failed to run: %s", AwsLambdaExecWrapperEnvVar, err)) - } else { - for _, line := range strings.Split(string(stdout), "\n") { - parts := strings.SplitN(line, "=", 2) - if len(parts) != 2 { - // Forward what the wrapper script prints to the standard output... - fmt.Println(line) - continue - } - name := parts[0] - val := parts[1] - - if os.Getenv(name) == val { - // Not changed, nothing to do. - continue - } - - if !strings.HasPrefix(name, "DD_") && !strings.HasPrefix(name, "AWS_LAMBDA_") { - logger.Debug(fmt.Sprintf("[%s] Skip %s=", AwsLambdaExecWrapperEnvVar, name)) - continue - } - - if err := os.Setenv(name, val); err != nil { - logger.Debug(fmt.Sprintf("[%s] Failed %s=%s: %s", AwsLambdaExecWrapperEnvVar, name, val, err)) - } else { - logger.Debug(fmt.Sprintf("[%s] Set %s=%s", AwsLambdaExecWrapperEnvVar, name, val)) - } - } + if val := os.Getenv(UniversalInstrumentation); val == "" { + _ = os.Setenv(UniversalInstrumentation, "1") } }