package jwt import ( "context" "git.diulo.com/mogfee/kit/errors" "git.diulo.com/mogfee/kit/middleware" "git.diulo.com/mogfee/kit/transport" "git.diulo.com/mogfee/kit/transport/http" "strings" ) const ( permissionNoCheck = "none" ) type userIdKey struct{} type authKey struct{} type needAuthKey struct{} type GetFromToken func(ctx context.Context) string type ParseFunc func(ctx context.Context, key string, tokenStr string) (*UserInfo, bool, error) type JwtOption func(o *options) func WithJwtKey(jwtKey string) JwtOption { return func(o *options) { o.jwtKey = jwtKey } } func WithValidate(validate func(authKey string) error) JwtOption { return func(o *options) { o.validate = validate } } func WithValidatePermission(validatePermission func(permissions []string, key string) bool) JwtOption { return func(o *options) { o.validatePermission = validatePermission } } func WithParseFunc(parseFunc ParseFunc) JwtOption { return func(o *options) { o.parseFunc = parseFunc } } func WithGetFromToken(fun GetFromToken) JwtOption { return func(o *options) { o.fromToken = fun } } type options struct { jwtKey string fromKey string //cookie:token header:key //validate func(authKey string) error //validatePermission func(validatePermission []string, key string) bool parseFunc ParseFunc //fromToken GetFromToken } func JWT(opts ...JwtOption) middleware.Middleware { var cfg = &options{ jwtKey: "JssLx22bjQwnyqby", validatePermission: InSlice, parseFunc: func(ctx context.Context, key string, tokenStr string) (*UserInfo, bool, error) { userInfo, err := Parse(key, tokenStr) if err != nil { return nil, false, err } return userInfo, userInfo.UserId > 0, nil }, fromToken: func(ctx context.Context) string { var tokenStr string if tr, ok := transport.FromServerContext(ctx); ok { tokenStr = tr.RequestHeader().Get("token") } return tokenStr }, } for _, o := range opts { o(cfg) } return func(handler middleware.Handler) middleware.Handler { return func(ctx context.Context, a any) (any, error) { authKey := FromAuthKeyContext(ctx) needAuth := FromNeedAuthContext(ctx) autoAuth := FromAutoAuthContext(ctx) tokenStr := getTokenStr(ctx, cfg.fromKey) userInfo, checkOk, err := cfg.parseFunc(ctx, cfg.jwtKey, tokenStr) if err != nil { return nil, err } //需要验证 if needAuth { if !checkOk { return nil, errors.Unauthorized("TOKEN_PERMISSION_BAD", "") } //if authKey != "" { // if !cfg.validatePermission(userInfo.Permissions, authKey) { // return nil, errors.Unauthorized("TOKEN_PERMISSION_BAD", "") // } //} //if cfg.validate != nil { // if err = cfg.validate(userInfo.UniqueId); err != nil { // return nil, err // } //} } else if autoAuth { } else { } if checkOk { ctx = SetUserContext(ctx, userInfo) } return handler(ctx, a) } } } func getTokenStr(ctx context.Context, fromKey string) string { arr := strings.Split(fromKey, ":") if len(arr) != 2 { return "" } switch arr[0] { case "cookie": if tr, ok := transport.FromServerContext(ctx); ok { if tr1, ok := tr.(http.Transporter); ok { if co, err := tr1.Request().Cookie(arr[1]); err == nil { return co.Value } } } case "header": if tr, ok := transport.FromServerContext(ctx); ok { return tr.RequestHeader().Get(arr[1]) } } return "" } func InSlice(validatePermission []string, key string) bool { for _, e := range validatePermission { if e == key { return true } } return false } func SetUserContext(ctx context.Context, user *UserInfo) context.Context { return context.WithValue(ctx, userIdKey{}, user) } func FromUserContext(ctx context.Context) (user *UserInfo, ok bool) { user, ok = ctx.Value(userIdKey{}).(*UserInfo) return } func SetAuthKeyContext(ctx context.Context, key string) context.Context { return context.WithValue(ctx, authKey{}, key) } func FromAuthKeyContext(ctx context.Context) string { v := ctx.Value(authKey{}) if v == nil { return "" } return v.(string) } func SetNeedAuthContext(ctx context.Context, auth bool) context.Context { return context.WithValue(ctx, needAuthKey{}, auth) } func FromNeedAuthContext(ctx context.Context) bool { v := ctx.Value(needAuthKey{}) if v == nil { return false } return v.(bool) } type autoAuthKey struct { } func SetAutoAuthContext(ctx context.Context, autoAuth bool) context.Context { return context.WithValue(ctx, autoAuthKey{}, autoAuth) } func FromAutoAuthContext(ctx context.Context) bool { v := ctx.Value(autoAuthKey{}) if v == nil { return false } return v.(bool) }