Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support http.Handler for RESPONSE_STREAM Lambda Function URLs #503

Merged
merged 8 commits into from
Apr 23, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
103 changes: 103 additions & 0 deletions lambdaurl/http_handler.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
//go:build go1.18
// +build go1.18

// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.

// Package lambdaurl serves requests from Lambda Function URLs using http.Handler.
package lambdaurl

import (
"context"
"encoding/base64"
"io"
"net/http"
"strings"
"sync"

"github.com/aws/aws-lambda-go/events"
"github.com/aws/aws-lambda-go/lambda"
)

type httpResponseWriter struct {
header http.Header
writer io.Writer
once sync.Once
status chan<- int
}

func (w *httpResponseWriter) Header() http.Header {
return w.header
}

func (w *httpResponseWriter) Write(p []byte) (int, error) {
w.once.Do(func() { w.status <- http.StatusOK })
return w.writer.Write(p)
}

func (w *httpResponseWriter) WriteHeader(statusCode int) {
w.once.Do(func() { w.status <- statusCode })
}

type requestContextKey struct{}

// RequestFromContext returns the *events.LambdaFunctionURLRequest from a context.
func RequestFromContext(ctx context.Context) (*events.LambdaFunctionURLRequest, bool) {
req, ok := ctx.Value(requestContextKey{}).(*events.LambdaFunctionURLRequest)
return req, ok
}

// Wrap converts an http.Handler into a lambda request handler.
// Only Lambda Function URLs configured with `InvokeMode: RESPONSE_STREAM` are supported with the returned handler.
// The response body of the handler will conform to the content-type `application/vnd.awslambda.http-integration-response`
func Wrap(handler http.Handler) func(context.Context, *events.LambdaFunctionURLRequest) (*events.LambdaFunctionURLStreamingResponse, error) {
return func(ctx context.Context, request *events.LambdaFunctionURLRequest) (*events.LambdaFunctionURLStreamingResponse, error) {
var body io.Reader = strings.NewReader(request.Body)
if request.IsBase64Encoded {
body = base64.NewDecoder(base64.StdEncoding, body)
}
url := "https://" + request.RequestContext.DomainName + request.RawPath
if request.RawQueryString != "" {
url += "?" + request.RawQueryString
}
ctx = context.WithValue(ctx, requestContextKey{}, request)
httpRequest, err := http.NewRequestWithContext(ctx, request.RequestContext.HTTP.Method, url, body)
if err != nil {
return nil, err
}
for k, v := range request.Headers {
httpRequest.Header.Add(k, v)
}
status := make(chan int) // Signals when it's OK to start returning the response body to Lambda
header := http.Header{}
r, w := io.Pipe()
go func() {
defer close(status)
defer w.Close() // TODO: recover and CloseWithError the any panic value once the runtime API client supports plumbing fatal errors through the reader
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will get back to this in the follow-up for all returned writers. As for now allowing the panic to crash the process and log is the "most correct" thing do do, at least compared to swallowing it without any logging / reporting!

handler.ServeHTTP(&httpResponseWriter{writer: w, header: header, status: status}, httpRequest)
}()
response := &events.LambdaFunctionURLStreamingResponse{
Body: r,
StatusCode: <-status,
}
if len(header) > 0 {
response.Headers = make(map[string]string, len(header))
for k, v := range header {
if k == "Set-Cookie" {
response.Cookies = v
} else {
response.Headers[k] = strings.Join(v, ",")
}
}
}
return response, nil
}
}

// Start wraps a http.Handler and calls lambda.StartHandlerFunc
// Only supports:
// - Lambda Function URLs configured with `InvokeMode: RESPONSE_STREAM`
// - Lambda Functions using the `provided` or `provided.al2` runtimes.
// - Lambda Functions using the `go1.x` runtime when compiled with `-tags lambda.norpc`
func Start(handler http.Handler, options ...lambda.Option) {
lambda.StartHandlerFunc(Wrap(handler), options...)
}
157 changes: 157 additions & 0 deletions lambdaurl/http_handler_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,157 @@
//go:build go1.18
// +build go1.18

// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
package lambdaurl

