Skip to content

Commit a78c49c

Browse files
valerenavalerena
authored
Pull upstream changes (#34)
1 parent 00d793b commit a78c49c

Some content is hidden

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

45 files changed

+1629
-312
lines changed

‎cmd/aws-lambda-rie/handlers.go

+3-3
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@
44
package main
55

66
import (
7+
"bytes"
78
"fmt"
8-
"io"
99
"io/ioutil"
1010
"math"
1111
"net/http"
@@ -24,7 +24,7 @@ import (
2424

2525
type Sandbox interface {
2626
Init(i *interop.Init, invokeTimeoutMs int64)
27-
Invoke(responseWriter io.Writer, invoke *interop.Invoke) error
27+
Invoke(responseWriter http.ResponseWriter, invoke *interop.Invoke) error
2828
}
2929

3030
var initDone bool
@@ -98,7 +98,7 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox) {
9898
InvokedFunctionArn: fmt.Sprintf("arn:aws:lambda:us-east-1:012345678912:function:%s", GetenvWithDefault("AWS_LAMBDA_FUNCTION_NAME", "test_function")),
9999
TraceID: r.Header.Get("X-Amzn-Trace-Id"),
100100
LambdaSegmentID: r.Header.Get("X-Amzn-Segment-Id"),
101-
Payload: bodyBytes,
101+
Payload: bytes.NewReader(bodyBytes),
102102
CorrelationID: "invokeCorrelationID",
103103
}
104104
fmt.Println("START RequestId: " + invokePayload.ID + " Version: " + functionVersion)

‎cmd/aws-lambda-rie/main.go

+7-10
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ import (
1515
log "github.com/sirupsen/logrus"
1616
)
1717

18-
1918
const (
2019
optBootstrap = "/opt/bootstrap"
2120
runtimeBootstrap = "/var/runtime/bootstrap"
@@ -58,24 +57,22 @@ func getCLIArgs() (options, []string) {
5857
}
5958

6059
func getBootstrap(args []string, opts options) (*rapidcore.Bootstrap, string) {
61-
var bootstrapLookupCmdList [][]string
60+
var bootstrapLookupCmd []string
6261
var handler string
6362
currentWorkingDir := "/var/task" // default value
6463

6564
if len(args) <= 1 {
66-
bootstrapLookupCmdList = [][]string{
67-
[]string{fmt.Sprintf("%s/bootstrap", currentWorkingDir)},
68-
[]string{optBootstrap},
69-
[]string{runtimeBootstrap},
65+
bootstrapLookupCmd = []string{
66+
fmt.Sprintf("%s/bootstrap", currentWorkingDir),
67+
optBootstrap,
68+
runtimeBootstrap,
7069
}
7170

7271
// handler is used later to set an env var for Lambda Image support
7372
handler = ""
7473
} else if len(args) > 1 {
7574

76-
bootstrapLookupCmdList = [][]string{
77-
args[1:],
78-
}
75+
bootstrapLookupCmd = args[1:]
7976

8077
if cwd, err := os.Getwd(); err == nil {
8178
currentWorkingDir = cwd
@@ -92,5 +89,5 @@ func getBootstrap(args []string, opts options) (*rapidcore.Bootstrap, string) {
9289
log.Panic("insufficient arguments: bootstrap not provided")
9390
}
9491

95-
return rapidcore.NewBootstrap(bootstrapLookupCmdList, currentWorkingDir), handler
92+
return rapidcore.NewBootstrapSingleCmd(bootstrapLookupCmd, currentWorkingDir), handler
9693
}

‎lambda/core/directinvoke/customerheaders.go

+41
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package directinvoke
5+
6+
import (
7+
"bytes"
8+
"encoding/base64"
9+
"encoding/json"
10+
)
11+
12+
type CustomerHeaders struct {
13+
CognitoIdentityID string `json:"Cognito-Identity-Id"`
14+
CognitoIdentityPoolID string `json:"Cognito-Identity-Pool-Id"`
15+
ClientContext string `json:"Client-Context"`
16+
}
17+
18+
func (s CustomerHeaders) Dump() string {
19+
if (s == CustomerHeaders{}) {
20+
return ""
21+
}
22+
23+
custHeadersJSON, err := json.Marshal(&s)
24+
if err != nil {
25+
panic(err)
26+
}
27+
28+
return base64.StdEncoding.EncodeToString(custHeadersJSON)
29+
}
30+
31+
func (s *CustomerHeaders) Load(in string) error {
32+
*s = CustomerHeaders{}
33+
34+
if in == "" {
35+
return nil
36+
}
37+
38+
base64Decoder := base64.NewDecoder(base64.StdEncoding, bytes.NewReader([]byte(in)))
39+
40+
return json.NewDecoder(base64Decoder).Decode(s)
41+
}

‎lambda/core/directinvoke/customerheaders_test.go

+25
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package directinvoke
5+
6+
import (
7+
"github.com/stretchr/testify/require"
8+
"testing"
9+
)
10+
11+
func TestCustomerHeadersEmpty(t *testing.T) {
12+
in := CustomerHeaders{}
13+
out := CustomerHeaders{}
14+
15+
require.NoError(t, out.Load(in.Dump()))
16+
require.Equal(t, in, out)
17+
}
18+
19+
func TestCustomerHeaders(t *testing.T) {
20+
in := CustomerHeaders{CognitoIdentityID: "asd"}
21+
out := CustomerHeaders{}
22+
23+
require.NoError(t, out.Load(in.Dump()))
24+
require.Equal(t, in, out)
25+
}

‎lambda/core/directinvoke/directinvoke.go

+106
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package directinvoke
5+
6+
import (
7+
"io"
8+
"net/http"
9+
10+
"github.com/go-chi/chi"
11+
"go.amzn.com/lambda/interop"
12+
)
13+
14+
const (
15+
InvokeIDHeader = "Invoke-Id"
16+
InvokedFunctionArnHeader = "Invoked-Function-Arn"
17+
VersionIDHeader = "Invoked-Function-Version"
18+
ReservationTokenHeader = "Reservation-Token"
19+
CustomerHeadersHeader = "Customer-Headers"
20+
ContentTypeHeader = "Content-Type"
21+
22+
ErrorTypeHeader = "Error-Type"
23+
24+
EndOfResponseTrailer = "End-Of-Response"
25+
26+
SandboxErrorType = "Error.Sandbox"
27+
)
28+
29+
const (
30+
EndOfResponseComplete = "Complete"
31+
EndOfResponseTruncated = "Truncated"
32+
EndOfResponseOversized = "Oversized"
33+
)
34+
35+
var MaxDirectResponseSize int64 = interop.MaxPayloadSize // this is intentionally not a constant so we can configure it via CLI
36+
37+
func renderBadRequest(w http.ResponseWriter, r *http.Request, errorType string) {
38+
w.Header().Set(ErrorTypeHeader, errorType)
39+
w.WriteHeader(http.StatusBadRequest)
40+
w.Header().Set(EndOfResponseTrailer, EndOfResponseComplete)
41+
}
42+
43+
// ReceiveDirectInvoke parses invoke and verifies it against Token message. Uses deadline provided by Token
44+
// Renders BadRequest in case of error
45+
func ReceiveDirectInvoke(w http.ResponseWriter, r *http.Request, token interop.Token) (*interop.Invoke, error) {
46+
w.Header().Set("Trailer", EndOfResponseTrailer)
47+
48+
custHeaders := CustomerHeaders{}
49+
if err := custHeaders.Load(r.Header.Get(CustomerHeadersHeader)); err != nil {
50+
renderBadRequest(w, r, interop.ErrMalformedCustomerHeaders.Error())
51+
return nil, interop.ErrMalformedCustomerHeaders
52+
}
53+
54+
inv := &interop.Invoke{
55+
ID: r.Header.Get(InvokeIDHeader),
56+
ReservationToken: chi.URLParam(r, "reservationtoken"),
57+
InvokedFunctionArn: r.Header.Get(InvokedFunctionArnHeader),
58+
VersionID: r.Header.Get(VersionIDHeader),
59+
ContentType: r.Header.Get(ContentTypeHeader),
60+
CognitoIdentityID: custHeaders.CognitoIdentityID,
61+
CognitoIdentityPoolID: custHeaders.CognitoIdentityPoolID,
62+
TraceID: token.TraceID,
63+
LambdaSegmentID: token.LambdaSegmentID,
64+
ClientContext: custHeaders.ClientContext,
65+
Payload: r.Body,
66+
CorrelationID: "invokeCorrelationID",
67+
DeadlineNs: token.DeadlineNs,
68+
}
69+
70+
if inv.ID != token.InvokeID {
71+
renderBadRequest(w, r, interop.ErrInvalidInvokeID.Error())
72+
return nil, interop.ErrInvalidInvokeID
73+
}
74+
75+
if inv.ReservationToken != token.ReservationToken {
76+
renderBadRequest(w, r, interop.ErrInvalidReservationToken.Error())
77+
return nil, interop.ErrInvalidReservationToken
78+
}
79+
80+
if inv.VersionID != token.VersionID {
81+
renderBadRequest(w, r, interop.ErrInvalidFunctionVersion.Error())
82+
return nil, interop.ErrInvalidFunctionVersion
83+
}
84+
85+
w.Header().Set(VersionIDHeader, token.VersionID)
86+
w.Header().Set(ReservationTokenHeader, token.ReservationToken)
87+
w.Header().Set(InvokeIDHeader, token.InvokeID)
88+
89+
return inv, nil
90+
}
91+
92+
func SendDirectInvokeResponse(additionalHeaders map[string]string, payload io.Reader, w http.ResponseWriter) error {
93+
for k, v := range additionalHeaders {
94+
w.Header().Add(k, v)
95+
}
96+
97+
n, err := io.Copy(w, io.LimitReader(payload, MaxDirectResponseSize+1)) // +1 because we do allow 10MB but not 10MB + 1 byte
98+
if err != nil {
99+
w.Header().Set(EndOfResponseTrailer, EndOfResponseTruncated)
100+
} else if n == MaxDirectResponseSize+1 {
101+
w.Header().Set(EndOfResponseTrailer, EndOfResponseOversized)
102+
} else {
103+
w.Header().Set(EndOfResponseTrailer, EndOfResponseComplete)
104+
}
105+
return err
106+
}

‎lambda/core/registrations.go

+8-1
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,13 @@ func (s *registrationServiceImpl) getInternalStateDescription(appCtx appctx.Appl
158158
}
159159

160160
func (s *registrationServiceImpl) CountAgents() int {
161+
s.mutex.Lock()
162+
defer s.mutex.Unlock()
163+
164+
return s.countAgentsUnsafe()
165+
}
166+
167+
func (s *registrationServiceImpl) countAgentsUnsafe() int {
161168
res := 0
162169
s.externalAgents.Visit(func(a *ExternalAgent) {
163170
res++
@@ -237,7 +244,7 @@ func (s *registrationServiceImpl) CreateInternalAgent(agentName string) (*Intern
237244
return nil, ErrRegistrationServiceOff
238245
}
239246

240-
if s.CountAgents() >= MaxAgentsAllowed {
247+
if s.countAgentsUnsafe() >= MaxAgentsAllowed {
241248
return nil, ErrTooManyExtensions
242249
}
243250

‎lambda/core/statejson/description.go

+16-2
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,9 @@ import (
1010

1111
// StateDescription ...
1212
type StateDescription struct {
13-
Name string `json:"name"`
14-
LastModified int64 `json:"lastModified"`
13+
Name string `json:"name"`
14+
LastModified int64 `json:"lastModified"`
15+
ResponseTimeNs int64 `json:"responseTimeNs"`
1516
}
1617

1718
// RuntimeDescription ...
@@ -34,10 +35,23 @@ type InternalStateDescription struct {
3435
FirstFatalError string `json:"firstFatalError"`
3536
}
3637

38+
// ResetDescription describes fields of the response to an INVOKE API request
39+
type ResetDescription struct {
40+
ExtensionsResetMs int64 `json:"extensionsResetMs"`
41+
}
42+
3743
func (s *InternalStateDescription) AsJSON() []byte {
3844
bytes, err := json.Marshal(s)
3945
if err != nil {
4046
log.Panicf("Failed to marshall internal states: %s", err)
4147
}
4248
return bytes
4349
}
50+
51+
func (s *ResetDescription) AsJSON() []byte {
52+
bytes, err := json.Marshal(s)
53+
if err != nil {
54+
log.Panicf("Failed to marshall reset description: %s", err)
55+
}
56+
return bytes
57+
}

‎lambda/core/states.go

+11-2
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ type Runtime struct {
8585
currentState RuntimeState
8686
stateLastModified time.Time
8787
Pid int
88+
responseTime time.Time
8889

8990
RuntimeStartedState RuntimeState
9091
RuntimeInitErrorState RuntimeState
@@ -150,19 +151,27 @@ func (s *Runtime) InitError() error {
150151
func (s *Runtime) ResponseSent() error {
151152
s.ManagedThread.Lock()
152153
defer s.ManagedThread.Unlock()
153-
return s.currentState.ResponseSent()
154+
err := s.currentState.ResponseSent()
155+
if err == nil {
156+
s.responseTime = time.Now()
157+
}
158+
return err
154159
}
155160

156161
// GetRuntimeDescription returns runtime description object for debugging purposes
157162
func (s *Runtime) GetRuntimeDescription() statejson.RuntimeDescription {
158163
s.ManagedThread.Lock()
159164
defer s.ManagedThread.Unlock()
160-
return statejson.RuntimeDescription{
165+
res := statejson.RuntimeDescription{
161166
State: statejson.StateDescription{
162167
Name: s.currentState.Name(),
163168
LastModified: s.stateLastModified.UnixNano() / int64(time.Millisecond),
164169
},
165170
}
171+
if !s.responseTime.IsZero() {
172+
res.State.ResponseTimeNs = s.responseTime.UnixNano()
173+
}
174+
return res
166175
}
167176

168177
// NewRuntime returns new Runtime instance.

‎lambda/core/states_test.go

+1
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,7 @@ func TestRuntimeStateTransitionsFromInvocationResponseState(t *testing.T) {
110110
runtime.SetState(runtime.RuntimeInvocationResponseState)
111111
assert.NoError(t, runtime.ResponseSent())
112112
assert.Equal(t, runtime.RuntimeResponseSentState, runtime.GetState())
113+
assert.NotEqual(t, 0, runtime.GetRuntimeDescription().State.ResponseTimeNs)
113114
// InvocationResponse-> InvocationResponse
114115
runtime.SetState(runtime.RuntimeInvocationResponseState)
115116
assert.Equal(t, ErrNotAllowed, runtime.InvocationResponse())

‎lambda/fatalerror/fatalerror.go

+1
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ const (
1616
AgentLaunchError ErrorType = "Extension.LaunchError" // agent could not be launched
1717
RuntimeExit ErrorType = "Runtime.ExitError"
1818
InvalidEntrypoint ErrorType = "Runtime.InvalidEntrypoint"
19+
InvalidWorkingDir ErrorType = "Runtime.InvalidWorkingDir"
1920
InvalidTaskConfig ErrorType = "Runtime.InvalidTaskConfig"
2021
Unknown ErrorType = "Unknown"
2122
)

0 commit comments

Comments
 (0)