From e2671fe96385cac1e87b55c7572bb2fcf4ba00a3 Mon Sep 17 00:00:00 2001
From: Das Jott <das.jott@gmail.com>
Date: Fri, 17 May 2019 16:45:49 +0200
Subject: [PATCH] Provide possibility to use key ids (#1289)

* provide possibility to use key ids
* kid tests
---
 middleware/jwt.go      | 19 +++++++--
 middleware/jwt_test.go | 90 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 106 insertions(+), 3 deletions(-)

diff --git a/middleware/jwt.go b/middleware/jwt.go
index 861d3142..279cb6f2 100644
--- a/middleware/jwt.go
+++ b/middleware/jwt.go
@@ -26,10 +26,14 @@ type (
 		// It may be used to define a custom JWT error.
 		ErrorHandler JWTErrorHandler
 
-		// Signing key to validate token.
-		// Required.
+		// Signing key to validate token. Used as fallback if SigningKeys has length 0.
+		// Required. This or SigningKeys.
 		SigningKey interface{}
 
+		// Map of signing keys to validate token with kid field usage.
+		// Required. This or SigningKey.
+		SigningKeys map[string]interface{}
+
 		// Signing method, used to check token signing method.
 		// Optional. Default value HS256.
 		SigningMethod string
@@ -110,7 +114,7 @@ func JWTWithConfig(config JWTConfig) echo.MiddlewareFunc {
 	if config.Skipper == nil {
 		config.Skipper = DefaultJWTConfig.Skipper
 	}
-	if config.SigningKey == nil {
+	if config.SigningKey == nil && len(config.SigningKeys) == 0 {
 		panic("echo: jwt middleware requires signing key")
 	}
 	if config.SigningMethod == "" {
@@ -133,6 +137,15 @@ func JWTWithConfig(config JWTConfig) echo.MiddlewareFunc {
 		if t.Method.Alg() != config.SigningMethod {
 			return nil, fmt.Errorf("unexpected jwt signing method=%v", t.Header["alg"])
 		}
+		if len(config.SigningKeys) > 0 {
+			if kid, ok := t.Header["kid"].(string); ok {
+				if key, ok := config.SigningKeys[kid]; ok {
+					return key, nil
+				}
+			}
+			return nil, fmt.Errorf("unexpected jwt key id=%v", t.Header["kid"])
+		}
+
 		return config.SigningKey, nil
 	}
 
diff --git a/middleware/jwt_test.go b/middleware/jwt_test.go
index fc90073a..8165c2d6 100644
--- a/middleware/jwt_test.go
+++ b/middleware/jwt_test.go
@@ -224,3 +224,93 @@ func TestJWT(t *testing.T) {
 		}
 	}
 }
+
+func TestJWTwithKID(t *testing.T) {
+	test := assert.New(t)
+
+	e := echo.New()
+	handler := func(c echo.Context) error {
+		return c.String(http.StatusOK, "test")
+	}
+	firstToken := "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiIsImtpZCI6ImZpcnN0T25lIn0.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9.w5VGpHOe0jlNgf7jMVLHzIYH_XULmpUlreJnilwSkWk"
+	secondToken := "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiIsImtpZCI6InNlY29uZE9uZSJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9.sdghDYQ85jdh0hgQ6bKbMguLI_NSPYWjkhVJkee-yZM"
+	wrongToken := "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiIsImtpZCI6InNlY29uZE9uZSJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9.RyhLybtVLpoewF6nz9YN79oXo32kAtgUxp8FNwTkb90"
+	staticToken := "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9.1_-XFYUPpJfgsaGwYhgZEt7hfySMg-a3GN-nfZmbW7o"
+	validKeys := map[string]interface{}{"firstOne": []byte("first_secret"), "secondOne": []byte("second_secret")}
+	invalidKeys := map[string]interface{}{"thirdOne": []byte("third_secret")}
+	staticSecret := []byte("static_secret")
+	invalidStaticSecret := []byte("invalid_secret")
+
+	for _, tc := range []struct {
+		expErrCode int // 0 for Success
+		config     JWTConfig
+		hdrAuth    string
+		info       string
+	}{
+		{
+			hdrAuth: DefaultJWTConfig.AuthScheme + " " + firstToken,
+			config:  JWTConfig{SigningKeys: validKeys},
+			info:    "First token valid",
+		},
+		{
+			hdrAuth: DefaultJWTConfig.AuthScheme + " " + secondToken,
+			config:  JWTConfig{SigningKeys: validKeys},
+			info:    "Second token valid",
+		},
+		{
+			expErrCode: http.StatusUnauthorized,
+			hdrAuth:    DefaultJWTConfig.AuthScheme + " " + wrongToken,
+			config:     JWTConfig{SigningKeys: validKeys},
+			info:       "Wrong key id token",
+		},
+		{
+			hdrAuth: DefaultJWTConfig.AuthScheme + " " + staticToken,
+			config:  JWTConfig{SigningKey: staticSecret},
+			info:    "Valid static secret token",
+		},
+		{
+			expErrCode: http.StatusUnauthorized,
+			hdrAuth:    DefaultJWTConfig.AuthScheme + " " + staticToken,
+			config:     JWTConfig{SigningKey: invalidStaticSecret},
+			info:       "Invalid static secret",
+		},
+		{
+			expErrCode: http.StatusUnauthorized,
+			hdrAuth:    DefaultJWTConfig.AuthScheme + " " + firstToken,
+			config:     JWTConfig{SigningKeys: invalidKeys},
+			info:       "Invalid keys first token",
+		},
+		{
+			expErrCode: http.StatusUnauthorized,
+			hdrAuth:    DefaultJWTConfig.AuthScheme + " " + secondToken,
+			config:     JWTConfig{SigningKeys: invalidKeys},
+			info:       "Invalid keys second token",
+		},
+	} {
+		req := httptest.NewRequest(http.MethodGet, "/", nil)
+		res := httptest.NewRecorder()
+		req.Header.Set(echo.HeaderAuthorization, tc.hdrAuth)
+		c := e.NewContext(req, res)
+
+		if tc.expErrCode != 0 {
+			h := JWTWithConfig(tc.config)(handler)
+			he := h(c).(*echo.HTTPError)
+			test.Equal(tc.expErrCode, he.Code, tc.info)
+			continue
+		}
+
+		h := JWTWithConfig(tc.config)(handler)
+		if test.NoError(h(c), tc.info) {
+			user := c.Get("user").(*jwt.Token)
+			switch claims := user.Claims.(type) {
+			case jwt.MapClaims:
+				test.Equal(claims["name"], "John Doe", tc.info)
+			case *jwtCustomClaims:
+				test.Equal(claims.Name, "John Doe", tc.info)
+				test.Equal(claims.Admin, true, tc.info)
+			default:
+				panic("unexpected type of claims")
+			}
+		}
+	}
+}