feat(Multi-tenancy): Add support for multi-tenancy (#11)

Makefile

@@ -10,7 +10,7 @@ DESTINATION_old:= bin/${BINARY_NAME}
 DESTINATION_x86_64 := bin/${BINARY_NAME}-x86_64
 DESTINATION_arm64 := bin/${BINARY_NAME}-arm64
 
-run_in_docker = docker run --env GOPROXY=direct -v $(shell pwd):/LambdaRuntimeLocal -w /LambdaRuntimeLocal golang:1.24 $(1)
+run_in_docker = docker run --env GOPROXY=direct -v $(shell pwd):/LambdaRuntimeLocal -w /LambdaRuntimeLocal golang:1.25 $(1)
 
 compile-with-docker-all:
 	$(call run_in_docker, make compile-lambda-linux-all)

cmd/aws-lambda-rie/main.go

@@ -4,151 +4,9 @@
 package main
 
 import (
-	"context"
-	"fmt"
-	"net"
-	"os"
-	"runtime/debug"
-
-	"github.com/jessevdk/go-flags"
-	"go.amzn.com/lambda/interop"
-	"go.amzn.com/lambda/rapidcore"
-
-	log "github.com/sirupsen/logrus"
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/rie"
 )
 
-const (
-	optBootstrap     = "/opt/bootstrap"
-	runtimeBootstrap = "/var/runtime/bootstrap"
-)
-
-type options struct {
-	LogLevel           string `long:"log-level" description:"The level of AWS Lambda Runtime Interface Emulator logs to display. Can also be set by the environment variable 'LOG_LEVEL'. Defaults to the value 'info'."`
-	InitCachingEnabled bool   `long:"enable-init-caching" description:"Enable support for Init Caching"`
-	// Do not have a default value so we do not need to keep it in sync with the default value in lambda/rapidcore/sandbox_builder.go
-	RuntimeAPIAddress               string `long:"runtime-api-address" description:"The address of the AWS Lambda Runtime API to communicate with the Lambda execution environment."`
-	RuntimeInterfaceEmulatorAddress string `long:"runtime-interface-emulator-address" default:"0.0.0.0:8080" description:"The address for the AWS Lambda Runtime Interface Emulator to accept HTTP request upon."`
-}
-
 func main() {
-	// More frequent GC reduces the tail latencies, equivalent to export GOGC=33
-	debug.SetGCPercent(33)
-
-	opts, args := getCLIArgs()
-
-	logLevel := "info"
-
-	// If you specify an option by using a parameter on the CLI command line, it overrides any value from either the corresponding environment variable.
-	//
-	// https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-envvars.html
-	if opts.LogLevel != "" {
-		logLevel = opts.LogLevel
-	} else if envLogLevel, envLogLevelSet := os.LookupEnv("LOG_LEVEL"); envLogLevelSet {
-		logLevel = envLogLevel
-	}
-
-	rapidcore.SetLogLevel(logLevel)
-
-	if opts.RuntimeAPIAddress != "" {
-		_, _, err := net.SplitHostPort(opts.RuntimeAPIAddress)
-
-		if err != nil {
-			log.WithError(err).Fatalf("The command line value for \"--runtime-api-address\" is not a valid network address %q.", opts.RuntimeAPIAddress)
-		}
-	}
-
-	_, _, err := net.SplitHostPort(opts.RuntimeInterfaceEmulatorAddress)
-
-	if err != nil {
-		log.WithError(err).Fatalf("The command line value for \"--runtime-interface-emulator-address\" is not a valid network address %q.", opts.RuntimeInterfaceEmulatorAddress)
-	}
-
-	bootstrap, handler := getBootstrap(args, opts)
-	sandbox := rapidcore.
-		NewSandboxBuilder().
-		AddShutdownFunc(context.CancelFunc(func() { os.Exit(0) })).
-		SetExtensionsFlag(true).
-		SetInitCachingFlag(opts.InitCachingEnabled)
-
-	if len(handler) > 0 {
-		sandbox.SetHandler(handler)
-	}
-
-	if opts.RuntimeAPIAddress != "" {
-		sandbox.SetRuntimeAPIAddress(opts.RuntimeAPIAddress)
-	}
-
-	sandboxContext, internalStateFn := sandbox.Create()
-	// Since we have not specified a custom interop server for standalone, we can
-	// directly reference the default interop server, which is a concrete type
-	sandbox.DefaultInteropServer().SetSandboxContext(sandboxContext)
-	sandbox.DefaultInteropServer().SetInternalStateGetter(internalStateFn)
-
-	startHTTPServer(opts.RuntimeInterfaceEmulatorAddress, sandbox, bootstrap)
-}
-
-func getCLIArgs() (options, []string) {
-	var opts options
-	parser := flags.NewParser(&opts, flags.IgnoreUnknown)
-	args, err := parser.ParseArgs(os.Args)
-
-	if err != nil {
-		log.WithError(err).Fatal("Failed to parse command line arguments:", os.Args)
-	}
-
-	return opts, args
-}
-
-func isBootstrapFileExist(filePath string) bool {
-	file, err := os.Stat(filePath)
-	return !os.IsNotExist(err) && !file.IsDir()
-}
-
-func getBootstrap(args []string, opts options) (interop.Bootstrap, string) {
-	var bootstrapLookupCmd []string
-	var handler string
-	currentWorkingDir := "/var/task" // default value
-
-	if len(args) <= 1 {
-		// set default value to /var/task/bootstrap, but switch to the other options if it doesn't exist
-		bootstrapLookupCmd = []string{
-			fmt.Sprintf("%s/bootstrap", currentWorkingDir),
-		}
-
-		if !isBootstrapFileExist(bootstrapLookupCmd[0]) {
-			var bootstrapCmdCandidates = []string{
-				optBootstrap,
-				runtimeBootstrap,
-			}
-
-			for i, bootstrapCandidate := range bootstrapCmdCandidates {
-				if isBootstrapFileExist(bootstrapCandidate) {
-					bootstrapLookupCmd = []string{bootstrapCmdCandidates[i]}
-					break
-				}
-			}
-		}
-
-		// handler is used later to set an env var for Lambda Image support
-		handler = ""
-	} else if len(args) > 1 {
-
-		bootstrapLookupCmd = args[1:]
-
-		if cwd, err := os.Getwd(); err == nil {
-			currentWorkingDir = cwd
-		}
-
-		if len(args) > 2 {
-			// Assume last arg is the handler
-			handler = args[len(args)-1]
-		}
-
-		log.Infof("exec '%s' (cwd=%s, handler=%s)", args[1], currentWorkingDir, handler)
-
-	} else {
-		log.Panic("insufficient arguments: bootstrap not provided")
-	}
-
-	return NewSimpleBootstrap(bootstrapLookupCmd, currentWorkingDir), handler
+	rie.Run()
 }

go.mod

@@ -1,12 +1,15 @@
-module go.amzn.com
+module github.com/aws/aws-lambda-runtime-interface-emulator
 
-go 1.24
+go 1.25
 
 require (
 	github.com/aws/aws-lambda-go v1.46.0
+	github.com/go-chi/chi v1.5.5
 	github.com/go-chi/chi/v5 v5.2.2
 	github.com/google/uuid v1.6.0
 	github.com/jessevdk/go-flags v1.5.0
+	github.com/orcaman/concurrent-map v1.0.0
+	github.com/santhosh-tekuri/jsonschema/v5 v5.3.1
 	github.com/sirupsen/logrus v1.9.3
 	github.com/stretchr/testify v1.9.0
 	golang.org/x/sync v0.6.0

go.sum

@@ -3,14 +3,20 @@ github.com/aws/aws-lambda-go v1.46.0/go.mod h1:dpMpZgvWx5vuQJfBt0zqBha60q7Dd7Rfg
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/go-chi/chi v1.5.5 h1:vOB/HbEMt9QqBqErz07QehcOKHaWFtuj87tTDVz2qXE=
+github.com/go-chi/chi v1.5.5/go.mod h1:C9JqLr3tIYjDOZpzn+BCuxY8z8vmca43EeMgyZt7irw=
 github.com/go-chi/chi/v5 v5.2.2 h1:CMwsvRVTbXVytCk1Wd72Zy1LAsAh9GxMmSNWLHCG618=
 github.com/go-chi/chi/v5 v5.2.2/go.mod h1:L2yAIGWB3H+phAw1NxKwWM+7eUH/lU8pOMm5hHcoops=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/jessevdk/go-flags v1.5.0 h1:1jKYvbxEjfUl0fmqTCOfonvskHHXMjBySTLW4y9LFvc=
 github.com/jessevdk/go-flags v1.5.0/go.mod h1:Fw0T6WPc1dYxT4mKEZRfG5kJhaTDP9pj1c2EWnYs/m4=
+github.com/orcaman/concurrent-map v1.0.0 h1:I/2A2XPCb4IuQWcQhBhSwGfiuybl/J0ev9HDbW65HOY=
+github.com/orcaman/concurrent-map v1.0.0/go.mod h1:Lu3tH6HLW3feq74c2GC+jIMS/K2CFcDWnWD9XkenwhI=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/santhosh-tekuri/jsonschema/v5 v5.3.1 h1:lZUw3E0/J3roVtGQ+SCrUrg3ON6NgVqpn3+iol9aGu4=
+github.com/santhosh-tekuri/jsonschema/v5 v5.3.1/go.mod h1:uToXkOrWAZ6/Oc07xWQrPOhJotwFIyu2bBVN41fcDUY=
 github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
 github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=

lambda/appctx/appctxutil.go → internal/lambda/appctx/appctxutil.go

@@ -8,8 +8,8 @@ import (
 	"net/http"
 	"strings"
 
-	"go.amzn.com/lambda/fatalerror"
-	"go.amzn.com/lambda/interop"
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/fatalerror"
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/interop"
 
 	log "github.com/sirupsen/logrus"
 )

lambda/appctx/appctxutil_test.go → internal/lambda/appctx/appctxutil_test.go

@@ -8,11 +8,11 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/fatalerror"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"go.amzn.com/lambda/fatalerror"
 
-	"go.amzn.com/lambda/interop"
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/interop"
 )
 
 func runTestRequestWithUserAgent(t *testing.T, userAgent string, expectedRuntimeRelease string) {

lambda/core/agentsmap.go → internal/lambda/core/agentsmap.go

@@ -5,6 +5,7 @@ package core
 
 import (
 	"errors"
+
 	"github.com/google/uuid"
 )
 

lambda/core/agentsmap_test.go → internal/lambda/core/agentsmap_test.go

@@ -4,9 +4,10 @@
 package core
 
 import (
+	"testing"
+
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
-	"testing"
 )
 
 func TestExternalAgentsMapLookupByName(t *testing.T) {

lambda/core/bandwidthlimiter/bandwidthlimiter.go → internal/lambda/core/bandwidthlimiter/bandwidthlimiter.go

@@ -6,7 +6,7 @@ package bandwidthlimiter
 import (
 	"io"
 
-	"go.amzn.com/lambda/interop"
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/interop"
 )
 
 func BandwidthLimitingCopy(dst *BandwidthLimitingWriter, src io.Reader) (written int64, err error) {

lambda/core/bandwidthlimiter/throttler.go → internal/lambda/core/bandwidthlimiter/throttler.go

@@ -12,8 +12,8 @@ import (
 
 	log "github.com/sirupsen/logrus"
 
-	"go.amzn.com/lambda/interop"
-	"go.amzn.com/lambda/metering"
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/interop"
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/metering"
 )
 
 var ErrBufferSizeTooLarge = errors.New("buffer size cannot be greater than bucket size")

lambda/core/credentials_test.go → internal/lambda/core/credentials_test.go

@@ -4,9 +4,10 @@
 package core
 
 import (
-	"github.com/stretchr/testify/assert"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/assert"
 )
 
 const (

lambda/core/directinvoke/customerheaders_test.go → internal/lambda/core/directinvoke/customerheaders_test.go

@@ -4,8 +4,9 @@
 package directinvoke
 
 import (
-	"github.com/stretchr/testify/require"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestCustomerHeadersEmpty(t *testing.T) {

lambda/core/directinvoke/directinvoke.go → internal/lambda/core/directinvoke/directinvoke.go

@@ -11,17 +11,18 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/core/bandwidthlimiter"
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/fatalerror"
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/interop"
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/metering"
 	"github.com/go-chi/chi/v5"
-	"go.amzn.com/lambda/core/bandwidthlimiter"
-	"go.amzn.com/lambda/fatalerror"
-	"go.amzn.com/lambda/interop"
-	"go.amzn.com/lambda/metering"
 
 	log "github.com/sirupsen/logrus"
 )
 
 const (
 	InvokeIDHeader                   = "Invoke-Id"
+	TenantIDHeader                   = "Tenant-Id"
 	InvokedFunctionArnHeader         = "Invoked-Function-Arn"
 	VersionIDHeader                  = "Invoked-Function-Version"
 	ReservationTokenHeader           = "Reservation-Token"
@@ -182,6 +183,7 @@ func ReceiveDirectInvoke(w http.ResponseWriter, r *http.Request, token interop.T
 		InvokeResponseMode:       InvokeResponseMode,
 		RestoreDurationNs:        token.RestoreDurationNs,
 		RestoreStartTimeMonotime: token.RestoreStartTimeMonotime,
+		TenantID:                 token.TenantID,
 	}
 
 	if inv.ID != token.InvokeID {
@@ -208,6 +210,10 @@ func ReceiveDirectInvoke(w http.ResponseWriter, r *http.Request, token interop.T
 	w.Header().Set(ReservationTokenHeader, token.ReservationToken)
 	w.Header().Set(InvokeIDHeader, token.InvokeID)
 
+	if token.TenantID != "" {
+		w.Header().Set(TenantIDHeader, string(token.TenantID))
+	}
+
 	return inv, nil
 }
 

lambda/core/directinvoke/directinvoke_test.go → internal/lambda/core/directinvoke/directinvoke_test.go

@@ -17,12 +17,12 @@ import (
 	"testing"
 	"time"
 
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/fatalerror"
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/interop"
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/metering"
 	"github.com/go-chi/chi/v5"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"go.amzn.com/lambda/fatalerror"
-	"go.amzn.com/lambda/interop"
-	"go.amzn.com/lambda/metering"
 )
 
 func NewResponseWriterWithoutFlushMethod() *ResponseWriterWithoutFlushMethod {
@@ -530,6 +530,23 @@ func TestConvertToInvokeResponseMode(t *testing.T) {
 	require.Equal(t, interop.ErrInvalidInvokeResponseMode, err)
 }
 
+func TestReceiveDirectInvoke_TenantIDHeader(t *testing.T) {
+	responseRecorder := httptest.NewRecorder()
+
+	token := createDummyToken()
+	token.TenantID = "blue"
+
+	req := makeDirectInvokeRequest([]byte(""), token.ReservationToken, token.InvokeID,
+		"", token.VersionID, "", "", "",
+		"", "", "")
+
+	_, err := ReceiveDirectInvoke(responseRecorder, req, token)
+	assert.NoError(t, err)
+
+	actual := responseRecorder.Header().Get(TenantIDHeader)
+	assert.Equal(t, "blue", actual)
+}
+
 func FuzzReceiveDirectInvoke(f *testing.F) {
 	testCustHeaders := CustomerHeaders{
 		CognitoIdentityID:     "id1",

lambda/core/directinvoke/util.go → internal/lambda/core/directinvoke/util.go

@@ -6,11 +6,12 @@ package directinvoke
 import (
 	"context"
 	"errors"
-	"go.amzn.com/lambda/core/bandwidthlimiter"
 	"io"
 	"net/http"
 	"time"
 
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/core/bandwidthlimiter"
+
 	log "github.com/sirupsen/logrus"
 )
 

lambda/core/externalagent.go → internal/lambda/core/externalagent.go

@@ -7,7 +7,7 @@ import (
 	"fmt"
 	"time"
 
-	"go.amzn.com/lambda/core/statejson"
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/core/statejson"
 
 	"github.com/google/uuid"
 )