package jwt import ( "context" "git.diulo.com/mogfee/kit/errors" "git.diulo.com/mogfee/kit/internal/token" "git.diulo.com/mogfee/kit/middleware" "git.diulo.com/mogfee/kit/transport" ) type userIdKey struct{} type authKey struct { } 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 } } type options struct { jwtKey string validate func(authKey string) error validatePermission func(validatePermission []string, key string) bool } func JWT(opts ...JwtOption) middleware.Middleware { var cfg = &options{ jwtKey: "JssLx22bjQwnyqby", validatePermission: InSlice, } for _, o := range opts { o(cfg) } return func(handler middleware.Handler) middleware.Handler { return func(ctx context.Context, a any) (any, error) { var tokenStr string if tr, ok := transport.FromServerContext(ctx); ok { tokenStr = tr.RequestHeader().Get("token") } userInfo, err := token.Parse(cfg.jwtKey, tokenStr) if err != nil { return nil, err } permission := FromAuthKeyContext(ctx) if permission != "" { if !cfg.validatePermission(userInfo.Permissions, permission) { return nil, errors.Unauthorized("TOKEN_PERMISSION_BAD", "") } } if cfg.validate != nil { if err = cfg.validate(userInfo.UniqueId); err != nil { return nil, err } } 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 *token.UserInfo) context.Context { return context.WithValue(ctx, userIdKey{}, user) } func FromUserContext(ctx context.Context) (user *token.UserInfo, ok bool) { user, ok = ctx.Value(userIdKey{}).(*token.UserInfo) return } func SetAuthKeyContext(ctx context.Context, key string) context.Context { return context.WithValue(ctx, authKey{}, key) } func FromAuthKeyContext(ctx context.Context) string { return ctx.Value(authKey{}).(string) }