<|start_filename|>lambda/rapidcore/env/customer.go<|end_filename|>
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package env

import (
	"os"
	"strings"

	log "github.com/sirupsen/logrus"
)

func isInternalEnvVar(envKey string) bool {
	// the rule is no '_' prefixed env. variables will be propagated to the runtime but the ones explicitly exempted
	allowedKeys := map[string]bool{
		"_HANDLER":                 true,
		"_AWS_XRAY_DAEMON_ADDRESS": true,
		"_AWS_XRAY_DAEMON_PORT":    true,
		"_LAMBDA_TELEMETRY_LOG_FD": true,
	}
	return strings.HasPrefix(envKey, "_") && !allowedKeys[envKey]
}

// CustomerEnvironmentVariables parses all environment variables that are
// not internal/credential/platform, and must be called before agent bootstrap.
func CustomerEnvironmentVariables() map[string]string {
	internalKeys := predefinedInternalEnvVarKeys()
	platformKeys := predefinedPlatformEnvVarKeys()
	runtimeKeys := predefinedRuntimeEnvVarKeys()
	credentialKeys := predefinedCredentialsEnvVarKeys()
	platformUnreservedKeys := predefinedPlatformUnreservedEnvVarKeys()
	isCustomer := func(key string) bool {
		return !internalKeys[key] &&
			!runtimeKeys[key] &&
			!platformKeys[key] &&
			!credentialKeys[key] &&
			!platformUnreservedKeys[key] &&
			!isInternalEnvVar(key)
	}

	customerEnv := map[string]string{}
	for _, keyval := range os.Environ() {
		key, val, err := SplitEnvironmentVariable(keyval)
		if err != nil {
			log.Warnf("Customer environment variable with invalid format: %s", err)
			continue
		}

		if isCustomer(key) {
			customerEnv[key] = val
		}
	}

	return customerEnv
}


<|start_filename|>lambda/metering/time_test.go<|end_filename|>
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package metering

import (
	"math"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestMonoToEpochPrecision(t *testing.T) {
	a := time.Now().UnixNano()
	b := MonoToEpoch(Monotime())

	// Conversion error is less than a millisecond.
	assert.True(t, math.Abs(float64(a-b)) < float64(time.Millisecond))
}

func TestExtensionsResetDurationProfilerForExtensionsResetWithNoExtensions(t *testing.T) {
	mono := Monotime()
	profiler := ExtensionsResetDurationProfiler{}

	profiler.extensionsResetStartTimeNs = mono
	profiler.extensionsResetEndTimeNs = mono + time.Second.Nanoseconds()
	profiler.AvailableNs = 3 * time.Second.Nanoseconds()
	profiler.NumAgentsRegisteredForShutdown = 0
	extensionsResetMs, resetTimeout := profiler.CalculateExtensionsResetMs()

	assert.Equal(t, int64(0), extensionsResetMs)
	assert.Equal(t, false, resetTimeout)
}

func TestExtensionsResetDurationProfilerForExtensionsResetWithinDeadline(t *testing.T) {
	mono := Monotime()
	profiler := ExtensionsResetDurationProfiler{}

	profiler.extensionsResetStartTimeNs = mono
	profiler.extensionsResetEndTimeNs = mono + time.Second.Nanoseconds()
	profiler.AvailableNs = 3 * time.Second.Nanoseconds()
	profiler.NumAgentsRegisteredForShutdown = 1
	extensionsResetMs, resetTimeout := profiler.CalculateExtensionsResetMs()

	assert.Equal(t, time.Second.Milliseconds(), extensionsResetMs)
	assert.Equal(t, false, resetTimeout)
}

func TestExtensionsResetDurationProfilerForExtensionsResetTimeout(t *testing.T) {
	mono := Monotime()
	profiler := ExtensionsResetDurationProfiler{}

	profiler.extensionsResetStartTimeNs = mono
	profiler.extensionsResetEndTimeNs = mono + 3*time.Second.Nanoseconds()
	profiler.AvailableNs = time.Second.Nanoseconds()
	profiler.NumAgentsRegisteredForShutdown = 1
	extensionsResetMs, resetTimeout := profiler.CalculateExtensionsResetMs()

	assert.Equal(t, time.Second.Milliseconds(), extensionsResetMs)
	assert.Equal(t, true, resetTimeout)
}

func TestExtensionsResetDurationProfilerEndToEnd(t *testing.T) {
	profiler := ExtensionsResetDurationProfiler{}

	profiler.Start()
	time.Sleep(time.Second)
	profiler.Stop()

	profiler.AvailableNs = 2 * time.Second.Nanoseconds()
	profiler.NumAgentsRegisteredForShutdown = 1
	extensionsResetMs, _ := profiler.CalculateExtensionsResetMs()

	assert.GreaterOrEqual(t, 2*time.Second.Milliseconds(), extensionsResetMs)
	assert.LessOrEqual(t, time.Second.Milliseconds(), extensionsResetMs)
}


<|start_filename|>lambda/core/directinvoke/customerheaders_test.go<|end_filename|>
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package directinvoke

import (
	"github.com/stretchr/testify/require"
	"testing"
)

func TestCustomerHeadersEmpty(t *testing.T) {
	in := CustomerHeaders{}
	out := CustomerHeaders{}

	require.NoError(t, out.Load(in.Dump()))
	require.Equal(t, in, out)
}

func TestCustomerHeaders(t *testing.T) {
	in := CustomerHeaders{CognitoIdentityID: "asd"}
	out := CustomerHeaders{}

	require.NoError(t, out.Load(in.Dump()))
	require.Equal(t, in, out)
}


<|start_filename|>lambda/rapi/rendering/rendering.go<|end_filename|>
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package rendering

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strconv"
	"sync"
	"time"

	"go.amzn.com/lambda/interop"
	"go.amzn.com/lambda/metering"
	"go.amzn.com/lambda/rapi/model"

	"github.com/go-chi/render"
	"github.com/google/uuid"
	log "github.com/sirupsen/logrus"
)

const (
	// ErrorTypeInternalServerError error type for internal server error
	ErrorTypeInternalServerError = "InternalServerError"
	// ErrorTypeInvalidStateTransition error type for invalid state transition
	ErrorTypeInvalidStateTransition = "InvalidStateTransition"
	// ErrorTypeInvalidRequestID error type for invalid request ID error
	ErrorTypeInvalidRequestID = "InvalidRequestID"
	// ErrorTypeRequestEntityTooLarge error type for payload too large
	ErrorTypeRequestEntityTooLarge = "RequestEntityTooLarge"
)

// ErrRenderingServiceStateNotSet returned when state not set
var ErrRenderingServiceStateNotSet = errors.New("EventRenderingService state not set")

// RendererState is renderer object state.
type RendererState interface {
	RenderAgentEvent(w http.ResponseWriter, r *http.Request) error
	RenderRuntimeEvent(w http.ResponseWriter, r *http.Request) error
}

// EventRenderingService is a state machine for rendering runtime and agent API responses.
type EventRenderingService struct {
	mutex        *sync.RWMutex
	currentState RendererState
}

// SetRenderer set current state
func (s *EventRenderingService) SetRenderer(state RendererState) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	s.currentState = state
}

// RenderAgentEvent delegates to state implementation.
func (s *EventRenderingService) RenderAgentEvent(w http.ResponseWriter, r *http.Request) error {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	if s.currentState == nil {
		return ErrRenderingServiceStateNotSet
	}
	return s.currentState.RenderAgentEvent(w, r)
}

// RenderRuntimeEvent delegates to state implementation.
func (s *EventRenderingService) RenderRuntimeEvent(w http.ResponseWriter, r *http.Request) error {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	if s.currentState == nil {
		return ErrRenderingServiceStateNotSet
	}
	return s.currentState.RenderRuntimeEvent(w, r)
}

// NewRenderingService returns new EventRenderingService.
func NewRenderingService() *EventRenderingService {
	return &EventRenderingService{
		mutex: &sync.RWMutex{},
	}
}

