Pulling upstream changes 2021-11

Author: valerena

cmd/aws-lambda-rie/main.go

@@ -21,7 +21,8 @@ const (
 )
 
 type options struct {
-	LogLevel string `long:"log-level" default:"info" description:"log level"`
+	LogLevel           string `long:"log-level" default:"info" description:"log level"`
+	InitCachingEnabled bool   `long:"enable-init-caching" description:"Enable support for Init Caching"`
 }
 
 func main() {
@@ -32,7 +33,11 @@ func main() {
 	rapidcore.SetLogLevel(opts.LogLevel)
 
 	bootstrap, handler := getBootstrap(args, opts)
-	sandbox := rapidcore.NewSandboxBuilder(bootstrap).AddShutdownFunc(context.CancelFunc(func() { os.Exit(0) })).SetExtensionsFlag(true)
+	sandbox := rapidcore.
+		NewSandboxBuilder(bootstrap).
+		AddShutdownFunc(context.CancelFunc(func() { os.Exit(0) })).
+		SetExtensionsFlag(true).
+		SetInitCachingFlag(opts.InitCachingEnabled)
 
 	if len(handler) > 0 {
 		sandbox.SetHandler(handler)
@@ -72,7 +77,7 @@ func getBootstrap(args []string, opts options) (*rapidcore.Bootstrap, string) {
 			fmt.Sprintf("%s/bootstrap", currentWorkingDir),
 		}
 
-		if !isBootstrapFileExist(bootstrapLookupCmd[0])  {
+		if !isBootstrapFileExist(bootstrapLookupCmd[0]) {
 			var bootstrapCmdCandidates = []string{
 				optBootstrap,
 				runtimeBootstrap,

lambda/agents/agent.go

@@ -25,13 +25,12 @@ type ExternalAgentProcess struct {
 }
 
 // NewExternalAgentProcess returns a new external agent process
-func NewExternalAgentProcess(path string, env []string, logWriter io.Writer) ExternalAgentProcess {
+func NewExternalAgentProcess(path string, env []string, stdoutWriter io.Writer, stderrWriter io.Writer) ExternalAgentProcess {
 	command := exec.Command(path)
 	command.Env = env
 
-	w := NewNewlineSplitWriter(logWriter)
-	command.Stdout = w
-	command.Stderr = w
+	command.Stdout = NewNewlineSplitWriter(stdoutWriter)
+	command.Stderr = NewNewlineSplitWriter(stderrWriter)
 	command.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
 
 	return ExternalAgentProcess{

lambda/agents/agent_test.go

@@ -204,30 +204,76 @@ func TestFindAgentMixed(t *testing.T) {
 // Test our ability to start agents
 func TestAgentStart(t *testing.T) {
 	assert := assert.New(t)
-	agent := NewExternalAgentProcess("../testdata/agents/bash_true.sh", []string{}, &mockWriter{})
+	agent := NewExternalAgentProcess("../testdata/agents/bash_true.sh", []string{}, &mockWriter{}, &mockWriter{})
 	assert.Nil(agent.Start())
 	assert.Nil(agent.Wait())
 }
 
 // Test that execution of invalid agents is correctly reported
 func TestInvalidAgentStart(t *testing.T) {
 	assert := assert.New(t)
-	agent := NewExternalAgentProcess("/bin/none", []string{}, &mockWriter{})
+	agent := NewExternalAgentProcess("/bin/none", []string{}, &mockWriter{}, &mockWriter{})
 	assert.True(os.IsNotExist(agent.Start()))
 }
 
-// Test that execution of invalid agents is correctly reported
-func TestAgentTelemetry(t *testing.T) {
+func TestAgentStdoutWriter(t *testing.T) {
+	// Given
+	assert := assert.New(t)
+
+	stdout := &mockWriter{}
+	stderr := &mockWriter{}
+	expectedStdout := "stdout line 1\nstdout line 2\nstdout line 3\n"
+	expectedStderr := ""
+
+	agent := NewExternalAgentProcess("../testdata/agents/bash_stdout.sh", []string{}, stdout, stderr)
+
+	// When
+	assert.NoError(agent.Start())
+	assert.NoError(agent.Wait())
+
+	// Then
+	assert.Equal(expectedStdout, string(bytes.Join(stdout.bytesReceived, []byte(""))))
+	assert.Equal(expectedStderr, string(bytes.Join(stderr.bytesReceived, []byte(""))))
+}
+
+func TestAgentStderrWriter(t *testing.T) {
+	// Given
 	assert := assert.New(t)
-	buffer := &mockWriter{}
 
-	agent := NewExternalAgentProcess("../testdata/agents/bash_echo.sh", []string{}, buffer)
+	stdout := &mockWriter{}
+	stderr := &mockWriter{}
+	expectedStdout := ""
+	expectedStderr := "stderr line 1\nstderr line 2\nstderr line 3\n"
+
+	agent := NewExternalAgentProcess("../testdata/agents/bash_stderr.sh", []string{}, stdout, stderr)
+
+	// When
+	assert.NoError(agent.Start())
+	assert.NoError(agent.Wait())
+
+	// Then
+	assert.Equal(expectedStdout, string(bytes.Join(stdout.bytesReceived, []byte(""))))
+	assert.Equal(expectedStderr, string(bytes.Join(stderr.bytesReceived, []byte(""))))
+}
+
+func TestAgentStdoutAndStderrSeperateWriters(t *testing.T) {
+	// Given
+	assert := assert.New(t)
+
+	stdout := &mockWriter{}
+	stderr := &mockWriter{}
+	expectedStdout := "stdout line 1\nstdout line 2\nstdout line 3\n"
+	expectedStderr := "stderr line 1\nstderr line 2\nstderr line 3\n"
+
+	agent := NewExternalAgentProcess("../testdata/agents/bash_stdout_and_stderr.sh", []string{}, stdout, stderr)
 
+	// When
 	assert.NoError(agent.Start())
 	assert.NoError(agent.Wait())
 
-	message := "hello world\n|barbaz\n|hello world\n|barbaz2"
-	assert.Equal(message, string(bytes.Join(buffer.bytesReceived, []byte("|"))))
+	// Then
+	assert.Equal(expectedStdout, string(bytes.Join(stdout.bytesReceived, []byte(""))))
+	assert.Equal(expectedStderr, string(bytes.Join(stderr.bytesReceived, []byte(""))))
 }
 
 type mockWriter struct {

lambda/appctx/appctxutil.go

@@ -5,11 +5,10 @@ package appctx
 
 import (
 	"context"
-	"net/http"
-	"strings"
-
 	"go.amzn.com/lambda/fatalerror"
 	"go.amzn.com/lambda/interop"
+	"net/http"
+	"strings"
 
 	log "github.com/sirupsen/logrus"
 )
@@ -24,6 +23,9 @@ type ReqCtxKey int
 // context object into request context.
 const ReqCtxApplicationContextKey ReqCtxKey = iota
 
+// MaxRuntimeReleaseLength Max length for user agent string.
+const MaxRuntimeReleaseLength = 128
+
 // FromRequest retrieves application context from the request context.
 func FromRequest(request *http.Request) ApplicationContext {
 	return request.Context().Value(ReqCtxApplicationContextKey).(ApplicationContext)
@@ -39,24 +41,78 @@ func GetRuntimeRelease(appCtx ApplicationContext) string {
 	return appCtx.GetOrDefault(AppCtxRuntimeReleaseKey, "").(string)
 }
 
-// UpdateAppCtxWithRuntimeRelease extracts runtime release info from user agent header and put it into appCtx.
+// GetUserAgentFromRequest Returns the first token -seperated by a space-
+// from request header 'User-Agent'.
+func GetUserAgentFromRequest(request *http.Request) string {
+	runtimeRelease := ""
+	userAgent := request.Header.Get("User-Agent")
+	// Split around spaces and use only the first token.
+	if fields := strings.Fields(userAgent); len(fields) > 0 && len(fields[0]) > 0 {
+		runtimeRelease = fields[0]
+	}
+	return runtimeRelease
+}
+
+// CreateRuntimeReleaseFromRequest Gets runtime features from request header
+// 'Lambda-Runtime-Features', and append it to the given runtime release.
+func CreateRuntimeReleaseFromRequest(request *http.Request, runtimeRelease string) string {
+	lambdaRuntimeFeaturesHeader := request.Header.Get("Lambda-Runtime-Features")
+
+	// "(", ")" are not valid token characters, and potentially could invalidate runtime_release
+	lambdaRuntimeFeaturesHeader = strings.ReplaceAll(lambdaRuntimeFeaturesHeader, "(", "")
+	lambdaRuntimeFeaturesHeader = strings.ReplaceAll(lambdaRuntimeFeaturesHeader, ")", "")
+
+	numberOfAppendedFeatures := 0
+	// Available length is a maximum length available for runtime features (including delimiters). From maximal runtime
+	// release length we subtract what we already have plus 3 additional bytes for a space and a pair of brackets for
+	// list of runtime features that is added later.
+	runtimeReleaseLength := len(runtimeRelease)
+	if runtimeReleaseLength == 0 {
+		runtimeReleaseLength = len("Unknown")
+	}
+	availableLength := MaxRuntimeReleaseLength - runtimeReleaseLength - 3
+	var lambdaRuntimeFeatures []string
+
+	for _, feature := range strings.Fields(lambdaRuntimeFeaturesHeader) {
+		featureLength := len(feature)
+		// If featureLength <= availableLength - numberOfAppendedFeatures
+		// (where numberOfAppendedFeatures is equal to number of delimiters needed).
+		if featureLength <= availableLength-numberOfAppendedFeatures {
+			availableLength -= featureLength
+			lambdaRuntimeFeatures = append(lambdaRuntimeFeatures, feature)
+			numberOfAppendedFeatures++
+		}
+	}
+	// Append valid features to runtime release.
+	if len(lambdaRuntimeFeatures) > 0 {
+		if runtimeRelease == "" {
+			runtimeRelease = "Unknown"
+		}
+		runtimeRelease += " (" + strings.Join(lambdaRuntimeFeatures, " ") + ")"
+	}
+
+	return runtimeRelease
+}
+
+// UpdateAppCtxWithRuntimeRelease extracts runtime release info from user agent & lambda runtime features
+// headers and update it into appCtx.
 // Sample UA:
 // Mozilla/5.0 (Windows NT 6.1; Win64; x64; rv:47.0) Gecko/20100101 Firefox/47.0
 func UpdateAppCtxWithRuntimeRelease(request *http.Request, appCtx ApplicationContext) bool {
-	// If appCtx has runtime release value already, skip updating for consistency.
-	if len(GetRuntimeRelease(appCtx)) > 0 {
-		return false
-	}
-
-	userAgent := request.Header.Get("User-Agent")
-	if len(userAgent) == 0 {
+	// If appCtx has runtime release value already, just append the runtime features.
+	if appCtxRuntimeRelease := GetRuntimeRelease(appCtx); len(appCtxRuntimeRelease) > 0 {
+		// if the runtime features are not appended before append them, otherwise ignore
+		if runtimeReleaseWithFeatures := CreateRuntimeReleaseFromRequest(request, appCtxRuntimeRelease); len(runtimeReleaseWithFeatures) > len(appCtxRuntimeRelease) &&
+			appCtxRuntimeRelease[len(appCtxRuntimeRelease)-1] != ')' {
+			appCtx.Store(AppCtxRuntimeReleaseKey, runtimeReleaseWithFeatures)
+			return true
+		}
 		return false
 	}
-
-	// Split around spaces and use only the first token.
-	if fields := strings.Fields(userAgent); len(fields) > 0 && len(fields[0]) > 0 {
-		appCtx.Store(AppCtxRuntimeReleaseKey,
-			fields[0])
+	// If appCtx doesn't have runtime release value, update it with user agent and runtime features.
+	if runtimeReleaseWithFeatures := CreateRuntimeReleaseFromRequest(request,
+		GetUserAgentFromRequest(request)); runtimeReleaseWithFeatures != "" {
+		appCtx.Store(AppCtxRuntimeReleaseKey, runtimeReleaseWithFeatures)
 		return true
 	}
 	return false

lambda/appctx/appctxutil_test.go

@@ -5,6 +5,7 @@ package appctx
 
 import (
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -30,6 +31,63 @@ func runTestRequestWithUserAgent(t *testing.T, userAgent string, expectedRuntime
 	assert.Equal(t, expectedRuntimeRelease, ctxRuntimeRelease, "failed to extract runtime_release token")
 }
 
+func TestCreateRuntimeReleaseFromRequest(t *testing.T) {
+	tests := map[string]struct {
+		userAgentHeader             string
+		lambdaRuntimeFeaturesHeader string
+		expectedRuntimeRelease      string
+	}{
+		"No User-Agent header": {
+			userAgentHeader:             "",
+			lambdaRuntimeFeaturesHeader: "httpcl/2.0 execwr",
+			expectedRuntimeRelease:      "Unknown (httpcl/2.0 execwr)",
+		},
+		"No Lambda-Runtime-Features header": {
+			userAgentHeader:             "Node.js/14.16.0",
+			lambdaRuntimeFeaturesHeader: "",
+			expectedRuntimeRelease:      "Node.js/14.16.0",
+		},
+		"Lambda-Runtime-Features header with additional spaces": {
+			userAgentHeader:             "Node.js/14.16.0",
+			lambdaRuntimeFeaturesHeader: "httpcl/2.0    execwr",
+			expectedRuntimeRelease:      "Node.js/14.16.0 (httpcl/2.0 execwr)",
+		},
+		"Lambda-Runtime-Features header with special characters": {
+			userAgentHeader:             "Node.js/14.16.0",
+			lambdaRuntimeFeaturesHeader: "httpcl/2.0@execwr-1 abcd?efg nodewr/(4.33)) nodewr/4.3",
+			expectedRuntimeRelease:      "Node.js/14.16.0 (httpcl/2.0@execwr-1 abcd?efg nodewr/4.33 nodewr/4.3)",
+		},
+		"Lambda-Runtime-Features header with long Lambda-Runtime-Features header": {
+			userAgentHeader:             "Node.js/14.16.0",
+			lambdaRuntimeFeaturesHeader: strings.Repeat("abcdef ", MaxRuntimeReleaseLength/7),
+			expectedRuntimeRelease:      "Node.js/14.16.0 (" + strings.Repeat("abcdef ", (MaxRuntimeReleaseLength-18-6)/7) + "abcdef)",
+		},
+		"Lambda-Runtime-Features header with long Lambda-Runtime-Features header with UTF-8 characters": {
+			userAgentHeader:             "Node.js/14.16.0",
+			lambdaRuntimeFeaturesHeader: strings.Repeat("我爱亚马逊 ", MaxRuntimeReleaseLength/16),
+			expectedRuntimeRelease:      "Node.js/14.16.0 (" + strings.Repeat("我爱亚马逊 ", (MaxRuntimeReleaseLength-18-15)/16) + "我爱亚马逊)",
+		},
+	}
+
+	for _, tc := range tests {
+		req := httptest.NewRequest("", "/", nil)
+		if tc.userAgentHeader != "" {
+			req.Header.Set("User-Agent", tc.userAgentHeader)
+		}
+		if tc.lambdaRuntimeFeaturesHeader != "" {
+			req.Header.Set("Lambda-Runtime-Features", tc.lambdaRuntimeFeaturesHeader)
+		}
+		appCtx := NewApplicationContext()
+		request := RequestWithAppCtx(req, appCtx)
+
+		UpdateAppCtxWithRuntimeRelease(request, appCtx)
+		runtimeRelease := GetRuntimeRelease(appCtx)
+
+		assert.LessOrEqual(t, len(runtimeRelease), MaxRuntimeReleaseLength)
+		assert.Equal(t, tc.expectedRuntimeRelease, runtimeRelease)
+	}
+}
+
 func TestUpdateAppCtxWithRuntimeRelease(t *testing.T) {
 	type pair struct {
 		in, wanted string
@@ -74,6 +132,25 @@ func TestUpdateAppCtxWithRuntimeReleaseWithBlankUserAgent(t *testing.T) {
 	assert.False(t, ok)
 }
 
+func TestUpdateAppCtxWithRuntimeReleaseWithLambdaRuntimeFeatures(t *testing.T) {
+	// GIVEN
+	// Simple LambdaRuntimeFeatures passed.
+	req := httptest.NewRequest("", "/", nil)
+	req.Header.Set("User-Agent", "Node.js/14.16.0")
+	req.Header.Set("Lambda-Runtime-Features", "httpcl/2.0 execwr nodewr/4.3")
+	request := RequestWithAppCtx(req, NewApplicationContext())
+	appCtx := request.Context().Value(ReqCtxApplicationContextKey).(ApplicationContext)
+
+	// DO
+	ok := UpdateAppCtxWithRuntimeRelease(request, appCtx)
+
+	//ASSERT
+	assert.True(t, ok, "runtime_release updated based only on User-Agent and valid features")
+	ctxRuntimeRelease, ok := appCtx.Load(AppCtxRuntimeReleaseKey)
+	assert.True(t, ok)
+	assert.Equal(t, "Node.js/14.16.0 (httpcl/2.0 execwr nodewr/4.3)", ctxRuntimeRelease)
+}
+
 // Test that RAPID allows updating runtime_release only once
 func TestUpdateAppCtxWithRuntimeReleaseMultipleTimes(t *testing.T) {
 	// GIVEN