Added multi-key support in jwt mode

Author: 01393707

example/server/server.go

@@ -26,7 +26,7 @@ func main() {
 	manager.MustTokenStorage(store.NewMemoryTokenStore())
 
 	// generate jwt access token
-	manager.MapAccessGenerate(generates.NewJWTAccessGenerate([]byte("00000000"), jwt.SigningMethodHS512))
+	manager.MapAccessGenerate(generates.NewJWTAccessGenerate("", []byte("00000000"), jwt.SigningMethodHS512))
 
 	clientStore := store.NewClientStore()
 	clientStore.Set("222222", &models.Client{

generates/jwt_access.go

@@ -1,100 +1,105 @@
-package generates
-
-import (
-	"context"
-	"encoding/base64"
-	"strings"
-	"time"
-
-	errs "errors"
-
-	"github.com/dgrijalva/jwt-go"
-	"gopkg.in/oauth2.v4"
-	"gopkg.in/oauth2.v4/errors"
-	"gopkg.in/oauth2.v4/utils/uuid"
-)
-
-// JWTAccessClaims jwt claims
-type JWTAccessClaims struct {
-	jwt.StandardClaims
-}
-
-// Valid claims verification
-func (a *JWTAccessClaims) Valid() error {
-	if time.Unix(a.ExpiresAt, 0).Before(time.Now()) {
-		return errors.ErrInvalidAccessToken
-	}
-	return nil
-}
-
-// NewJWTAccessGenerate create to generate the jwt access token instance
-func NewJWTAccessGenerate(key []byte, method jwt.SigningMethod) *JWTAccessGenerate {
-	return &JWTAccessGenerate{
-		SignedKey:    key,
-		SignedMethod: method,
-	}
-}
-
-// JWTAccessGenerate generate the jwt access token
-type JWTAccessGenerate struct {
-	SignedKey    []byte
-	SignedMethod jwt.SigningMethod
-}
-
-// Token based on the UUID generated token
-func (a *JWTAccessGenerate) Token(ctx context.Context, data *oauth2.GenerateBasic, isGenRefresh bool) (string, string, error) {
-	claims := &JWTAccessClaims{
-		StandardClaims: jwt.StandardClaims{
-			Audience:  data.Client.GetID(),
-			Subject:   data.UserID,
-			ExpiresAt: data.TokenInfo.GetAccessCreateAt().Add(data.TokenInfo.GetAccessExpiresIn()).Unix(),
-		},
-	}
-
-	token := jwt.NewWithClaims(a.SignedMethod, claims)
-	var key interface{}
-	if a.isEs() {
-		v, err := jwt.ParseECPrivateKeyFromPEM(a.SignedKey)
-		if err != nil {
-			return "", "", err
-		}
-		key = v
-	} else if a.isRsOrPS() {
-		v, err := jwt.ParseRSAPrivateKeyFromPEM(a.SignedKey)
-		if err != nil {
-			return "", "", err
-		}
-		key = v
-	} else if a.isHs() {
-		key = a.SignedKey
-	} else {
-		return "", "", errs.New("unsupported sign method")
-	}
-
-	access, err := token.SignedString(key)
-	if err != nil {
-		return "", "", err
-	}
-	refresh := ""
-
-	if isGenRefresh {
-		refresh = base64.URLEncoding.EncodeToString(uuid.NewSHA1(uuid.Must(uuid.NewRandom()), []byte(access)).Bytes())
-		refresh = strings.ToUpper(strings.TrimRight(refresh, "="))
-	}
-
-	return access, refresh, nil
-}
-
-func (a *JWTAccessGenerate) isEs() bool {
-	return strings.HasPrefix(a.SignedMethod.Alg(), "ES")
-}
-
-func (a *JWTAccessGenerate) isRsOrPS() bool {
-	isRs := strings.HasPrefix(a.SignedMethod.Alg(), "RS")
-	isPs := strings.HasPrefix(a.SignedMethod.Alg(), "PS")
-	return isRs || isPs
-}
-
-func (a *JWTAccessGenerate) isHs() bool {
-	return strings.HasPrefix(a.SignedMethod.Alg(), "HS")
-}
+package generates
+
+import (
+	"context"
+	"encoding/base64"
+	"strings"
+	"time"
+
+	errs "errors"
+
+	"github.com/dgrijalva/jwt-go"
+	"gopkg.in/oauth2.v4"
+	"gopkg.in/oauth2.v4/errors"
+	"gopkg.in/oauth2.v4/utils/uuid"
+)
+
+// JWTAccessClaims jwt claims
+type JWTAccessClaims struct {
+	jwt.StandardClaims
+}
+
+// Valid claims verification
+func (a *JWTAccessClaims) Valid() error {
+	if time.Unix(a.ExpiresAt, 0).Before(time.Now()) {
+		return errors.ErrInvalidAccessToken
+	}
+	return nil
+}
+
+// NewJWTAccessGenerate create to generate the jwt access token instance
+func NewJWTAccessGenerate(kid string, key []byte, method jwt.SigningMethod) *JWTAccessGenerate {
+	return &JWTAccessGenerate{
+		SignedKeyId:  kid,
+		SignedKey:    key,
+		SignedMethod: method,
+	}
+}
+
+// JWTAccessGenerate generate the jwt access token
+type JWTAccessGenerate struct {
+	SignedKeyId  string
+	SignedKey    []byte
+	SignedMethod jwt.SigningMethod
+}
+
+// Token based on the UUID generated token
+func (a *JWTAccessGenerate) Token(ctx context.Context, data *oauth2.GenerateBasic, isGenRefresh bool) (string, string, error) {
+	claims := &JWTAccessClaims{
+		StandardClaims: jwt.StandardClaims{
+			Audience:  data.Client.GetID(),
+			Subject:   data.UserID,
+			ExpiresAt: data.TokenInfo.GetAccessCreateAt().Add(data.TokenInfo.GetAccessExpiresIn()).Unix(),
+		},
+	}
+
+	token := jwt.NewWithClaims(a.SignedMethod, claims)
+	if a.SignedKeyId != "" {
+		token.Header["kid"] = a.SignedKeyId
+	}
+	var key interface{}
+	if a.isEs() {
+		v, err := jwt.ParseECPrivateKeyFromPEM(a.SignedKey)
+		if err != nil {
+			return "", "", err
+		}
+		key = v
+	} else if a.isRsOrPS() {
+		v, err := jwt.ParseRSAPrivateKeyFromPEM(a.SignedKey)
+		if err != nil {
+			return "", "", err
+		}
+		key = v
+	} else if a.isHs() {
+		key = a.SignedKey
+	} else {
+		return "", "", errs.New("unsupported sign method")
+	}
+
+	access, err := token.SignedString(key)
+	if err != nil {
+		return "", "", err
+	}
+	refresh := ""
+
+	if isGenRefresh {
+		refresh = base64.URLEncoding.EncodeToString(uuid.NewSHA1(uuid.Must(uuid.NewRandom()), []byte(access)).Bytes())
+		refresh = strings.ToUpper(strings.TrimRight(refresh, "="))
+	}
+
+	return access, refresh, nil
+}
+
+func (a *JWTAccessGenerate) isEs() bool {
+	return strings.HasPrefix(a.SignedMethod.Alg(), "ES")
+}
+
+func (a *JWTAccessGenerate) isRsOrPS() bool {
+	isRs := strings.HasPrefix(a.SignedMethod.Alg(), "RS")
+	isPs := strings.HasPrefix(a.SignedMethod.Alg(), "PS")
+	return isRs || isPs
+}
+
+func (a *JWTAccessGenerate) isHs() bool {
+	return strings.HasPrefix(a.SignedMethod.Alg(), "HS")
+}

generates/jwt_access_test.go

@@ -28,7 +28,7 @@ func TestJWTAccess(t *testing.T) {
 			},
 		}
 
-		gen := generates.NewJWTAccessGenerate([]byte("00000000"), jwt.SigningMethodHS512)
+		gen := generates.NewJWTAccessGenerate("", []byte("00000000"), jwt.SigningMethodHS512)
 		access, refresh, err := gen.Token(context.Background(), data, true)
 		So(err, ShouldBeNil)
 		So(access, ShouldNotBeEmpty)