// InvokeRendererMetrics contains metrics of invoke request
type InvokeRendererMetrics struct {
	ReadTime  time.Duration
	SizeBytes int
}

// InvokeRenderer knows how to render invoke event.
type InvokeRenderer struct {
	ctx                 context.Context
	invoke              *interop.Invoke
	tracingHeaderParser func(context.Context, *interop.Invoke) string
	requestBuffer       []byte
	requestMutex        sync.Mutex
	metrics             InvokeRendererMetrics
}

// NewAgentInvokeEvent forms a new AgentInvokeEvent from INVOKE request
func NewAgentInvokeEvent(req *interop.Invoke) (*model.AgentInvokeEvent, error) {
	deadlineMono, err := strconv.ParseInt(req.DeadlineNs, 10, 64)
	if err != nil {
		return nil, err
	}
	deadline := metering.MonoToEpoch(deadlineMono) / int64(time.Millisecond)
	return &model.AgentInvokeEvent{
		AgentEvent: &model.AgentEvent{
			EventType:  "INVOKE",
			DeadlineMs: deadline,
		},
		RequestID:          req.ID,
		InvokedFunctionArn: req.InvokedFunctionArn,
		Tracing:            model.NewXRayTracing(req.TraceID),
	}, nil
}

// RenderAgentEvent renders invoke event json for agent.
func (s *InvokeRenderer) RenderAgentEvent(writer http.ResponseWriter, request *http.Request) error {
	event, err := NewAgentInvokeEvent(s.invoke)
	if err != nil {
		return err
	}

	bytes, err := json.Marshal(event)
	if err != nil {
		return err
	}

	renderAgentInvokeHeaders(writer, uuid.New()) // TODO: check this thing

	if _, err := writer.Write(bytes); err != nil {
		return err
	}
	return nil
}

func (s *InvokeRenderer) bufferInvokeRequest() error {
	s.requestMutex.Lock()
	defer s.requestMutex.Unlock()
	var err error = nil
	if nil == s.requestBuffer {
		reader := io.LimitReader(s.invoke.Payload, interop.MaxPayloadSize)
		start := time.Now()
		s.requestBuffer, err = ioutil.ReadAll(reader)
		s.metrics = InvokeRendererMetrics{
			ReadTime:  time.Since(start),
			SizeBytes: len(s.requestBuffer),
		}
	}
	return err
}

// RenderRuntimeEvent renders invoke payload for runtime.
func (s *InvokeRenderer) RenderRuntimeEvent(writer http.ResponseWriter, request *http.Request) error {
	invoke := s.invoke
	customerTraceID := s.tracingHeaderParser(s.ctx, s.invoke)

	cognitoIdentityJSON := ""
	if len(invoke.CognitoIdentityID) != 0 || len(invoke.CognitoIdentityPoolID) != 0 {
		cognitoJSON, err := json.Marshal(model.CognitoIdentity{
			CognitoIdentityID:     invoke.CognitoIdentityID,
			CognitoIdentityPoolID: invoke.CognitoIdentityPoolID,
		})
		if err != nil {
			return err
		}

		cognitoIdentityJSON = string(cognitoJSON)
	}

	var deadlineHeader string
	if t, err := strconv.ParseInt(invoke.DeadlineNs, 10, 64); err == nil {
		deadlineHeader = strconv.FormatInt(metering.MonoToEpoch(t)/int64(time.Millisecond), 10)
	} else {
		log.WithError(err).Warn("Failed to compute deadline header")
	}

	renderInvokeHeaders(writer, invoke.ID, customerTraceID, invoke.ClientContext,
		cognitoIdentityJSON, invoke.InvokedFunctionArn, deadlineHeader, invoke.ContentType)

	if invoke.Payload != nil {
		if err := s.bufferInvokeRequest(); err != nil {
			return err
		}
		_, err := writer.Write(s.requestBuffer)
		return err
	}

	return nil
}

// NewInvokeRenderer returns new invoke event renderer
func NewInvokeRenderer(ctx context.Context, invoke *interop.Invoke, traceParser func(context.Context, *interop.Invoke) string) *InvokeRenderer {
	return &InvokeRenderer{
		invoke:              invoke,
		ctx:                 ctx,
		tracingHeaderParser: traceParser,
		requestBuffer:       nil,
		requestMutex:        sync.Mutex{},
	}
}

func (s *InvokeRenderer) GetMetrics() InvokeRendererMetrics {
	s.requestMutex.Lock()
	defer s.requestMutex.Unlock()
	return s.metrics
}

// ShutdownRenderer renderer for shutdown event.
type ShutdownRenderer struct {
	AgentEvent model.AgentShutdownEvent
}

// RenderAgentEvent renders shutdown event for agent.
func (s *ShutdownRenderer) RenderAgentEvent(w http.ResponseWriter, r *http.Request) error {
	bytes, err := json.Marshal(s.AgentEvent)
	if err != nil {
		return err
	}
	if _, err := w.Write(bytes); err != nil {
		return err
	}
	return nil
}

// RenderRuntimeEvent renders shutdown event for runtime.
func (s *ShutdownRenderer) RenderRuntimeEvent(w http.ResponseWriter, r *http.Request) error {
	panic("We should SIGTERM runtime")
}

func setHeaderIfNotEmpty(headers http.Header, key string, value string) {
	if len(value) != 0 {
		headers.Set(key, value)
	}
}

func setHeaderOrDefault(headers http.Header, key, val, defaultVal string) {
	if val == "" {
		headers.Set(key, defaultVal)
		return
	}
	headers.Set(key, val)
}

func renderInvokeHeaders(writer http.ResponseWriter, invokeID string, customerTraceID string, clientContext string,
	cognitoIdentity string, invokedFunctionArn string, deadlineMs string, contentType string) {
	headers := writer.Header()
	setHeaderIfNotEmpty(headers, "Lambda-Runtime-Aws-Request-Id", invokeID)
	setHeaderIfNotEmpty(headers, "Lambda-Runtime-Trace-Id", customerTraceID)
	setHeaderIfNotEmpty(headers, "Lambda-Runtime-Client-Context", clientContext)
	setHeaderIfNotEmpty(headers, "Lambda-Runtime-Cognito-Identity", cognitoIdentity)
	setHeaderIfNotEmpty(headers, "Lambda-Runtime-Invoked-Function-Arn", invokedFunctionArn)
	setHeaderIfNotEmpty(headers, "Lambda-Runtime-Deadline-Ms", deadlineMs)
	setHeaderOrDefault(headers, "Content-Type", contentType, "application/json")
	writer.WriteHeader(http.StatusOK)
}

// RenderRuntimeLogsResponse renders response from Telemetry API
func RenderRuntimeLogsResponse(w http.ResponseWriter, respBody []byte, status int, headers map[string][]string) error {
	respHeaders := w.Header()
	for k, vals := range headers {
		for _, v := range vals {
			respHeaders.Add(k, v)
		}
	}

	w.WriteHeader(status)

	_, err := w.Write(respBody)
	return err
}

func renderAgentInvokeHeaders(writer http.ResponseWriter, eventID uuid.UUID) {
	headers := writer.Header()
	headers.Set("Lambda-Extension-Event-Identifier", eventID.String())
	headers.Set("Content-Type", "application/json")
	writer.WriteHeader(http.StatusOK)
}

// RenderForbiddenWithTypeMsg method for rendering error response
func RenderForbiddenWithTypeMsg(w http.ResponseWriter, r *http.Request, errorType string, format string, args ...interface{}) {
	render.Status(r, http.StatusForbidden)
	render.JSON(w, r, &model.ErrorResponse{
		ErrorType:    errorType,
		ErrorMessage: fmt.Sprintf(format, args...),
	})
}

// RenderInternalServerError method for rendering error response
func RenderInternalServerError(w http.ResponseWriter, r *http.Request) {
	render.Status(r, http.StatusInternalServerError)
	render.JSON(w, r, &model.ErrorResponse{
		ErrorMessage: "Internal Server Error",
		ErrorType:    ErrorTypeInternalServerError,
	})
}

