Merge pull request #137 from krakend/unknown_key_list

add a configurable unknown key list

jose.go

@@ -45,6 +45,7 @@ func NewValidator(signatureConfig *SignatureConfig, cookieEf, headerEf Extractor
 		SecretURL:           signatureConfig.SecretURL,
 		CipherKey:           signatureConfig.CipherKey,
 		KeyIdentifyStrategy: signatureConfig.KeyIdentifyStrategy,
+		UnknownKeysTTL:      signatureConfig.UnknownKeysTTL,
 	}
 
 	sp, err := SecretProvider(cfg, te)

jwk.go

@@ -37,6 +37,7 @@ type SecretProviderConfig struct {
 	SecretURL           string
 	CipherKey           []byte
 	KeyIdentifyStrategy string
+	UnknownKeysTTL      string
 }
 
 var (
@@ -200,6 +201,7 @@ func newJWKClientOptions(cfg SecretProviderConfig) (JWKClientOptions, error) {
 			},
 		},
 		KeyIdentifyStrategy: cfg.KeyIdentifyStrategy,
+		UnknownKeysTTL:      cfg.UnknownKeysTTL,
 	}, nil
 }
 

jwk_client.go

@@ -2,7 +2,10 @@ package jose
 
 import (
 	"net/http"
+	"sync"
+	"time"
 
+	"github.com/go-jose/go-jose/v3"
 	"github.com/go-jose/go-jose/v3/jwt"
 	"github.com/krakend/go-auth0/v2"
 )
