Decode and pass ClientContext during dispatch

cmd/localstack/custom_interop.go

@@ -5,6 +5,7 @@ package main
 
 import (
 	"bytes"
+	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -104,6 +105,47 @@ func (l *LocalStackAdapter) SendResult(invokeId string, body []byte, isError boo
 	return l.post(endpoint, body)
 }
 
+func decodeInvokeRequest(body []byte) (*interop.Invoke, error) {
+	var request *lsapi.InvokeRequest
+	if err := json.Unmarshal(body, &request); err != nil {
+		return nil, err
+	}
+	if request == nil {
+		return nil, errors.New("invocation request must be an object")
+	}
+	clientContext, err := base64.StdEncoding.DecodeString(request.ClientContext)
+	if err != nil {
+		return nil, err
+	}
+	return &interop.Invoke{
+		ID:                 request.InvokeId,
+		InvokedFunctionArn: request.InvokedFunctionArn,
+		Payload:            strings.NewReader(request.Payload),
+		NeedDebugLogs:      true,
+		TraceID:            request.TraceId,
+		ClientContext:      string(clientContext),
+	}, nil
+}
+
+// newInvokeHandler validates requests before acknowledging and dispatching them asynchronously
+func newInvokeHandler(dispatch func(*interop.Invoke)) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		body, err := io.ReadAll(r.Body)
+		if err != nil {
+			http.Error(w, "Failed to read invocation request", http.StatusInternalServerError)
+			return
+		}
+		invoke, err := decodeInvokeRequest(body)
+		if err != nil {
+			http.Error(w, "Invalid invocation request", http.StatusBadRequest)
+			return
+		}
+		go dispatch(invoke)
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("OK"))
+	}
+}
+
 func NewCustomInteropServer(lsOpts *LsOpts, delegate interop.Server, logCollector *LogCollector, eventsAPI *LocalStackEventsAPI) (server *CustomInteropServer) {
 	server = &CustomInteropServer{
 		delegate:         delegate.(*rapidcore.Server),
@@ -117,114 +159,80 @@ func NewCustomInteropServer(lsOpts *LsOpts, delegate interop.Server, logCollecto
 		resetDone: make(chan struct{}, 1),
 	}
 
-	// TODO: extract this
 	go func() {
 		r := chi.NewRouter()
-		r.Post("/invoke", func(w http.ResponseWriter, r *http.Request) {
-			invokeR := lsapi.InvokeRequest{}
-			bytess, err := io.ReadAll(r.Body)
+		r.Post("/invoke", newInvokeHandler(func(invoke *interop.Invoke) {
+			invokeResp := &standalone.ResponseWriterProxy{}
+			// The synthetic START and INIT_REPORT lines are emitted via LocalStackEventsAPI
+			// from rapid's lifecycle events, so they land at the AWS-faithful points (e.g.
+			// after an inline suppressed init's own logs) — see events.go.
+
+			// First invocation into a successfully initialized on-demand environment: REPORT
+			// carries the Init phase duration as measured by rapid (take-once; absent on warm
+			// starts, failed/timed-out inits, and non-on-demand environments).
+			initDurationMS, hasInitDuration := server.eventsAPI.TakeColdStartInitDuration()
+
+			invokeStart := time.Now()
+			err := server.Invoke(invokeResp, invoke)
+			timeout := int(server.delegate.GetInvokeTimeout().Seconds())
+			isErr := false
+			status := ""
+			errorType := ""
 			if err != nil {
-				log.Error(err)
-			}
-
-			go func() {
-				err = json.Unmarshal(bytess, &invokeR)
-				if err != nil {
-					log.Error(err)
-				}
-
-				invokeResp := &standalone.ResponseWriterProxy{}
-				// The synthetic START and INIT_REPORT lines are emitted via LocalStackEventsAPI
-				// from rapid's lifecycle events, so they land at the AWS-faithful points (e.g.
-				// after an inline suppressed init's own logs) — see events.go.
-
-				// First invocation into a successfully initialized on-demand environment: REPORT
-				// carries the Init phase duration as measured by rapid (take-once; absent on warm
-				// starts, failed/timed-out inits, and non-on-demand environments).
-				initDurationMS, hasInitDuration := server.eventsAPI.TakeColdStartInitDuration()
-
-				invokeStart := time.Now()
-				err = server.Invoke(invokeResp, &interop.Invoke{
-					ID:                 invokeR.InvokeId,
-					InvokedFunctionArn: invokeR.InvokedFunctionArn,
-					Payload:            strings.NewReader(invokeR.Payload), // r.Body,
-					NeedDebugLogs:      true,
-					TraceID:            invokeR.TraceId,
-					// TODO: set correct segment ID from request
-					//LambdaSegmentID:    "LambdaSegmentID", // r.Header.Get("X-Amzn-Segment-Id"),
-					//CognitoIdentityID:     "",
-					//CognitoIdentityPoolID: "",
-					//DeadlineNs:            "",
-					//ClientContext:         "",
-					//ContentType:           "",
-					//ReservationToken:      "",
-					//VersionID:             "",
-					//InvokeReceivedTime:    0,
-					//ResyncState:           interop.Resync{},
-				})
-				timeout := int(server.delegate.GetInvokeTimeout().Seconds())
-				isErr := false
-				status := ""
-				errorType := ""
-				if err != nil {
-					switch {
-					case errors.Is(err, rapidcore.ErrInvokeTimeout):
-						log.Debugf("Got invoke timeout")
-						isErr = true
-						status = "timeout"
-						errorResponse := lsapi.ErrorResponse{
-							ErrorType: "Sandbox.Timedout",
-							ErrorMessage: fmt.Sprintf(
-								"RequestId: %s Error: Task timed out after %d.00 seconds",
-								invokeR.InvokeId,
-								timeout,
-							),
-						}
-						jsonErrorResponse, err := json.Marshal(errorResponse)
-						if err != nil {
-							log.Fatalln("unable to marshall json timeout response")
-						}
-						_, err = invokeResp.Write(jsonErrorResponse)
-						if err != nil {
-							log.Fatalln("unable to write to response")
-						}
-					case errors.Is(err, rapidcore.ErrInvokeDoneFailed):
-						// The error response body was already written by rapid and is sent below.
-						// When an init failure was folded into this invocation (AWS suppressed
-						// init), the REPORT additionally carries the failure status and the
-						// scrubbed fatal error type (e.g. Runtime.Unknown).
-						if errType := server.eventsAPI.InitErrorType(); errType != "" {
-							isErr = true
-							status = "error"
-							errorType = errType
-						}
-					default:
-						log.Fatalln(err)
+				switch {
+				case errors.Is(err, rapidcore.ErrInvokeTimeout):
+					log.Debugf("Got invoke timeout")
+					isErr = true
+					status = "timeout"
+					errorResponse := lsapi.ErrorResponse{
+						ErrorType: "Sandbox.Timedout",
+						ErrorMessage: fmt.Sprintf(
+							"RequestId: %s Error: Task timed out after %d.00 seconds",
+							invoke.ID,
+							timeout,
+						),
 					}
-				}
-				// optional sleep. can be used for debugging purposes
-				if lsOpts.PostInvokeWaitMS != "" {
-					waitMS, err := strconv.Atoi(lsOpts.PostInvokeWaitMS)
+					jsonErrorResponse, err := json.Marshal(errorResponse)
 					if err != nil {
-						log.Fatalln(err)
+						log.Fatalln("unable to marshall json timeout response")
 					}
-					time.Sleep(time.Duration(waitMS) * time.Millisecond)
-				}
-				timeoutDuration := time.Duration(timeout) * time.Second
-				memorySize := GetEnvOrDie("AWS_LAMBDA_FUNCTION_MEMORY_SIZE")
-				PrintEndReports(invokeR.InvokeId, initDurationMS, hasInitDuration, status, errorType, memorySize, invokeStart, timeoutDuration, logCollector)
-
-				if err2 := server.localStackAdapter.SendLogs(invokeR.InvokeId, logCollector.getLogs()); err2 != nil {
-					log.Error("failed to send logs to LocalStack: ", err2)
+					_, err = invokeResp.Write(jsonErrorResponse)
+					if err != nil {
+						log.Fatalln("unable to write to response")
+					}
+				case errors.Is(err, rapidcore.ErrInvokeDoneFailed):
+					// The error response body was already written by rapid and is sent below.
+					// When an init failure was folded into this invocation (AWS suppressed
+					// init), the REPORT additionally carries the failure status and the
+					// scrubbed fatal error type (e.g. Runtime.Unknown).
+					if errType := server.eventsAPI.InitErrorType(); errType != "" {
+						isErr = true
+						status = "error"
+						errorType = errType
+					}
+				default:
+					log.Fatalln(err)
 				}
-				if err2 := server.localStackAdapter.SendResult(invokeR.InvokeId, invokeResp.Body, isErr); err2 != nil {
-					log.Error("failed to send result to LocalStack: ", err2)
+			}
+			// optional sleep. can be used for debugging purposes
+			if lsOpts.PostInvokeWaitMS != "" {
+				waitMS, err := strconv.Atoi(lsOpts.PostInvokeWaitMS)
+				if err != nil {
+					log.Fatalln(err)
 				}
-			}()
+				time.Sleep(time.Duration(waitMS) * time.Millisecond)
+			}
+			timeoutDuration := time.Duration(timeout) * time.Second
+			memorySize := GetEnvOrDie("AWS_LAMBDA_FUNCTION_MEMORY_SIZE")
+			PrintEndReports(invoke.ID, initDurationMS, hasInitDuration, status, errorType, memorySize, invokeStart, timeoutDuration, logCollector)
 
-			w.WriteHeader(200)
-			_, _ = w.Write([]byte("OK"))
-		})
+			if err2 := server.localStackAdapter.SendLogs(invoke.ID, logCollector.getLogs()); err2 != nil {
+				log.Error("failed to send logs to LocalStack: ", err2)
+			}
+			if err2 := server.localStackAdapter.SendResult(invoke.ID, invokeResp.Body, isErr); err2 != nil {
+				log.Error("failed to send result to LocalStack: ", err2)
+			}
+		}))
 		err := http.ListenAndServe(":"+server.port, r)
 		if err != nil {
 			log.Error(err)

cmd/localstack/custom_interop_test.go

@@ -1,12 +1,17 @@
 package main
 
 import (
+	"encoding/base64"
 	"encoding/json"
+	"errors"
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
+	"time"
 
+	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/interop"
 	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lsapi"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -36,6 +41,137 @@ func TestInvokeRequestContract(t *testing.T) {
 	assert.Equal(t, "arn:aws:lambda:us-east-1:000000000000:function:my-fn", req.InvokedFunctionArn)
 	assert.Equal(t, `{"key":"value"}`, req.Payload)
 	assert.Equal(t, "Root=1-abc;Parent=def;Sampled=1", req.TraceId)
+	assert.Empty(t, req.ClientContext)
+}
+
+// TestInvokeRequestClientContextContract checks that JSON encoding preserves the base64
+// context verbatim and omits an empty context. Decoding belongs to the invocation handler.
+func TestInvokeRequestClientContextContract(t *testing.T) {
+	raw := `{"client-context":"eyJjdXN0b20iOnsiZm9vIjoiYmFyIn19"}`
+	var req lsapi.InvokeRequest
+	require.NoError(t, json.Unmarshal([]byte(raw), &req))
+	assert.Equal(t, "eyJjdXN0b20iOnsiZm9vIjoiYmFyIn19", req.ClientContext)
+	encoded, err := json.Marshal(req)
+	require.NoError(t, err)
+	assert.Contains(t, string(encoded), `"client-context":"eyJjdXN0b20iOnsiZm9vIjoiYmFyIn19"`)
+	encoded, err = json.Marshal(lsapi.InvokeRequest{})
+	require.NoError(t, err)
+	assert.NotContains(t, string(encoded), "client-context")
+}
+
+// TestInvokeHandlerForwardsContextPerRequest reuses one handler across successive requests
+// to check decoding, context isolation, and preservation of the remaining invocation fields.
+func TestInvokeHandlerForwardsContextPerRequest(t *testing.T) {
+	dispatched := make(chan *interop.Invoke, 1)
+	handler := newInvokeHandler(func(invoke *interop.Invoke) { dispatched <- invoke })
+	contexts := []struct {
+		name    string
+		context string
+		present bool
+	}{
+		{"full", `{"custom":{"foo":"bar"},"client":{"app_title":"test"},"env":{"platform":"test"}}`, true},
+		{"nested", `{"custom":{"nested":[1,true,null],"unicode":"\u00e9"}}`, true},
+		// Neither form of an absent context may inherit the preceding request's value.
+		{"omitted", "", false},
+		{"empty", "", true},
+		// Base64-looking data inside the decoded JSON must remain untouched.
+		{"decoded-once", `"eyJmb28iOiJiYXIifQ=="`, true},
+	}
+	for _, test := range contexts {
+		t.Run(test.name, func(t *testing.T) {
+			wire := map[string]string{
+				"invoke-id":            test.name,
+				"invoked-function-arn": "arn:aws:lambda:us-east-1:000000000000:function:my-fn",
+				"payload":              `{"value":"hello"}`,
+				"trace-id":             "Root=1-abc;Parent=def;Sampled=1",
+			}
+			if test.present {
+				wire["client-context"] = base64.StdEncoding.EncodeToString([]byte(test.context))
+			}
+			body, err := json.Marshal(wire)
+			require.NoError(t, err)
+			recorder := httptest.NewRecorder()
+			handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/invoke", strings.NewReader(string(body))))
+			assert.Equal(t, http.StatusOK, recorder.Code)
+			assert.Equal(t, "OK", recorder.Body.String())
+			select {
+			case invoke := <-dispatched:
+				assert.Equal(t, test.context, invoke.ClientContext)
+				assert.Equal(t, test.name, invoke.ID)
+				assert.Equal(t, wire["invoked-function-arn"], invoke.InvokedFunctionArn)
+				assert.Equal(t, wire["trace-id"], invoke.TraceID)
+				assert.True(t, invoke.NeedDebugLogs)
+				payload, err := io.ReadAll(invoke.Payload)
+				require.NoError(t, err)
+				assert.Equal(t, wire["payload"], string(payload))
+			case <-time.After(time.Second):
+				t.Fatal("valid invocation was not dispatched")
+			}
+		})
+	}
+}
+
+type failedInvocationBody struct{}
+
+func (failedInvocationBody) Read([]byte) (int, error) {
+	return 0, errors.New("body read failed")
+}
+
+// TestInvokeHandlerRejectsInvalidRequests checks that bad input gets HTTP 400 and unreadable
+// bodies get HTTP 500, with neither case starting a Lambda invocation in the background.
+func TestInvokeHandlerRejectsInvalidRequests(t *testing.T) {
+	for _, test := range []struct {
+		name   string
+		body   io.Reader
+		status int
+	}{
+		{"malformed-json", strings.NewReader(`{invalid`), http.StatusBadRequest},
+		{"null", strings.NewReader(`null`), http.StatusBadRequest},
+		{"array", strings.NewReader(`[]`), http.StatusBadRequest},
+		{"invalid-base64", strings.NewReader(`{"client-context":"not base64!"}`), http.StatusBadRequest},
+		{"wrong-context-type", strings.NewReader(`{"client-context":42}`), http.StatusBadRequest},
+		{"body-read-failure", failedInvocationBody{}, http.StatusInternalServerError},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			dispatched := make(chan *interop.Invoke, 1)
+			handler := newInvokeHandler(func(invoke *interop.Invoke) { dispatched <- invoke })
+			recorder := httptest.NewRecorder()
+			handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/invoke", test.body))
+			assert.Equal(t, test.status, recorder.Code)
+			select {
+			case <-dispatched:
+				t.Fatal("invalid invocation was dispatched")
+			case <-time.After(20 * time.Millisecond):
+			}
+		})
+	}
+}
+
+// TestInvokeHandlerAcknowledgesBeforeCompletion holds the invocation open while checking
+// that /invoke returns 200 OK. Execution results are delivered later through callbacks.
+func TestInvokeHandlerAcknowledgesBeforeCompletion(t *testing.T) {
+	release := make(chan struct{})
+	completed := make(chan struct{})
+	defer close(release)
+	handler := newInvokeHandler(func(*interop.Invoke) { <-release; close(completed) })
+	returned := make(chan struct{})
+	recorder := httptest.NewRecorder()
+	go func() {
+		handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/invoke", strings.NewReader(`{"invoke-id":"async"}`)))
+		close(returned)
+	}()
+	select {
+	case <-returned:
+		assert.Equal(t, http.StatusOK, recorder.Code)
+		assert.Equal(t, "OK", recorder.Body.String())
+	case <-time.After(time.Second):
+		t.Fatal("handler waited for invocation completion")
+	}
+	select {
+	case <-completed:
+		t.Fatal("invocation completed before release")
+	default:
+	}
 }
 
 // TestLogResponseContract verifies that LogResponse uses the "logs" JSON key expected by

internal/lsapi/types.go

@@ -6,6 +6,7 @@ type InvokeRequest struct {
 	InvokedFunctionArn string `json:"invoked-function-arn"`
 	Payload            string `json:"payload"`
 	TraceId            string `json:"trace-id"`
+	ClientContext      string `json:"client-context,omitempty"`
 }
 
 // LogResponse is sent by the runtime to report logs for a completed invocation.