-
Notifications
You must be signed in to change notification settings - Fork 2
/
context.go
142 lines (126 loc) · 3.78 KB
/
context.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
package context
import (
"context"
"fmt"
"github.com/gin-gonic/gin"
"github.com/grpc-ecosystem/go-grpc-middleware/util/metautils"
"github.com/woaijssss/tros/constants"
"github.com/woaijssss/tros/pkg/utils/encrypt"
"google.golang.org/grpc/metadata"
"strings"
)
type TrContext struct {
context.Context
TraceId string `json:"traceId,omitempty"`
UserId int64 `json:"userId,omitempty"`
OpId int64 `json:"opId,omitempty"`
RunAs int64 `json:"runAs,omitempty"`
Roles string `json:"roles,omitempty"`
BizTypes int `json:"bizTypes,omitempty"`
GroupId int64 `json:"groupId,omitempty"`
Platform string `json:"platform,omitempty"`
UserAgent string `json:"userAgent,omitempty"`
Lang string `json:"lang,omitempty"`
GoId uint64 `json:"goId,omitempty"`
Token string `json:"token,omitempty"`
ShareToken string `json:"shareToken,omitempty"`
RemoteIp string `json:"remoteIp,omitempty"`
CompanyId int64 `json:"companyId,omitempty"`
Product int `json:"product,omitempty"`
Extra map[string]any `json:"extra,omitempty"`
}
const (
headerRealIP = "X-Real-Ip"
headerUserAgent = "User-Agent"
headerRefer = "Referer"
metadataTrace = "Tr-Trace"
)
func (ctx *TrContext) SetExtraKeyValue(key string, val any) {
if ctx.Extra == nil {
ctx.Extra = map[string]any{key: val}
} else {
ctx.Extra[key] = val
}
}
func (ctx *TrContext) GetExtraValue(key string) any {
if ctx.Extra == nil {
return nil
} else {
return ctx.Extra[key]
}
}
func GenTraceID() string {
return encrypt.EncodeMD5AsEmpty()
}
func GetContextWithTraceId() context.Context {
ctx := context.Background()
return AddTraceID(ctx, GenTraceID())
}
func GetTraceIdFromContext(ctx context.Context) string {
traceId, ok := ctx.Value(constants.TraceId).(string)
if !ok {
return ""
}
return traceId
}
func GetUserIdFromContext(ctx context.Context) string {
if c, ok := ctx.(*gin.Context); ok {
return c.GetString(constants.UserId)
}
return ctx.Value(constants.UserId).(string)
}
func InsertTraceID(ctx context.Context) context.Context {
md := metautils.ExtractIncoming(ctx)
s := md.Get(metadataTrace)
if s == "" {
s = GenTraceID()
}
return AddTraceID(ctx, s)
}
func InsertRemoteIp(ctx context.Context) context.Context {
md, ok := metadata.FromIncomingContext(ctx)
if ok {
if ipHeaders, ok := md["x-forwarded-for"]; ok && len(ipHeaders) > 0 {
/*
The IP passed in by grpc gateway may still have "127.0.0.1" or other IPs after the real IP, so a replacement is needed
*/
ips := strings.Split(ipHeaders[0], ",")
ip := ""
if len(ips) >= 1 {
ip = ips[0]
}
fmt.Println("Remote IP:", ip)
return context.WithValue(ctx, constants.RemoteIp, ip)
}
}
if c, ok := ctx.(*gin.Context); ok {
c.Set(constants.RemoteIp, c.RemoteIP())
}
return ctx
}
func AddTraceID(ctx context.Context, traceId string) context.Context {
if traceId == "" {
return ctx
}
if c, ok := ctx.(*gin.Context); ok {
c.Set(constants.TraceId, traceId)
if c.Request != nil {
newCtx := context.WithValue(c.Request.Context(), constants.TraceId, traceId)
c.Request = c.Request.WithContext(newCtx)
c.Request.Header.Set(metadataTrace, traceId)
}
return c
}
return context.WithValue(ctx, constants.TraceId, traceId)
}
func AddUserID(ctx context.Context, userId string) context.Context {
if c, ok := ctx.(*gin.Context); ok { // 兼容纯gin模式的请求
c.Set(constants.UserId, userId)
if c.Request != nil {
newCtx := context.WithValue(c.Request.Context(), constants.UserId, userId)
c.Request = c.Request.WithContext(newCtx)
}
return c
}
return context.WithValue(ctx, constants.UserId, userId) // 兼容grpc-gateway模式的请求
}