// RenderRequestEntityTooLarge method for rendering error response
func RenderRequestEntityTooLarge(w http.ResponseWriter, r *http.Request) {
	render.Status(r, http.StatusRequestEntityTooLarge)
	render.JSON(w, r, &model.ErrorResponse{
		ErrorMessage: fmt.Sprintf("Exceeded maximum allowed payload size (%d bytes).", interop.MaxPayloadSize),
		ErrorType:    ErrorTypeRequestEntityTooLarge,
	})
}

// RenderInvalidRequestID renders invalid request ID error response
func RenderInvalidRequestID(w http.ResponseWriter, r *http.Request) {
	render.Status(r, http.StatusBadRequest)
	render.JSON(w, r, &model.ErrorResponse{
		ErrorMessage: "Invalid request ID",
		ErrorType:    "InvalidRequestID",
	})
}

// RenderAccepted method for rendering accepted status response
func RenderAccepted(w http.ResponseWriter, r *http.Request) {
	render.Status(r, http.StatusAccepted)
	render.JSON(w, r, &model.StatusResponse{
		Status: "OK",
	})
}

// RenderInteropError is a convenience method for interpreting interop errors
func RenderInteropError(writer http.ResponseWriter, request *http.Request, err error) {
	if err == interop.ErrInvalidInvokeID || err == interop.ErrResponseSent {
		RenderInvalidRequestID(writer, request)
	} else {
		log.Panic(err)
	}
}


<|start_filename|>lambda/rapidcore/standalone/reserveHandler.go<|end_filename|>
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package standalone

import (
	"net/http"

	log "github.com/sirupsen/logrus"
	"go.amzn.com/lambda/core/directinvoke"
	"go.amzn.com/lambda/interop"
	"go.amzn.com/lambda/rapidcore"
)

const (
	ReservationTokenHeader = "Reservation-Token"
	InvokeIDHeader         = "Invoke-ID"
	VersionIDHeader        = "Version-ID"
)

func tokenToHeaders(w http.ResponseWriter, token interop.Token) {
	w.Header().Set(ReservationTokenHeader, token.ReservationToken)
	w.Header().Set(directinvoke.InvokeIDHeader, token.InvokeID)
	w.Header().Set(directinvoke.VersionIDHeader, token.VersionID)
}

func ReserveHandler(w http.ResponseWriter, r *http.Request, s rapidcore.InteropServer) {
	reserveResp, err := s.Reserve("", r.Header.Get("X-Amzn-Trace-Id"), r.Header.Get("X-Amzn-Segment-Id"))

	if err != nil {
		switch err {
		case rapidcore.ErrInitAlreadyDone:
			// init already happened before, just provide internal state and return
			tokenToHeaders(w, reserveResp.Token)
			InternalStateHandler(w, r, s)
		case rapidcore.ErrReserveReservationDone:
			// TODO use http.StatusBadGateway
			w.WriteHeader(http.StatusGatewayTimeout)
		case rapidcore.ErrInitDoneFailed, rapidcore.ErrInitError:
			w.WriteHeader(DoneFailedHTTPCode)
			w.Write(reserveResp.InternalState.AsJSON())
		case rapidcore.ErrTerminated:
			w.WriteHeader(DoneFailedHTTPCode)
			w.Write(reserveResp.InternalState.AsJSON())
		default:
			log.Errorf("Failed to reserve: %s", err)
			w.WriteHeader(400)
		}
		return
	}

	tokenToHeaders(w, reserveResp.Token)
	w.Write(reserveResp.InternalState.AsJSON())
}


<|start_filename|>cmd/aws-lambda-rie/handlers.go<|end_filename|>
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"math"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"go.amzn.com/lambda/interop"
	"go.amzn.com/lambda/rapidcore"

	"github.com/google/uuid"

	log "github.com/sirupsen/logrus"
)

type Sandbox interface {
	Init(i *interop.Init, invokeTimeoutMs int64)
	Invoke(responseWriter http.ResponseWriter, invoke *interop.Invoke) error
}

var initDone bool

func GetenvWithDefault(key string, defaultValue string) string {
	envValue := os.Getenv(key)

	if envValue == "" {
		return defaultValue
	}

	return envValue
}

func printEndReports(invokeId string, initDuration string, memorySize string, invokeStart time.Time, timeoutDuration time.Duration) {
	// Calcuation invoke duration
	invokeDuration := math.Min(float64(time.Now().Sub(invokeStart).Nanoseconds()),
		float64(timeoutDuration.Nanoseconds())) / float64(time.Millisecond)

	fmt.Println("END RequestId: " + invokeId)
	// We set the Max Memory Used and Memory Size to be the same (whatever it is set to) since there is
	// not a clean way to get this information from rapidcore
	fmt.Printf(
		"REPORT RequestId: %s\t"+
			initDuration+
			"Duration: %.2f ms\t"+
			"Billed Duration: %.f ms\t"+
			"Memory Size: %s MB\t"+
			"Max Memory Used: %s MB\t\n",
		invokeId, invokeDuration, math.Ceil(invokeDuration), memorySize, memorySize)
}

func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox) {
	log.Debugf("invoke: -> %s %s %v", r.Method, r.URL, r.Header)
	bodyBytes, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.Errorf("Failed to read invoke body: %s", err)
		w.WriteHeader(500)
		return
	}

	initDuration := ""
	inv := GetenvWithDefault("AWS_LAMBDA_FUNCTION_TIMEOUT", "300")
	timeoutDuration, _ := time.ParseDuration(inv + "s")
	// Default
	timeout, err := strconv.ParseInt(inv, 10, 64)
	if err != nil {
		panic(err)
	}

	functionVersion := GetenvWithDefault("AWS_LAMBDA_FUNCTION_VERSION", "$LATEST")
	memorySize := GetenvWithDefault("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "3008")

	if !initDone {

		initStart, initEnd := InitHandler(sandbox, functionVersion, timeout)

		// Calculate InitDuration
		initTimeMS := math.Min(float64(initEnd.Sub(initStart).Nanoseconds()),
			float64(timeoutDuration.Nanoseconds())) / float64(time.Millisecond)

		initDuration = fmt.Sprintf("Init Duration: %.2f ms\t", initTimeMS)

		// Set initDone so next invokes do not try to Init the function again
		initDone = true
	}

	invokeStart := time.Now()
	invokePayload := &interop.Invoke{
		ID:                 uuid.New().String(),
		InvokedFunctionArn: fmt.Sprintf("arn:aws:lambda:us-east-1:012345678912:function:%s", GetenvWithDefault("AWS_LAMBDA_FUNCTION_NAME", "test_function")),
		TraceID:            r.Header.Get("X-Amzn-Trace-Id"),
		LambdaSegmentID:    r.Header.Get("X-Amzn-Segment-Id"),
		Payload:            bytes.NewReader(bodyBytes),
		CorrelationID:      "invokeCorrelationID",
	}
	fmt.Println("START RequestId: " + invokePayload.ID + " Version: " + functionVersion)

	// If we write to 'w' directly and waitUntilRelease fails, we won't be able to propagate error anymore
	invokeResp := &ResponseWriterProxy{}
	if err := sandbox.Invoke(invokeResp, invokePayload); err != nil {
		switch err {

		// Reserve errors:
		case rapidcore.ErrAlreadyReserved:
			log.Errorf("Failed to reserve: %s", err)
			w.WriteHeader(http.StatusBadRequest)
			return
		case rapidcore.ErrInternalServerError:
			w.WriteHeader(http.StatusInternalServerError)
			return
		case rapidcore.ErrInitDoneFailed:
			w.WriteHeader(http.StatusBadGateway)
			w.Write(invokeResp.Body)
			return
		case rapidcore.ErrReserveReservationDone:
			// TODO use http.StatusBadGateway
			w.WriteHeader(http.StatusGatewayTimeout)
			return

		// Invoke errors:
		case rapidcore.ErrNotReserved:
		case rapidcore.ErrAlreadyReplied:
		case rapidcore.ErrAlreadyInvocating:
			log.Errorf("Failed to set reply stream: %s", err)
			w.WriteHeader(http.StatusBadRequest)
			return
		case rapidcore.ErrInvokeReservationDone:
			// TODO use http.StatusBadGateway
			w.WriteHeader(http.StatusGatewayTimeout)
			return
		case rapidcore.ErrInvokeResponseAlreadyWritten:
			return
		// AwaitRelease errors:
		case rapidcore.ErrInvokeDoneFailed:
			w.WriteHeader(http.StatusBadGateway)
			w.Write(invokeResp.Body)
			return
		case rapidcore.ErrReleaseReservationDone:
			// TODO return sandbox status when we implement async reset handling
			// TODO use http.StatusOK
			w.WriteHeader(http.StatusGatewayTimeout)
			return
		case rapidcore.ErrInvokeTimeout:
			printEndReports(invokePayload.ID, initDuration, memorySize, invokeStart, timeoutDuration)

			w.Write([]byte(fmt.Sprintf("Task timed out after %d.00 seconds", timeout)))
			time.Sleep(100 * time.Millisecond)
			//initDone = false
			return
		}
	}

	printEndReports(invokePayload.ID, initDuration, memorySize, invokeStart, timeoutDuration)

	if invokeResp.StatusCode != 0 {
		w.WriteHeader(invokeResp.StatusCode)
	}
	w.Write(invokeResp.Body)
}

