-
Notifications
You must be signed in to change notification settings - Fork 0
/
alb.go
151 lines (144 loc) · 4.08 KB
/
alb.go
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
// Package alb provides adapter enabling usage of http.Handler inside AWS Lambda
// running behind AWS ALB as described here:
// https://docs.aws.amazon.com/lambda/latest/dg/services-alb.html
//
// Usage example:
//
// package main
//
// import (
// "fmt"
// "net/http"
//
// "github.com/artyom/alb"
// "github.com/aws/aws-lambda-go/lambda"
// )
//
// func main() { lambda.Start(alb.Handler(http.HandlerFunc(hello))) }
//
// func hello(w http.ResponseWriter, r *http.Request) {
// fmt.Fprintln(w, "Hello from AWS Lambda behind ALB")
// }
//
// Note: since both request and reply to/from AWS Lambda are passed as
// json-encoded payloads, their sizes are limited. AWS documentation states
// that: "The maximum size of the request body that you can send to a Lambda
// function is 1 MB. [...] The maximum size of the response JSON that the Lambda
// function can send is 1 MB." Exact limit of response size also depends on
// whether its body is valid utf8 or not, as non-utf8 payloads are transparently
// base64-encoded, which adds some overhead.
//
// For further details see
// https://docs.aws.amazon.com/elasticloadbalancing/latest/application/lambda-functions.html
package alb
import (
"bytes"
"context"
"encoding/base64"
"io/ioutil"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"unicode/utf8"
)
// Handler returns function suitable to use as an AWS Lambda handler with
// github.com/aws/aws-lambda-go/lambda package.
//
// Note that request is fully cached in memory.
func Handler(h http.Handler) func(context.Context, request) (*response, error) {
if h == nil {
panic("Wrap called with nil handler")
}
hh := lambdaHandler{handler: h}
return hh.Run
}
type request struct {
Method string `json:"httpMethod"`
Path string `json:"path"` // escaped
Query map[string]string `json:"queryStringParameters"` // escaped
Headers map[string]string `json:"headers"`
Body string `json:"body"`
BodyEncoded bool `json:"isBase64Encoded"`
}
type response struct {
StatusCode int `json:"statusCode"`
Status string `json:"statusDescription"`
Headers map[string]string `json:"headers"`
Body string `json:"body"`
BodyEncoded bool `json:"isBase64Encoded"`
}
type lambdaHandler struct {
handler http.Handler
}
func (h *lambdaHandler) Run(ctx context.Context, req request) (*response, error) {
u, err := buildURL(req.Path, req.Query)
if err != nil {
return nil, err
}
headers := make(http.Header, len(req.Headers))
for k, v := range req.Headers {
headers.Set(k, v)
}
r := &http.Request{
ProtoMajor: 1,
ProtoMinor: 1,
Proto: "HTTP/1.1",
Method: req.Method,
URL: u,
Header: headers,
Host: headers.Get("Host"),
}
r = r.WithContext(ctx)
switch {
case req.BodyEncoded:
b, err := base64.StdEncoding.DecodeString(req.Body)
if err != nil {
return nil, err
}
r.Body = ioutil.NopCloser(bytes.NewReader(b))
r.ContentLength = int64(len(b))
default:
r.Body = ioutil.NopCloser(strings.NewReader(req.Body))
r.ContentLength = int64(len(req.Body))
}
recorder := httptest.NewRecorder()
h.handler.ServeHTTP(recorder, r)
res := recorder.Result()
out := &response{
StatusCode: res.StatusCode,
Status: res.Status,
Headers: make(map[string]string, len(res.Header)),
}
for k, vv := range res.Header {
out.Headers[k] = strings.Join(vv, ",")
}
if b := recorder.Body.Bytes(); utf8.Valid(b) {
out.Body = recorder.Body.String()
} else {
out.Body = base64.StdEncoding.EncodeToString(b)
out.BodyEncoded = true
}
return out, nil
}
// buildURL constructs url from already escaped path and query string parameters
// minimizing allocations and escaping overhead.
func buildURL(path string, query map[string]string) (*url.URL, error) {
if len(query) == 0 {
return url.Parse(path)
}
var b strings.Builder
b.WriteString(path)
b.WriteByte('?')
var i int
for k, v := range query {
if i != 0 {
b.WriteByte('&')
}
b.WriteString(k)
b.WriteByte('=')
b.WriteString(v)
i++
}
return url.Parse(b.String())
}