fix: custom KeyFunc doesn't useful

hertz-contrib · Oct 3, 2023 · 3f069c4 · 3f069c4
1 parent 8a8e45c
commit 3f069c4
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 3 deletions.
diff --git a/auth_jwt.go b/auth_jwt.go
@@ -572,11 +572,20 @@ func (mw *HertzJWTMiddleware) LogoutHandler(ctx context.Context, c *app.RequestC
 func (mw *HertzJWTMiddleware) signedString(token *jwt.Token) (string, error) {
 	var tokenString string
 	var err error
-	if mw.usingPublicKeyAlgo() {
-		tokenString, err = token.SignedString(mw.privKey)
+	if mw.KeyFunc != nil {
+		key, innerErr := mw.KeyFunc(token)
+		if innerErr != nil {
+			return "", innerErr
+		}
+		tokenString, err = token.SignedString(key)
 	} else {
-		tokenString, err = token.SignedString(mw.Key)
+		if mw.usingPublicKeyAlgo() {
+			tokenString, err = token.SignedString(mw.privKey)
+		} else {
+			tokenString, err = token.SignedString(mw.Key)
+		}
 	}
+
 	return tokenString, err
 }
 

diff --git a/auth_jwt_test.go b/auth_jwt_test.go
@@ -1113,3 +1113,30 @@ func TestLogout(t *testing.T) {
 	assert.DeepEqual(t, http.StatusOK, w.Code)
 	assert.DeepEqual(t, fmt.Sprintf("%s=; domain=%s; path=/", cookieName, cookieDomain), w.Header().Get("Set-Cookie"))
 }
+
+func TestSignedStringWithKeyFunc(t *testing.T) {
+	var used string
+	mwWithoutKeyFunc, _ := New(&HertzJWTMiddleware{
+		Realm:         "test zone",
+		Key:           key,
+		Timeout:       time.Hour,
+		Authenticator: defaultAuthenticator,
+	})
+	mw, _ := New(&HertzJWTMiddleware{
+		Realm: "test zone",
+		Key:   key,
+		KeyFunc: func(*jwt.Token) (interface{}, error) {
+			used = "keyFunc"
+			return []byte("secret"), nil
+		},
+		Timeout:       time.Hour,
+		Authenticator: defaultAuthenticator,
+	})
+	token := jwt.New(jwt.SigningMethodHS256)
+
+	// 调用要测试的函数
+	tokenString, _ := mw.signedString(token)
+	tokenStringWithoutKeyFunc, _ := mwWithoutKeyFunc.signedString(token)
+	assert.NotEqual(t, tokenString, tokenStringWithoutKeyFunc)
+	assert.DeepEqual(t, "keyFunc", used)
+}