func InitHandler(sandbox Sandbox, functionVersion string, timeout int64) (time.Time, time.Time) {
	additionalFunctionEnvironmentVariables := map[string]string{}

	// Add default Env Vars if they were not defined. This is a required otherwise 1p Python2.7, Python3.6, and
	// possibly others pre runtime API runtimes will fail. This will be overwritten if they are defined on the system.
	additionalFunctionEnvironmentVariables["AWS_LAMBDA_LOG_GROUP_NAME"] = "/aws/lambda/Functions"
	additionalFunctionEnvironmentVariables["AWS_LAMBDA_LOG_STREAM_NAME"] = "$LATEST"
	additionalFunctionEnvironmentVariables["AWS_LAMBDA_FUNCTION_VERSION"] = "$LATEST"
	additionalFunctionEnvironmentVariables["AWS_LAMBDA_FUNCTION_MEMORY_SIZE"] = "3008"
	additionalFunctionEnvironmentVariables["AWS_LAMBDA_FUNCTION_NAME"] = "test_function"

	// Forward Env Vars from the running system (container) to what the function can view. Without this, Env Vars will
	// not be viewable when the function runs.
	for _, env := range os.Environ() {
		// Split the env into by the first "=". This will account for if the env var's value has a '=' in it
		envVar := strings.SplitN(env, "=", 2)
		additionalFunctionEnvironmentVariables[envVar[0]] = envVar[1]
	}

	initStart := time.Now()
	// pass to rapid
	sandbox.Init(&interop.Init{
		Handler:           GetenvWithDefault("AWS_LAMBDA_FUNCTION_HANDLER", os.Getenv("_HANDLER")),
		CorrelationID:     "initCorrelationID",
		AwsKey:            os.Getenv("AWS_ACCESS_KEY_ID"),
		AwsSecret:         os.Getenv("AWS_SECRET_ACCESS_KEY"),
		AwsSession:        os.Getenv("AWS_SESSION_TOKEN"),
		XRayDaemonAddress: "0.0.0.0:0", // TODO
		FunctionName:      GetenvWithDefault("AWS_LAMBDA_FUNCTION_NAME", "test_function"),
		FunctionVersion:   functionVersion,

		CustomerEnvironmentVariables: additionalFunctionEnvironmentVariables,
	}, timeout*1000)
	initEnd := time.Now()
	return initStart, initEnd
}


<|start_filename|>lambda/rapid/bootstrap.go<|end_filename|>
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package rapid

import (
	"os"

	"go.amzn.com/lambda/fatalerror"
)

type Bootstrap interface {
	Cmd() ([]string, error)              // returns the args of bootstrap, where args[0] is the path to executable
	Env(e EnvironmentVariables) []string // returns the environment variables to be passed to the bootstrapped process
	Cwd() (string, error)                // returns the working directory of the bootstrap process
	ExtraFiles() []*os.File              // returns the extra file descriptors apart from 1 & 2 to be passed to runtime
	CachedFatalError(err error) (fatalerror.ErrorType, string, bool)
}


<|start_filename|>lambda/rapidcore/standalone/router.go<|end_filename|>
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package standalone

import (
	"context"
	"net/http"

	"go.amzn.com/lambda/rapidcore"
	"go.amzn.com/lambda/rapidcore/telemetry"

	"github.com/go-chi/chi"
)

func NewHTTPRouter(sandbox rapidcore.Sandbox, eventLog *telemetry.EventLog, shutdownFunc context.CancelFunc) *chi.Mux {
	ipcSrv := sandbox.InteropServer()
	r := chi.NewRouter()
	r.Use(standaloneAccessLogDecorator)
	r.Post("/2015-03-31/functions/*/invocations", func(w http.ResponseWriter, r *http.Request) { Execute(w, r, sandbox) })
	r.Post("/test/init", func(w http.ResponseWriter, r *http.Request) { InitHandler(w, r, sandbox) })
	r.Post("/test/reserve", func(w http.ResponseWriter, r *http.Request) { ReserveHandler(w, r, ipcSrv) })
	r.Post("/test/invoke", func(w http.ResponseWriter, r *http.Request) { InvokeHandler(w, r, ipcSrv) })
	r.Post("/test/waitUntilRelease", func(w http.ResponseWriter, r *http.Request) { WaitUntilReleaseHandler(w, r, ipcSrv) })
	r.Post("/test/reset", func(w http.ResponseWriter, r *http.Request) { ResetHandler(w, r, ipcSrv) })
	r.Post("/test/shutdown", func(w http.ResponseWriter, r *http.Request) { ShutdownHandler(w, r, ipcSrv, shutdownFunc) })
	r.Post("/test/directInvoke/{reservationtoken}", func(w http.ResponseWriter, r *http.Request) { DirectInvokeHandler(w, r, ipcSrv) })
	r.Get("/test/internalState", func(w http.ResponseWriter, r *http.Request) { InternalStateHandler(w, r, ipcSrv) })
	r.Get("/test/eventLog", func(w http.ResponseWriter, r *http.Request) { EventLogHandler(w, r, eventLog) })

	return r
}


<|start_filename|>lambda/rapidcore/standalone/resetHandler.go<|end_filename|>
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package standalone

import (
	"net/http"

	"go.amzn.com/lambda/rapidcore"
)

type resetAPIRequest struct {
	Reason    string `json:"reason"`
	TimeoutMs int64  `json:"timeoutMs"`
}

func ResetHandler(w http.ResponseWriter, r *http.Request, s rapidcore.InteropServer) {
	reset := resetAPIRequest{}
	if lerr := readBodyAndUnmarshalJSON(r, &reset); lerr != nil {
		lerr.Send(w, r)
		return
	}

	resetDescription, err := s.Reset(reset.Reason, reset.TimeoutMs)
	if err != nil {
		(&FailureReply{}).Send(w, r)
		return
	}

	w.Write(resetDescription.AsJSON())
}


<|start_filename|>lambda/core/directinvoke/customerheaders.go<|end_filename|>
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package directinvoke

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
)

type CustomerHeaders struct {
	CognitoIdentityID     string `json:"Cognito-Identity-Id"`
	CognitoIdentityPoolID string `json:"Cognito-Identity-Pool-Id"`
	ClientContext         string `json:"Client-Context"`
}

func (s CustomerHeaders) Dump() string {
	if (s == CustomerHeaders{}) {
		return ""
	}

	custHeadersJSON, err := json.Marshal(&s)
	if err != nil {
		panic(err)
	}

	return base64.StdEncoding.EncodeToString(custHeadersJSON)
}

