kit/middleware/jwt/jwt.go

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)
}