feat: apply $AWS_LAMBDA_EXEC_WRAPPER if set

lambda/wrapper.go

@@ -0,0 +1,46 @@
+// 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 2016-present Datadog, Inc.
+
+// Specify the noexecwrapper build tag to remove the wrapper tampoline from
+// this library if it is undesirable.
+//go:build unix && !noexecwrapper
+// +build unix,!noexecwrapper
+
+package lambda
+
+import (
+	"log"
+	"os"
+	"syscall"
+)
+
+const awsLambdaExecWrapper = "AWS_LAMBDA_EXEC_WRAPPER"
+
+func init() {
+	// Honor the AWS_LAMBDA_EXEC_WRAPPER configuration at startup, trying to emulate
+	// the behavior of managed runtimes, as this configuration is otherwise not applied
+	// by provided runtimes (or go1.x).
+	execAWSLambdaExecWrapper(os.Getenv, syscall.Exec)
+}
+
+// If AWS_LAMBDA_EXEC_WRAPPER is defined, replace the current process by spawning
+// it with the current process' arguments (including the program name). If the call
+// to syscall.Exec fails, this aborts the process with a fatal error.
+func execAWSLambdaExecWrapper(
+	getenv func(key string) string,
+	sysExec func(argv0 string, argv []string, envv []string) error,
+) {
+	wrapper := getenv(awsLambdaExecWrapper)
+	if wrapper == "" {
+		return
+	}
+
+	// The AWS_LAMBDA_EXEC_WRAPPER variable is blanked before replacing the process
+	// in order to avoid endlessly restarting the process.
+	env := append(os.Environ(), awsLambdaExecWrapper+"=")
+	if err := sysExec(wrapper, append([]string{wrapper}, os.Args...), env); err != nil {
+		log.Fatalf("failed to sysExec() %s=%s: %v", awsLambdaExecWrapper, wrapper, err)
+	}
+}

lambda/wrapper_test.go

@@ -0,0 +1,58 @@
+// 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 2016-present Datadog, Inc.
+
+//go:build unix && !noexecwrapper
+// +build unix,!noexecwrapper
+
+package lambda
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestExecAwsLambdaExecWrapperNotSet(t *testing.T) {
+	exec, execCalled := mockExec(t, "<nope>")
+	execAWSLambdaExecWrapper(
+		mockedGetenv(t, ""),
+		exec,
+	)
+	require.False(t, *execCalled)
+}
+
+func TestExecAwsLambdaExecWrapperSet(t *testing.T) {
+	wrapper := "/path/to/wrapper/entry/point"
+	exec, execCalled := mockExec(t, wrapper)
+	execAWSLambdaExecWrapper(
+		mockedGetenv(t, wrapper),
+		exec,
+	)
+	require.True(t, *execCalled)
+}
+
+func mockExec(t *testing.T, value string) (mock func(string, []string, []string) error, called *bool) {
+	mock = func(argv0 string, argv []string, envv []string) error {
+		*called = true
+		require.Equal(t, value, argv0)
+		require.Equal(t, append([]string{value}, os.Args...), argv)
+		require.Equal(t, awsLambdaExecWrapper+"=", envv[len(envv)-1])
+		return nil
+	}
+	called = ptrTo(false)
+	return
+}
+
+func mockedGetenv(t *testing.T, value string) func(string) string {
+	return func(key string) string {
+		require.Equal(t, awsLambdaExecWrapper, key)
+		return value
+	}
+}
+
+func ptrTo[T any](val T) *T {
+	return &val
+}