x

middleware/jwt/jwt.go

@@ -15,6 +15,8 @@ const (
 type userIdKey struct{}
 type authKey struct {
 }
+
+type ParseFunc func(key string, tokenStr string) (*token.UserInfo, error)
 type JwtOption func(o *options)
 
 func WithJwtKey(jwtKey string) JwtOption {
@@ -35,16 +37,24 @@ func WithValidatePermission(validatePermission func(permissions []string, key st
 	}
 }
 
+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:          token.Parse,
 	}
 	for _, o := range opts {
 		o(cfg)
@@ -62,7 +72,7 @@ func JWT(opts ...JwtOption) middleware.Middleware {
 			if tr, ok := transport.FromServerContext(ctx); ok {
 				tokenStr = tr.RequestHeader().Get("token")
 			}
-			userInfo, err := token.Parse(cfg.jwtKey, tokenStr)
+			userInfo, err := cfg.parseFunc(cfg.jwtKey, tokenStr)
 			if err != nil {
 				return nil, err
 			}