feat: apply $AWS_LAMBDA_EXEC_WRAPPER if set

Managed AWS Lambda runtimes wrap the function's runtime entry using a script
specified by the `AWS_LAMBDA_EXEC_WRAPPER` environment variable, however this is
not performed by provided runtimes (`provided`, `provided.al2`, `go1.x`).

This adds an `init` function that re-startes the current process after wrapping
it in the wrapper to emulate the behavior of other lambda runtimes, although
this will cause initialization of some other go packages to be run twice, which
will have an impact on cold start times.

See also: https://docs.aws.amazon.com/lambda/latest/dg/runtimes-modify.html#runtime-wrapper

Fixes aws#523

Author: RomainMuller

lambda/wrapper.go

@@ -0,0 +1,45 @@
+// 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
+
+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,57 @@
+// 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
+
+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
+}