aws-lambda-go/lambdaurl/http_handler.go at 53f3f37eaa0fe9adc7312835ec182cede33b6d22 · aws/aws-lambda-go · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
//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 detectContentTypeContextKey struct{}

// WithDetectContentType sets the behavior of content type detection when the Content-Type header is not already provided.
// When true, the first Write call will pass the intial bytes to http.DetectContentType.
// When false, and if no Content-Type is provided, no Content-Type will be sent back to Lambda,
// and the Lambda Function URL will fallback to it's default.
//
// Note: The http.ResponseWriter passed to the handler is unbuffered.
// This may result in different Content-Type headers in the Function URL response when compared to http.ListenAndServe.
//
// Usage:
//
//	lambdaurl.Start(
//	        http.HandlerFunc(func (w http.ResponseWriter, r *http.Request) {
//	                w.Write("<!DOCTYPE html><html></html>")
//	        }),
//	        lambdaurl.WithDetectContentType(true)
//	)
func WithDetectContentType(detectContentType bool) lambda.Option {
	return lambda.WithContextValue(detectContentTypeContextKey{}, detectContentType)
}

type httpResponseWriter struct {
	detectContentType bool
	header            http.Header
	writer            io.Writer
	once              sync.Once
	ready             chan<- header
}

type header struct {
	code   int
	header http.Header
}

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

func (w *httpResponseWriter) Write(p []byte) (int, error) {
	w.writeHeader(http.StatusOK, p)
	return w.writer.Write(p)
}

func (w *httpResponseWriter) WriteHeader(statusCode int) {
	w.writeHeader(statusCode, nil)
}

func (w *httpResponseWriter) writeHeader(statusCode int, initialPayload []byte) {
	w.once.Do(func() {
		if w.detectContentType {
			if w.Header().Get("Content-Type") == "" {
				w.Header().Set("Content-Type", detectContentType(initialPayload))
			}
		}
		w.ready <- header{code: statusCode, header: w.header}
	})
}

func detectContentType(p []byte) string {
	// http.DetectContentType returns "text/plain; charset=utf-8" for nil and zero-length byte slices.
	// This is a weird behavior, since otherwise it defaults to "application/octet-stream"! So we'll do that.
	// This differs from http.ListenAndServe, which set no Content-Type when the initial Flush body is empty.
	if len(p) == 0 {
		return "application/octet-stream"
	}
	return http.DetectContentType(p)
}

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
		}
		httpRequest.RemoteAddr = request.RequestContext.HTTP.SourceIP
		hasCookie := false
		for k, v := range request.Headers {
			httpRequest.Header.Add(k, v)
			if strings.ToLower(k) == "cookie" {
				hasCookie = true
			}
		}
		if !hasCookie {
			for _, cookie := range request.Cookies {
				httpRequest.Header.Add("Cookie", cookie)
			}
		}

		ready := make(chan header) // Signals when it's OK to start returning the response body to Lambda
		r, w := io.Pipe()
		responseWriter := &httpResponseWriter{writer: w, ready: ready}
		if detectContentType, ok := ctx.Value(detectContentTypeContextKey{}).(bool); ok {
			responseWriter.detectContentType = detectContentType
		}
		go func() {
			defer close(ready)
			defer w.Close() // TODO: recover and CloseWithError the any panic value once the runtime API client supports plumbing fatal errors through the reader
			//nolint:errcheck
			defer responseWriter.Write(nil) // force default status, headers, content type detection, if none occurred during the execution of the handler
			handler.ServeHTTP(responseWriter, httpRequest)
		}()
		header := <-ready
		response := &events.LambdaFunctionURLStreamingResponse{
			Body:       r,
			StatusCode: header.code,
		}
		if len(header.header) > 0 {
			response.Headers = make(map[string]string, len(header.header))
			for k, v := range header.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...)
}