Skip to content

Commit b97db3d

Browse files
viren-nadkarniviren-nadkarni
authored
Decode and pass ClientContext during dispatch (#142)
1 parent 1250973 commit b97db3d

3 files changed

Lines changed: 244 additions & 99 deletions

File tree

‎cmd/localstack/custom_interop.go‎

Lines changed: 107 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ package main
55

66
import (
77
"bytes"
8+
"encoding/base64"
89
"encoding/json"
910
"errors"
1011
"fmt"
@@ -104,6 +105,47 @@ func (l *LocalStackAdapter) SendResult(invokeId string, body []byte, isError boo
104105
return l.post(endpoint, body)
105106
}
106107

108+
func decodeInvokeRequest(body []byte) (*interop.Invoke, error) {
109+
var request *lsapi.InvokeRequest
110+
if err := json.Unmarshal(body, &request); err != nil {
111+
return nil, err
112+
}
113+
if request == nil {
114+
return nil, errors.New("invocation request must be an object")
115+
}
116+
clientContext, err := base64.StdEncoding.DecodeString(request.ClientContext)
117+
if err != nil {
118+
return nil, err
119+
}
120+
return &interop.Invoke{
121+
ID: request.InvokeId,
122+
InvokedFunctionArn: request.InvokedFunctionArn,
123+
Payload: strings.NewReader(request.Payload),
124+
NeedDebugLogs: true,
125+
TraceID: request.TraceId,
126+
ClientContext: string(clientContext),
127+
}, nil
128+
}
129+
130+
// newInvokeHandler validates requests before acknowledging and dispatching them asynchronously
131+
func newInvokeHandler(dispatch func(*interop.Invoke)) http.HandlerFunc {
132+
return func(w http.ResponseWriter, r *http.Request) {
133+
body, err := io.ReadAll(r.Body)
134+
if err != nil {
135+
http.Error(w, "Failed to read invocation request", http.StatusInternalServerError)
136+
return
137+
}
138+
invoke, err := decodeInvokeRequest(body)
139+
if err != nil {
140+
http.Error(w, "Invalid invocation request", http.StatusBadRequest)
141+
return
142+
}
143+
go dispatch(invoke)
144+
w.WriteHeader(http.StatusOK)
145+
_, _ = w.Write([]byte("OK"))
146+
}
147+
}
148+
107149
func NewCustomInteropServer(lsOpts *LsOpts, delegate interop.Server, logCollector *LogCollector, eventsAPI *LocalStackEventsAPI) (server *CustomInteropServer) {
108150
server = &CustomInteropServer{
109151
delegate: delegate.(*rapidcore.Server),
@@ -117,114 +159,80 @@ func NewCustomInteropServer(lsOpts *LsOpts, delegate interop.Server, logCollecto
117159
resetDone: make(chan struct{}, 1),
118160
}
119161

120-
// TODO: extract this
121162
go func() {
122163
r := chi.NewRouter()
123-
r.Post("/invoke", func(w http.ResponseWriter, r *http.Request) {
124-
invokeR := lsapi.InvokeRequest{}
125-
bytess, err := io.ReadAll(r.Body)
164+
r.Post("/invoke", newInvokeHandler(func(invoke *interop.Invoke) {
165+
invokeResp := &standalone.ResponseWriterProxy{}
166+
// The synthetic START and INIT_REPORT lines are emitted via LocalStackEventsAPI
167+
// from rapid's lifecycle events, so they land at the AWS-faithful points (e.g.
168+
// after an inline suppressed init's own logs) — see events.go.
169+
170+
// First invocation into a successfully initialized on-demand environment: REPORT
171+
// carries the Init phase duration as measured by rapid (take-once; absent on warm
172+
// starts, failed/timed-out inits, and non-on-demand environments).
173+
initDurationMS, hasInitDuration := server.eventsAPI.TakeColdStartInitDuration()
174+
175+
invokeStart := time.Now()
176+
err := server.Invoke(invokeResp, invoke)
177+
timeout := int(server.delegate.GetInvokeTimeout().Seconds())
178+
isErr := false
179+
status := ""
180+
errorType := ""
126181
if err != nil {
127-
log.Error(err)
128-
}
129-
130-
go func() {
131-
err = json.Unmarshal(bytess, &invokeR)
132-
if err != nil {
133-
log.Error(err)
134-
}
135-
136-
invokeResp := &standalone.ResponseWriterProxy{}
137-
// The synthetic START and INIT_REPORT lines are emitted via LocalStackEventsAPI
138-
// from rapid's lifecycle events, so they land at the AWS-faithful points (e.g.
139-
// after an inline suppressed init's own logs) — see events.go.
140-
141-
// First invocation into a successfully initialized on-demand environment: REPORT
142-
// carries the Init phase duration as measured by rapid (take-once; absent on warm
143-
// starts, failed/timed-out inits, and non-on-demand environments).
144-
initDurationMS, hasInitDuration := server.eventsAPI.TakeColdStartInitDuration()
145-
146-
invokeStart := time.Now()
147-
err = server.Invoke(invokeResp, &interop.Invoke{
148-
ID: invokeR.InvokeId,
149-
InvokedFunctionArn: invokeR.InvokedFunctionArn,
150-
Payload: strings.NewReader(invokeR.Payload), // r.Body,
151-
NeedDebugLogs: true,
152-
TraceID: invokeR.TraceId,
153-
// TODO: set correct segment ID from request
154-
//LambdaSegmentID: "LambdaSegmentID", // r.Header.Get("X-Amzn-Segment-Id"),
155-
//CognitoIdentityID: "",
156-
//CognitoIdentityPoolID: "",
157-
//DeadlineNs: "",
158-
//ClientContext: "",
159-
//ContentType: "",
160-
//ReservationToken: "",
161-
//VersionID: "",
162-
//InvokeReceivedTime: 0,
163-
//ResyncState: interop.Resync{},
164-
})
165-
timeout := int(server.delegate.GetInvokeTimeout().Seconds())
166-
isErr := false
167-
status := ""
168-
errorType := ""
169-
if err != nil {
170-
switch {
171-
case errors.Is(err, rapidcore.ErrInvokeTimeout):
172-
log.Debugf("Got invoke timeout")
173-
isErr = true
174-
status = "timeout"
175-
errorResponse := lsapi.ErrorResponse{
176-
ErrorType: "Sandbox.Timedout",
177-
ErrorMessage: fmt.Sprintf(
178-
"RequestId: %s Error: Task timed out after %d.00 seconds",
179-
invokeR.InvokeId,
180-
timeout,
181-
),
182-
}
183-
jsonErrorResponse, err := json.Marshal(errorResponse)
184-
if err != nil {
185-
log.Fatalln("unable to marshall json timeout response")
186-
}
187-
_, err = invokeResp.Write(jsonErrorResponse)
188-
if err != nil {
189-
log.Fatalln("unable to write to response")
190-
}
191-
case errors.Is(err, rapidcore.ErrInvokeDoneFailed):
192-
// The error response body was already written by rapid and is sent below.
193-
// When an init failure was folded into this invocation (AWS suppressed
194-
// init), the REPORT additionally carries the failure status and the
195-
// scrubbed fatal error type (e.g. Runtime.Unknown).
196-
if errType := server.eventsAPI.InitErrorType(); errType != "" {
197-
isErr = true
198-
status = "error"
199-
errorType = errType
200-
}
201-
default:
202-
log.Fatalln(err)
182+
switch {
183+
case errors.Is(err, rapidcore.ErrInvokeTimeout):
184+
log.Debugf("Got invoke timeout")
185+
isErr = true
186+
status = "timeout"
187+
errorResponse := lsapi.ErrorResponse{
188+
ErrorType: "Sandbox.Timedout",
189+
ErrorMessage: fmt.Sprintf(
190+
"RequestId: %s Error: Task timed out after %d.00 seconds",
191+
invoke.ID,
192+
timeout,
193+
),
203194
}
204-
}
205-
// optional sleep. can be used for debugging purposes
206-
if lsOpts.PostInvokeWaitMS != "" {
207-
waitMS, err := strconv.Atoi(lsOpts.PostInvokeWaitMS)
195+
jsonErrorResponse, err := json.Marshal(errorResponse)
208196
if err != nil {
209-
log.Fatalln(err)
197+
log.Fatalln("unable to marshall json timeout response")
210198
}
211-
time.Sleep(time.Duration(waitMS) * time.Millisecond)
212-
}
213-
timeoutDuration := time.Duration(timeout) * time.Second
214-
memorySize := GetEnvOrDie("AWS_LAMBDA_FUNCTION_MEMORY_SIZE")
215-
PrintEndReports(invokeR.InvokeId, initDurationMS, hasInitDuration, status, errorType, memorySize, invokeStart, timeoutDuration, logCollector)
216-
217-
if err2 := server.localStackAdapter.SendLogs(invokeR.InvokeId, logCollector.getLogs()); err2 != nil {
218-
log.Error("failed to send logs to LocalStack: ", err2)
199+
_, err = invokeResp.Write(jsonErrorResponse)
200+
if err != nil {
201+
log.Fatalln("unable to write to response")
202+
}
203+
case errors.Is(err, rapidcore.ErrInvokeDoneFailed):
204+
// The error response body was already written by rapid and is sent below.
205+
// When an init failure was folded into this invocation (AWS suppressed
206+
// init), the REPORT additionally carries the failure status and the
207+
// scrubbed fatal error type (e.g. Runtime.Unknown).
208+
if errType := server.eventsAPI.InitErrorType(); errType != "" {
209+
isErr = true
210+
status = "error"
211+
errorType = errType
212+
}
213+
default:
214+
log.Fatalln(err)
219215
}
220-
if err2 := server.localStackAdapter.SendResult(invokeR.InvokeId, invokeResp.Body, isErr); err2 != nil {
221-
log.Error("failed to send result to LocalStack: ", err2)
216+
}
217+
// optional sleep. can be used for debugging purposes
218+
if lsOpts.PostInvokeWaitMS != "" {
219+
waitMS, err := strconv.Atoi(lsOpts.PostInvokeWaitMS)
220+
if err != nil {
221+
log.Fatalln(err)
222222
}
223-
}()
223+
time.Sleep(time.Duration(waitMS) * time.Millisecond)
224+
}
225+
timeoutDuration := time.Duration(timeout) * time.Second
226+
memorySize := GetEnvOrDie("AWS_LAMBDA_FUNCTION_MEMORY_SIZE")
227+
PrintEndReports(invoke.ID, initDurationMS, hasInitDuration, status, errorType, memorySize, invokeStart, timeoutDuration, logCollector)
224228

225-
w.WriteHeader(200)
226-
_, _ = w.Write([]byte("OK"))
227-
})
229+
if err2 := server.localStackAdapter.SendLogs(invoke.ID, logCollector.getLogs()); err2 != nil {
230+
log.Error("failed to send logs to LocalStack: ", err2)
231+
}
232+
if err2 := server.localStackAdapter.SendResult(invoke.ID, invokeResp.Body, isErr); err2 != nil {
233+
log.Error("failed to send result to LocalStack: ", err2)
234+
}
235+
}))
228236
err := http.ListenAndServe(":"+server.port, r)
229237
if err != nil {
230238
log.Error(err)

‎cmd/localstack/custom_interop_test.go‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,17 @@
11
package main
22

33
import (
4+
"encoding/base64"
45
"encoding/json"
6+
"errors"
57
"io"
68
"net/http"
79
"net/http/httptest"
10+
"strings"
811
"testing"
12+
"time"
913

14+
"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/interop"
1015
"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lsapi"
1116
"github.com/stretchr/testify/assert"
1217
"github.com/stretchr/testify/require"
@@ -36,6 +41,137 @@ func TestInvokeRequestContract(t *testing.T) {
3641
assert.Equal(t, "arn:aws:lambda:us-east-1:000000000000:function:my-fn", req.InvokedFunctionArn)
3742
assert.Equal(t, `{"key":"value"}`, req.Payload)
3843
assert.Equal(t, "Root=1-abc;Parent=def;Sampled=1", req.TraceId)
44+
assert.Empty(t, req.ClientContext)
45+
}
46+
47+
// TestInvokeRequestClientContextContract checks that JSON encoding preserves the base64
48+
// context verbatim and omits an empty context. Decoding belongs to the invocation handler.
49+
func TestInvokeRequestClientContextContract(t *testing.T) {
50+
raw := `{"client-context":"eyJjdXN0b20iOnsiZm9vIjoiYmFyIn19"}`
51+
var req lsapi.InvokeRequest
52+
require.NoError(t, json.Unmarshal([]byte(raw), &req))
53+
assert.Equal(t, "eyJjdXN0b20iOnsiZm9vIjoiYmFyIn19", req.ClientContext)
54+
encoded, err := json.Marshal(req)
55+
require.NoError(t, err)
56+
assert.Contains(t, string(encoded), `"client-context":"eyJjdXN0b20iOnsiZm9vIjoiYmFyIn19"`)
57+
encoded, err = json.Marshal(lsapi.InvokeRequest{})
58+
require.NoError(t, err)
59+
assert.NotContains(t, string(encoded), "client-context")
60+
}
61+
62+
// TestInvokeHandlerForwardsContextPerRequest reuses one handler across successive requests
63+
// to check decoding, context isolation, and preservation of the remaining invocation fields.
64+
func TestInvokeHandlerForwardsContextPerRequest(t *testing.T) {
65+
dispatched := make(chan *interop.Invoke, 1)
66+
handler := newInvokeHandler(func(invoke *interop.Invoke) { dispatched <- invoke })
67+
contexts := []struct {
68+
name string
69+
context string
70+
present bool
71+
}{
72+
{"full", `{"custom":{"foo":"bar"},"client":{"app_title":"test"},"env":{"platform":"test"}}`, true},
73+
{"nested", `{"custom":{"nested":[1,true,null],"unicode":"\u00e9"}}`, true},
74+
// Neither form of an absent context may inherit the preceding request's value.
75+
{"omitted", "", false},
76+
{"empty", "", true},
77+
// Base64-looking data inside the decoded JSON must remain untouched.
78+
{"decoded-once", `"eyJmb28iOiJiYXIifQ=="`, true},
79+
}
80+
for _, test := range contexts {
81+
t.Run(test.name, func(t *testing.T) {
82+
wire := map[string]string{
83+
"invoke-id": test.name,
84+
"invoked-function-arn": "arn:aws:lambda:us-east-1:000000000000:function:my-fn",
85+
"payload": `{"value":"hello"}`,
86+
"trace-id": "Root=1-abc;Parent=def;Sampled=1",
87+
}
88+
if test.present {
89+
wire["client-context"] = base64.StdEncoding.EncodeToString([]byte(test.context))
90+
}
91+
body, err := json.Marshal(wire)
92+
require.NoError(t, err)
93+
recorder := httptest.NewRecorder()
94+
handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/invoke", strings.NewReader(string(body))))
95+
assert.Equal(t, http.StatusOK, recorder.Code)
96+
assert.Equal(t, "OK", recorder.Body.String())
97+
select {
98+
case invoke := <-dispatched:
99+
assert.Equal(t, test.context, invoke.ClientContext)
100+
assert.Equal(t, test.name, invoke.ID)
101+
assert.Equal(t, wire["invoked-function-arn"], invoke.InvokedFunctionArn)
102+
assert.Equal(t, wire["trace-id"], invoke.TraceID)
103+
assert.True(t, invoke.NeedDebugLogs)
104+
payload, err := io.ReadAll(invoke.Payload)
105+
require.NoError(t, err)
106+
assert.Equal(t, wire["payload"], string(payload))
107+
case <-time.After(time.Second):
108+
t.Fatal("valid invocation was not dispatched")
109+
}
110+
})
111+
}
112+
}
113+
114+
type failedInvocationBody struct{}
115+
116+
func (failedInvocationBody) Read([]byte) (int, error) {
117+
return 0, errors.New("body read failed")
118+
}
119+
120+
// TestInvokeHandlerRejectsInvalidRequests checks that bad input gets HTTP 400 and unreadable
121+
// bodies get HTTP 500, with neither case starting a Lambda invocation in the background.
122+
func TestInvokeHandlerRejectsInvalidRequests(t *testing.T) {
123+
for _, test := range []struct {
124+
name string
125+
body io.Reader
126+
status int
127+
}{
128+
{"malformed-json", strings.NewReader(`{invalid`), http.StatusBadRequest},
129+
{"null", strings.NewReader(`null`), http.StatusBadRequest},
130+
{"array", strings.NewReader(`[]`), http.StatusBadRequest},
131+
{"invalid-base64", strings.NewReader(`{"client-context":"not base64!"}`), http.StatusBadRequest},
132+
{"wrong-context-type", strings.NewReader(`{"client-context":42}`), http.StatusBadRequest},
133+
{"body-read-failure", failedInvocationBody{}, http.StatusInternalServerError},
134+
} {
135+
t.Run(test.name, func(t *testing.T) {
136+
dispatched := make(chan *interop.Invoke, 1)
137+
handler := newInvokeHandler(func(invoke *interop.Invoke) { dispatched <- invoke })
138+
recorder := httptest.NewRecorder()
139+
handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/invoke", test.body))
140+
assert.Equal(t, test.status, recorder.Code)
141+
select {
142+
case <-dispatched:
143+
t.Fatal("invalid invocation was dispatched")
144+
case <-time.After(20 * time.Millisecond):
145+
}
146+
})
147+
}
148+
}
149+
150+
// TestInvokeHandlerAcknowledgesBeforeCompletion holds the invocation open while checking
151+
// that /invoke returns 200 OK. Execution results are delivered later through callbacks.
152+
func TestInvokeHandlerAcknowledgesBeforeCompletion(t *testing.T) {
153+
release := make(chan struct{})
154+
completed := make(chan struct{})
155+
defer close(release)
156+
handler := newInvokeHandler(func(*interop.Invoke) { <-release; close(completed) })
157+
returned := make(chan struct{})
158+
recorder := httptest.NewRecorder()
159+
go func() {
160+
handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/invoke", strings.NewReader(`{"invoke-id":"async"}`)))
161+
close(returned)
162+
}()
163+
select {
164+
case <-returned:
165+
assert.Equal(t, http.StatusOK, recorder.Code)
166+
assert.Equal(t, "OK", recorder.Body.String())
167+
case <-time.After(time.Second):
168+
t.Fatal("handler waited for invocation completion")
169+
}
170+
select {
171+
case <-completed:
172+
t.Fatal("invocation completed before release")
173+
default:
174+
}
39175
}
40176

41177
// TestLogResponseContract verifies that LogResponse uses the "logs" JSON key expected by

‎internal/lsapi/types.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ type InvokeRequest struct {
66
InvokedFunctionArn string `json:"invoked-function-arn"`
77
Payload string `json:"payload"`
88
TraceId string `json:"trace-id"`
9+
ClientContext string `json:"client-context,omitempty"`
910
}
1011

1112
// LogResponse is sent by the runtime to report logs for a completed invocation.

0 commit comments

Comments
 (0)