kit/middleware/jwt/jwt.go

package jwt

import (
	"context"
	"git.diulo.com/mogfee/kit/errors"
	"git.diulo.com/mogfee/kit/middleware"
)

type JwtOption func(o *options)

func WithJwtKey(jwtKey string) JwtOption {
	return func(o *options) {
		o.jwtKey = jwtKey
	}
}

func WithFromKey(fromKey string) JwtOption {
	return func(o *options) {
		o.fromKey = fromKey
	}
}
func WithValidate(val JwtValidate) JwtOption {
	return func(o *options) {
		o.validate = val
	}
}

type options struct {
	jwtKey   string
	fromKey  string //cookie:token header:key
	validate JwtValidate
}
type JwtValidate interface {
	//GetToken 获取token
	GetToken(ctx context.Context, key string) (tokenStr string)
	//ParseToken 解析token获取用户信息
	ParseToken(ctx context.Context, key string, token string) (*UserInfo, error)
	//Validate 校验权限
	Validate(ctx context.Context, permission string, permissions []string) error
}

func JWT(opts ...JwtOption) middleware.Middleware {
	var cfg = &options{
		jwtKey:   "JssLx22bjQwnyqby",
		fromKey:  "header:token",
		validate: &JwtDefault{},
	}
	for _, o := range opts {
		o(cfg)
	}

	return func(handler middleware.Handler) middleware.Handler {
		return func(ctx context.Context, a any) (any, error) {
			//1. 解析token
			//2. 获取用户信息
			//3. 校验权限
			//4. 设置ctx
			authKey := FromAuthKeyContext(ctx)
			needAuth := FromNeedAuthContext(ctx)

			// 解析token
			tokenStr := cfg.validate.GetToken(ctx, cfg.fromKey)
			if tokenStr == "" && needAuth {
				return nil, errors.Unauthorized("NO_TOKEN", "")
			}

			if tokenStr != "" {
				if err := func() error {
					userInfo, err := cfg.validate.ParseToken(ctx, cfg.jwtKey, tokenStr)
					if err != nil {
						return err
					}
					if needAuth && userInfo.UserId == "" {
						return errors.Unauthorized("TOKEN_BAD", "")
					}
					if authKey != "" {
						if err = cfg.validate.Validate(ctx, authKey, userInfo.Permissions); err != nil {
							return err
						}
					}
					if userInfo.UserId != "" {
						ctx = SetUserContext(ctx, userInfo)
					}
					return nil
				}(); err != nil {
					if needAuth {
						return nil, err
					}
				}
			}
			return handler(ctx, a)
		}
	}
}