func (s *CustomerHeaders) Load(in string) error {
	*s = CustomerHeaders{}

	if in == "" {
		return nil
	}

	base64Decoder := base64.NewDecoder(base64.StdEncoding, bytes.NewReader([]byte(in)))

	return json.NewDecoder(base64Decoder).Decode(s)
}


<|start_filename|>lambda/core/registrations.go<|end_filename|>
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package core

import (
	"errors"
	"os"
	"sync"

	"go.amzn.com/lambda/appctx"
	"go.amzn.com/lambda/core/statejson"

	"github.com/google/uuid"
)

type registrationServiceState int

const (
	registrationServiceOn registrationServiceState = iota
	registrationServiceOff
)

const MaxAgentsAllowed = 10

// Event represents a platform event which agent can subscribe to
type Event string

const (
	// InvokeEvent is dispatched when INVOKE happens
	InvokeEvent Event = "INVOKE"
	// ShutdownEvent is dispatched when SHUTDOWN or RESET happen
	ShutdownEvent Event = "SHUTDOWN"
)

// ErrRegistrationServiceOff returned on attempt to register after registration service has been turned off.
var ErrRegistrationServiceOff = errors.New("ErrRegistrationServiceOff")

// ErrTooManyExtensions means MaxAgentsAllowed limit is exceeded
var ErrTooManyExtensions = errors.New("ErrTooManyExtensions")

type AgentInfoErrorType string

const (
	PermissionDenied  AgentInfoErrorType = "PermissionDenied"
	TooManyExtensions AgentInfoErrorType = "TooManyExtensions"
	UnknownError      AgentInfoErrorType = "UnknownError"
)

func MapErrorToAgentInfoErrorType(err error) AgentInfoErrorType {
	if os.IsPermission(err) {
		return PermissionDenied
	}
	if err == ErrTooManyExtensions {
		return TooManyExtensions
	}
	return UnknownError
}

// AgentInfo holds information about an agent renderable in customer logs
type AgentInfo struct {
	Name          string
	State         string
	Subscriptions []string
	ErrorType     string
}

// FunctionMetadata holds static information regarding the function (Name, Version, Handler)
type FunctionMetadata struct {
	FunctionName    string
	FunctionVersion string
	Handler         string
}

// RegistrationService keeps track of registered parties, including external agents, threads, and runtime.
type RegistrationService interface {
	CreateExternalAgent(agentName string) (*ExternalAgent, error)
	CreateInternalAgent(agentName string) (*InternalAgent, error)
	PreregisterRuntime(r *Runtime) error
	SetFunctionMetadata(metadata FunctionMetadata)
	GetFunctionMetadata() FunctionMetadata
	GetRuntime() *Runtime
	GetRegisteredAgentsSize() uint16
	FindExternalAgentByName(agentName string) (*ExternalAgent, bool)
	FindInternalAgentByName(agentName string) (*InternalAgent, bool)
	FindExternalAgentByID(agentID uuid.UUID) (*ExternalAgent, bool)
	FindInternalAgentByID(agentID uuid.UUID) (*InternalAgent, bool)
	TurnOff()
	InitFlow() InitFlowSynchronization
	GetInternalStateDescriptor(appCtx appctx.ApplicationContext) func() statejson.InternalStateDescription
	GetExternalAgents() []*ExternalAgent
	GetSubscribedExternalAgents(eventType Event) []*ExternalAgent
	GetSubscribedInternalAgents(eventType Event) []*InternalAgent
	CountAgents() int
	Clear()
	AgentsInfo() []AgentInfo
}

type registrationServiceImpl struct {
	runtime          *Runtime
	internalAgents   InternalAgentsMap
	externalAgents   ExternalAgentsMap
	state            registrationServiceState
	mutex            *sync.Mutex
	initFlow         InitFlowSynchronization
	invokeFlow       InvokeFlowSynchronization
	functionMetadata FunctionMetadata
}

func (s *registrationServiceImpl) Clear() {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	s.runtime = nil
	s.internalAgents.Clear()
	s.externalAgents.Clear()
	s.state = registrationServiceOn
}

func (s *registrationServiceImpl) InitFlow() InitFlowSynchronization {
	return s.initFlow
}

// GetInternalStateDescriptor returns function that returns internal state description for debugging purposes
func (s *registrationServiceImpl) GetInternalStateDescriptor(appCtx appctx.ApplicationContext) func() statejson.InternalStateDescription {
	return func() statejson.InternalStateDescription {
		return s.getInternalStateDescription(appCtx)
	}
}

func (s *registrationServiceImpl) getInternalStateDescription(appCtx appctx.ApplicationContext) statejson.InternalStateDescription {
	isd := statejson.InternalStateDescription{
		Extensions: []statejson.ExtensionDescription{},
	}

	if s.runtime != nil {
		// we use pointer here so that 'runtime' json field is nil if runtime is not set (as opposed to filled with default values)
		rtdesc := s.runtime.GetRuntimeDescription()
		isd.Runtime = &rtdesc
	}

	s.mutex.Lock()
	defer s.mutex.Unlock()

	s.internalAgents.Visit(func(agent *InternalAgent) {
		isd.Extensions = append(isd.Extensions, agent.GetAgentDescription())
	})

	s.externalAgents.Visit(func(agent *ExternalAgent) {
		isd.Extensions = append(isd.Extensions, agent.GetAgentDescription())
	})

	if fatalerror, found := appctx.LoadFirstFatalError(appCtx); found {
		isd.FirstFatalError = string(fatalerror)
	}

	return isd
}

func (s *registrationServiceImpl) CountAgents() int {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	return s.countAgentsUnsafe()
}

func (s *registrationServiceImpl) countAgentsUnsafe() int {
	res := 0
	s.externalAgents.Visit(func(a *ExternalAgent) {
		res++
	})
	s.internalAgents.Visit(func(a *InternalAgent) {
		res++
	})
	return res
}

func (s *registrationServiceImpl) GetExternalAgents() []*ExternalAgent {
	agents := []*ExternalAgent{}
	s.externalAgents.Visit(func(a *ExternalAgent) {
		agents = append(agents, a)
	})
	return agents
}

func (s *registrationServiceImpl) GetInternalAgents() []*InternalAgent {
	agents := []*InternalAgent{}
	s.internalAgents.Visit(func(a *InternalAgent) {
		agents = append(agents, a)
	})
	return agents
}

func (s *registrationServiceImpl) GetSubscribedExternalAgents(eventType Event) []*ExternalAgent {
	agents := []*ExternalAgent{}
	s.externalAgents.Visit(func(a *ExternalAgent) {
		if a.IsSubscribed(eventType) {
			agents = append(agents, a)
		}
	})
	return agents
}

func (s *registrationServiceImpl) GetSubscribedInternalAgents(eventType Event) []*InternalAgent {
	agents := []*InternalAgent{}
	s.internalAgents.Visit(func(a *InternalAgent) {
		if a.IsSubscribed(eventType) {
			agents = append(agents, a)
		}
	})
	return agents
}

// CreateExternalAgent creates agent in agent collection
func (s *registrationServiceImpl) CreateExternalAgent(agentName string) (*ExternalAgent, error) {
	agent := NewExternalAgent(agentName, s.initFlow, s.invokeFlow)

	s.mutex.Lock()
	defer s.mutex.Unlock()

	if s.state != registrationServiceOn {
		return nil, ErrRegistrationServiceOff
	}

	if _, found := s.internalAgents.FindByName(agentName); found {
		return nil, ErrAgentNameCollision
	}

	if err := s.externalAgents.Insert(agent); err != nil {
		return nil, err
	}

	return agent, nil
}

