Skip to content

Commit 5c6579e

Browse files
bmoffattaidansteele
bmoffatt
and
aidansteele
authored
Support http.Handler for RESPONSE_STREAM Lambda Function URLs (#503)
* initial * typo in events redaction * 1.18+ * remove the panic recover for now - the runtime api client code does not yet re-propogate the crash * Fix typo Co-authored-by: Aidan Steele <[email protected]> * Update http_handler.go * base64 decode branch coverage * cover RequestFromContext --------- Co-authored-by: Aidan Steele <[email protected]>
1 parent dc78417 commit 5c6579e

13 files changed

+606
-0
lines changed

‎lambdaurl/http_handler.go

+103
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
//go:build go1.18
2+
// +build go1.18
3+
4+
// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
5+
6+
// Package lambdaurl serves requests from Lambda Function URLs using http.Handler.
7+
package lambdaurl
8+
9+
import (
10+
"context"
11+
"encoding/base64"
12+
"io"
13+
"net/http"
14+
"strings"
15+
"sync"
16+
17+
"github.com/aws/aws-lambda-go/events"
18+
"github.com/aws/aws-lambda-go/lambda"
19+
)
20+
21+
type httpResponseWriter struct {
22+
header http.Header
23+
writer io.Writer
24+
once sync.Once
25+
status chan<- int
26+
}
27+
28+
func (w *httpResponseWriter) Header() http.Header {
29+
return w.header
30+
}
31+
32+
func (w *httpResponseWriter) Write(p []byte) (int, error) {
33+
w.once.Do(func() { w.status <- http.StatusOK })
34+
return w.writer.Write(p)
35+
}
36+
37+
func (w *httpResponseWriter) WriteHeader(statusCode int) {
38+
w.once.Do(func() { w.status <- statusCode })
39+
}
40+
41+
type requestContextKey struct{}
42+
43+
// RequestFromContext returns the *events.LambdaFunctionURLRequest from a context.
44+
func RequestFromContext(ctx context.Context) (*events.LambdaFunctionURLRequest, bool) {
45+
req, ok := ctx.Value(requestContextKey{}).(*events.LambdaFunctionURLRequest)
46+
return req, ok
47+
}
48+
49+
// Wrap converts an http.Handler into a lambda request handler.
50+
// Only Lambda Function URLs configured with `InvokeMode: RESPONSE_STREAM` are supported with the returned handler.
51+
// The response body of the handler will conform to the content-type `application/vnd.awslambda.http-integration-response`
52+
func Wrap(handler http.Handler) func(context.Context, *events.LambdaFunctionURLRequest) (*events.LambdaFunctionURLStreamingResponse, error) {
53+
return func(ctx context.Context, request *events.LambdaFunctionURLRequest) (*events.LambdaFunctionURLStreamingResponse, error) {
54+
var body io.Reader = strings.NewReader(request.Body)
55+
if request.IsBase64Encoded {
56+
body = base64.NewDecoder(base64.StdEncoding, body)
57+
}
58+
url := "https://" + request.RequestContext.DomainName + request.RawPath
59+
if request.RawQueryString != "" {
60+
url += "?" + request.RawQueryString
61+
}
62+
ctx = context.WithValue(ctx, requestContextKey{}, request)
63+
httpRequest, err := http.NewRequestWithContext(ctx, request.RequestContext.HTTP.Method, url, body)
64+
if err != nil {
65+
return nil, err
66+
}
67+
for k, v := range request.Headers {
68+
httpRequest.Header.Add(k, v)
69+
}
70+
status := make(chan int) // Signals when it's OK to start returning the response body to Lambda
71+
header := http.Header{}
72+
r, w := io.Pipe()
73+
go func() {
74+
defer close(status)
75+
defer w.Close() // TODO: recover and CloseWithError the any panic value once the runtime API client supports plumbing fatal errors through the reader
76+
handler.ServeHTTP(&httpResponseWriter{writer: w, header: header, status: status}, httpRequest)
77+
}()
78+
response := &events.LambdaFunctionURLStreamingResponse{
79+
Body: r,
80+
StatusCode: <-status,
81+
}
82+
if len(header) > 0 {
83+
response.Headers = make(map[string]string, len(header))
84+
for k, v := range header {
85+
if k == "Set-Cookie" {
86+
response.Cookies = v
87+
} else {
88+
response.Headers[k] = strings.Join(v, ",")
89+
}
90+
}
91+
}
92+
return response, nil
93+
}
94+
}
95+
96+
// Start wraps a http.Handler and calls lambda.StartHandlerFunc
97+
// Only supports:
98+
// - Lambda Function URLs configured with `InvokeMode: RESPONSE_STREAM`
99+
// - Lambda Functions using the `provided` or `provided.al2` runtimes.
100+
// - Lambda Functions using the `go1.x` runtime when compiled with `-tags lambda.norpc`
101+
func Start(handler http.Handler, options ...lambda.Option) {
102+
lambda.StartHandlerFunc(Wrap(handler), options...)
103+
}

‎lambdaurl/http_handler_test.go

+157
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
//go:build go1.18
2+
// +build go1.18
3+
4+
// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
5+
package lambdaurl
6+
7+
import (
8+
"bytes"
9+
"context"
10+
_ "embed"
11+
"encoding/json"
12+
"io"
13+
"io/ioutil"
14+
"log"
15+
"net/http"
16+
"testing"
17+
"time"
18+
19+
"github.com/aws/aws-lambda-go/events"
20+
"github.com/stretchr/testify/assert"
21+
"github.com/stretchr/testify/require"
22+
)
23+
24+
//go:embed testdata/function-url-request-with-headers-and-cookies-and-text-body.json
25+
var helloRequest []byte
26+
27+
//go:embed testdata/function-url-domain-only-get-request.json
28+
var domainOnlyGetRequest []byte
29+
30+
//go:embed testdata/function-url-domain-only-get-request-trailing-slash.json
31+
var domainOnlyWithSlashGetRequest []byte
32+
33+
//go:embed testdata/function-url-domain-only-request-with-base64-encoded-body.json
34+
var base64EncodedBodyRequest []byte
35+
36+
func TestWrap(t *testing.T) {
37+
for name, params := range map[string]struct {
38+
input []byte
39+
handler http.HandlerFunc
40+
expectStatus int
41+
expectBody string
42+
expectHeaders map[string]string
43+
expectCookies []string
44+
}{
45+
"hello": {
46+
input: helloRequest,
47+
handler: func(w http.ResponseWriter, r *http.Request) {
48+
w.Header().Add("Hello", "world1")
49+
w.Header().Add("Hello", "world2")
50+
http.SetCookie(w, &http.Cookie{Name: "yummy", Value: "cookie"})
51+
http.SetCookie(w, &http.Cookie{Name: "yummy", Value: "cake"})
52+
http.SetCookie(w, &http.Cookie{Name: "fruit", Value: "banana", Expires: time.Date(2000, time.January, 0, 0, 0, 0, 0, time.UTC)})
53+
for _, c := range r.Cookies() {
54+
http.SetCookie(w, c)
55+
}
56+
57+
w.WriteHeader(http.StatusTeapot)
58+
encoder := json.NewEncoder(w)
59+
_ = encoder.Encode(struct{ RequestQueryParams, Method any }{r.URL.Query(), r.Method})
60+
},
61+
expectStatus: http.StatusTeapot,
62+
expectHeaders: map[string]string{
63+
"Hello": "world1,world2",
64+
},
65+
expectCookies: []string{
66+
"yummy=cookie",
67+
"yummy=cake",
68+
"fruit=banana; Expires=Fri, 31 Dec 1999 00:00:00 GMT",
69+
"foo=bar",
70+
"hello=hello",
71+
},
72+
expectBody: `{"RequestQueryParams":{"foo":["bar"],"hello":["world"]},"Method":"POST"}` + "\n",
73+
},
74+
"mux": {
75+
input: helloRequest,
76+
handler: func(w http.ResponseWriter, r *http.Request) {
77+
log.Println(r.URL)
78+
mux := http.NewServeMux()
79+
mux.HandleFunc("/hello", func(w http.ResponseWriter, r *http.Request) {
80+
w.WriteHeader(200)
81+
_, _ = w.Write([]byte("Hello World!"))
82+
})
83+
mux.ServeHTTP(w, r)
84+
},
85+
expectStatus: 200,
86+
expectBody: "Hello World!",
87+
},
88+
"get-implicit-trailing-slash": {
89+
input: domainOnlyGetRequest,
90+
handler: func(w http.ResponseWriter, r *http.Request) {
91+
encoder := json.NewEncoder(w)
92+
_ = encoder.Encode(r.Method)
93+
_ = encoder.Encode(r.URL.String())
94+
},
95+
expectStatus: http.StatusOK,
96+
expectBody: "\"GET\"\n\"https://lambda-url-id.lambda-url.us-west-2.on.aws/\"\n",
97+
},
98+
"get-explicit-trailing-slash": {
99+
input: domainOnlyWithSlashGetRequest,
100+
handler: func(w http.ResponseWriter, r *http.Request) {
101+
encoder := json.NewEncoder(w)
102+
_ = encoder.Encode(r.Method)
103+
_ = encoder.Encode(r.URL.String())
104+
},
105+
expectStatus: http.StatusOK,
106+
expectBody: "\"GET\"\n\"https://lambda-url-id.lambda-url.us-west-2.on.aws/\"\n",
107+
},
108+
"empty handler": {
109+
input: helloRequest,
110+
handler: func(w http.ResponseWriter, r *http.Request) {},
111+
expectStatus: http.StatusOK,
112+
},
113+
"base64request": {
114+
input: base64EncodedBodyRequest,
115+
handler: func(w http.ResponseWriter, r *http.Request) {
116+
_, _ = io.Copy(w, r.Body)
117+
},
118+
expectStatus: http.StatusOK,
119+
expectBody: "<idk/>",
120+
},
121+
} {
122+
t.Run(name, func(t *testing.T) {
123+
handler := Wrap(params.handler)
124+
var req events.LambdaFunctionURLRequest
125+
require.NoError(t, json.Unmarshal(params.input, &req))
126+
res, err := handler(context.Background(), &req)
127+
require.NoError(t, err)
128+
resultBodyBytes, err := ioutil.ReadAll(res)
129+
require.NoError(t, err)
130+
resultHeaderBytes, resultBodyBytes, ok := bytes.Cut(resultBodyBytes, []byte{0, 0, 0, 0, 0, 0, 0, 0})
131+
require.True(t, ok)
132+
var resultHeader struct {
133+
StatusCode int
134+
Headers map[string]string
135+
Cookies []string
136+
}
137+
require.NoError(t, json.Unmarshal(resultHeaderBytes, &resultHeader))
138+
assert.Equal(t, params.expectBody, string(resultBodyBytes))
139+
assert.Equal(t, params.expectStatus, resultHeader.StatusCode)
140+
assert.Equal(t, params.expectHeaders, resultHeader.Headers)
141+
assert.Equal(t, params.expectCookies, resultHeader.Cookies)
142+
})
143+
}
144+
}
145+
146+
func TestRequestContext(t *testing.T) {
147+
var req *events.LambdaFunctionURLRequest
148+
require.NoError(t, json.Unmarshal(helloRequest, &req))
149+
handler := Wrap(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
150+
reqFromContext, exists := RequestFromContext(r.Context())
151+
require.True(t, exists)
152+
require.NotNil(t, reqFromContext)
153+
assert.Equal(t, req, reqFromContext)
154+
}))
155+
_, err := handler(context.Background(), req)
156+
require.NoError(t, err)
157+
}

