From 8bacf5f39d7a63bcb666b178d8ce9b2f54d406a3 Mon Sep 17 00:00:00 2001 From: lambda-tooling+rie Date: Fri, 31 Jul 2026 09:09:13 +0000 Subject: [PATCH] Update from upstream - 2026-07-31 --- .github/workflows/check-binaries.yml | 46 ++++++++++--------- .github/workflows/integ-tests.yml | 22 +++------ .github/workflows/release.yml | 22 +++++---- .gitignore | 1 + .../internal/invoke/rie_invoke_request.go | 6 +++ .../invoke/rie_invoke_request_test.go | 1 + .../interop/mock_invoke_request.go | 17 +++++++ .../interop/sandbox_model.go | 2 + .../lambda-managed-instances/invoke/consts.go | 2 + .../invoke/invoke_router.go | 4 ++ .../invoke/mock_runtime_error_request.go | 17 +++++++ .../invoke/mock_runtime_response_request.go | 17 +++++++ .../invoke/running_invoke.go | 21 +++++++++ .../invoke/running_invoke_test.go | 4 ++ .../invoke/runtime_error_request.go | 6 +++ .../invoke/runtime_response_request.go | 14 ++++-- .../invoke/runtime_response_sender.go | 3 ++ .../invoke/runtime_response_sender_test.go | 2 + .../functional/process_supervisor.go | 14 ++++-- .../testutils/functional/runtime_actions.go | 30 ++++++++---- .../testutils/functional/runtime_client.go | 7 +++ 21 files changed, 196 insertions(+), 62 deletions(-) diff --git a/.github/workflows/check-binaries.yml b/.github/workflows/check-binaries.yml index fd12637c..66fa2ba3 100644 --- a/.github/workflows/check-binaries.yml +++ b/.github/workflows/check-binaries.yml @@ -15,19 +15,19 @@ jobs: report_contents: ${{ steps.save-output.outputs.report_contents }} steps: - name: Setup python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: '3.11' - name: Checkout code - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: ref: main - name: Download latest release - uses: robinraju/release-downloader@v1.10 - with: - latest: true - fileName: 'aws-lambda-rie*' - out-file-path: "bin" + run: | + mkdir -p bin + gh release download --pattern 'aws-lambda-rie*' --dir bin + env: + GH_TOKEN: ${{ github.token }} - name: Run check for vulnerabilities id: check-binaries run: | @@ -61,25 +61,29 @@ jobs: name: Save outputs for the check with the latest build id: save-new-version run: | - if [ "${{ steps.check-new-version.outcome }}" == "failure" ]; then + if [ "${CHECK_OUTCOME}" == "failure" ]; then fixed="No" else fixed="Yes" fi echo "fixed=$fixed" >> "$GITHUB_OUTPUT" + env: + CHECK_OUTCOME: ${{ steps.check-new-version.outcome }} - if: always() && steps.save-output.outputs.report_contents name: Create GitHub Issue indicating vulnerabilities id: create-issue - uses: dacbd/create-issue-action@main - with: - token: ${{ github.token }} - title: | - CVEs found in latest RIE release - body: | - ### CVEs found in latest RIE release - ``` - ${{ steps.save-output.outputs.report_contents }} - ``` - - #### Are these resolved by building with the latest patch version of Go (${{ steps.check-new-version.outputs.latest_version }})?: - > **${{ steps.save-new-version.outputs.fixed }}** + run: | + gh issue create \ + --title "CVEs found in latest RIE release" \ + --body "### CVEs found in latest RIE release + \`\`\` + ${REPORT_CONTENTS} + \`\`\` + + #### Are these resolved by building with the latest patch version of Go (${LATEST_VERSION})?: + > **${FIXED}**" + env: + GH_TOKEN: ${{ github.token }} + REPORT_CONTENTS: ${{ steps.save-output.outputs.report_contents }} + LATEST_VERSION: ${{ steps.check-new-version.outputs.latest_version }} + FIXED: ${{ steps.save-new-version.outputs.fixed }} diff --git a/.github/workflows/integ-tests.yml b/.github/workflows/integ-tests.yml index e97b0c1f..5fce4bca 100644 --- a/.github/workflows/integ-tests.yml +++ b/.github/workflows/integ-tests.yml @@ -13,41 +13,33 @@ permissions: jobs: go-tests: runs-on: ubuntu-latest - environment: - name: integ-tests steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 - name: run go tests run: make tests-with-docker integ-tests-x86: runs-on: ubuntu-latest - environment: - name: integ-tests steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 with: python-version: '3.11' - name: run integration tests run: make integ-tests-with-docker-x86-64 integ-tests-arm64: runs-on: ubuntu-latest - environment: - name: integ-tests steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 with: python-version: '3.11' - name: run integration tests run: make integ-tests-with-docker-arm64 integ-tests-old: runs-on: ubuntu-latest - environment: - name: integ-tests steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 with: python-version: '3.11' - name: run integration tests diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index bb3c5146..1b809a2e 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -19,11 +19,11 @@ jobs: environment: Release runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 with: ref: main - name: Set up python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: '3.11' - name: Build @@ -33,12 +33,14 @@ jobs: make tests-with-docker make integ-tests - name: Release - uses: softprops/action-gh-release@v2 - with: - name: Release ${{ github.event.inputs.releaseVersion }} - tag_name: v${{ github.event.inputs.releaseVersion }} - body: ${{ github.event.inputs.releaseBody }} - files: | - bin/aws-lambda-rie - bin/aws-lambda-rie-arm64 + run: | + gh release create "v${RELEASE_VERSION}" \ + --title "Release ${RELEASE_VERSION}" \ + --notes "${RELEASE_BODY}" \ + bin/aws-lambda-rie \ + bin/aws-lambda-rie-arm64 \ bin/aws-lambda-rie-x86_64 + env: + GH_TOKEN: ${{ github.token }} + RELEASE_VERSION: ${{ github.event.inputs.releaseVersion }} + RELEASE_BODY: ${{ github.event.inputs.releaseBody }} diff --git a/.gitignore b/.gitignore index 85b35a21..ab34db8f 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,4 @@ tags .idea .DS_Store .venv +build.log diff --git a/internal/lambda-managed-instances/aws-lambda-rie/internal/invoke/rie_invoke_request.go b/internal/lambda-managed-instances/aws-lambda-rie/internal/invoke/rie_invoke_request.go index 608d8e42..a5d1d3d1 100644 --- a/internal/lambda-managed-instances/aws-lambda-rie/internal/invoke/rie_invoke_request.go +++ b/internal/lambda-managed-instances/aws-lambda-rie/internal/invoke/rie_invoke_request.go @@ -40,6 +40,7 @@ type rieInvokeRequest struct { cognitoIdentityPoolId string clientContext string responseMode string + internalInvocationID string functionVersionID string } @@ -98,6 +99,7 @@ func NewRieInvokeRequest(request *http.Request, writer http.ResponseWriter) (*ri cognitoIdentityPoolId: cognitoIdentityPoolId, clientContext: clientContext, responseMode: request.Header.Get(invoke.ResponseModeHeader), + internalInvocationID: uuid.New().String(), } return req, nil @@ -184,3 +186,7 @@ func (r *rieInvokeRequest) UpdateFromInitData(initData interop.InitStaticDataPro func (r *rieInvokeRequest) FunctionVersionID() string { return r.functionVersionID } + +func (r *rieInvokeRequest) InternalInvocationID() string { + return r.internalInvocationID +} diff --git a/internal/lambda-managed-instances/aws-lambda-rie/internal/invoke/rie_invoke_request_test.go b/internal/lambda-managed-instances/aws-lambda-rie/internal/invoke/rie_invoke_request_test.go index 7ead546c..c4265441 100644 --- a/internal/lambda-managed-instances/aws-lambda-rie/internal/invoke/rie_invoke_request_test.go +++ b/internal/lambda-managed-instances/aws-lambda-rie/internal/invoke/rie_invoke_request_test.go @@ -143,6 +143,7 @@ func TestNewRieInvokeRequest(t *testing.T) { if tt.want.invokeID == "" { tt.want.invokeID = got.invokeID } + tt.want.internalInvocationID = got.internalInvocationID assert.Equal(t, tt.want, got) }) diff --git a/internal/lambda-managed-instances/interop/mock_invoke_request.go b/internal/lambda-managed-instances/interop/mock_invoke_request.go index 84e0a1ac..b6aef276 100644 --- a/internal/lambda-managed-instances/interop/mock_invoke_request.go +++ b/internal/lambda-managed-instances/interop/mock_invoke_request.go @@ -143,6 +143,23 @@ func (_m *MockInvokeRequest) FunctionVersionID() string { return r0 } +func (_m *MockInvokeRequest) InternalInvocationID() string { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for InternalInvocationID") + } + + var r0 string + if rf, ok := ret.Get(0).(func() string); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(string) + } + + return r0 +} + func (_m *MockInvokeRequest) InvokeID() string { ret := _m.Called() diff --git a/internal/lambda-managed-instances/interop/sandbox_model.go b/internal/lambda-managed-instances/interop/sandbox_model.go index d88e441c..8f50ea45 100644 --- a/internal/lambda-managed-instances/interop/sandbox_model.go +++ b/internal/lambda-managed-instances/interop/sandbox_model.go @@ -302,6 +302,8 @@ type InvokeRequest interface { UpdateFromInitData(InitStaticDataProvider) model.AppError FunctionVersionID() string + + InternalInvocationID() string } type InitStaticDataProvider interface { diff --git a/internal/lambda-managed-instances/invoke/consts.go b/internal/lambda-managed-instances/invoke/consts.go index 1632a4f4..d6a140dc 100644 --- a/internal/lambda-managed-instances/invoke/consts.go +++ b/internal/lambda-managed-instances/invoke/consts.go @@ -9,6 +9,8 @@ const ( FunctionErrorBodyTrailer = "lambda-runtime-function-error-body" ResponseModeHeader = "invoke-response-mode" TraceIdHeader = "x-amzn-trace-id" + + RuntimeInvocationIdHeader = "lambda-runtime-invocation-id" ) type InvokeBodyResponseStatus string diff --git a/internal/lambda-managed-instances/invoke/invoke_router.go b/internal/lambda-managed-instances/invoke/invoke_router.go index 6625bd2f..17be7f86 100644 --- a/internal/lambda-managed-instances/invoke/invoke_router.go +++ b/internal/lambda-managed-instances/invoke/invoke_router.go @@ -35,6 +35,8 @@ type RuntimeResponseRequest interface { BodyReader() io.Reader TrailerError() ErrorForInvoker + + InvocationID() string } type RuntimeErrorRequest interface { @@ -47,6 +49,8 @@ type RuntimeErrorRequest interface { ReturnCode() int ErrorDetails() string GetXrayErrorCause() json.RawMessage + + InvocationID() string } type runningInvoke interface { diff --git a/internal/lambda-managed-instances/invoke/mock_runtime_error_request.go b/internal/lambda-managed-instances/invoke/mock_runtime_error_request.go index 09d09cf1..38b39284 100644 --- a/internal/lambda-managed-instances/invoke/mock_runtime_error_request.go +++ b/internal/lambda-managed-instances/invoke/mock_runtime_error_request.go @@ -171,6 +171,23 @@ func (_m *MockRuntimeErrorRequest) ReturnCode() int { return r0 } +func (_m *MockRuntimeErrorRequest) InvocationID() string { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for InvocationID") + } + + var r0 string + if rf, ok := ret.Get(0).(func() string); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(string) + } + + return r0 +} + func NewMockRuntimeErrorRequest(t interface { mock.TestingT Cleanup(func()) diff --git a/internal/lambda-managed-instances/invoke/mock_runtime_response_request.go b/internal/lambda-managed-instances/invoke/mock_runtime_response_request.go index 5da7b503..bc550bc8 100644 --- a/internal/lambda-managed-instances/invoke/mock_runtime_response_request.go +++ b/internal/lambda-managed-instances/invoke/mock_runtime_response_request.go @@ -122,6 +122,23 @@ func (_m *MockRuntimeResponseRequest) TrailerError() ErrorForInvoker { return r0 } +func (_m *MockRuntimeResponseRequest) InvocationID() string { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for InvocationID") + } + + var r0 string + if rf, ok := ret.Get(0).(func() string); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(string) + } + + return r0 +} + func NewMockRuntimeResponseRequest(t interface { mock.TestingT Cleanup(func()) diff --git a/internal/lambda-managed-instances/invoke/running_invoke.go b/internal/lambda-managed-instances/invoke/running_invoke.go index feec3496..e08b12c7 100644 --- a/internal/lambda-managed-instances/invoke/running_invoke.go +++ b/internal/lambda-managed-instances/invoke/running_invoke.go @@ -68,6 +68,8 @@ type runningInvokeImpl struct { invokeRespSender InvokeResponseSender runtimeNext http.ResponseWriter + internalInvocationID string + responderFactoryFunc ResponderFactoryFunc sendInvokeToRuntime func(context.Context, interop.InitStaticDataProvider, interop.InvokeRequest, http.ResponseWriter, string) (int64, time.Duration, time.Duration, model.AppError) createTracingData func(traceId string, tracingMode intmodel.XrayTracingMode, segmentIDGenerator func() string) (downstreamTraceId string, tracingCtx *interop.TracingCtx) @@ -99,6 +101,8 @@ func newRunningInvoke( func (r *runningInvokeImpl) RunInvokeAndSendResult(ctx context.Context, initData interop.InitStaticDataProvider, invokeReq interop.InvokeRequest, metrics interop.InvokeMetrics) model.AppError { downstreamTraceId, tracingCtx := r.createTracingData(invokeReq.TraceId(), initData.XRayTracingMode(), xray.GenerateSegmentID) + r.internalInvocationID = invokeReq.InternalInvocationID() + metrics.TriggerStartRequest() if err := metrics.SendInvokeStartEvent(tracingCtx); err != nil { logging.Error(ctx, "Failed to send InvokeStartEvent", "err", err) @@ -290,11 +294,28 @@ func (r *runningInvokeImpl) RuntimeResponse(ctx context.Context, runtimeRespReq return model.NewCustomerError(model.ErrorRuntimeInvokeResponseInProgress) } + if echoedID := runtimeRespReq.InvocationID(); echoedID != "" && r.internalInvocationID != "" { + if echoedID != r.internalInvocationID { + logging.Warn(ctx, "Cross-wiring detected: invocation ID mismatch on response", + "expected", r.internalInvocationID, "received", echoedID) + r.responseState.CompareAndSwap(stateGotResponse, stateNoResponse) + return model.NewCustomerError(model.ErrorRuntimeInvokeTimeout) + } + } + r.runtimeResponseChan <- runtimeRespReq return <-r.responseSentChan } func (r *runningInvokeImpl) RuntimeError(ctx context.Context, runtimeErrReq RuntimeErrorRequest) model.AppError { + if echoedID := runtimeErrReq.InvocationID(); echoedID != "" && r.internalInvocationID != "" { + if echoedID != r.internalInvocationID { + logging.Warn(ctx, "Cross-wiring detected: invocation ID mismatch on error", + "expected", r.internalInvocationID, "received", echoedID) + return model.NewCustomerError(model.ErrorRuntimeInvokeTimeout) + } + } + oldState := r.responseState.Swap(stateGotError) if oldState == stateGotError { logging.Warn(ctx, "Invalid invoke state : error in progress") diff --git a/internal/lambda-managed-instances/invoke/running_invoke_test.go b/internal/lambda-managed-instances/invoke/running_invoke_test.go index 951f650e..4f69ae68 100644 --- a/internal/lambda-managed-instances/invoke/running_invoke_test.go +++ b/internal/lambda-managed-instances/invoke/running_invoke_test.go @@ -66,6 +66,10 @@ func hijackRunningInvokeDeps(ri *runningInvokeImpl, mocks *runningInvokeMocks) { func createMocksAndInitRunningInvoke(t *testing.T) (*runningInvokeMocks, *runningInvokeImpl) { mocks := newRunningInvokeMocks(t) + mocks.eaInvokeRequest.On("InternalInvocationID").Return("").Maybe() + mocks.runtimeRespReq.On("InvocationID").Return("").Maybe() + mocks.runtimeErrorReq.On("InvocationID").Return("").Maybe() + ri := newRunningInvoke( mocks.runtimeNextRequest, func(ctx context.Context, ir interop.InvokeRequest) InvokeResponseSender { diff --git a/internal/lambda-managed-instances/invoke/runtime_error_request.go b/internal/lambda-managed-instances/invoke/runtime_error_request.go index afe5edcd..990bf105 100644 --- a/internal/lambda-managed-instances/invoke/runtime_error_request.go +++ b/internal/lambda-managed-instances/invoke/runtime_error_request.go @@ -29,6 +29,7 @@ type runtimeError struct { invokeID interop.InvokeID errorType model.ErrorType errorCategory model.ErrorCategory + invocationID string errorDetails string xrayErrorCause json.RawMessage @@ -42,6 +43,7 @@ func NewRuntimeError(ctx context.Context, request *http.Request, invokeID intero invokeID: invokeID, errorType: model.GetValidRuntimeOrFunctionErrorType(request.Header.Get(RuntimeErrorTypeHeader)), errorCategory: RuntimeErrorCategory, + invocationID: request.Header.Get(RuntimeInvocationIdHeader), errorDetails: errorDetails, xrayErrorCause: getValidatedErrorCause(ctx, request.Header.Get(LambdaXRayErrorCauseHeader)), } @@ -89,6 +91,10 @@ func (r *runtimeError) GetXrayErrorCause() json.RawMessage { return r.xrayErrorCause } +func (r *runtimeError) InvocationID() string { + return r.invocationID +} + func (r *runtimeError) ReturnCode() int { return http.StatusOK } diff --git a/internal/lambda-managed-instances/invoke/runtime_response_request.go b/internal/lambda-managed-instances/invoke/runtime_response_request.go index 3d3c2ba8..7e51d52a 100644 --- a/internal/lambda-managed-instances/invoke/runtime_response_request.go +++ b/internal/lambda-managed-instances/invoke/runtime_response_request.go @@ -31,6 +31,7 @@ type runtimeResponse struct { contentType string invokeID interop.InvokeID responseMode string + invocationID string } func NewRuntimeResponse(ctx context.Context, request *http.Request, writer http.ResponseWriter, invokeID interop.InvokeID) runtimeResponse { @@ -40,10 +41,11 @@ func NewRuntimeResponse(ctx context.Context, request *http.Request, writer http. contentType = "application/octet-stream" } resp := runtimeResponse{ - request: request, - rc: http.NewResponseController(writer), - contentType: contentType, - invokeID: invokeID, + request: request, + rc: http.NewResponseController(writer), + contentType: contentType, + invokeID: invokeID, + invocationID: request.Header.Get(RuntimeInvocationIdHeader), } switch mode := request.Header.Get(RuntimeResponseModeHeader); mode { @@ -87,6 +89,10 @@ func (r *runtimeResponse) ResponseMode() string { return r.responseMode } +func (r *runtimeResponse) InvocationID() string { + return r.invocationID +} + func (r *runtimeResponse) TrailerError() ErrorForInvoker { typ := r.request.Trailer.Get(FunctionErrorTypeTrailer) if typ == "" { diff --git a/internal/lambda-managed-instances/invoke/runtime_response_sender.go b/internal/lambda-managed-instances/invoke/runtime_response_sender.go index 813282df..88a8d72b 100644 --- a/internal/lambda-managed-instances/invoke/runtime_response_sender.go +++ b/internal/lambda-managed-instances/invoke/runtime_response_sender.go @@ -42,6 +42,9 @@ func sendInvokeToRuntime(ctx context.Context, initData interop.InitStaticDataPro if cogId := buildCognitoIdentifyHeader(invokeReq); cogId != "" { runtimeReq.Header().Set(RuntimeCognitoIdentifyHeader, cogId) } + if internalInvocationID := invokeReq.InternalInvocationID(); internalInvocationID != "" { + runtimeReq.Header().Set(RuntimeInvocationIdHeader, internalInvocationID) + } runtimeReq.WriteHeader(http.StatusOK) timedReader := &utils.TimedReader{ diff --git a/internal/lambda-managed-instances/invoke/runtime_response_sender_test.go b/internal/lambda-managed-instances/invoke/runtime_response_sender_test.go index 3d527b42..a8796876 100644 --- a/internal/lambda-managed-instances/invoke/runtime_response_sender_test.go +++ b/internal/lambda-managed-instances/invoke/runtime_response_sender_test.go @@ -62,6 +62,7 @@ func buildInvokeReqMocks(invokeReq *interop.MockInvokeRequest) { invokeReq.On("ClientContext").Return("client-context-example") invokeReq.On("CognitoId").Return("cognito_id_12345") invokeReq.On("CognitoPoolId").Return("cognito_pool_id_6789") + invokeReq.On("InternalInvocationID").Return("") } func buildInitDataMocks(initData *interop.MockInitStaticDataProvider) { @@ -147,6 +148,7 @@ func TestSendResponse_OmitsEmptyOptionalHeaders(t *testing.T) { mocks.invokeReq.On("ClientContext").Return("") mocks.invokeReq.On("CognitoId").Return("") mocks.invokeReq.On("CognitoPoolId").Return("") + mocks.invokeReq.On("InternalInvocationID").Return("") mocks.invokeReq.On("BodyReader").Return(mocks.reader) recorder := httptest.NewRecorder() diff --git a/internal/lambda-managed-instances/testutils/functional/process_supervisor.go b/internal/lambda-managed-instances/testutils/functional/process_supervisor.go index 8fb4c22e..072470f8 100644 --- a/internal/lambda-managed-instances/testutils/functional/process_supervisor.go +++ b/internal/lambda-managed-instances/testutils/functional/process_supervisor.go @@ -40,9 +40,10 @@ type RuntimeEnv struct { } type RuntimeExecutionEnvironment struct { - Actions []ExecutionEnvironmentAction - InvokeID interop.InvokeID - RuntimeEnv *RuntimeEnv + Actions []ExecutionEnvironmentAction + InvokeID interop.InvokeID + InvocationID string + RuntimeEnv *RuntimeEnv } func (r *RuntimeEnv) Exec(request *model.ExecRequest) (<-chan struct{}, error) { @@ -134,6 +135,7 @@ func (r *RuntimeExecutionEnvironment) executeEnvActions(client *Client, t *testi if resp != nil { r.InvokeID = resp.Header.Get(invoke.RuntimeRequestIdHeader) + r.InvocationID = resp.Header.Get(invoke.RuntimeInvocationIdHeader) a.ValidateStatus(t, resp) } case StdoutAction: @@ -146,6 +148,9 @@ func (r *RuntimeExecutionEnvironment) executeEnvActions(client *Client, t *testi if a.InvokeID == "" { a.InvokeID = r.InvokeID } + if a.ResponseHeaders == nil && r.InvocationID != "" { + a.ResponseHeaders = map[string]string{invoke.RuntimeInvocationIdHeader: r.InvocationID} + } executeAndValidateAction(a, client, t) case InvocationStreamingResponseAction: if a.InvokeID == "" { @@ -156,6 +161,9 @@ func (r *RuntimeExecutionEnvironment) executeEnvActions(client *Client, t *testi if a.InvokeID == "" { a.InvokeID = r.InvokeID } + if a.ResponseHeaders == nil && r.InvocationID != "" { + a.ResponseHeaders = map[string]string{invoke.RuntimeInvocationIdHeader: r.InvocationID} + } executeAndValidateAction(a, client, t) case ExitAction: a.exitProcessOnce = r.RuntimeEnv.ExitProcessOnce diff --git a/internal/lambda-managed-instances/testutils/functional/runtime_actions.go b/internal/lambda-managed-instances/testutils/functional/runtime_actions.go index bea7a1f6..503fcc92 100644 --- a/internal/lambda-managed-instances/testutils/functional/runtime_actions.go +++ b/internal/lambda-managed-instances/testutils/functional/runtime_actions.go @@ -157,13 +157,14 @@ type InvocationResponseAction struct { ContentType string InvokeID interop.InvokeID ResponseModeHeader string + ResponseHeaders map[string]string ExpectedStatus int ExpectedBody string Trailers map[string]string } func (a InvocationResponseAction) Execute(t *testing.T, client *Client) (*http.Response, error) { - return client.Response(a.InvokeID, a.Payload, a.ContentType, a.ResponseModeHeader, a.Trailers) + return client.ResponseWithHeaders(a.InvokeID, a.Payload, a.ContentType, a.ResponseModeHeader, a.Trailers, a.ResponseHeaders) } func (a InvocationResponseAction) ValidateStatus(t *testing.T, resp *http.Response) { @@ -185,17 +186,28 @@ func (a InvocationResponseAction) String() string { } type InvocationResponseErrorAction struct { - Payload string - ContentType string - InvokeID interop.InvokeID - ErrorType string - ErrorCause string - ExpectedStatus int - ExpectedBody string + Payload string + ContentType string + InvokeID interop.InvokeID + ErrorType string + ErrorCause string + ResponseHeaders map[string]string + ExpectedStatus int + ExpectedBody string } func (a InvocationResponseErrorAction) Execute(t *testing.T, client *Client) (*http.Response, error) { - return client.ResponseError(a.InvokeID, a.Payload, a.ContentType, a.ErrorType, a.ErrorCause) + headers := map[string]string{ContentTypeHeader: a.ContentType} + if len(a.ErrorType) > 0 { + headers[LambdaErrorTypeHeader] = a.ErrorType + } + if len(a.ErrorCause) > 0 { + headers[LambdaXRayErrorCauseHeader] = a.ErrorCause + } + for k, v := range a.ResponseHeaders { + headers[k] = v + } + return client.ResponseErrorWithHeaders(a.InvokeID, a.Payload, headers) } func (a InvocationResponseErrorAction) ValidateStatus(t *testing.T, resp *http.Response) { diff --git a/internal/lambda-managed-instances/testutils/functional/runtime_client.go b/internal/lambda-managed-instances/testutils/functional/runtime_client.go index c8660169..34014b5b 100644 --- a/internal/lambda-managed-instances/testutils/functional/runtime_client.go +++ b/internal/lambda-managed-instances/testutils/functional/runtime_client.go @@ -77,12 +77,19 @@ func (client *Client) Next(body io.Reader) *http.Response { } func (client *Client) Response(invokeID interop.InvokeID, payload io.Reader, contentType string, responseModeHeader string, trailers map[string]string) (*http.Response, error) { + return client.ResponseWithHeaders(invokeID, payload, contentType, responseModeHeader, trailers, nil) +} + +func (client *Client) ResponseWithHeaders(invokeID interop.InvokeID, payload io.Reader, contentType string, responseModeHeader string, trailers map[string]string, extraHeaders map[string]string) (*http.Response, error) { url := fmt.Sprintf("%s/runtime/invocation/%s/response", client.baseurl, invokeID) headers := map[string]string{ContentTypeHeader: contentType} if responseModeHeader != "" { headers[LambdaResponseModeHeader] = responseModeHeader } + for k, v := range extraHeaders { + headers[k] = v + } return client.postBufferedResponse(url, payload, headers, trailers) }