added basic provisioner refresh on jwks cache miss

authority/provisioner/keystore.go

@@ -2,6 +2,8 @@ package provisioner
 
 import (
 	"encoding/json"
+	"io"
+	"log"
 	"math/rand"
 	"regexp"
 	"strconv"
@@ -10,22 +12,28 @@ import (
 
 	"github.com/pkg/errors"
 	"go.step.sm/crypto/jose"
+	"golang.org/x/sync/singleflight"
 )
 
 const (
 	defaultCacheAge    = 12 * time.Hour
 	defaultCacheJitter = 1 * time.Hour
+	// minReloadInterval is the minimum time between cache refreshes on cache
+	// misses.
+	minReloadInterval = 1 * time.Minute
 )
 
 var maxAgeRegex = regexp.MustCompile(`max-age=(\d+)`)
 
 type keyStore struct {
 	sync.RWMutex
-	client HTTPClient
-	uri    string
-	keySet jose.JSONWebKeySet
-	expiry time.Time
-	jitter time.Duration
+	client     HTTPClient
+	uri        string
+	keySet     jose.JSONWebKeySet
+	expiry     time.Time
+	jitter     time.Duration
+	nextReload time.Time
+	group      singleflight.Group
 }
 
 func newKeyStore(client HTTPClient, uri string) (*keyStore, error) {
@@ -46,24 +54,44 @@ func newKeyStore(client HTTPClient, uri string) (*keyStore, error) {
 func (ks *keyStore) Get(kid string) (keys []jose.JSONWebKey) {
 	ks.RLock()
 	// Force reload if expiration has passed
-	if time.Now().After(ks.expiry) {
+	if time.Now().After(ks.expiry) && time.Now().After(ks.nextReload) {
 		ks.RUnlock()
 		ks.reload()
 		ks.RLock()
 	}
 	keys = ks.keySet.Key(kid)
+	// Reload on unknown kid in case the provider has rotated their keys without
+	// us knowing about it.
+	if len(keys) == 0 && time.Now().After(ks.nextReload) {
+		ks.RUnlock()
+		ks.reload()
+		ks.RLock()
+		keys = ks.keySet.Key(kid)
+	}
 	ks.RUnlock()
 	return
 }
 
 func (ks *keyStore) reload() {
-	if keys, age, err := getKeysFromJWKsURI(ks.client, ks.uri); err == nil {
+	// Coalesce concurrent reloads into a single request, so that a burst of
+	// unknown key ids will only result in one fetch.
+	_, _, _ = ks.group.Do(ks.uri, func() (any, error) {
+		keys, age, err := getKeysFromJWKsURI(ks.client, ks.uri)
+
 		ks.Lock()
-		ks.keySet = keys
-		ks.jitter = getCacheJitter(age)
-		ks.expiry = getExpirationTime(age, ks.jitter)
+		ks.nextReload = time.Now().Add(minReloadInterval)
+		if err == nil {
+			ks.keySet = keys
+			ks.jitter = getCacheJitter(age)
+			ks.expiry = getExpirationTime(age, ks.jitter)
+		}
 		ks.Unlock()
-	}
+
+		if err != nil {
+			log.Printf("Failed to reload the JWK Set from %s: %v", ks.uri, err)
+		}
+		return nil, nil
+	})
 }
 
 func getKeysFromJWKsURI(client HTTPClient, uri string) (jose.JSONWebKeySet, time.Duration, error) {
@@ -73,6 +101,10 @@ func getKeysFromJWKsURI(client HTTPClient, uri string) (jose.JSONWebKeySet, time
 		return keys, 0, errors.Wrapf(err, "failed to connect to %s", uri)
 	}
 	defer resp.Body.Close()
+	if resp.StatusCode >= 400 {
+		b, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
+		return keys, 0, errors.Errorf("error reading %s: status=%d, response=%s", uri, resp.StatusCode, b)
+	}
 	if err := json.NewDecoder(resp.Body).Decode(&keys); err != nil {
 		return keys, 0, errors.Wrapf(err, "error reading %s", uri)
 	}

authority/provisioner/keystore_test.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"reflect"
+	"sync"
 	"testing"
 	"time"
 
@@ -35,6 +36,7 @@ func Test_newKeyStore(t *testing.T) {
 	}{
 		{"ok", args{srv.Client(), srv.URL}, ks.keySet, false},
 		{"fail", args{srv.Client(), srv.URL + "/error"}, jose.JSONWebKeySet{}, true},
+		{"fail json error body", args{srv.Client(), srv.URL + "/error-json"}, jose.JSONWebKeySet{}, true},
 		{"fail client", args{http.DefaultClient, srv.URL}, jose.JSONWebKeySet{}, true},
 	}
 	for _, tt := range tests {
@@ -153,6 +155,123 @@ func Test_keyStore_Get(t *testing.T) {
 	}
 }
 
+func Test_keyStore_Get_unknownKeyID(t *testing.T) {
+	srv := generateJWKServer(2)
+	defer srv.Close()
+
+	ks, err := newKeyStore(srv.Client(), srv.URL+"/rotate")
+	assert.FatalError(t, err)
+	ks.RLock()
+	cached := ks.keySet
+	ks.RUnlock()
+
+	rotated := rotateJWKServer(t, srv)
+	assert.Len(t, 1, ks.Get(cached.Keys[0].KeyID))
+
+	assert.Len(t, 1, ks.Get(rotated.Keys[0].KeyID))
+	assert.Len(t, 1, ks.Get(rotated.Keys[1].KeyID))
+}
+
+func Test_keyStore_Get_unknownKeyIDIsRateLimited(t *testing.T) {
+	srv := generateJWKServer(2)
+	defer srv.Close()
+
+	client := &countingClient{HTTPClient: srv.Client()}
+	ks, err := newKeyStore(client, srv.URL+"/rotate")
+	assert.FatalError(t, err)
+
+	rotated := rotateJWKServer(t, srv)
+	assert.Len(t, 1, ks.Get(rotated.Keys[0].KeyID))
+	assert.Equals(t, 2, client.gets) // one on init, one on the unknown key id
+
+	for range 10 {
+		assert.Len(t, 0, ks.Get("foobar"))
+	}
+	assert.Equals(t, 2, client.gets)
+}
+
+func Test_keyStore_Get_unknownKeyIDCoalescesReloads(t *testing.T) {
+	srv := generateJWKServer(2)
+	defer srv.Close()
+
+	client := &countingClient{HTTPClient: srv.Client()}
+	ks, err := newKeyStore(client, srv.URL+"/rotate")
+	assert.FatalError(t, err)
+
+	rotated := rotateJWKServer(t, srv)
+
+	// nextReload only advances once a reload finishes, so a simultaneous burst
+	// all passes the check in Get.
+	found := make([]int, 50)
+	var wg sync.WaitGroup
+	for i := range found {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			found[i] = len(ks.Get(rotated.Keys[0].KeyID))
+		}()
+	}
+	wg.Wait()
+
+	for _, n := range found {
+		assert.Equals(t, 1, n)
+	}
+	assert.Equals(t, 2, client.gets) // one on init, one shared by the burst
+}
+
+func Test_keyStore_Get_failedReloadKeepsCachedKeys(t *testing.T) {
+	srv := generateJWKServer(2)
+	defer srv.Close()
+
+	ks, err := newKeyStore(srv.Client(), srv.URL+"/rotate")
+	assert.FatalError(t, err)
+	ks.RLock()
+	cached := ks.keySet
+	ks.RUnlock()
+
+	// The endpoint starts answering with a JSON error body, which decodes into
+	// an empty key set unless the status code is checked.
+	failJWKServer(t, srv)
+
+	// An unknown key id triggers a reload, and that reload fails.
+	assert.Len(t, 0, ks.Get("foobar"))
+
+	assert.Len(t, 1, ks.Get(cached.Keys[0].KeyID))
+	assert.Len(t, 1, ks.Get(cached.Keys[1].KeyID))
+}
+
+// rotateJWKServer rotates the keys on srv and returns the new set.
+func rotateJWKServer(t *testing.T, srv *httptest.Server) jose.JSONWebKeySet {
+	t.Helper()
+	var keySet jose.JSONWebKeySet
+	resp, err := srv.Client().Get(srv.URL + "/rotate/next")
+	assert.FatalError(t, err)
+	defer resp.Body.Close()
+	assert.FatalError(t, json.NewDecoder(resp.Body).Decode(&keySet))
+	return keySet
+}
+
+// failJWKServer makes the /rotate endpoint of srv start returning an error.
+func failJWKServer(t *testing.T, srv *httptest.Server) {
+	t.Helper()
+	resp, err := srv.Client().Get(srv.URL + "/rotate/fail")
+	assert.FatalError(t, err)
+	assert.FatalError(t, resp.Body.Close())
+}
+
+type countingClient struct {
+	HTTPClient
+	mu   sync.Mutex
+	gets int
+}
+
+func (c *countingClient) Get(uri string) (*http.Response, error) {
+	c.mu.Lock()
+	c.gets++
+	c.mu.Unlock()
+	return c.HTTPClient.Get(uri)
+}
+
 func Test_abs(t *testing.T) {
 	maxInt64 := time.Duration(1<<63 - 1)
 	minInt64 := time.Duration(-1 << 63)

authority/provisioner/utils_test.go

@@ -341,6 +341,7 @@ func generateOIDC() (*OIDC, error) {
 			JWKSetURI: "https://example.com/.well-known/jwks",
 		},
 		keyStore: &keyStore{
+			client: erroringHTTPClient{},
 			keySet: jose.JSONWebKeySet{Keys: []jose.JSONWebKey{*jwk}},
 			expiry: time.Now().Add(24 * time.Hour),
 		},
@@ -373,6 +374,7 @@ func generateGCP() (*GCP, error) {
 		DisableSSHCAUser: &DefaultDisableSSHCAUser,
 		config:           newGCPConfig(),
 		keyStore: &keyStore{
+			client: erroringHTTPClient{},
 			keySet: jose.JSONWebKeySet{Keys: []jose.JSONWebKey{*jwk}},
 			expiry: time.Now().Add(24 * time.Hour),
 		},
@@ -610,6 +612,7 @@ func generateAzure() (*Azure, error) {
 			JWKSetURI: "https://login.microsoftonline.com/common/discovery/keys",
 		},
 		keyStore: &keyStore{
+			client: erroringHTTPClient{},
 			keySet: jose.JSONWebKeySet{Keys: []jose.JSONWebKey{*jwk}},
 			expiry: time.Now().Add(24 * time.Hour),
 		},
@@ -1127,11 +1130,31 @@ func generateJWKServerHandler(n int, srv *httptest.Server) http.Handler {
 	}
 
 	defaultKeySet := must(generateJSONWebKeySet(n))[0].(jose.JSONWebKeySet)
+	rotatingKeySet := defaultKeySet
+	rotateFails := false
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		hits.Hits++
 		switch r.RequestURI {
 		case "/error":
 			http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
+		case "/error-json":
+			w.Header().Add("Content-Type", "application/json")
+			w.WriteHeader(http.StatusNotFound)
+			w.Write([]byte(`{"message":"Not Found"}`))
+		case "/rotate":
+			if rotateFails {
+				w.Header().Add("Content-Type", "application/json")
+				w.WriteHeader(http.StatusNotFound)
+				w.Write([]byte(`{"message":"Not Found"}`))
+				return
+			}
+			w.Header().Add("Cache-Control", "max-age=604800")
+			writeJSON(w, getPublic(rotatingKeySet))
+		case "/rotate/next":
+			rotatingKeySet = must(generateJSONWebKeySet(n))[0].(jose.JSONWebKeySet)
+			writeJSON(w, getPublic(rotatingKeySet))
+		case "/rotate/fail":
+			rotateFails = true
 		case "/hits":
 			writeJSON(w, hits)
 		case "/.well-known/openid-configuration":