You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

140 lines
3.3 KiB

package jwt
import (
"context"
"git.diulo.com/mogfee/kit/errors"
"git.diulo.com/mogfee/kit/middleware"
"git.diulo.com/mogfee/kit/transport"
)
const (
permissionNoCheck = "none"
)
type userIdKey struct{}
type authKey struct{}
type needAuthKey struct{}
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
}
}
type options struct {
jwtKey string
validate func(authKey string) error
validatePermission func(validatePermission []string, key string) bool
parseFunc ParseFunc
}
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
},
}
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)
var tokenStr string
if tr, ok := transport.FromServerContext(ctx); ok {
tokenStr = tr.RequestHeader().Get("token")
}
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
}
}
}
if checkOk {
ctx = SetUserContext(ctx, userInfo)
}
return handler(ctx, a)
}
}
}
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)
}