From 9ad8b647c2c4024c56ba7805af96400e6cad854a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=8E=E4=BC=9F=E4=B9=90?= <mogfee@163.com>
Date: Thu, 16 Mar 2023 10:51:14 +0800
Subject: [PATCH] x

---
 middleware/jwt/jwt.go | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/middleware/jwt/jwt.go b/middleware/jwt/jwt.go
index a01b013..12d5a63 100644
--- a/middleware/jwt/jwt.go
+++ b/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
 			}