/
ctx.go
107 lines (84 loc) · 2.26 KB
/
ctx.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
package ctx
import (
"context"
"net/http"
)
type key int
const (
HeaderUUIDAPI = "X-CTX-CarbonAPI-UUID"
HeaderUUIDZipper = "X-CTX-CarbonZipper-UUID"
uuidKey key = iota
headersToPassKey
headersToLogKey
maxDataPoints
)
func ifaceToString(v interface{}) string {
if v != nil {
return v.(string)
}
return ""
}
func getCtxString(ctx context.Context, k key) string {
return ifaceToString(ctx.Value(k))
}
func getCtxInt64(ctx context.Context, k key) int64 {
v := ctx.Value(k)
if v != nil {
return v.(int64)
}
return 0
}
func getCtxMapString(ctx context.Context, k key) map[string]string {
v := ctx.Value(k)
if v != nil {
vv, ok := v.(map[string]string)
if !ok {
return map[string]string{}
}
return vv
}
return map[string]string{}
}
func GetPassHeaders(ctx context.Context) map[string]string {
return getCtxMapString(ctx, headersToPassKey)
}
func SetPassHeaders(ctx context.Context, h map[string]string) context.Context {
return context.WithValue(ctx, headersToPassKey, h)
}
func GetLogHeaders(ctx context.Context) map[string]string {
return getCtxMapString(ctx, headersToLogKey)
}
func SetLogHeaders(ctx context.Context, h map[string]string) context.Context {
return context.WithValue(ctx, headersToLogKey, h)
}
func GetUUID(ctx context.Context) string {
return getCtxString(ctx, uuidKey)
}
func SetUUID(ctx context.Context, v string) context.Context {
return context.WithValue(ctx, uuidKey, v)
}
func SetMaxDatapoints(ctx context.Context, h int64) context.Context {
return context.WithValue(ctx, maxDataPoints, h)
}
func GetMaxDatapoints(ctx context.Context) int64 {
return getCtxInt64(ctx, maxDataPoints)
}
func ParseCtx(h http.HandlerFunc, uuidKey string) http.HandlerFunc {
return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
uuid := req.Header.Get(uuidKey)
ctx := req.Context()
ctx = SetUUID(ctx, uuid)
h.ServeHTTP(rw, req.WithContext(ctx))
})
}
func MarshalPassHeaders(ctx context.Context, response *http.Request) *http.Request {
headers := GetPassHeaders(ctx)
for name, value := range headers {
response.Header.Add(name, value)
}
return response
}
func MarshalCtx(ctx context.Context, response *http.Request, uuidKey string) *http.Request {
response.Header.Add(uuidKey, GetUUID(ctx))
return response
}