‎lambdaurl/testdata/function-url-domain-only-get-request-trailing-slash.json

+44
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
{
2+
"headers": {
3+
"accept": "application/xml",
4+
"accept-encoding": "gzip, deflate",
5+
"content-type": "application/json",
6+
"host": "lambda-url-id.lambda-url.us-west-2.on.aws",
7+
"user-agent": "python-requests/2.28.2",
8+
"x-amz-content-sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
9+
"x-amz-date": "20230418T170147Z",
10+
"x-amz-security-token": "security-token",
11+
"x-amzn-tls-cipher-suite": "ECDHE-RSA-AES128-GCM-SHA256",
12+
"x-amzn-tls-version": "TLSv1.2",
13+
"x-amzn-trace-id": "Root=1-643eccfb-7c4d3f09749a95a044db997a",
14+
"x-forwarded-for": "127.0.0.1",
15+
"x-forwarded-port": "443",
16+
"x-forwarded-proto": "https"
17+
},
18+
"isBase64Encoded": false,
19+
"rawPath": "/",
20+
"rawQueryString": "",
21+
"requestContext": {
22+
"accountId": "aws-account-id",
23+
"apiId": "lambda-url-id",
24+
"authorizer": {
25+
"iam": {}
26+
},
27+
"domainName": "lambda-url-id.lambda-url.us-west-2.on.aws",
28+
"domainPrefix": "lambda-url-id",
29+
"http": {
30+
"method": "GET",
31+
"path": "/",
32+
"protocol": "HTTP/1.1",
33+
"sourceIp": "127.0.0.1",
34+
"userAgent": "python-requests/2.28.2"
35+
},
36+
"requestId": "3a72f39b-d6bd-4a4f-b040-f94d09b4daa3",
37+
"routeKey": "$default",
38+
"stage": "$default",
39+
"time": "18/Apr/2023:17:01:47 +0000",
40+
"timeEpoch": 1681837307717
41+
},
42+
"routeKey": "$default",
43+
"version": "2.0"
44+
}

