/
middleware.go
130 lines (115 loc) · 3.78 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
package server
import (
"bytes"
"encoding/json"
"net/http"
"strings"
)
// 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)
}
})
}
// 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")
if r.Method == http.MethodOptions {
w.WriteHeader(http.StatusOK)
return
}
}
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)
}