fix: IdentityHandler arg (#159)

fix #158

Author: appleboy

README.md

@@ -43,11 +43,15 @@ type login struct {
 	Password string `form:"password" json:"password" binding:"required"`
 }
 
+var identityKey = "id"
+
 func helloHandler(c *gin.Context) {
 	claims := jwt.ExtractClaims(c)
+	user, _ := c.Get(identityKey)
 	c.JSON(200, gin.H{
-		"userID": claims["id"],
-		"text":   "Hello World.",
+		"userID":   claims["id"],
+		"userName": user.(*User).UserName,
+		"text":     "Hello World.",
 	})
 }
 
@@ -74,15 +78,21 @@ func main() {
 		Key:         []byte("secret key"),
 		Timeout:     time.Hour,
 		MaxRefresh:  time.Hour,
-		IdentityKey: "id",
+		IdentityKey: identityKey,
 		PayloadFunc: func(data interface{}) jwt.MapClaims {
 			if v, ok := data.(*User); ok {
 				return jwt.MapClaims{
-					"id": v.UserName,
+					identityKey: v.UserName,
 				}
 			}
 			return jwt.MapClaims{}
 		},
+		IdentityHandler: func(c *gin.Context) interface{} {
+			claims := jwt.ExtractClaims(c)
+			return &User{
+				UserName: claims["id"].(string),
+			}
+		},
 		Authenticator: func(c *gin.Context) (interface{}, error) {
 			var loginVals login
 			if err := c.ShouldBind(&loginVals); err != nil {
@@ -102,7 +112,7 @@ func main() {
 			return nil, jwt.ErrFailedAuthentication
 		},
 		Authorizator: func(data interface{}, c *gin.Context) bool {
-			if v, ok := data.(string); ok && v == "admin" {
+			if v, ok := data.(*User); ok && v.UserName == "admin" {
 				return true
 			}
 

auth_jwt.go

@@ -69,7 +69,7 @@ type GinJWTMiddleware struct {
 	RefreshResponse func(*gin.Context, int, string, time.Time)
 
 	// Set the identity handler function
-	IdentityHandler func(jwt.MapClaims) interface{}
+	IdentityHandler func(*gin.Context) interface{}
 
 	// Set the identity key
 	IdentityKey string
@@ -166,6 +166,9 @@ var (
 
 	// ErrInvalidPubKey indicates the the given public key is invalid
 	ErrInvalidPubKey = errors.New("public key invalid")
+
+	// IdentityKey default identity key
+	IdentityKey = "identity"
 )
 
 // New for check error with GinJWTMiddleware
@@ -283,11 +286,12 @@ func (mw *GinJWTMiddleware) MiddlewareInit() error {
 	}
 
 	if mw.IdentityKey == "" {
-		mw.IdentityKey = "identity"
+		mw.IdentityKey = IdentityKey
 	}
 
 	if mw.IdentityHandler == nil {
-		mw.IdentityHandler = func(claims jwt.MapClaims) interface{} {
+		mw.IdentityHandler = func(c *gin.Context) interface{} {
+			claims := ExtractClaims(c)
 			return claims[mw.IdentityKey]
 		}
 	}
@@ -334,9 +338,9 @@ func (mw *GinJWTMiddleware) middlewareImpl(c *gin.Context) {
 	}
 
 	claims := token.Claims.(jwt.MapClaims)
-
-	identity := mw.IdentityHandler(claims)
 	c.Set("JWT_PAYLOAD", claims)
+	identity := mw.IdentityHandler(c)
+
 	if identity != nil {
 		c.Set(mw.IdentityKey, identity)
 	}

example/basic/server.go

@@ -15,11 +15,15 @@ type login struct {
 	Password string `form:"password" json:"password" binding:"required"`
 }
 
+var identityKey = "id"
+
 func helloHandler(c *gin.Context) {
 	claims := jwt.ExtractClaims(c)
+	user, _ := c.Get(identityKey)
 	c.JSON(200, gin.H{
-		"userID": claims["id"],
-		"text":   "Hello World.",
+		"userID":   claims["id"],
+		"userName": user.(*User).UserName,
+		"text":     "Hello World.",
 	})
 }
 
@@ -46,15 +50,21 @@ func main() {
 		Key:         []byte("secret key"),
 		Timeout:     time.Hour,
 		MaxRefresh:  time.Hour,
-		IdentityKey: "id",
+		IdentityKey: identityKey,
 		PayloadFunc: func(data interface{}) jwt.MapClaims {
 			if v, ok := data.(*User); ok {
 				return jwt.MapClaims{
-					"id": v.UserName,
+					identityKey: v.UserName,
 				}
 			}
 			return jwt.MapClaims{}
 		},
+		IdentityHandler: func(c *gin.Context) interface{} {
+			claims := jwt.ExtractClaims(c)
+			return &User{
+				UserName: claims["id"].(string),
+			}
+		},
 		Authenticator: func(c *gin.Context) (interface{}, error) {
 			var loginVals login
 			if err := c.ShouldBind(&loginVals); err != nil {
@@ -74,7 +84,7 @@ func main() {
 			return nil, jwt.ErrFailedAuthentication
 		},
 		Authorizator: func(data interface{}, c *gin.Context) bool {
-			if v, ok := data.(string); ok && v == "admin" {
+			if v, ok := data.(*User); ok && v.UserName == "admin" {
 				return true
 			}