forked from nytimes/gizmo
-
Notifications
You must be signed in to change notification settings - Fork 0
/
middleware.go
191 lines (168 loc) · 5.49 KB
/
middleware.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
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
package server
import (
"bytes"
"encoding/json"
"net/http"
"strings"
"golang.org/x/net/context"
)
// JSONContextToHTTP is a middleware func to convert a ContextHandler an http.Handler.
func JSONContextToHTTP(ep JSONContextEndpoint) ContextHandler {
return ContextHandlerFunc(func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
if r.Body != nil {
defer func() {
if err := r.Body.Close(); err != nil {
}
}()
}
// it's JSON, so always set that content type
w.Header().Set("Content-Type", jsonContentType)
// prepare to grab the response from the ep
var b bytes.Buffer
encoder := json.NewEncoder(&b)
// call the func and return err or not
code, res, err := ep(ctx, r)
w.WriteHeader(code)
if err != nil {
res = err
}
err = encoder.Encode(res)
if err != nil {
LogWithFields(r).Error("unable to JSON encode response: ", err)
}
if _, err := w.Write(b.Bytes()); err != nil {
LogWithFields(r).Warn("unable to write response: ", err)
}
})
}
// JSONToHTTP is the middleware func to convert a JSONEndpoint to
// an http.HandlerFunc.
func JSONToHTTP(ep JSONEndpoint) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Body != nil {
defer func() {
if err := r.Body.Close(); err != nil {
Log.Warn("unable to close request body: ", err)
}
}()
}
// it's JSON, so always set that content type
w.Header().Set("Content-Type", jsonContentType)
// prepare to grab the response from the ep
var b bytes.Buffer
encoder := json.NewEncoder(&b)
// call the func and return err or not
code, res, err := ep(r)
w.WriteHeader(code)
if err != nil {
res = err
}
err = encoder.Encode(res)
if err != nil {
LogWithFields(r).Error("unable to JSON encode response: ", err)
}
if _, err := w.Write(b.Bytes()); err != nil {
LogWithFields(r).Warn("unable to write response: ", err)
}
})
}
// ContextToHTTP is a middleware func to convert a ContextHandler an http.Handler.
func ContextToHTTP(ctx context.Context, ep ContextHandler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ep.ServeHTTPContext(ctx, w, r)
})
}
// WithCloseHandler returns a Handler cancelling the context when the client
// connection close unexpectedly.
func WithCloseHandler(h ContextHandler) ContextHandler {
return ContextHandlerFunc(func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
// Cancel the context if the client closes the connection
if wcn, ok := w.(http.CloseNotifier); ok {
var cancel context.CancelFunc
ctx, cancel = context.WithCancel(ctx)
defer cancel()
notify := wcn.CloseNotify()
go func() {
select {
case <-notify:
cancel()
case <-ctx.Done():
}
}()
}
h.ServeHTTPContext(ctx, w, r)
})
}
// CORSHandler is a middleware func for setting all headers that enable CORS.
// If an originSuffix is provided, a strings.HasSuffix check will be performed
// before adding any CORS header. If an empty string is provided, any Origin
// header found will be placed into the CORS header. If no Origin header is
// found, no headers will be added.
func CORSHandler(f http.Handler, originSuffix string) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
origin := r.Header.Get("Origin")
if origin != "" &&
(originSuffix == "" || strings.HasSuffix(origin, originSuffix)) {
w.Header().Set("Access-Control-Allow-Origin", origin)
w.Header().Set("Access-Control-Allow-Credentials", "true")
w.Header().Set("Access-Control-Allow-Headers", "Content-Type, x-requested-by, *")
w.Header().Set("Access-Control-Allow-Methods", "GET, PUT, POST, DELETE, OPTIONS")
}
f.ServeHTTP(w, r)
})
}
// NoCacheHandler is a middleware func for setting the Cache-Control to no-cache.
func NoCacheHandler(f http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
w.Header().Set("Pragma", "no-cache")
w.Header().Set("Expires", "0")
f.ServeHTTP(w, r)
})
}
// JSONPHandler is a middleware func for wrapping response body with JSONP.
func JSONPHandler(f http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// using a custom ResponseWriter so we can
// capture the response of the main request,
// wrap our JSONP stuff around it
// and only write to the actual response once
jw := &jsonpResponseWriter{w: w}
f.ServeHTTP(jw, r)
// add the JSONP only if the callback exists
callbackLabel := r.FormValue("callback")
if callbackLabel != "" {
var result []byte
result = append(jsonpStart, []byte(callbackLabel)...)
result = append(result, jsonpSecond...)
result = append(result, jw.buf.Bytes()...)
result = append(result, jsonpEnd...)
if _, err := w.Write(result); err != nil {
LogWithFields(r).Warn("unable to write JSONP response: ", err)
}
} else {
// if no callback, just write the bytes
if _, err := w.Write(jw.buf.Bytes()); err != nil {
LogWithFields(r).Warn("unable to write response: ", err)
}
}
})
}
var (
jsonpStart = []byte("/**/")
jsonpSecond = []byte("(")
jsonpEnd = []byte(");")
)
type jsonpResponseWriter struct {
w http.ResponseWriter
buf bytes.Buffer
}
func (w *jsonpResponseWriter) Header() http.Header {
return w.w.Header()
}
func (w *jsonpResponseWriter) WriteHeader(h int) {
w.w.WriteHeader(h)
}
func (w *jsonpResponseWriter) Write(b []byte) (int, error) {
return w.buf.Write(b)
}