82 lines
1.8 KiB
Go
82 lines
1.8 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"time"
|
||
|
||
"google.golang.org/grpc"
|
||
"google.golang.org/grpc/metadata"
|
||
"myschools.me/wyh/grpcservice/model"
|
||
)
|
||
|
||
func AuthGrpcInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
|
||
md, ok := metadata.FromIncomingContext(ctx)
|
||
if !ok {
|
||
return handler(ctx, req)
|
||
}
|
||
// md 是一个 map[string][]string 类型的
|
||
val, ok := md["token"]
|
||
if !ok {
|
||
return handler(ctx, req)
|
||
}
|
||
token, err := grpcAuth(val[0])
|
||
if err != nil {
|
||
return handler(ctx, req)
|
||
}
|
||
ctx = context.WithValue(ctx, "token", token)
|
||
return handler(ctx, req)
|
||
}
|
||
|
||
func grpcAuth(token string) (*model.Token, error) {
|
||
if token == "1111" {
|
||
return &model.Token{
|
||
Token: "1111",
|
||
ExpireIn: 7200,
|
||
}, nil
|
||
}
|
||
// 从缓存获取token中存储的结构体返回，
|
||
|
||
return nil, errors.New("token不合法")
|
||
}
|
||
|
||
func ContextWithToken(token *string, reqid *string, timeout int) (context.Context, context.CancelFunc) {
|
||
if timeout == 0 {
|
||
timeout = 1
|
||
}
|
||
requestid := ""
|
||
if reqid != nil {
|
||
requestid = *reqid
|
||
}
|
||
accesstoken := ""
|
||
if token != nil {
|
||
accesstoken = *token
|
||
}
|
||
md := metadata.MD{}
|
||
md.Set("token", accesstoken) //这里的TOKENNAME无论大小写，获取时均为小写
|
||
md.Set("requestid", requestid)
|
||
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
|
||
// ctx = context.WithValue(ctx, "requestid", requestid)
|
||
ctx = metadata.NewOutgoingContext(ctx, md)
|
||
return ctx, cancel
|
||
}
|
||
|
||
// ContextRequestID 从context中获取请求ID，任何出错时均返回""
|
||
func ContextRequestID(ctx context.Context) string {
|
||
md, ok := metadata.FromIncomingContext(ctx)
|
||
if !ok {
|
||
return ""
|
||
}
|
||
|
||
val, ok := md["requestid"]
|
||
if !ok {
|
||
return ""
|
||
}
|
||
|
||
if len(val) == 0 {
|
||
return ""
|
||
}
|
||
|
||
return val[0]
|
||
}
|