@@ -57,23 +60,36 @@ func TokenIDGetterFactory(keyIdentifyStrategy string) TokenIDGetter {
 type JWKClientOptions struct {
 	auth0.JWKClientOptions
 	KeyIdentifyStrategy string
+	UnknownKeysTTL      string
 }
 
 type JWKClient struct {
 	*auth0.JWKClient
 	extractor     auth0.RequestTokenExtractor
 	tokenIDGetter TokenIDGetter
+	misses        missTracker
 }
 
 // NewJWKClientWithCache creates a new JWKClient instance from the provided options and custom extractor and keycacher.
 // Passing nil to keyCacher will create a persistent key cacher.
 // the extractor is also saved in the extended JWKClient.
 func NewJWKClientWithCache(options JWKClientOptions, extractor auth0.RequestTokenExtractor, keyCacher auth0.KeyCacher) *JWKClient {
-	return &JWKClient{
+	c := &JWKClient{
 		JWKClient:     auth0.NewJWKClientWithCache(options.JWKClientOptions, extractor, keyCacher),
 		extractor:     extractor,
 		tokenIDGetter: TokenIDGetterFactory(options.KeyIdentifyStrategy),
+		misses:        noTracker,
 	}
+
+	if ttl, err := time.ParseDuration(options.UnknownKeysTTL); err == nil && ttl >= time.Second {
+		c.misses = &memoryMissTracker{
+			keys: []unknownKey{},
+			mu:   new(sync.Mutex),
+			ttl:  ttl,
+		}
+	}
+
+	return c
 }
 
 // GetSecret implements the GetSecret method of the SecretProvider interface.
@@ -93,3 +109,88 @@ func (j *JWKClient) SecretFromToken(token *jwt.JSONWebToken) (interface{}, error
 	keyID := j.tokenIDGetter.Get(token)
 	return j.GetKey(keyID)
 }
+
+// GetKey wraps the internal key getter so it can manage the misses and avoid smashing the JWK
+// provider looking for unknown keys
+func (j *JWKClient) GetKey(keyID string) (jose.JSONWebKey, error) {
+	if j.misses.Exists(keyID) {
+		return jose.JSONWebKey{}, ErrNoKeyFound
+	}
+
+	k, err := j.JWKClient.GetKey(keyID)
+	if err != nil {
+		j.misses.Add(keyID)
+	}
+	return k, err
+}
+
+// missTracker is an interface defining the required signatures for tracking
+// keys missing from the received jwk
+type missTracker interface {
+	Exists(string) bool
+	Add(string)
+}
+
+// noopMissTracker is a missTracker that does nothing and always allows the client
+// to contact the jwk provider
+type noopMissTracker struct{}
+
+func (noopMissTracker) Exists(_ string) bool { return false }
+func (noopMissTracker) Add(_ string)         {}
+
+var noTracker = noopMissTracker{}
+
+// memoryMissTracker is a missTracker that keeps a list of missed keys in the last TTL period.
+// When the Exists method is called, it maintain the size of the list, removing all the entries
+// stored for more than the defined TTL.
+type memoryMissTracker struct {
+	keys []unknownKey
+	mu   *sync.Mutex
+	ttl  time.Duration
+}
+
+type unknownKey struct {
+	name string
+	time time.Time
+}
+
+// Exists looks for the key in the list and removes all evicted entries found before the required one. If the required is evicted,
+// it removes it and returns false, so the client can try to fetch it again.
+func (u *memoryMissTracker) Exists(key string) bool {
+	u.mu.Lock()
+	defer u.mu.Unlock()
+
+	now := time.Now()
+	cutPosition := -1
+	var found bool
+
+	for i, uk := range u.keys {
+		evicted := now.Sub(uk.time) >= u.ttl
+		if evicted {
+			cutPosition = i
+		}
+		if uk.name == key {
+			found = !evicted
+			break
+		}
+	}
+
+	if cutPosition == -1 {
+		return found
+	}
+
+	if len(u.keys) > cutPosition+1 {
+		u.keys = u.keys[cutPosition+1:]
+	} else {
+		u.keys = []unknownKey{}
+	}
+
+	return found
+}
+
+// Add appends a key and a timestamp to the end of the list of keys
+func (u *memoryMissTracker) Add(key string) {
+	u.mu.Lock()
+	u.keys = append(u.keys, unknownKey{name: key, time: time.Now()})
+	u.mu.Unlock()
+}

jwk_client_test.go

@@ -3,6 +3,7 @@ package jose
 import (
 	"net/http"
 	"net/http/httptest"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -79,3 +80,65 @@ func TestJWKClient_globalCache(t *testing.T) {
 		t.Errorf("invalid count %d", count)
 	}
 }
+
+func Test_memoryMissTracker(t *testing.T) {
+	now := time.Now()
+	uks := &memoryMissTracker{
+		mu: new(sync.Mutex),
+		keys: []unknownKey{
+			{
+				name: "key1",
+				time: now.Add(-time.Hour),
+			},
+			{
+				name: "key2",
+				time: now.Add(-2 * time.Minute),
+			},
+			{
+				name: "key3",
+				time: now.Add(-time.Second),
+			},
+			{
+				name: "key4",
+				time: now.Add(-time.Millisecond),
+			},
+		},
+		ttl: time.Minute,
+	}
+
+	if uks.Exists("key1") {
+		t.Errorf("key1 should not be present in list of misses %+v", uks)
+	}
+
+	if len(uks.keys) != 3 {
+		t.Errorf("wrong size %+v", uks)
+	}
+
+	if !uks.Exists("key3") {
+		t.Errorf("key3 should be present in list of misses %+v", uks)
+	}
+
+	if uks.Exists("key2") {
+		t.Errorf("key2 should not be present in list of misses %+v", uks)
+	}
+
+	if len(uks.keys) != 2 {
+		t.Errorf("wrong size %+v", uks)
+	}
+
+	if uks.Exists("key1") {
+		t.Errorf("key1 should not be present in list of misses %+v", uks)
+	}
+
+	if !uks.Exists("key4") {
+		t.Errorf("key4 should be present in list of misses %+v", uks)
+	}
+
+	if !uks.Exists("key3") {
+		t.Errorf("key3 should be present in list of misses %+v", uks)
+	}
+
+	if len(uks.keys) != 2 {
+		t.Errorf("wrong size %+v", uks)
+	}
+}

jws.go

@@ -44,6 +44,7 @@ type SignatureConfig struct {
 	KeyIdentifyStrategy     string     `json:"key_identify_strategy"`
 	OperationDebug          bool       `json:"operation_debug,omitempty"`
 	Leeway                  string     `json:"leeway"`
+	UnknownKeysTTL          string     `json:"failed_jwk_key_cooldown"`
 }
 
 type SignerConfig struct {