import (
"bytes"
"context"
_ "embed"
"encoding/json"
"io"
"io/ioutil"
"log"
"net/http"
"testing"
"time"

"github.com/aws/aws-lambda-go/events"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

//go:embed testdata/function-url-request-with-headers-and-cookies-and-text-body.json
var helloRequest []byte

//go:embed testdata/function-url-domain-only-get-request.json
var domainOnlyGetRequest []byte

//go:embed testdata/function-url-domain-only-get-request-trailing-slash.json
var domainOnlyWithSlashGetRequest []byte

//go:embed testdata/function-url-domain-only-request-with-base64-encoded-body.json
var base64EncodedBodyRequest []byte

func TestWrap(t *testing.T) {
for name, params := range map[string]struct {
input []byte
handler http.HandlerFunc
expectStatus int
expectBody string
expectHeaders map[string]string
expectCookies []string
}{
"hello": {
input: helloRequest,
handler: func(w http.ResponseWriter, r *http.Request) {
w.Header().Add("Hello", "world1")
w.Header().Add("Hello", "world2")
http.SetCookie(w, &http.Cookie{Name: "yummy", Value: "cookie"})
http.SetCookie(w, &http.Cookie{Name: "yummy", Value: "cake"})
http.SetCookie(w, &http.Cookie{Name: "fruit", Value: "banana", Expires: time.Date(2000, time.January, 0, 0, 0, 0, 0, time.UTC)})
for _, c := range r.Cookies() {
http.SetCookie(w, c)
}

w.WriteHeader(http.StatusTeapot)
encoder := json.NewEncoder(w)
_ = encoder.Encode(struct{ RequestQueryParams, Method any }{r.URL.Query(), r.Method})
},
expectStatus: http.StatusTeapot,
expectHeaders: map[string]string{
"Hello": "world1,world2",
},
expectCookies: []string{
"yummy=cookie",
"yummy=cake",
"fruit=banana; Expires=Fri, 31 Dec 1999 00:00:00 GMT",
"foo=bar",
"hello=hello",
},
expectBody: `{"RequestQueryParams":{"foo":["bar"],"hello":["world"]},"Method":"POST"}` + "\n",
},
"mux": {
input: helloRequest,
handler: func(w http.ResponseWriter, r *http.Request) {
log.Println(r.URL)
mux := http.NewServeMux()
mux.HandleFunc("/hello", func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(200)
_, _ = w.Write([]byte("Hello World!"))
})
mux.ServeHTTP(w, r)
},
expectStatus: 200,
expectBody: "Hello World!",
},
"get-implicit-trailing-slash": {
input: domainOnlyGetRequest,
handler: func(w http.ResponseWriter, r *http.Request) {
encoder := json.NewEncoder(w)
_ = encoder.Encode(r.Method)
_ = encoder.Encode(r.URL.String())
},
expectStatus: http.StatusOK,
expectBody: "\"GET\"\n\"https://lambda-url-id.lambda-url.us-west-2.on.aws/\"\n",
},
"get-explicit-trailing-slash": {
input: domainOnlyWithSlashGetRequest,
handler: func(w http.ResponseWriter, r *http.Request) {
encoder := json.NewEncoder(w)
_ = encoder.Encode(r.Method)
_ = encoder.Encode(r.URL.String())
},
expectStatus: http.StatusOK,
expectBody: "\"GET\"\n\"https://lambda-url-id.lambda-url.us-west-2.on.aws/\"\n",
},
"empty handler": {
input: helloRequest,
handler: func(w http.ResponseWriter, r *http.Request) {},
expectStatus: http.StatusOK,
},
"base64request": {
input: base64EncodedBodyRequest,
handler: func(w http.ResponseWriter, r *http.Request) {
_, _ = io.Copy(w, r.Body)
},
expectStatus: http.StatusOK,
expectBody: "<idk/>",
},
} {
t.Run(name, func(t *testing.T) {
handler := Wrap(params.handler)
var req events.LambdaFunctionURLRequest
require.NoError(t, json.Unmarshal(params.input, &req))
res, err := handler(context.Background(), &req)
require.NoError(t, err)
resultBodyBytes, err := ioutil.ReadAll(res)
require.NoError(t, err)
resultHeaderBytes, resultBodyBytes, ok := bytes.Cut(resultBodyBytes, []byte{0, 0, 0, 0, 0, 0, 0, 0})
require.True(t, ok)
var resultHeader struct {
StatusCode int
Headers map[string]string
Cookies []string
}
require.NoError(t, json.Unmarshal(resultHeaderBytes, &resultHeader))
assert.Equal(t, params.expectBody, string(resultBodyBytes))
assert.Equal(t, params.expectStatus, resultHeader.StatusCode)
assert.Equal(t, params.expectHeaders, resultHeader.Headers)
assert.Equal(t, params.expectCookies, resultHeader.Cookies)
})
}
}