// CreateInternalAgent creates agent in agent collection
func (s *registrationServiceImpl) CreateInternalAgent(agentName string) (*InternalAgent, error) {
	agent := NewInternalAgent(agentName, s.initFlow, s.invokeFlow)

	s.mutex.Lock()
	defer s.mutex.Unlock()

	if s.state != registrationServiceOn {
		return nil, ErrRegistrationServiceOff
	}

	if s.countAgentsUnsafe() >= MaxAgentsAllowed {
		return nil, ErrTooManyExtensions
	}

	if _, found := s.externalAgents.FindByName(agentName); found {
		return nil, ErrAgentNameCollision
	}

	if err := s.internalAgents.Insert(agent); err != nil {
		return nil, err
	}

	return agent, nil
}

// PreregisterRuntime allows to preregister a runtime.
func (s *registrationServiceImpl) PreregisterRuntime(r *Runtime) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if s.state != registrationServiceOn {
		return ErrRegistrationServiceOff
	}

	s.runtime = r

	// RUNTIME IS NOT PART OF SUBSCRIPTIONS

	return nil
}

// SetFunctionMetadata sets the static function metadata object
func (s *registrationServiceImpl) SetFunctionMetadata(metadata FunctionMetadata) {
	s.functionMetadata = metadata
}

// GetFunctionMetadata returns the static function metadata object
func (s *registrationServiceImpl) GetFunctionMetadata() FunctionMetadata {
	return s.functionMetadata
}

// GetRuntime retrieves runtime.
func (s *registrationServiceImpl) GetRuntime() *Runtime {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	return s.runtime
}

// GetRegisteredAgentsSize retrieves the number of agents registered with the platform.
func (s *registrationServiceImpl) GetRegisteredAgentsSize() uint16 {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	return uint16(s.externalAgents.Size()) + uint16(s.internalAgents.Size())
}

// FindExternalAgentByName
func (s *registrationServiceImpl) FindExternalAgentByName(name string) (agent *ExternalAgent, found bool) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if agent, found = s.externalAgents.FindByName(name); found {
		return
	}
	return
}

// FindInternalAgentByName
func (s *registrationServiceImpl) FindInternalAgentByName(name string) (agent *InternalAgent, found bool) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if agent, found = s.internalAgents.FindByName(name); found {
		return
	}
	return
}

// FindExternalAgentByID
func (s *registrationServiceImpl) FindExternalAgentByID(agentID uuid.UUID) (agent *ExternalAgent, found bool) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if agent, found = s.externalAgents.FindByID(agentID); found {
		return
	}
	return
}

// FindInternalAgentByID
func (s *registrationServiceImpl) FindInternalAgentByID(agentID uuid.UUID) (agent *InternalAgent, found bool) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if agent, found = s.internalAgents.FindByID(agentID); found {
		return
	}
	return
}

// ReportAgentsInfo returns information about all agents
func (s *registrationServiceImpl) AgentsInfo() []AgentInfo {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	agentsInfo := []AgentInfo{}
	for _, agent := range s.GetExternalAgents() {
		agentsInfo = append(agentsInfo, AgentInfo{
			agent.Name,
			agent.GetState().Name(),
			agent.SubscribedEvents(),
			agent.ErrorType(),
		})
	}

	for _, agent := range s.GetInternalAgents() {
		agentsInfo = append(agentsInfo, AgentInfo{
			agent.Name,
			agent.GetState().Name(),
			agent.SubscribedEvents(),
			agent.ErrorType(),
		})
	}

	return agentsInfo
}

// TurnOff turns off registration service.
func (s *registrationServiceImpl) TurnOff() {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	s.state = registrationServiceOff
}

// NewRegistrationService returns new RegistrationService instance.
func NewRegistrationService(initFlow InitFlowSynchronization, invokeFlow InvokeFlowSynchronization) RegistrationService {
	return &registrationServiceImpl{
		mutex:          &sync.Mutex{},
		state:          registrationServiceOn,
		internalAgents: NewInternalAgentsMap(),
		externalAgents: NewExternalAgentsMap(),
		initFlow:       initFlow,
		invokeFlow:     invokeFlow,
	}
}


<|start_filename|>lambda/rapid/exit.go<|end_filename|>
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package rapid

import (
	"fmt"
	"os"

	"go.amzn.com/lambda/extensions"
	"go.amzn.com/lambda/fatalerror"
	"go.amzn.com/lambda/interop"
	"go.amzn.com/lambda/rapi/model"

	log "github.com/sirupsen/logrus"
)

func checkInteropError(format string, err error) {
	if err == interop.ErrInvalidInvokeID || err == interop.ErrResponseSent {
		log.Warnf(format, err)
	} else {
		log.Panicf(format, err)
	}
}

func trySendDefaultErrorResponse(interopServer interop.Server, invokeID string, errorType fatalerror.ErrorType, err error) {
	resp := model.ErrorResponse{
		ErrorType:    string(errorType),
		ErrorMessage: fmt.Sprintf("Error: %v", err),
	}

	if invokeID != "" {
		resp.ErrorMessage = fmt.Sprintf("RequestId: %s Error: %v", invokeID, err)
	}

	if err := interopServer.SendErrorResponse(invokeID, resp.AsInteropError()); err != nil {
		checkInteropError("Failed to send default error response: %s", err)
	}
}

func reportErrorAndExit(doneFailMsg *interop.DoneFail, invokeID string, interopServer interop.Server, err error) {
	// This function maintains compatibility of exit sequence behaviour
	// with Sandbox Factory in the absence of extensions

	// NOTE this check will prevent us from sending FAULT message in case
	// response (positive or negative) has already been sent. This is done
	// to maintain legacy behavior of RAPID.
	// ALSO NOTE, this works in case of positive response because this will
	// be followed by RAPID exit.
	if !interopServer.IsResponseSent() {
		trySendDefaultErrorResponse(interopServer, invokeID, doneFailMsg.ErrorType, err)
	}

	if err := interopServer.CommitResponse(); err != nil {
		checkInteropError("Failed to commit error response: %s", err)
	}

	// old behavior: no DoneFails
	doneMsg := &interop.Done{
		WaitForExit:   true,
		CorrelationID: doneFailMsg.CorrelationID, // required for standalone mode
		Meta:          doneFailMsg.Meta,
	}

	if err := interopServer.SendDone(doneMsg); err != nil {
		checkInteropError("Failed to send DONE during exit: %s", err)
	}

	os.Exit(1)
}

func reportErrorAndRequestReset(doneFailMsg *interop.DoneFail, invokeID string, interopServer interop.Server, err error) {
	if err == errResetReceived {
		// errResetReceived is returned when execution flow was interrupted by the Reset message,
		// hence this error deserves special handling and we yield to main receive loop to handle it
		return
	}

	trySendDefaultErrorResponse(interopServer, invokeID, doneFailMsg.ErrorType, err)

	if err := interopServer.CommitResponse(); err != nil {
		checkInteropError("Failed to commit error response: %s", err)
	}

	if err := interopServer.SendDoneFail(doneFailMsg); err != nil {
		checkInteropError("Failed to send DONEFAIL: %s", err)
	}
}

func handleInitError(doneFailMsg *interop.DoneFail, execCtx *rapidContext, invokeID string, interopServer interop.Server, err error) {
	if execCtx.standaloneMode {
		reportErrorAndRequestReset(doneFailMsg, invokeID, interopServer, err)
		return
	}

	if !execCtx.HasActiveExtensions() {
		// we don't expect Slicer to send RESET during INIT, that's why we Exit here
		reportErrorAndExit(doneFailMsg, invokeID, interopServer, err)
	}

	reportErrorAndRequestReset(doneFailMsg, invokeID, interopServer, err)
}

func handleInvokeError(doneFailMsg *interop.DoneFail, execCtx *rapidContext, invokeID string, interopServer interop.Server, err error) {
	if execCtx.standaloneMode {
		reportErrorAndRequestReset(doneFailMsg, invokeID, interopServer, err)
		return
	}

	// Invoke with extensions disabled maintains behaviour parity with pre-extensions rapid
	if !extensions.AreEnabled() {
		reportErrorAndExit(doneFailMsg, invokeID, interopServer, err)
	}

	reportErrorAndRequestReset(doneFailMsg, invokeID, interopServer, err)
}


