Skip to content

Commit 21ce1f0

Browse files
skyrpexskyrpex
committed
wip
1 parent c91a64a commit 21ce1f0

56 files changed

Lines changed: 1963 additions & 130 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -193,4 +193,4 @@ See [CONTRIBUTING](CONTRIBUTING.md#security-issue-notifications) for more inform
193193
194194
## License
195195
196-
This project is licensed under the Apache-2.0 License.
196+
This project is licensed under the Apache-2.0 License.

‎internal/lambda-managed-instances/aws-lambda-rie/internal/app.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,11 @@ func (h *HTTPHandler) invoke(w http.ResponseWriter, r *http.Request) {
6060
return
6161
}
6262

63-
invokeReq := rieinvoke.NewRieInvokeRequest(r, w)
63+
invokeReq, err := rieinvoke.NewRieInvokeRequest(r, w)
64+
if err != nil {
65+
h.respondWithError(w, err)
66+
return
67+
}
6468
ctx := logging.WithInvokeID(r.Context(), invokeReq.InvokeID())
6569

6670
metrics := invoke.NewInvokeMetrics(nil, &noOpCounter{})

‎internal/lambda-managed-instances/aws-lambda-rie/internal/init.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,8 +61,9 @@ func GetInitRequestMessage(fileUtil utils.FileUtil, args []string) (intmodel.Ini
6161
XrayTracingMode: intmodel.XRayTracingModePassThrough,
6262
CurrentWorkingDir: cwd,
6363
RuntimeBinaryCommand: cmd,
64-
AvailabilityZoneId: "",
65-
AmiId: "",
64+
65+
AvailabilityZoneId: "use1-az1",
66+
AmiId: "",
6667
}, nil
6768
}
6869

‎internal/lambda-managed-instances/aws-lambda-rie/internal/init_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ func Test_getInitRequestMessage(t *testing.T) {
6262
XrayTracingMode: intmodel.XRayTracingModePassThrough,
6363
CurrentWorkingDir: "REPLACE",
6464
RuntimeBinaryCommand: []string{"/path/to/bootstrap"},
65-
AvailabilityZoneId: "",
65+
AvailabilityZoneId: "use1-az1",
6666
AmiId: "",
6767
},
6868
},
@@ -116,7 +116,7 @@ func Test_getInitRequestMessage(t *testing.T) {
116116
XrayTracingMode: intmodel.XRayTracingModePassThrough,
117117
CurrentWorkingDir: "/var/task",
118118
RuntimeBinaryCommand: []string{"/custom/bootstrap", "custom_handler"},
119-
AvailabilityZoneId: "",
119+
AvailabilityZoneId: "use1-az1",
120120
AmiId: "",
121121
},
122122
},

‎internal/lambda-managed-instances/aws-lambda-rie/internal/invoke/consts.go‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package invoke
5+
6+
const (
7+
RequestIdHeader = "X-Amzn-RequestId"
8+
9+
ClientContextHeader = "X-Amz-Client-Context"
10+
11+
CognitoIdentityHeader = "X-Amz-Cognito-Identity"
12+
)

‎internal/lambda-managed-instances/aws-lambda-rie/internal/invoke/rie_invoke_request.go‎

Lines changed: 44 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,12 @@
44
package invoke
55