func TestRequestContext(t *testing.T) {
var req *events.LambdaFunctionURLRequest
require.NoError(t, json.Unmarshal(helloRequest, &req))
handler := Wrap(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
reqFromContext, exists := RequestFromContext(r.Context())
require.True(t, exists)
require.NotNil(t, reqFromContext)
assert.Equal(t, req, reqFromContext)
}))
_, err := handler(context.Background(), req)
require.NoError(t, err)
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
{
"headers": {
"accept": "application/xml",
"accept-encoding": "gzip, deflate",
"content-type": "application/json",
"host": "lambda-url-id.lambda-url.us-west-2.on.aws",
"user-agent": "python-requests/2.28.2",
"x-amz-content-sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
"x-amz-date": "20230418T170147Z",
"x-amz-security-token": "security-token",
"x-amzn-tls-cipher-suite": "ECDHE-RSA-AES128-GCM-SHA256",
"x-amzn-tls-version": "TLSv1.2",
"x-amzn-trace-id": "Root=1-643eccfb-7c4d3f09749a95a044db997a",
"x-forwarded-for": "127.0.0.1",
"x-forwarded-port": "443",
"x-forwarded-proto": "https"
},
"isBase64Encoded": false,
"rawPath": "/",
"rawQueryString": "",
"requestContext": {
"accountId": "aws-account-id",
"apiId": "lambda-url-id",
"authorizer": {
"iam": {}
},
"domainName": "lambda-url-id.lambda-url.us-west-2.on.aws",
"domainPrefix": "lambda-url-id",
"http": {
"method": "GET",
"path": "/",
"protocol": "HTTP/1.1",
"sourceIp": "127.0.0.1",
"userAgent": "python-requests/2.28.2"
},
"requestId": "3a72f39b-d6bd-4a4f-b040-f94d09b4daa3",
"routeKey": "$default",
"stage": "$default",
"time": "18/Apr/2023:17:01:47 +0000",
"timeEpoch": 1681837307717
},
"routeKey": "$default",
"version": "2.0"
}
44 changes: 44 additions & 0 deletions lambdaurl/testdata/function-url-domain-only-get-request.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
{
"headers": {
"accept": "application/xml",
"accept-encoding": "gzip, deflate",
"content-type": "application/json",
"host": "lambda-url-id.lambda-url.us-west-2.on.aws",
"user-agent": "python-requests/2.28.2",
"x-amz-content-sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
"x-amz-date": "20230418T170147Z",
"x-amz-security-token": "security-token",
"x-amzn-tls-cipher-suite": "ECDHE-RSA-AES128-GCM-SHA256",
"x-amzn-tls-version": "TLSv1.2",
"x-amzn-trace-id": "Root=1-643eccfb-4c9be61972302fa41111a443",
"x-forwarded-for": "127.0.0.1",
"x-forwarded-port": "443",
"x-forwarded-proto": "https"
},
"isBase64Encoded": false,
"rawPath": "/",
"rawQueryString": "",
"requestContext": {
"accountId": "aws-account-id",
"apiId": "lambda-url-id",
"authorizer": {
"iam": {}
},
"domainName": "lambda-url-id.lambda-url.us-west-2.on.aws",
"domainPrefix": "lambda-url-id",
"http": {
"method": "GET",
"path": "/",
"protocol": "HTTP/1.1",
"sourceIp": "127.0.0.1",
"userAgent": "python-requests/2.28.2"
},
"requestId": "deeb7e49-a9a8-4a8f-bcd1-5482231e2087",
"routeKey": "$default",
"stage": "$default",
"time": "18/Apr/2023:17:01:47 +0000",
"timeEpoch": 1681837307545
},
"routeKey": "$default",
"version": "2.0"
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
{
"body": "PGlkay8+",
"headers": {
"accept": "*/*",
"accept-encoding": "gzip, deflate",
"content-length": "6",
"content-type": "idk",
"host": "lambda-url-id.lambda-url.us-west-2.on.aws",
"user-agent": "python-requests/2.28.2",
"x-amz-content-sha256": "0ab2082273499eaa495f2196e32d8c794745e58a20a0c93182c59d2165432839",
"x-amz-date": "20230418T170147Z",
"x-amz-security-token": "security-token",
"x-amzn-tls-cipher-suite": "ECDHE-RSA-AES128-GCM-SHA256",
"x-amzn-tls-version": "TLSv1.2",
"x-amzn-trace-id": "Root=1-643eccfb-7fdecb844a12b4b45645132d",
"x-forwarded-for": "127.0.0.1",
"x-forwarded-port": "443",
"x-forwarded-proto": "https"
},
"isBase64Encoded": true,
"rawPath": "/",
"rawQueryString": "",
"requestContext": {
"accountId": "aws-account-id",
"apiId": "lambda-url-id",
"authorizer": {
"iam": {}
},
"domainName": "lambda-url-id.lambda-url.us-west-2.on.aws",
"domainPrefix": "lambda-url-id",
"http": {
"method": "POST",
"path": "/",
"protocol": "HTTP/1.1",
"sourceIp": "127.0.0.1",
"userAgent": "python-requests/2.28.2"
},
"requestId": "9701a3d4-36ad-40bd-bf0b-a525c987d27f",
"routeKey": "$default",
"stage": "$default",
"time": "18/Apr/2023:17:01:47 +0000",
"timeEpoch": 1681837307386
},
"routeKey": "$default",
"version": "2.0"
}