From e6a814b115b2aeedcdc001f87364d9598237093e Mon Sep 17 00:00:00 2001 From: Viren Nadkarni Date: Tue, 13 Oct 2026 19:05:24 +0530 Subject: [PATCH 1/2] Decode and pass ClientContext during dispatch --- cmd/localstack/custom_interop.go | 206 +++++++++++++------------- cmd/localstack/custom_interop_test.go | 126 ++++++++++++++++ internal/lsapi/types.go | 1 + 3 files changed, 234 insertions(+), 99 deletions(-) diff --git a/cmd/localstack/custom_interop.go b/cmd/localstack/custom_interop.go index 69d88dc..87e4c13 100644 --- a/cmd/localstack/custom_interop.go +++ b/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) diff --git a/cmd/localstack/custom_interop_test.go b/cmd/localstack/custom_interop_test.go index d089106..3634f7b 100644 --- a/cmd/localstack/custom_interop_test.go +++ b/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,127 @@ 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) +} + +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") +} + +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}, + {"omitted", "", false}, + {"empty", "", true}, + {"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") +} + +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): + } + }) + } +} + +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 diff --git a/internal/lsapi/types.go b/internal/lsapi/types.go index 7eaaf6b..5f32d7b 100644 --- a/internal/lsapi/types.go +++ b/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. From f57a27fca16b9fd39171217175ec321eaea9a813 Mon Sep 17 00:00:00 2001 From: Viren Nadkarni Date: Tue, 13 Oct 2026 19:55:58 +0530 Subject: [PATCH 2/2] Add comments --- cmd/localstack/custom_interop_test.go | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/cmd/localstack/custom_interop_test.go b/cmd/localstack/custom_interop_test.go index 3634f7b..dfdeefb 100644 --- a/cmd/localstack/custom_interop_test.go +++ b/cmd/localstack/custom_interop_test.go @@ -44,6 +44,8 @@ func TestInvokeRequestContract(t *testing.T) { 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 @@ -57,6 +59,8 @@ func TestInvokeRequestClientContextContract(t *testing.T) { 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 }) @@ -67,8 +71,10 @@ func TestInvokeHandlerForwardsContextPerRequest(t *testing.T) { }{ {"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 { @@ -111,6 +117,8 @@ 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 @@ -139,6 +147,8 @@ func TestInvokeHandlerRejectsInvalidRequests(t *testing.T) { } } +// 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{})