99 lines
2.3 KiB
Go
99 lines
2.3 KiB
Go
package trace
|
|
|
|
import (
|
|
"context"
|
|
"gitlab.33.cn/chat/dtalk/pkg/api"
|
|
"gitlab.33.cn/chat/dtalk/pkg/api/trace"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/metadata"
|
|
)
|
|
|
|
const (
|
|
// todo rename, oa那里也要一起改 X-Gateway-*
|
|
OAOpeId string = "X-OA-Ope-Id"
|
|
OATraceId string = "X-OA-Trace-Id"
|
|
)
|
|
|
|
// NewTraceIdFromMD 从grpc metadata获取traceId
|
|
func NewTraceIdFromMD(md metadata.MD) (string, bool) {
|
|
if md == nil {
|
|
return "", false
|
|
}
|
|
|
|
var t string
|
|
|
|
if traceId := md.Get(OATraceId); len(traceId) > 0 {
|
|
t = traceId[0]
|
|
} else {
|
|
return "", false
|
|
}
|
|
|
|
return t, true
|
|
}
|
|
|
|
// NewAddressFromMD 从grpc metadata获取操作者 ID
|
|
func NewAddressFromMD(md metadata.MD) (string, bool) {
|
|
if md == nil {
|
|
return "", false
|
|
}
|
|
|
|
var t string
|
|
|
|
if opeId := md.Get(OAOpeId); len(opeId) > 0 {
|
|
t = opeId[0]
|
|
} else {
|
|
return "", false
|
|
}
|
|
|
|
return t, true
|
|
}
|
|
|
|
// ServerUnaryInterceptor 默认令牌服务端一元拦截器
|
|
func ServerUnaryInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
|
|
return handler(wrapServerContext(ctx), req)
|
|
}
|
|
|
|
// wrapServerContext 包装服务端上下文
|
|
func wrapServerContext(ctx context.Context) context.Context {
|
|
md, ok := metadata.FromIncomingContext(ctx)
|
|
if !ok {
|
|
return ctx
|
|
}
|
|
|
|
token, ok := NewTraceIdFromMD(md)
|
|
if !ok {
|
|
// 如果没有就自己创一个 ID
|
|
ctx = context.WithValue(ctx, trace.DtalkTraceId, trace.NewTraceIdWithContext(ctx))
|
|
} else {
|
|
ctx = context.WithValue(ctx, trace.DtalkTraceId, token)
|
|
}
|
|
|
|
address, ok := NewAddressFromMD(md)
|
|
if !ok {
|
|
return ctx
|
|
}
|
|
ctx = context.WithValue(ctx, api.Address, address)
|
|
|
|
return ctx
|
|
}
|
|
|
|
// UnaryClientInterceptor 默认令牌客户端一元拦截器
|
|
func UnaryClientInterceptor(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
|
|
return invoker(wrapClientContext(ctx), method, req, reply, cc, opts...)
|
|
}
|
|
|
|
// wrapClientContext 包装客户端上下文
|
|
func wrapClientContext(ctx context.Context) context.Context {
|
|
t, ok := ctx.Value(trace.DtalkTraceId).(string)
|
|
if ok {
|
|
ctx = metadata.AppendToOutgoingContext(ctx, OATraceId, t)
|
|
}
|
|
|
|
token := api.NewAddrWithContext(ctx)
|
|
if ok {
|
|
ctx = metadata.AppendToOutgoingContext(ctx, OAOpeId, token)
|
|
}
|
|
|
|
return ctx
|
|
}
|