Prevent possible race condition (#42)

go.mod

@@ -4,4 +4,7 @@ go 1.21
 
 require golang.org/x/time v0.9.0
 
-retract [v0.5.0, v0.5.15] // HTTP client only overwrites and appends JWK to local cache during refresh: https://github.com/MicahParks/jwkset/issues/40
+retract (
+	v0.6.0 // Potential race condition in refresh goroutine: https://github.com/MicahParks/jwkset/pull/42
+	[v0.5.0, v0.5.15] // HTTP client only overwrites and appends JWK to local cache during refresh: https://github.com/MicahParks/jwkset/issues/40
+)

storage.go

@@ -72,11 +72,6 @@ func (m *MemoryJWKSet) KeyDelete(_ context.Context, keyID string) (ok bool, err
 	}
 	return ok, nil
 }
-func (m *MemoryJWKSet) KeyDeleteAll() {
-	m.mux.Lock()
-	defer m.mux.Unlock()
-	m.set = make([]JWK, 0)
-}
 func (m *MemoryJWKSet) KeyRead(_ context.Context, keyID string) (JWK, error) {
 	m.mux.RLock()
 	defer m.mux.RUnlock()
@@ -95,16 +90,9 @@ func (m *MemoryJWKSet) KeyReadAll(_ context.Context) ([]JWK, error) {
 func (m *MemoryJWKSet) KeyWrite(_ context.Context, jwk JWK) error {
 	m.mux.Lock()
 	defer m.mux.Unlock()
-	for i, j := range m.set {
-		if j.Marshal().KID == jwk.Marshal().KID {
-			m.set[i] = jwk
-			return nil
-		}
-	}
 	m.set = append(m.set, jwk)
 	return nil
 }
-
 func (m *MemoryJWKSet) JSON(ctx context.Context) (json.RawMessage, error) {
 	jwks, err := m.Marshal(ctx)
 	if err != nil {
@@ -259,19 +247,18 @@ func NewStorageFromHTTP(u *url.URL, options HTTPClientStorageOptions) (Storage,
 		if err != nil {
 			return fmt.Errorf("failed to decode JWK Set response: %w", err)
 		}
-		store.KeyDeleteAll() // Clear local cache in case of key revocation.
-		for _, marshal := range jwks.Keys {
+		store.mux.Lock()
+		defer store.mux.Unlock()
+		store.set = make([]JWK, len(jwks.Keys)) // Clear local cache in case of key revocation.
+		for i, marshal := range jwks.Keys {
 			marshalOptions := JWKMarshalOptions{
 				Private: true,
 			}
 			jwk, err := NewJWKFromMarshal(marshal, marshalOptions, options.ValidateOptions)
 			if err != nil {
 				return fmt.Errorf("failed to create JWK from JWK Marshal: %w", err)
 			}
-			err = store.KeyWrite(options.Ctx, jwk)
-			if err != nil {
-				return fmt.Errorf("failed to write JWK to memory storage: %w", err)
-			}
+			store.set[i] = jwk
 		}
 		return nil
 	}

storage_test.go

@@ -77,6 +77,13 @@ func TestMemoryKeyRead(t *testing.T) {
 	if !bytes.Equal(key.Key().([]byte), hmacKey1) {
 		t.Fatalf("Read key does not match written key.")
 	}
+	ok, err := store.KeyDelete(params.ctx, kidWritten)
+	if err != nil {
+		t.Fatalf("Failed to delete written key. %s", err)
+	}
+	if !ok {
+		t.Fatalf("Failed to delete written key.")
+	}
 
 	jwk = newStorageTestJWK(t, hmacKey2, kidWritten)
 	err = store.KeyWrite(params.ctx, jwk)
@@ -93,10 +100,13 @@ func TestMemoryKeyRead(t *testing.T) {
 		t.Fatalf("Read key does not match written key.")
 	}
 
-	_, err = store.KeyDelete(params.ctx, kidWritten)
+	ok, err = store.KeyDelete(params.ctx, kidWritten)
 	if err != nil {
 		t.Fatalf("Failed to delete written key. %s", err)
 	}
+	if !ok {
+		t.Fatalf("Failed to delete written key.")
+	}
 
 	_, err = store.KeyRead(params.ctx, kidWritten)
 	if !errors.Is(err, ErrKeyNotFound) {