<|start_filename|>lambda/rapi/model/agentevent.go<|end_filename|>
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package model

// AgentEvent is one of INVOKE, SHUTDOWN agent events
type AgentEvent struct {
	EventType  string `json:"eventType"`
	DeadlineMs int64  `json:"deadlineMs"`
}

// AgentInvokeEvent is the response to agent's get next request
type AgentInvokeEvent struct {
	*AgentEvent
	RequestID          string   `json:"requestId"`
	InvokedFunctionArn string   `json:"invokedFunctionArn"`
	Tracing            *Tracing `json:"tracing,omitempty"`
}

// AgentShutdownEvent is the response to agent's get next request
type AgentShutdownEvent struct {
	*AgentEvent
	ShutdownReason string `json:"shutdownReason"`
}


<|start_filename|>lambda/rapidcore/standalone/directInvokeHandler.go<|end_filename|>
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package standalone

import (
	log "github.com/sirupsen/logrus"
	"go.amzn.com/lambda/core/directinvoke"
	"go.amzn.com/lambda/rapidcore"
	"net/http"
)

func DirectInvokeHandler(w http.ResponseWriter, r *http.Request, s rapidcore.InteropServer) {
	tok := s.CurrentToken()
	if tok == nil {
		log.Errorf("Attempt to call directInvoke without Reserve")
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	invoke, err := directinvoke.ReceiveDirectInvoke(w, r, *tok)
	if err != nil {
		log.Errorf("direct invoke error: %s", err)
		return
	}

	if err := s.FastInvoke(w, invoke, true); err != nil {
		switch err {
		case rapidcore.ErrNotReserved:
		case rapidcore.ErrAlreadyReplied:
		case rapidcore.ErrAlreadyInvocating:
			log.Errorf("Failed to set reply stream: %s", err)
			w.WriteHeader(http.StatusBadRequest)
			return
		case rapidcore.ErrInvokeReservationDone:
			w.WriteHeader(http.StatusBadGateway)
		}
	}
}


<|start_filename|>lambda/rapi/router_test.go<|end_filename|>
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package rapi

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"

	"go.amzn.com/lambda/interop"
	"go.amzn.com/lambda/rapi/model"

	"go.amzn.com/lambda/testdata"
)

func createInvoke(id string) *interop.Invoke {
	return &interop.Invoke{
		ID:                 id,
		InvokedFunctionArn: "arn::dummy:Function",
		Payload:            strings.NewReader("{\"invoke\":\"" + id + "\"}"),
		DeadlineNs:         "123456",
	}
}

// Make a test request
func makeTestRequest(t *testing.T, router http.Handler, request *http.Request) *httptest.ResponseRecorder {
	responseRecorder := httptest.NewRecorder()
	router.ServeHTTP(responseRecorder, request)
	t.Logf("test(%v) = %v", request.URL, responseRecorder.Code)
	return responseRecorder
}

// Make a test request in a benchmark
func makeBenchRequest(b *testing.B, router http.Handler, request *http.Request) *httptest.ResponseRecorder {
	responseRecorder := httptest.NewRecorder()
	b.StartTimer()
	router.ServeHTTP(responseRecorder, request)
	b.StopTimer()
	return responseRecorder
}

// Verify response error type
func assertResponseErrorType(t *testing.T, expectedErrorType string, response *httptest.ResponseRecorder) {
	errResp := model.ErrorResponse{}
	err := json.Unmarshal(response.Body.Bytes(), &errResp)
	assert.Nil(t, err)
	assert.Equal(t, expectedErrorType, errResp.ErrorType)
}

// TestAcceptXML tests that server response is always
// rendered as JSON, regardless of the value provided
// in "Accept" header.
//
// When using render.Render(...), chi rendering library
// would attempt to render response using content type
// specified in the "Accept" header.
//
// The purpose of this test is to confirm that RAPID
// renders all server responses as application/json.
func TestAcceptXML(t *testing.T) {
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	responseRecorder := httptest.NewRecorder()
	request := httptest.NewRequest("POST", "/runtime/invocation/x-y-z/error", bytes.NewReader([]byte("")))
	// Tell server that client side accepts "application/xml".
	request.Header.Add("Accept", "application/xml")
	// Serve request.
	router.ServeHTTP(responseRecorder, request)
	v := &model.ErrorResponse{}
	// Expected response is 403 transition is not allowed, rendered as JSON.
	err := json.Unmarshal(responseRecorder.Body.Bytes(), v)
	if err != nil {
		t.Error("Expected JSON document, received: ", responseRecorder.Body.String())
	}
	assert.Equal(t, "InvalidRequestID", v.ErrorType)
	assert.Equal(t, http.StatusBadRequest, responseRecorder.Code)
}

// Verify that unsupported methods return 404
func Test404PageNotFound(t *testing.T) {
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	responseRecorder := makeTestRequest(t, router, httptest.NewRequest("POST", "/runtime/unsupported", bytes.NewReader([]byte(""))))
	assert.Equal(t, http.StatusNotFound, responseRecorder.Code)
	assert.Equal(t, "404 page not found\n", responseRecorder.Body.String())
}

func Test405MethodNotAllowed(t *testing.T) {
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	responseRecorder := makeTestRequest(t, router, httptest.NewRequest("DELETE", "/runtime/invocation/ABC/error", bytes.NewReader([]byte(""))))
	assert.Equal(t, http.StatusMethodNotAllowed, responseRecorder.Code)
}

func TestInitErrorAccepted(t *testing.T) {
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	responseRecorder := makeTestRequest(t, router, httptest.NewRequest("POST", "/runtime/init/error", bytes.NewReader([]byte("{}"))))
	assert.Equal(t, http.StatusAccepted, responseRecorder.Code)
}

func TestInitErrorForbidden(t *testing.T) {
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	flowTest.ConfigureForInvoke(context.Background(), createInvoke("InvokeA"))
	responseRecorder := makeTestRequest(t, router, httptest.NewRequest("GET", "/runtime/invocation/next", nil))
	assert.Equal(t, http.StatusOK, responseRecorder.Code)
	responseRecorder = makeTestRequest(t, router, httptest.NewRequest("POST", "/runtime/init/error", bytes.NewReader([]byte("{}"))))
	assert.Equal(t, http.StatusForbidden, responseRecorder.Code)
}

func TestInvokeResponseAccepted(t *testing.T) {
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	flowTest.ConfigureForInvoke(context.Background(), createInvoke("InvokeA"))
	responseRecorder := makeTestRequest(t, router, httptest.NewRequest("GET", "/runtime/invocation/next", nil))
	assert.Equal(t, http.StatusOK, responseRecorder.Code)
	responseRecorder = makeTestRequest(t, router, httptest.NewRequest("POST", "/runtime/invocation/InvokeA/response", bytes.NewReader([]byte("{}"))))
	assert.Equal(t, http.StatusAccepted, responseRecorder.Code)
}

func TestInvokeErrorResponseAccepted(t *testing.T) {
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	flowTest.ConfigureForInvoke(context.Background(), createInvoke("InvokeA"))
	responseRecorder := makeTestRequest(t, router, httptest.NewRequest("GET", "/runtime/invocation/next", nil))
	assert.Equal(t, http.StatusOK, responseRecorder.Code)
	responseRecorder = makeTestRequest(t, router, httptest.NewRequest("POST", "/runtime/invocation/InvokeA/error", bytes.NewReader([]byte("{}"))))
	assert.Equal(t, http.StatusAccepted, responseRecorder.Code)
}

func TestInvokeNextTwice(t *testing.T) {
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	flowTest.ConfigureForInvoke(context.Background(), createInvoke("InvokeA"))
	responseRecorder := makeTestRequest(t, router, httptest.NewRequest("GET", "/runtime/invocation/next", nil))
	assert.Equal(t, http.StatusOK, responseRecorder.Code)
	responseRecorder = makeTestRequest(t, router, httptest.NewRequest("GET", "/runtime/invocation/next", nil))
	assert.Equal(t, http.StatusOK, responseRecorder.Code)
}