66
import (
7+
"encoding/base64"
8+
"encoding/json"
79
"errors"
10+
"fmt"
811
"io"
12+
"log/slog"
913
"net/http"
1014
"time"
1115

@@ -16,6 +20,11 @@ import (
1620
"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/rapid/model"
1721
)
1822

23+
type cognitoIdentity struct {
24+
CognitoIdentityID string `json:"cognitoIdentityId"`
25+
CognitoIdentityPoolID string `json:"cognitoIdentityPoolId"`
26+
}
27+
1928
type rieInvokeRequest struct {
2029
request *http.Request
2130
writer http.ResponseWriter
@@ -35,18 +44,47 @@ type rieInvokeRequest struct {
3544
functionVersionID string
3645
}
3746

38-
func NewRieInvokeRequest(request *http.Request, writer http.ResponseWriter) *rieInvokeRequest {
47+
func NewRieInvokeRequest(request *http.Request, writer http.ResponseWriter) (*rieInvokeRequest, model.AppError) {
3948

4049
contentType := request.Header.Get(invoke.СontentTypeHeader)
4150
if contentType == "" {
4251
contentType = "application/json"
4352
}
4453

45-
invokeID := request.Header.Get("X-Amzn-RequestId")
54+
invokeID := request.Header.Get(RequestIdHeader)
4655
if invokeID == "" {
4756
invokeID = uuid.New().String()
4857
}
4958

59+
clientContext := ""
60+
if encodedClientContext := request.Header.Get(ClientContextHeader); encodedClientContext != "" {
61+
decodedClientContext, err := base64.StdEncoding.DecodeString(encodedClientContext)
62+
if err != nil {
63+
slog.Warn("Failed to decode X-Amz-Client-Context header", "err", err)
64+
return nil, model.NewClientError(
65+
fmt.Errorf("X-Amz-Client-Context must be a valid base64 encoded string: %w", err),
66+
model.ErrorSeverityInvalid,
67+
model.ErrorMalformedRequest,
68+
)
69+
}
70+
clientContext = string(decodedClientContext)
71+
}
72+
73+
var cognitoIdentityId, cognitoIdentityPoolId string
74+
if cognitoIdentityHeader := request.Header.Get(CognitoIdentityHeader); cognitoIdentityHeader != "" {
75+
var cognito cognitoIdentity
76+
if err := json.Unmarshal([]byte(cognitoIdentityHeader), &cognito); err != nil {
77+
slog.Warn("Failed to parse X-Amz-Cognito-Identity header", "err", err)
78+
return nil, model.NewClientError(
79+
fmt.Errorf("X-Amz-Cognito-Identity must be a valid JSON string: %w", err),
80+
model.ErrorSeverityInvalid,
81+
model.ErrorMalformedRequest,
82+
)
83+
}
84+
cognitoIdentityId = cognito.CognitoIdentityID
85+
cognitoIdentityPoolId = cognito.CognitoIdentityPoolID
86+
}
87+
5088
req := &rieInvokeRequest{
5189
request: request,
5290
writer: writer,
@@ -56,13 +94,13 @@ func NewRieInvokeRequest(request *http.Request, writer http.ResponseWriter) *rie
5694
responseBandwidthRate: 2 * 1024 * 1024,
5795
responseBandwidthBurstSize: 6 * 1024 * 1024,
5896
traceId: request.Header.Get(invoke.TraceIdHeader),
59-
cognitoIdentityId: "",
60-
cognitoIdentityPoolId: "",
61-
clientContext: request.Header.Get("X-Amz-Client-Context"),
97+
cognitoIdentityId: cognitoIdentityId,
98+
cognitoIdentityPoolId: cognitoIdentityPoolId,
99+
clientContext: clientContext,
62100
responseMode: request.Header.Get(invoke.ResponseModeHeader),
63101
}
64102

65-
return req
103+
return req, nil
66104
}
67105

68106
func (r *rieInvokeRequest) ContentType() string {

‎internal/lambda-managed-instances/aws-lambda-rie/internal/invoke/rie_invoke_request_test.go‎

Lines changed: 77 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,14 +10,18 @@ import (
1010

1111
"github.com/stretchr/testify/assert"
1212
"github.com/stretchr/testify/require"
13+
14+
"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/rapid/model"
1315
)
1416

1517
func TestNewRieInvokeRequest(t *testing.T) {
1618
tests := []struct {
17-
name string
18-
request func() *http.Request
19-
writer http.ResponseWriter
20-
want *rieInvokeRequest
19+
name string
20+
request func() *http.Request
21+
writer http.ResponseWriter
22+
want *rieInvokeRequest
23+
wantError bool
24+
wantErrorContain string
2125
}{
2226
{
2327
name: "no_headers_in_request",
@@ -37,6 +41,7 @@ func TestNewRieInvokeRequest(t *testing.T) {
3741
cognitoIdentityPoolId: "",
3842
clientContext: "",
3943
},
44+
wantError: false,
4045
},
4146
{
4247
name: "all_headers_present_in_request",
@@ -46,6 +51,7 @@ func TestNewRieInvokeRequest(t *testing.T) {
4651
r.Header.Set("X-Amzn-Trace-Id", "Root=1-5e1b4151-5ac6c58f3375aa3c7c6b73c9")
4752
r.Header.Set("X-Amz-Client-Context", "eyJjdXN0b20iOnsidGVzdCI6InZhbHVlIn19")
4853
r.Header.Set("X-Amzn-RequestId", "test-invoke-id")
54+
r.Header.Set("X-Amz-Cognito-Identity", `{"cognitoIdentityId":"us-east-1:12345678-1234-1234-1234-123456789012","cognitoIdentityPoolId":"us-east-1:87654321-4321-4321-4321-210987654321"}`)
4955
require.NoError(t, err)
5056
return r
5157
},
@@ -57,16 +63,80 @@ func TestNewRieInvokeRequest(t *testing.T) {
5763
responseBandwidthRate: 2 * 1024 * 1024,
5864
responseBandwidthBurstSize: 6 * 1024 * 1024,
5965
traceId: "Root=1-5e1b4151-5ac6c58f3375aa3c7c6b73c9",
60-
cognitoIdentityId: "",
66+
cognitoIdentityId: "us-east-1:12345678-1234-1234-1234-123456789012",
67+
cognitoIdentityPoolId: "us-east-1:87654321-4321-4321-4321-210987654321",
68+
clientContext: `{"custom":{"test":"value"}}`,
69+
},
70+
wantError: false,
71+
},
72+
{
73+
name: "malformed_cognito_identity_header",
74+
request: func() *http.Request {
75+
r, err := http.NewRequest("GET", "http://localhost/", nil)
76+
r.Header.Set("X-Amzn-RequestId", "test-invoke-id")
77+
r.Header.Set("X-Amz-Cognito-Identity", "not-valid-json{")
78+
require.NoError(t, err)
79+
return r
80+
},
81+
writer: httptest.NewRecorder(),
82+
want: nil,
83+
wantError: true,
84+
wantErrorContain: "X-Amz-Cognito-Identity must be a valid JSON string",
85+
},
86+
{
87+
name: "malformed_client_context_header",
88+
request: func() *http.Request {
89+
r, err := http.NewRequest("GET", "http://localhost/", nil)
90+
r.Header.Set("X-Amzn-RequestId", "test-invoke-id")
91+
r.Header.Set("X-Amz-Client-Context", "not-valid-base64!!!")
92+
require.NoError(t, err)
93+
return r
94+
},
95+
writer: httptest.NewRecorder(),
96+
want: nil,
97+
wantError: true,
98+
wantErrorContain: "X-Amz-Client-Context must be a valid base64 encoded string",
99+
},
100+
{
101+
name: "partial_cognito_identity_header",
102+
request: func() *http.Request {
103+
r, err := http.NewRequest("GET", "http://localhost/", nil)
104+
r.Header.Set("X-Amzn-RequestId", "test-invoke-id")
105+
r.Header.Set("X-Amz-Cognito-Identity", `{"cognitoIdentityId":"us-east-1:only-id"}`)
106+
require.NoError(t, err)
107+
return r
108+
},
109+
writer: httptest.NewRecorder(),
110+
want: &rieInvokeRequest{
111+
invokeID: "test-invoke-id",
112+
contentType: "application/json",
113+
maxPayloadSize: 6*1024*1024 + 100,
114+
responseBandwidthRate: 2 * 1024 * 1024,
115+
responseBandwidthBurstSize: 6 * 1024 * 1024,
116+
traceId: "",
117+
cognitoIdentityId: "us-east-1:only-id",
61118
cognitoIdentityPoolId: "",
62-
clientContext: "eyJjdXN0b20iOnsidGVzdCI6InZhbHVlIn19",
119+
clientContext: "",
63120
},
121+
wantError: false,
64122
},
65123
}
66124
for _, tt := range tests {
67125
t.Run(tt.name, func(t *testing.T) {
68126
r := tt.request()
69-
got := NewRieInvokeRequest(r, tt.writer)
127+
got, err := NewRieInvokeRequest(r, tt.writer)
128+
129+
if tt.wantError {
130+
assert.NotNil(t, err)
131+
assert.Nil(t, got)
132+
assert.Equal(t, model.ErrorMalformedRequest, err.ErrorType())
133+
assert.Equal(t, http.StatusBadRequest, err.ReturnCode())
134+
assert.Contains(t, err.Error(), tt.wantErrorContain)
135+
return
136+
}
137+
138+
assert.Nil(t, err)
139+
require.NotNil(t, got)
70140

71141
tt.want.request = r
72142
tt.want.writer = tt.writer

‎internal/lambda-managed-instances/aws-lambda-rie/internal/run.go‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,9 @@ import (
1010
"os"
1111
"time"
1212

13+
"github.com/google/uuid"
14+
"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lmds"
15+
1316
rieinvoke "github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/aws-lambda-rie/internal/invoke"
1417
"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/aws-lambda-rie/internal/telemetry"
1518
"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/interop"
@@ -47,8 +50,9 @@ func Run(supv supvmodel.ProcessSupervisor, args []string, fileUtil utils.FileUti
4750
responderFactoryFunc := func(_ context.Context, invokeReq interop.InvokeRequest) invoke.InvokeResponseSender {
4851
return rieinvoke.NewResponder(invokeReq)
4952
}
50-
invokeRouter := invoke.NewInvokeRouter(rapid.MaxIdleRuntimesQueueSize, eventsAPI, responderFactoryFunc, timeout.NewRecentCache())
53+
invokeRouter := invoke.NewInvokeRouter(rapid.RuntimePoolSize, eventsAPI, responderFactoryFunc, timeout.NewRecentCache())
5154

55+
metadataToken := uuid.NewString()
5256
deps := rapid.Dependencies{
5357
EventsAPI: eventsAPI,
5458
LogsEgressAPI: telemetry.NewLogsEgress(telemetryAPIRelay, os.Stdout),
@@ -57,9 +61,10 @@ func Run(supv supvmodel.ProcessSupervisor, args []string, fileUtil utils.FileUti
5761
RuntimeAPIAddrPort: runtimeAPIAddr,
5862
FileUtils: fileUtil,
5963
InvokeRouter: invokeRouter,
64+
MetadataService: lmds.NewService(metadataToken),
6065
}
6166

62-
raptorApp, err := raptor.StartApp(deps, "", noOpLogger{})
67+
raptorApp, err := raptor.StartApp(deps, "", metadataToken, noOpLogger{})
6368
if err != nil {
6469
return nil, nil, nil, fmt.Errorf("could not start runtime api server: %w", err)
6570
}

0 commit comments

Comments
 (0)