-
Notifications
You must be signed in to change notification settings - Fork 276
/
http.go
48 lines (39 loc) · 1.24 KB
/
http.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
package requestid
import "net/http"
type transport struct {
base http.RoundTripper
}
// NewRoundTripper creates a new RoundTripper which adds the request id to the outgoing headers.
func NewRoundTripper(base http.RoundTripper) http.RoundTripper {
return &transport{base: base}
}
func (t *transport) RoundTrip(req *http.Request) (res *http.Response, err error) {
requestID := FromContext(req.Context())
if requestID != "" && req.Header.Get(headerName) == "" {
req.Header.Set(headerName, requestID)
}
return t.base.RoundTrip(req)
}
type httpMiddleware struct {
next http.Handler
}
// HTTPMiddleware creates a new http middleware that populates the request id.
func HTTPMiddleware() func(next http.Handler) http.Handler {
return func(next http.Handler) http.Handler {
return httpMiddleware{next: next}
}
}
func (h httpMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
requestID := FromHTTPHeader(r.Header)
if requestID == "" {
requestID = New()
}
ctx := WithValue(r.Context(), requestID)
r = r.WithContext(ctx)
h.next.ServeHTTP(w, r)
}
// FromHTTPHeader returns the request id in the HTTP header. If no request id exists,
// an empty string is returned.
func FromHTTPHeader(hdr http.Header) string {
return hdr.Get(headerName)
}