‎lambdaurl/testdata/function-url-domain-only-get-request.json

+44
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
{
2+
"headers": {
3+
"accept": "application/xml",
4+
"accept-encoding": "gzip, deflate",
5+
"content-type": "application/json",
6+
"host": "lambda-url-id.lambda-url.us-west-2.on.aws",
7+
"user-agent": "python-requests/2.28.2",
8+
"x-amz-content-sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
9+
"x-amz-date": "20230418T170147Z",
10+
"x-amz-security-token": "security-token",
11+
"x-amzn-tls-cipher-suite": "ECDHE-RSA-AES128-GCM-SHA256",
12+
"x-amzn-tls-version": "TLSv1.2",
13+
"x-amzn-trace-id": "Root=1-643eccfb-4c9be61972302fa41111a443",
14+
"x-forwarded-for": "127.0.0.1",
15+
"x-forwarded-port": "443",
16+
"x-forwarded-proto": "https"
17+
},
18+
"isBase64Encoded": false,
19+
"rawPath": "/",
20+
"rawQueryString": "",
21+
"requestContext": {
22+
"accountId": "aws-account-id",
23+
"apiId": "lambda-url-id",
24+
"authorizer": {
25+
"iam": {}
26+
},
27+
"domainName": "lambda-url-id.lambda-url.us-west-2.on.aws",
28+
"domainPrefix": "lambda-url-id",
29+
"http": {
30+
"method": "GET",
31+
"path": "/",
32+
"protocol": "HTTP/1.1",
33+
"sourceIp": "127.0.0.1",
34+
"userAgent": "python-requests/2.28.2"
35+
},
36+
"requestId": "deeb7e49-a9a8-4a8f-bcd1-5482231e2087",
37+
"routeKey": "$default",
38+
"stage": "$default",
39+
"time": "18/Apr/2023:17:01:47 +0000",
40+
"timeEpoch": 1681837307545
41+
},
42+
"routeKey": "$default",
43+
"version": "2.0"
44+
}

