feat: allow refreshing keys once via jwks URL when current jwt kid is not found

Author: alexzhangreslv

README.md

@@ -51,6 +51,7 @@ OpaDebugMode | Set the opa response in the http response body when the request i
 PayloadFields | The field-name in the JWT payload that are required (e.g. `exp`). Multiple field names may be specified (string array)
 Required | Is `Authorization` header with JWT token required for every request.
 Keys | Used to validate JWT signature. Multiple keys are supported. Allowed values include certificates, public keys, symmetric keys. In case the value is a valid URL, the plugin will fetch keys from the JWK endpoint.
+ForceRefreshKeys | Force fetching keys from JWKS service when the key of current JWT token is not found. If set false, keys will only be refreshed every 15 minutes by default.
 Alg | Used to verify which PKI algorithm is used in the JWT.
 JwksHeaders | Map used to add headers to a JWKS request (e.g. credentials for a 3rd party JWKS service).
 JwtHeaders | Map used to inject JWT payload fields as HTTP request headers.

jwt.go

@@ -23,6 +23,7 @@ import (
 	"net/url"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 )
 
@@ -35,6 +36,7 @@ type Config struct {
 	PayloadFields      []string
 	Required           bool
 	Keys               []string
+	ForceRefreshKeys   bool
 	Alg                string
 	OpaHeaders         map[string]string
 	JwtHeaders         map[string]string
@@ -74,6 +76,9 @@ type JwtPlugin struct {
 	opaHttpStatusField string
 	jwtCookieKey       string
 	jwtQueryKey        string
+
+	keysLock        sync.RWMutex
+	forceRefreshCmd chan chan<- struct{}
 }
 
 // LogEvent contains a single log entry
@@ -189,20 +194,43 @@ func New(_ context.Context, next http.Handler, config *Config, _ string) (http.H
 			return nil, err
 		}
 		if len(jwtPlugin.jwkEndpoints) > 0 {
+			if config.ForceRefreshKeys {
+				jwtPlugin.forceRefreshCmd = make(chan chan<- struct{})
+			}
 			go jwtPlugin.BackgroundRefresh()
 		}
 	}
 	return jwtPlugin, nil
 }
 
 func (jwtPlugin *JwtPlugin) BackgroundRefresh() {
+	jwtPlugin.FetchKeys()
 	for {
-		jwtPlugin.FetchKeys()
-		time.Sleep(15 * time.Minute) // 15 min
+		select {
+		case keysFetchedChan := <-jwtPlugin.forceRefreshCmd:
+			jwtPlugin.FetchKeys()
+			keysFetchedChan <- struct{}{}
+		case <-time.After(15 * time.Minute):
+			jwtPlugin.FetchKeys()
+		}
+	}
+}
+
+func (jwtPlugin *JwtPlugin) forceRefreshKeys() (refreshed bool) {
+	if jwtPlugin.forceRefreshCmd == nil || len(jwtPlugin.jwkEndpoints) == 0 {
+		return
 	}
+	refreshedCh := make(chan struct{})
+	jwtPlugin.forceRefreshCmd <- refreshedCh
+	<-refreshedCh
+	refreshed = true
+	return
 }
 
 func (jwtPlugin *JwtPlugin) ParseKeys(certificates []string) error {
+	jwtPlugin.keysLock.Lock()
+	defer jwtPlugin.keysLock.Unlock()
+
 	for _, certificate := range certificates {
 		if block, rest := pem.Decode([]byte(certificate)); block != nil {
 			if len(rest) > 0 {
@@ -235,6 +263,7 @@ func (jwtPlugin *JwtPlugin) ParseKeys(certificates []string) error {
 func (jwtPlugin *JwtPlugin) FetchKeys() {
 	logInfo(fmt.Sprintf("FetchKeys - #%d jwkEndpoints to fetch", len(jwtPlugin.jwkEndpoints))).
 		print()
+	fetchedKeys := map[string]interface{}{}
 	for _, u := range jwtPlugin.jwkEndpoints {
 		req, err := http.NewRequest("GET", u.String(), nil)
 		if err != nil {
@@ -281,7 +310,7 @@ func (jwtPlugin *JwtPlugin) FetchKeys() {
 					ptr := new(rsa.PublicKey)
 					ptr.N = new(big.Int).SetBytes(nBytes)
 					ptr.E = int(new(big.Int).SetBytes(eBytes).Uint64())
-					jwtPlugin.keys[key.Kid] = ptr
+					fetchedKeys[key.Kid] = ptr
 				}
 			case "EC":
 				{
@@ -323,7 +352,7 @@ func (jwtPlugin *JwtPlugin) FetchKeys() {
 					ptr.Curve = crv
 					ptr.X = new(big.Int).SetBytes(xBytes)
 					ptr.Y = new(big.Int).SetBytes(yBytes)
-					jwtPlugin.keys[key.Kid] = ptr
+					fetchedKeys[key.Kid] = ptr
 				}
 			case "oct":
 				{
@@ -337,11 +366,18 @@ func (jwtPlugin *JwtPlugin) FetchKeys() {
 							break
 						}
 					}
-					jwtPlugin.keys[key.Kid] = kBytes
+					fetchedKeys[key.Kid] = kBytes
 				}
 			}
 		}
 	}
+
+	jwtPlugin.keysLock.Lock()
+	defer jwtPlugin.keysLock.Unlock()
+
+	for k, v := range fetchedKeys {
+		jwtPlugin.keys[k] = v
+	}
 }
 
 func (jwtPlugin *JwtPlugin) ServeHTTP(rw http.ResponseWriter, request *http.Request) {
@@ -373,7 +409,7 @@ func (jwtPlugin *JwtPlugin) CheckToken(request *http.Request, rw http.ResponseWr
 	if jwtToken != nil {
 		sub = fmt.Sprint(jwtToken.Payload["sub"])
 		// only verify jwt tokens if keys are configured
-		if len(jwtPlugin.keys) > 0 || len(jwtPlugin.jwkEndpoints) > 0 {
+		if len(jwtPlugin.getKeysSync()) > 0 || len(jwtPlugin.jwkEndpoints) > 0 {
 			if err = jwtPlugin.VerifyToken(jwtToken); err != nil {
 				logError(fmt.Sprintf("Token is invalid - err: %s", err.Error())).
 					withSub(sub).
@@ -550,6 +586,12 @@ func (jwtPlugin *JwtPlugin) remoteAddr(req *http.Request) Network {
 	}
 }
 
+func (jwtPlugin *JwtPlugin) getKeysSync() map[string]interface{} {
+	jwtPlugin.keysLock.RLock()
+	defer jwtPlugin.keysLock.RUnlock()
+	return jwtPlugin.keys
+}
+
 func (jwtPlugin *JwtPlugin) VerifyToken(jwtToken *JWT) error {
 	for _, h := range jwtToken.Header.Crit {
 		if _, ok := supportedHeaderNames[h]; !ok {
@@ -564,11 +606,14 @@ func (jwtPlugin *JwtPlugin) VerifyToken(jwtToken *JWT) error {
 	if jwtPlugin.alg != "" && jwtToken.Header.Alg != jwtPlugin.alg {
 		return fmt.Errorf("incorrect alg, expected %s got %s", jwtPlugin.alg, jwtToken.Header.Alg)
 	}
-	key, ok := jwtPlugin.keys[jwtToken.Header.Kid]
+	key, ok := jwtPlugin.getKeysSync()[jwtToken.Header.Kid]
+	if !ok && jwtPlugin.forceRefreshKeys() {
+		key, ok = jwtPlugin.getKeysSync()[jwtToken.Header.Kid]
+	}
 	if ok {
 		return a.verify(key, a.hash, jwtToken.Plaintext, jwtToken.Signature)
 	} else {
-		for _, key := range jwtPlugin.keys {
+		for _, key := range jwtPlugin.getKeysSync() {
 			err := a.verify(key, a.hash, jwtToken.Plaintext, jwtToken.Signature)
 			if err == nil {
 				return nil

jwt_test.go

@@ -618,6 +618,54 @@ func TestNewJWKEndpoint(t *testing.T) {
 	}
 }
 
+func TestForceRefreshKeys(t *testing.T) {
+	keys := `{"keys":[{"kty":"oct","kid":"57bd26a0-6209-4a93-a688-f8752be5d191","k":"eW91ci01MTItYml0LXNlY3JldA","alg":"HS512"}]}`
+	token := "Bearer eyJhbGciOiJIUzUxMiIsInR5cCI6IkpXVCIsImNyaXQiOlsia2lkIl0sImtpZCI6IjU3YmQyNmEwLTYyMDktNGE5My1hNjg4LWY4NzUyYmU1ZDE5MSJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWUsImlhdCI6MTUxNjIzOTAyMn0.573ixRAw4I4XUFJwJGpv5dHNOGaexX5zTtF0nOQTWuU2_JyZjD-7cuMPxQUHOv8RR0kQrS0uVdo_N1lzTCPFnA"
+	jwksCalledCounter := 0
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		defer func() { jwksCalledCounter++ }()
+		w.WriteHeader(http.StatusOK)
+		if jwksCalledCounter == 0 {
+			fmt.Fprintln(w, `{"keys":[]}`)
+			return
+		}
+		_, _ = fmt.Fprintln(w, keys)
+	}))
+	defer ts.Close()
+	cfg := Config{
+		Keys:             []string{ts.URL},
+		ForceRefreshKeys: true,
+	}
+	ctx := context.Background()
+	nextCalled := false
+	next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { nextCalled = true })
+	opa, err := New(ctx, next, &cfg, "test-traefik-jwt-plugin")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	recorder := httptest.NewRecorder()
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://localhost", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	req.Header.Add("Authorization", token)
+
+	opa.ServeHTTP(recorder, req)
+
+	resp := recorder.Result()
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("Expected status code %d, received %d", http.StatusOK, resp.StatusCode)
+	}
+	if !nextCalled {
+		t.Fatalf("next.ServeHTTP was called: %t, expected: %t", nextCalled, true)
+	}
+	if jwksCalledCounter != 2 {
+		t.Fatalf("jwks was called: %d times, expected: %d", jwksCalledCounter, 2)
+	}
+}
+
 func TestIssue3(t *testing.T) {
 	cfg := Config{
 		JwtHeaders: map[string]string{"Subject": "sub", "User": "preferred_username"},