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]>
aws · Apr 23, 2023 · 5c6579e · 5c6579e
1 parent dc78417
commit 5c6579e
Show file tree

Hide file tree

Showing 13 changed files with 606 additions and 0 deletions.
diff --git a/lambdaurl/http_handler.go b/lambdaurl/http_handler.go
@@ -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...)
+}
diff --git a/lambdaurl/http_handler_test.go b/lambdaurl/http_handler_test.go
@@ -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)
+}
diff --git a/lambdaurl/testdata/function-url-domain-only-get-request-trailing-slash.json b/lambdaurl/testdata/function-url-domain-only-get-request-trailing-slash.json
@@ -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"
+}
diff --git a/lambdaurl/testdata/function-url-domain-only-get-request.json b/lambdaurl/testdata/function-url-domain-only-get-request.json
@@ -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"
+}
diff --git a/lambdaurl/testdata/function-url-domain-only-request-with-base64-encoded-body.json b/lambdaurl/testdata/function-url-domain-only-request-with-base64-encoded-body.json
@@ -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"
+}