Skip to content

Commit

Permalink
Support http.Handler for RESPONSE_STREAM Lambda Function URLs (#503)
Browse files Browse the repository at this point in the history
* 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]>
bmoffatt and aidansteele authored Apr 23, 2023

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.
1 parent dc78417 commit 5c6579e
Showing 13 changed files with 606 additions and 0 deletions.
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
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"
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
{
"body": "{\"hello\": \"world\"}",
"cookies": [
"foo=bar",
"hello=hello"
],
"headers": {
"accept": "*/*",
"accept-encoding": "gzip, deflate",
"content-length": "18",
"content-type": "application/json",
"cookie": "foo=bar; hello=hello",
"header1": "h1",
"header2": "h1,h2",
"header3": "h3",
"host": "lambda-url-id.lambda-url.us-west-2.on.aws",
"user-agent": "python-requests/2.28.2",
"x-amz-content-sha256": "5f8f04f6a3a892aaabbddb6cf273894493773960d4a325b105fee46eef4304f1",
"x-amz-date": "20230418T170146Z",
"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-643eccfa-2c6028925c2b249524664087",
"x-forwarded-for": "127.0.0.1",
"x-forwarded-port": "443",
"x-forwarded-proto": "https"
},
"isBase64Encoded": false,
"queryStringParameters": {
"foo": "bar",
"hello": "world"
},
"rawPath": "/hello",
"rawQueryString": "hello=world&foo=bar",
"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": "/hello",
"protocol": "HTTP/1.1",
"sourceIp": "127.0.0.1",
"userAgent": "python-requests/2.28.2"
},
"requestId": "5bbd0e3e-fe7a-4299-9076-32d4de45391b",
"routeKey": "$default",
"stage": "$default",
"time": "18/Apr/2023:17:01:46 +0000",
"timeEpoch": 1681837306806
},
"routeKey": "$default",
"version": "2.0"
}
54 changes: 54 additions & 0 deletions lambdaurl/testdata/gen-events.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
#!/bin/bash
set -euo pipefail

url_id="${1}" # should be the lambda function url domain prefix for an echo function
region=${AWS_REGION:-us-west-2}
url="https://${url_id}.lambda-url.${region}.on.aws"
account_id=$(aws sts get-caller-identity --output text --query "Account")

redact () {
#https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_identifiers.html#identifiers-unique-ids
sed "s/${url_id}/lambda-url-id/g" \
| sed 's/A[A-Z][A-Z]A[A-Z1-9]*\([":]\)/iam-unique-id\1/g' \
| sed "s/${account_id}/aws-account-id/g" \
| jq '.headers |= (.["x-amz-security-token"] = "security-token" )' \
| jq '.headers |= (.["x-forwarded-for"] = "127.0.0.1")' \
| jq '.requestContext.authorizer |= (.["iam"] = {})' \
| jq '.requestContext.http |= (.["sourceIp"] = "127.0.0.1")'
}

awscurl --service lambda --region $region \
-X POST \
-H 'Header1: h1' \
-H 'Header2: h1,h2' \
-H 'Header3: h3' \
-H 'Cookie: foo=bar; hello=hello' \
-H 'Content-Type: application/json' \
-d '{"hello": "world"}' \
"$url/hello?hello=world&foo=bar" \
| redact \
| tee function-url-request-with-headers-and-cookies-and-text-body.json \
| jq

awscurl --service lambda --region $region \
-X POST \
-d '<idk/>' \
"$url" \
| redact \
| tee function-url-domain-only-request-with-base64-encoded-body.json \
| jq

awscurl --service lambda --region $region \
-X GET \
"$url" \
| redact \
| tee function-url-domain-only-get-request.json \
| jq

awscurl --service lambda --region $region \
-X GET \
"$url/" \
| redact \
| tee function-url-domain-only-get-request-trailing-slash.json \
| jq

2 changes: 2 additions & 0 deletions lambdaurl/testdata/testfunc/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
.aws-sam/
samconfig.toml
11 changes: 11 additions & 0 deletions lambdaurl/testdata/testfunc/echo/main.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
package main

import (
"github.com/aws/aws-lambda-go/lambda"
)

func main() {
lambda.Start(func(req any) (any, error) {
return req, nil
})
}
7 changes: 7 additions & 0 deletions lambdaurl/testdata/testfunc/go.mod
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
module testfunc

require github.com/aws/aws-lambda-go v1.40.0

replace github.com/aws/aws-lambda-go => ../../../

go 1.20
4 changes: 4 additions & 0 deletions lambdaurl/testdata/testfunc/go.sum
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/stretchr/testify v1.7.2 h1:4jaiDzPyXQvSd7D0EjG45355tLlV3VOECpq10pLC+8s=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
39 changes: 39 additions & 0 deletions lambdaurl/testdata/testfunc/site/main.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
package main

import (
"context"
"encoding/json"
"io"
"net/http"
"os"
"strings"

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

func logLambdaRequest(ctx context.Context) {
req, ok := lambdaurl.RequestFromContext(ctx)
if ok {
enc := json.NewEncoder(os.Stdout)
enc.SetIndent("", " ")
enc.Encode(req)
}
}

func root(w http.ResponseWriter, r *http.Request) {
logLambdaRequest(r.Context())
}

func hello(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/html")
io.Copy(w, strings.NewReader(`<html><body>Hello World!</body></html>`))
}

func main() {
http.HandleFunc("/hello", hello)
http.HandleFunc("/", root)
if os.Getenv("AWS_LAMBDA_RUNTIME_API") != "" {
lambdaurl.Start(http.DefaultServeMux)
}
http.ListenAndServe(":9001", nil)
}
37 changes: 37 additions & 0 deletions lambdaurl/testdata/testfunc/template.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: >
Test handler for github.com/aws/aws-lambda-go/lambdaurl
Globals:
Function:
Timeout: 3
Runtime: provided.al2
Handler: bootstrap
Architectures: [ arm64 ]
Resources:
Site:
Type: AWS::Serverless::Function
Metadata:
BuildMethod: go1.x
Properties:
CodeUri: site
FunctionUrlConfig:
AuthType: AWS_IAM
InvokeMode: RESPONSE_STREAM
Echo:
Type: AWS::Serverless::Function
Metadata:
BuildMethod: go1.x
Properties:
CodeUri: echo
FunctionUrlConfig:
AuthType: AWS_IAM
InvokeMode: RESPONSE_STREAM

Outputs:
SiteUrl:
Description: "Site Lambda Function URL"
Value: !GetAtt SiteUrl.FunctionUrl
EchoUrl:
Description: "Echo Lambda Function URL"
Value: !GetAtt EchoUrl.FunctionUrl

0 comments on commit 5c6579e

Please sign in to comment.