|
1
|
package reqid |
|
2
|
|
|
3
|
import ( |
|
4
|
"context" |
|
5
|
"crypto/rand" |
|
6
|
"encoding/hex" |
|
7
|
"net/http" |
|
8
|
) |
|
9
|
|
|
10
|
type requestIDKey string |
|
11
|
|
|
12
|
const ( |
|
13
|
RequestID requestIDKey = "request-id" |
|
14
|
|
|
15
|
Header = "X-Request-ID" |
|
16
|
) |
|
17
|
|
|
18
|
// Middleware http middleware that sets random generated request id to each |
|
19
|
// request it get fed. |
|
20
|
func Middleware(next http.Handler) http.Handler { |
|
21
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
|
22
|
rid := r.Header.Get(Header) |
|
23
|
if rid == "" { |
|
24
|
rid = generateRequestID() |
|
25
|
r.Header.Set(Header, rid) |
|
26
|
} |
|
27
|
|
|
28
|
ctx := SetContext(r.Context(), rid) |
|
29
|
r = r.WithContext(ctx) |
|
30
|
|
|
31
|
w.Header().Add(Header, rid) |
|
32
|
next.ServeHTTP(w, r) |
|
33
|
}) |
|
34
|
} |
|
35
|
|
|
36
|
// Get returns the request ID of http request (looks up from headers) |
|
37
|
func Get(r *http.Request) string { |
|
38
|
return r.Header.Get(Header) |
|
39
|
} |
|
40
|
|
|
41
|
// GetContext returns the request ID from context. |
|
42
|
func GetContext(ctx context.Context) string { |
|
43
|
rid, ok := ctx.Value(RequestID).(string) |
|
44
|
if !ok { |
|
45
|
return "" |
|
46
|
} |
|
47
|
return rid |
|
48
|
} |
|
49
|
|
|
50
|
// SetContext gets a parent context and returns a child context with the set |
|
51
|
// provided request ID |
|
52
|
func SetContext(ctx context.Context, reqID string) context.Context { |
|
53
|
return context.WithValue(ctx, RequestID, reqID) |
|
54
|
} |
|
55
|
|
|
56
|
func generateRequestID() string { |
|
57
|
b := make([]byte, 13) |
|
58
|
_, err := rand.Read(b) |
|
59
|
if err != nil { |
|
60
|
return "unknown" |
|
61
|
} |
|
62
|
return hex.EncodeToString(b) |
|
63
|
} |