func TestInvokeResponseInvalidRequestID(t *testing.T) {
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	flowTest.ConfigureForInvoke(context.Background(), createInvoke("ABC"))
	responseRecorder := makeTestRequest(t, router, httptest.NewRequest("GET", "/runtime/invocation/next", nil))
	assert.Equal(t, http.StatusOK, responseRecorder.Code)
	responseRecorder = makeTestRequest(t, router, httptest.NewRequest("POST", "/runtime/invocation/XYZ/response", bytes.NewReader([]byte("{}"))))
	assert.Equal(t, http.StatusBadRequest, responseRecorder.Code)
	assertResponseErrorType(t, "InvalidRequestID", responseRecorder)
}

func TestInvokeErrorResponseInvalidRequestID(t *testing.T) {
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	flowTest.ConfigureForInvoke(context.Background(), createInvoke("ABC"))
	responseRecorder := makeTestRequest(t, router, httptest.NewRequest("GET", "/runtime/invocation/next", nil))
	assert.Equal(t, http.StatusOK, responseRecorder.Code)
	responseRecorder = makeTestRequest(t, router, httptest.NewRequest("POST", "/runtime/invocation/XYZ/error", bytes.NewReader([]byte("{}"))))
	assert.Equal(t, http.StatusBadRequest, responseRecorder.Code)
	assertResponseErrorType(t, "InvalidRequestID", responseRecorder)
}

func TestInvokeResponseTwice(t *testing.T) {
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	flowTest.ConfigureForInvoke(context.Background(), createInvoke("ABC"))
	responseRecorder := makeTestRequest(t, router, httptest.NewRequest("GET", "/runtime/invocation/next", nil))
	assert.Equal(t, http.StatusOK, responseRecorder.Code)
	responseRecorder = makeTestRequest(t, router, httptest.NewRequest("POST", "/runtime/invocation/ABC/response", bytes.NewReader([]byte("{}"))))
	assert.Equal(t, http.StatusAccepted, responseRecorder.Code)
	responseRecorder = makeTestRequest(t, router, httptest.NewRequest("POST", "/runtime/invocation/ABC/response", bytes.NewReader([]byte("{}"))))
	assert.Equal(t, http.StatusForbidden, responseRecorder.Code)
	assertResponseErrorType(t, "InvalidStateTransition", responseRecorder)
}

func TestInvokeErrorResponseTwice(t *testing.T) {
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	flowTest.ConfigureForInvoke(context.Background(), createInvoke("ABC"))
	responseRecorder := makeTestRequest(t, router, httptest.NewRequest("GET", "/runtime/invocation/next", nil))
	assert.Equal(t, http.StatusOK, responseRecorder.Code)
	responseRecorder = makeTestRequest(t, router, httptest.NewRequest("POST", "/runtime/invocation/ABC/error", bytes.NewReader([]byte("{}"))))
	assert.Equal(t, http.StatusAccepted, responseRecorder.Code)
	responseRecorder = makeTestRequest(t, router, httptest.NewRequest("POST", "/runtime/invocation/ABC/error", bytes.NewReader([]byte("{}"))))
	assert.Equal(t, http.StatusForbidden, responseRecorder.Code)
	assertResponseErrorType(t, "InvalidStateTransition", responseRecorder)
}

func TestInvokeResponseAfterErrorResponse(t *testing.T) {
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	flowTest.ConfigureForInvoke(context.Background(), createInvoke("ABC"))
	responseRecorder := makeTestRequest(t, router, httptest.NewRequest("GET", "/runtime/invocation/next", nil))
	assert.Equal(t, http.StatusOK, responseRecorder.Code)
	responseRecorder = makeTestRequest(t, router, httptest.NewRequest("POST", "/runtime/invocation/ABC/error", bytes.NewReader([]byte("{}"))))
	assert.Equal(t, http.StatusAccepted, responseRecorder.Code)
	responseRecorder = makeTestRequest(t, router, httptest.NewRequest("POST", "/runtime/invocation/ABC/response", bytes.NewReader([]byte("{}"))))
	assert.Equal(t, http.StatusForbidden, responseRecorder.Code)
	assertResponseErrorType(t, "InvalidStateTransition", responseRecorder)
}

func TestInvokeErrorResponseAfterResponse(t *testing.T) {
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	flowTest.ConfigureForInvoke(context.Background(), createInvoke("ABC"))
	responseRecorder := makeTestRequest(t, router, httptest.NewRequest("GET", "/runtime/invocation/next", nil))
	assert.Equal(t, http.StatusOK, responseRecorder.Code)
	responseRecorder = makeTestRequest(t, router, httptest.NewRequest("POST", "/runtime/invocation/ABC/response", bytes.NewReader([]byte("{}"))))
	assert.Equal(t, http.StatusAccepted, responseRecorder.Code)
	responseRecorder = makeTestRequest(t, router, httptest.NewRequest("POST", "/runtime/invocation/ABC/error", bytes.NewReader([]byte("{}"))))
	assert.Equal(t, http.StatusForbidden, responseRecorder.Code)
	assertResponseErrorType(t, "InvalidStateTransition", responseRecorder)
}

func TestMoreThanOneInvoke(t *testing.T) {
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	var responseRecorder *httptest.ResponseRecorder
	for _, id := range []string{"A", "B", "C"} {
		flowTest.ConfigureForInvoke(context.Background(), createInvoke(id))
		responseRecorder = makeTestRequest(t, router, httptest.NewRequest("GET", "/runtime/invocation/next", nil))
		assert.Equal(t, http.StatusOK, responseRecorder.Code)
		responseRecorder = makeTestRequest(t, router, httptest.NewRequest("POST", fmt.Sprintf("/runtime/invocation/%s/response", id), bytes.NewReader([]byte("{}"))))
		assert.Equal(t, http.StatusAccepted, responseRecorder.Code)
	}
}

func benchmarkInvoke(b *testing.B, payload []byte) {
	b.StopTimer()
	b.ReportAllocs()
	flowTest := testdata.NewFlowTest()
	flowTest.ConfigureForInit()
	router := NewRouter(flowTest.AppCtx, flowTest.RegistrationService, flowTest.RenderingService)
	for i := 0; i < b.N; i++ {
		id := uuid.New().String()
		flowTest.ConfigureForInvoke(context.Background(), createInvoke(id))
		makeBenchRequest(b, router, httptest.NewRequest("GET", "/runtime/invocation/next", nil))
		makeBenchRequest(b, router, httptest.NewRequest("POST", fmt.Sprintf("/runtime/invocation/%s/response", id), bytes.NewReader(payload)))
	}
}

func BenchmarkInvokeWithEmptyPayload(b *testing.B) {
	benchmarkInvoke(b, []byte(""))
}

func BenchmarkInvokeWith4KBPayload(b *testing.B) {
	benchmarkInvoke(b, bytes.Repeat([]byte("a"), 4*1024))
}

func BenchmarkInvokeWith512KBPayload(b *testing.B) {
	benchmarkInvoke(b, bytes.Repeat([]byte("a"), 512*1024))
}

func BenchmarkInvokeWith1MBPayload(b *testing.B) {
	benchmarkInvoke(b, bytes.Repeat([]byte("a"), 1*1024*1024))
}

func BenchmarkInvokeWith2MBPayload(b *testing.B) {
	benchmarkInvoke(b, bytes.Repeat([]byte("a"), 2*1024*1024))
}

func BenchmarkInvokeWith4MBPayload(b *testing.B) {
	benchmarkInvoke(b, bytes.Repeat([]byte("a"), 4*1024*1024))
}

func BenchmarkInvokeWith6MBPayload(b *testing.B) {
	benchmarkInvoke(b, bytes.Repeat([]byte("a"), 6*1024*1024))
}