‎lambdaurl/testdata/function-url-domain-only-request-with-base64-encoded-body.json

+46
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
{
2+
"body": "PGlkay8+",
3+
"headers": {
4+
"accept": "*/*",
5+
"accept-encoding": "gzip, deflate",
6+
"content-length": "6",
7+
"content-type": "idk",
8+
"host": "lambda-url-id.lambda-url.us-west-2.on.aws",
9+
"user-agent": "python-requests/2.28.2",
10+
"x-amz-content-sha256": "0ab2082273499eaa495f2196e32d8c794745e58a20a0c93182c59d2165432839",
11+
"x-amz-date": "20230418T170147Z",
12+
"x-amz-security-token": "security-token",
13+
"x-amzn-tls-cipher-suite": "ECDHE-RSA-AES128-GCM-SHA256",
14+
"x-amzn-tls-version": "TLSv1.2",
15+
"x-amzn-trace-id": "Root=1-643eccfb-7fdecb844a12b4b45645132d",
16+
"x-forwarded-for": "127.0.0.1",
17+
"x-forwarded-port": "443",
18+
"x-forwarded-proto": "https"
19+
},
20+
"isBase64Encoded": true,
21+
"rawPath": "/",
22+
"rawQueryString": "",
23+
"requestContext": {
24+
"accountId": "aws-account-id",
25+
"apiId": "lambda-url-id",
26+
"authorizer": {
27+
"iam": {}
28+
},
29+
"domainName": "lambda-url-id.lambda-url.us-west-2.on.aws",
30+
"domainPrefix": "lambda-url-id",
31+
"http": {
32+
"method": "POST",
33+
"path": "/",
34+
"protocol": "HTTP/1.1",
35+
"sourceIp": "127.0.0.1",
36+
"userAgent": "python-requests/2.28.2"
37+
},
38+
"requestId": "9701a3d4-36ad-40bd-bf0b-a525c987d27f",
39+
"routeKey": "$default",
40+
"stage": "$default",
41+
"time": "18/Apr/2023:17:01:47 +0000",
42+
"timeEpoch": 1681837307386
43+
},
44+
"routeKey": "$default",
45+
"version": "2.0"
46+
}

0 commit comments

Comments
 (0)