From 1766dc23e681fe80a83de20d1a72de1d8ffb6adc Mon Sep 17 00:00:00 2001 From: Romain Marcadier Date: Mon, 16 Oct 2023 10:57:59 +0200 Subject: [PATCH 1/6] 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 #523 --- lambda/wrapper.go | 45 +++++++++++++++++++++++++++++++++ lambda/wrapper_test.go | 57 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 102 insertions(+) create mode 100644 lambda/wrapper.go create mode 100644 lambda/wrapper_test.go diff --git a/lambda/wrapper.go b/lambda/wrapper.go new file mode 100644 index 00000000..5186d426 --- /dev/null +++ b/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) + } +} diff --git a/lambda/wrapper_test.go b/lambda/wrapper_test.go new file mode 100644 index 00000000..afe4502f --- /dev/null +++ b/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, "") + 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 +} From 460eea26527da2ee0b3c7a10ec72cff413ab9836 Mon Sep 17 00:00:00 2001 From: Romain Marcadier Date: Mon, 16 Oct 2023 11:29:45 +0200 Subject: [PATCH 2/6] lint fix --- lambda/wrapper.go | 4 ++-- lambda/wrapper_test.go | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/lambda/wrapper.go b/lambda/wrapper.go index 5186d426..134e0ed3 100644 --- a/lambda/wrapper.go +++ b/lambda/wrapper.go @@ -21,13 +21,13 @@ 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) + 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( +func execAWSLambdaExecWrapper( getenv func(key string) string, sysExec func(argv0 string, argv []string, envv []string) error, ) { diff --git a/lambda/wrapper_test.go b/lambda/wrapper_test.go index afe4502f..72d95cd0 100644 --- a/lambda/wrapper_test.go +++ b/lambda/wrapper_test.go @@ -16,7 +16,7 @@ import ( func TestExecAwsLambdaExecWrapperNotSet(t *testing.T) { exec, execCalled := mockExec(t, "") - execAwsLambdaExecWrapper( + execAWSLambdaExecWrapper( mockedGetenv(t, ""), exec, ) @@ -26,7 +26,7 @@ func TestExecAwsLambdaExecWrapperNotSet(t *testing.T) { func TestExecAwsLambdaExecWrapperSet(t *testing.T) { wrapper := "/path/to/wrapper/entry/point" exec, execCalled := mockExec(t, wrapper) - execAwsLambdaExecWrapper( + execAWSLambdaExecWrapper( mockedGetenv(t, wrapper), exec, ) From 31a466b6747241d34419e0fd40c2584be02f427e Mon Sep 17 00:00:00 2001 From: Romain Marcadier Date: Mon, 16 Oct 2023 14:06:51 +0200 Subject: [PATCH 3/6] add legacy build tagging --- lambda/wrapper.go | 1 + lambda/wrapper_test.go | 1 + 2 files changed, 2 insertions(+) diff --git a/lambda/wrapper.go b/lambda/wrapper.go index 134e0ed3..4721297c 100644 --- a/lambda/wrapper.go +++ b/lambda/wrapper.go @@ -6,6 +6,7 @@ // 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 diff --git a/lambda/wrapper_test.go b/lambda/wrapper_test.go index 72d95cd0..584234e4 100644 --- a/lambda/wrapper_test.go +++ b/lambda/wrapper_test.go @@ -4,6 +4,7 @@ // Copyright 2016-present Datadog, Inc. //go:build unix && !noexecwrapper +// +build unix,!noexecwrapper package lambda From 15f432beabc47321bdb43f4ebdfcd27dfd0f8c24 Mon Sep 17 00:00:00 2001 From: Romain Marcadier Date: Tue, 31 Oct 2023 10:18:05 +0100 Subject: [PATCH 4/6] switch from init to option enablement --- lambda/handler.go | 29 +++++++++++++++++++++++++++++ lambda/wrapper.go | 15 +++++++-------- lambda/wrapper_test.go | 10 ++++++++++ 3 files changed, 46 insertions(+), 8 deletions(-) diff --git a/lambda/handler.go b/lambda/handler.go index e4cfaf7a..a4960b0a 100644 --- a/lambda/handler.go +++ b/lambda/handler.go @@ -10,8 +10,10 @@ import ( "fmt" "io" "io/ioutil" // nolint:staticcheck + "os" "reflect" "strings" + "syscall" "github.com/aws/aws-lambda-go/lambda/handlertrace" ) @@ -28,6 +30,8 @@ type handlerOptions struct { jsonResponseIndentValue string enableSIGTERM bool sigtermCallbacks []func() + enableExecWrapper bool + execWrapperCallbacks []func() } type Option func(*handlerOptions) @@ -102,6 +106,28 @@ func WithEnableSIGTERM(callbacks ...func()) Option { }) } +// WithEnableExecWrapper enables applying the value of the AWS_LAMBDA_EXEC_WRAPPER environment +// variable. If this fariable is set, the current process will be re-started, wrapped under the +// specified wrapper script. Optionally, an array of callback functions to run before restarting +// the process may be provided. +// +// Usage: +// +// lambda.StartWithOptions( +// func (event any) (any, error) { +// return event, nil +// }, +// lambda.WithEnableExecWrapper(func(){ +// log.Print("[AWS_LAMBDA_EXEC_WRAPPER] process is about to be re-started...") +// }), +// ) +func WithEnableExecWrapper(callbacks ...func()) Option { + return Option(func(h *handlerOptions) { + h.execWrapperCallbacks = append(h.execWrapperCallbacks, callbacks...) + h.enableExecWrapper = true + }) +} + // handlerTakesContext returns whether the handler takes a context.Context as its first argument. func handlerTakesContext(handler reflect.Type) (bool, error) { switch handler.NumIn() { @@ -184,6 +210,9 @@ func newHandler(handlerFunc interface{}, options ...Option) *handlerOptions { for _, option := range options { option(h) } + if h.enableExecWrapper { + execAWSLambdaExecWrapper(os.Getenv, syscall.Exec, h.execWrapperCallbacks) + } if h.enableSIGTERM { enableSIGTERM(h.sigtermCallbacks) } diff --git a/lambda/wrapper.go b/lambda/wrapper.go index 4721297c..f40e5b9a 100644 --- a/lambda/wrapper.go +++ b/lambda/wrapper.go @@ -13,30 +13,29 @@ 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) -} - +// execAWSLambdaExecWrapper applies the AWS_LAMBDA_EXEC_WRAPPER environment variable. // 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, + callbacks []func(), ) { wrapper := getenv(awsLambdaExecWrapper) if wrapper == "" { return } + // Execute the provided callbacks before re-starting the process... + for _, callback := range callbacks { + callback() + } + // 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+"=") diff --git a/lambda/wrapper_test.go b/lambda/wrapper_test.go index 584234e4..b1ea317f 100644 --- a/lambda/wrapper_test.go +++ b/lambda/wrapper_test.go @@ -16,22 +16,32 @@ import ( ) func TestExecAwsLambdaExecWrapperNotSet(t *testing.T) { + var called bool + callback := func() { called = true } + exec, execCalled := mockExec(t, "") execAWSLambdaExecWrapper( mockedGetenv(t, ""), exec, + []func(){callback}, ) require.False(t, *execCalled) + require.False(t, called) } func TestExecAwsLambdaExecWrapperSet(t *testing.T) { + var called bool + callback := func() { called = true } + wrapper := "/path/to/wrapper/entry/point" exec, execCalled := mockExec(t, wrapper) execAWSLambdaExecWrapper( mockedGetenv(t, wrapper), exec, + []func(){callback}, ) require.True(t, *execCalled) + require.True(t, called) } func mockExec(t *testing.T, value string) (mock func(string, []string, []string) error, called *bool) { From 5953b5ce8bf8ef5841142cde69aa410f450197ed Mon Sep 17 00:00:00 2001 From: Romain Marcadier Date: Tue, 31 Oct 2023 10:21:49 +0100 Subject: [PATCH 5/6] remove build tag (no longer needed) --- lambda/wrapper.go | 5 ----- lambda/wrapper_test.go | 3 --- 2 files changed, 8 deletions(-) diff --git a/lambda/wrapper.go b/lambda/wrapper.go index f40e5b9a..854a239c 100644 --- a/lambda/wrapper.go +++ b/lambda/wrapper.go @@ -3,11 +3,6 @@ // 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 ( diff --git a/lambda/wrapper_test.go b/lambda/wrapper_test.go index b1ea317f..bd294e9e 100644 --- a/lambda/wrapper_test.go +++ b/lambda/wrapper_test.go @@ -3,9 +3,6 @@ // 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 ( From d4c6b0dd25eedddb900eadbcc6067116cd0ec81d Mon Sep 17 00:00:00 2001 From: Romain Marcadier Date: Tue, 31 Oct 2023 10:27:06 +0100 Subject: [PATCH 6/6] remove generics from test to make it compatible with go < 1.18 --- lambda/wrapper_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lambda/wrapper_test.go b/lambda/wrapper_test.go index bd294e9e..05ed1d85 100644 --- a/lambda/wrapper_test.go +++ b/lambda/wrapper_test.go @@ -60,6 +60,6 @@ func mockedGetenv(t *testing.T, value string) func(string) string { } } -func ptrTo[T any](val T) *T { +func ptrTo(val bool) *bool { return &val }