snapshot of https://github.com/centrifugal/centrifugo.git (master, 2026-09-16; blobs > 131072B omitted)

internal/jwks/cache.go

@@ -0,0 +1,20 @@
+package jwks
+
+import (
+	"errors"
+)
+
+var (
+	// ErrCacheNotFound returned when cache value not found.
+	ErrCacheNotFound = errors.New("cache: value not found")
+)
+
+// Cache works with cache layer. The cacheKey is opaque to the cache and chosen by
+// the caller — Manager uses a composite of the resolved JWKS endpoint URL and the
+// JWT header kid so that keys cached from one trust domain (templated URL) cannot
+// satisfy lookups for another. See Manager.FetchKey.
+type Cache interface {
+	Add(cacheKey string, key *JWK) error
+	Get(cacheKey string) (*JWK, error)
+	Len() (int, error)
+}

internal/jwks/cache_ttl.go

@@ -0,0 +1,125 @@
+package jwks
+
+import (
+	"sync"
+	"time"
+)
+
+type item struct {
+	sync.RWMutex
+	data       *JWK
+	expiration *time.Time
+}
+
+func (i *item) touch(d time.Duration) {
+	i.Lock()
+	exp := time.Now().Add(d)
+	i.expiration = &exp
+	i.Unlock()
+}
+
+func (i *item) expired() bool {
+	i.RLock()
+	res := true
+	if i.expiration != nil {
+		res = i.expiration.Before(time.Now())
+	}
+	i.RUnlock()
+	return res
+}
+
+// TTLCache is a TTL bases in-memory cache.
+type TTLCache struct {
+	mu       sync.RWMutex
+	ttl      time.Duration
+	stop     chan struct{}
+	stopOnce sync.Once
+	items    map[string]*item
+}
+
+// NewTTLCache returns a new instance of ttl cache.
+func NewTTLCache(ttl time.Duration) *TTLCache {
+	cache := &TTLCache{
+		ttl:   ttl,
+		stop:  make(chan struct{}),
+		items: make(map[string]*item),
+	}
+	cache.run()
+	return cache
+}
+
+func (tc *TTLCache) cleanup() {
+	tc.mu.Lock()
+	for key, item := range tc.items {
+		if item.expired() {
+			delete(tc.items, key)
+		}
+	}
+	tc.mu.Unlock()
+}
+
+func (tc *TTLCache) run() {
+	d := tc.ttl
+	if d < time.Second {
+		d = time.Second
+	}
+
+	ticker := time.NewTicker(d)
+	go func() {
+		for {
+			select {
+			case <-ticker.C:
+				tc.cleanup()
+			case <-tc.stop:
+				ticker.Stop()
+				return
+			}
+		}
+	}()
+}
+
+// Add item into cache under the given cacheKey.
+func (tc *TTLCache) Add(cacheKey string, key *JWK) error {
+	tc.mu.Lock()
+	item := &item{data: key}
+	item.touch(tc.ttl)
+	tc.items[cacheKey] = item
+	tc.mu.Unlock()
+	return nil
+}
+
+// Get item by cacheKey.
+func (tc *TTLCache) Get(cacheKey string) (*JWK, error) {
+	tc.mu.RLock()
+	item, ok := tc.items[cacheKey]
+	if !ok || item.expired() {
+		tc.mu.RUnlock()
+		return nil, ErrCacheNotFound
+	}
+	item.touch(tc.ttl)
+	tc.mu.RUnlock()
+	return item.data, nil
+}
+
+// Stop stops TTL cache.
+func (tc *TTLCache) Stop() error {
+	tc.stopOnce.Do(func() {
+		close(tc.stop)
+	})
+	return nil
+}
+
+func (tc *TTLCache) remove(cacheKey string) error {
+	tc.mu.Lock()
+	delete(tc.items, cacheKey)
+	tc.mu.Unlock()
+	return nil
+}
+
+// Len returns current size of cache.
+func (tc *TTLCache) Len() (int, error) {
+	tc.mu.RLock()
+	n := len(tc.items)
+	tc.mu.RUnlock()
+	return n, nil
+}

internal/jwks/cache_ttl_test.go

@@ -0,0 +1,161 @@
+package jwks
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestTTLCacheInit(t *testing.T) {
+	cache := NewTTLCache(5 * time.Minute)
+	require.NotNil(t, cache)
+}
+
+func TestTTLCacheAdd(t *testing.T) {
+	testCases := []struct {
+		Name string
+		TTL  time.Duration
+		Ops  int
+	}{
+		{
+			Name: "OK",
+			TTL:  5 * time.Second,
+			Ops:  100,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.Name, func(t *testing.T) {
+			cache := NewTTLCache(tc.TTL)
+			require.NotNil(t, cache)
+
+			for i := 0; i < tc.Ops; i++ {
+				kid := fmt.Sprintf("key-%d", i+1)
+				require.NoError(t, cache.Add(kid, &JWK{
+					Kid: kid,
+					Kty: "RSA",
+					Alg: "RS256",
+					Use: "sig",
+				}))
+			}
+		})
+	}
+}
+
+func TestTTLCacheGet(t *testing.T) {
+	testCases := []struct {
+		Name  string
+		Key   *JWK
+		Kid   string
+		Error error
+	}{
+		{
+			Name: "OK",
+			Key: &JWK{
+				Kid: "202101",
+				Kty: "RSA",
+				Alg: "RS256",
+				Use: "sig",
+			},
+			Kid: "202101",
+		},
+		{
+			Name: "NotFound",
+			Key: &JWK{
+				Kid: "202101",
+				Kty: "RSA",
+				Alg: "RS256",
+				Use: "sig",
+			},
+			Kid:   "202102",
+			Error: ErrCacheNotFound,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.Name, func(t *testing.T) {
+			cache := NewTTLCache(5 * time.Minute)
+			require.NotNil(t, cache)
+			require.NoError(t, cache.Add(tc.Key.Kid, tc.Key))
+
+			key, err := cache.Get(tc.Kid)
+			if tc.Error != nil {
+				require.Error(t, err)
+				require.ErrorIs(t, err, tc.Error)
+			} else {
+				require.NoError(t, err)
+				require.EqualValues(t, tc.Key, key)
+			}
+		})
+	}
+}
+
+func TestTTLCacheRemove(t *testing.T) {
+	testCases := []struct {
+		Name      string
+		NumAdd    int
+		NumDelete int
+		Len       int
+	}{
+		{
+			Name:      "OK",
+			NumAdd:    75,
+			NumDelete: 50,
+			Len:       25,
+		},
+		{
+			Name:      "RemoveUntilEmpty",
+			NumAdd:    75,
+			NumDelete: 100,
+			Len:       0,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.Name, func(t *testing.T) {
+			cache := NewTTLCache(5 * time.Minute)
+			require.NotNil(t, cache)
+
+			for i := 0; i < tc.NumAdd; i++ {
+				kid := fmt.Sprintf("key-%d", i+1)
+				require.NoError(t, cache.Add(kid, &JWK{
+					Kid: kid,
+					Kty: "RSA",
+					Alg: "RS256",
+					Use: "sig",
+				}))
+			}
+
+			for i := 0; i < tc.NumDelete; i++ {
+				kid := fmt.Sprintf("key-%d", i+1)
+				require.NoError(t, cache.remove(kid))
+			}
+
+			n, err := cache.Len()
+			require.NoError(t, err)
+			require.Equal(t, tc.Len, n)
+		})
+	}
+}
+
+func TestTTLCacheCleanup(t *testing.T) {
+	cache := NewTTLCache(1 * time.Millisecond)
+
+	for i := 0; i < 10; i++ {
+		kid := fmt.Sprintf("key-%d", i+1)
+		require.NoError(t, cache.Add(kid, &JWK{
+			Kid: kid,
+			Kty: "RSA",
+			Alg: "RS256",
+			Use: "sig",
+		}))
+	}
+
+	time.Sleep(2 * time.Second)
+
+	n, err := cache.Len()
+	require.NoError(t, err)
+	require.Equal(t, 0, n)
+}

internal/jwks/manager.go

@@ -0,0 +1,235 @@
+package jwks
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"path"
+	"time"
+
+	"github.com/rakutentech/jwk-go/jwk"
+	"github.com/rs/zerolog/log"
+	"github.com/valyala/fasttemplate"
+	"golang.org/x/sync/singleflight"
+)
+
+const (
+	_defaultRetries            = 2
+	_defaultTimeout            = 1 * time.Second
+	_defaultMaxIdleConnPerHost = 255
+	_defaultTTL                = 1 * time.Hour
+)
+
+// JWK represents an unparsed JSON Web Key (JWK) in its wire format.
+type JWK = jwk.JWK
+
+var (
+	// ErrInvalidURL returned when input url has invalid format.
+	ErrInvalidURL = errors.New("jwks: invalid url value or format")
+	// ErrInvalidNumRetries returned when number of retries is zero.
+	ErrInvalidNumRetries = errors.New("jwks: invalid number of retries")
+	// ErrKeyIDNotProvided returned when input kid is not present.
+	ErrKeyIDNotProvided = errors.New("jwks: kid is not provided")
+	// ErrPublicKeyNotFound returned when no public key is found.
+	ErrPublicKeyNotFound = errors.New("jwks: public key not found")
+
+	errUnexpectedStatusCode = errors.New("jwks: unexpected status code")
+	errUnmarshal            = errors.New("jwks: unmarshal error")
+	errConvert              = errors.New("jwks: convert error")
+)
+
+// Manager fetches and returns JWK from public source.
+type Manager struct {
+	url      *fasttemplate.Template
+	cache    Cache
+	client   *http.Client
+	useCache bool
+	retries  uint
+	group    singleflight.Group
+}
+
+func defaultHTTPClient() *http.Client {
+	return &http.Client{
+		Transport: &http.Transport{
+			MaxIdleConnsPerHost: _defaultMaxIdleConnPerHost,
+		},
+		Timeout: _defaultTimeout,
+	}
+}
+
+// NewManager returns a new instance of Manager.
+func NewManager(rawURL string, opts ...Option) (*Manager, error) {
+	u, err := url.Parse(rawURL)
+	if err != nil {
+		return nil, ErrInvalidURL
+	}
+	if u.Scheme != "http" && u.Scheme != "https" {
+		return nil, fmt.Errorf("endpoint must have http:// or https:// scheme, got: %s", rawURL)
+	}
+	urlTemplate := fasttemplate.New(rawURL, "{{", "}}")
+
+	mng := &Manager{
+		url: urlTemplate,
+
+		cache:    NewTTLCache(_defaultTTL),
+		client:   defaultHTTPClient(),
+		useCache: true,
+		retries:  _defaultRetries,
+	}
+
+	for _, opt := range opts {
+		opt(mng)
+	}
+
+	if mng.retries == 0 {
+		return nil, ErrInvalidNumRetries
+	}
+
+	return mng, nil
+}
+
+// FetchKey fetches JWKS from public source or cache.
+//
+// The cache and singleflight keys are scoped to the resolved JWKS endpoint URL,
+// not only the JWT header kid. This prevents a key cached from one trust domain
+// (e.g., tenant A's templated JWKS URL) from satisfying a verification request
+// for a different trust domain (tenant B) that happens to advertise the same kid.
+// JWT kid values are not globally unique by spec — common operational labels like
+// "default" or rotation IDs may collide across issuers.
+func (m *Manager) FetchKey(ctx context.Context, kid string, tokenVars map[string]any) (*JWK, error) {
+	if kid == "" {
+		return nil, ErrKeyIDNotProvided
+	}
+
+	jwkURL, err := m.resolveURL(tokenVars)
+	if err != nil {
+		return nil, err
+	}
+
+	cacheKey := cacheKey(jwkURL, kid)
+
+	// If useCache is true, first try to get key from cache.
+	if m.useCache {
+		key, err := m.cache.Get(cacheKey)
+		if err == nil {
+			return key, nil
+		}
+	}
+
+	// Otherwise fetch from public JWKS.
+	v, err, _ := m.group.Do(cacheKey, func() (any, error) {
+		return m.fetchKey(ctx, jwkURL, kid)
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return v.(*JWK), nil
+}
+
+// cacheKey builds a cache/singleflight key namespaced by the resolved JWKS URL.
+// The NUL byte is a delimiter that cannot appear in a valid URL or JWT kid, so
+// the encoding is unambiguous.
+func cacheKey(resolvedURL, kid string) string {
+	return resolvedURL + "\x00" + kid
+}
+
+func (m *Manager) loadData(req *http.Request) ([]byte, error) {
+	resp, err := m.client.Do(req) //nolint:gosec // URL is from server configuration, not user input.
+	if err != nil {
+		return nil, err
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("%w: %d", errUnexpectedStatusCode, resp.StatusCode)
+	}
+	return io.ReadAll(resp.Body)
+}
+
+// resolveURL applies tokenVars to the URL template and rejects results with
+// path traversal sequences. Runs on every FetchKey call (including cache hits)
+// so that cache lookups are scoped to the validated, fully-resolved URL.
+func (m *Manager) resolveURL(tokenVars map[string]any) (string, error) {
+	jwkURL := m.url.ExecuteString(tokenVars)
+
+	u, err := url.Parse(jwkURL)
+	if err != nil {
+		return "", fmt.Errorf("error parsing JWKS URL: %w", err)
+	}
+	if u.Path != "" {
+		if cleanPath := path.Clean(u.Path); cleanPath != u.Path {
+			log.Info().Str("path", u.Path).Str("clean_path", cleanPath).
+				Msg("JWKS URL path contains traversal sequences, request rejected")
+			return "", fmt.Errorf("JWKS URL path contains traversal sequences: %q", u.Path)
+		}
+	}
+	return jwkURL, nil
+}
+
+func (m *Manager) fetchKey(ctx context.Context, jwkURL, kid string) (*JWK, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, jwkURL, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	var set jwk.KeySpecSet
+	var data []byte
+	var lastError error
+
+	retries := m.retries
+	for {
+		if retries == 0 {
+			return nil, lastError
+		}
+		retries--
+		var err error
+		data, err = m.loadData(req)
+		if err != nil {
+			lastError = err
+			continue
+		}
+		break
+	}
+
+	if err := json.Unmarshal(data, &set); err != nil {
+		return nil, fmt.Errorf("%w: %v", errUnmarshal, err)
+	}
+
+	if len(set.Keys) == 0 {
+		return nil, ErrPublicKeyNotFound
+	}
+
+	var res *JWK
+
+	// Save new set into cache.
+	for _, spec := range set.Keys {
+		key, err := spec.ToJWK()
+		if err != nil {
+			return nil, fmt.Errorf("%w: %v", errConvert, err)
+		}
+
+		if key.Use != "sig" {
+			// Not interested in other types of Use in Centrifugo.
+			continue
+		}
+
+		if m.useCache {
+			_ = m.cache.Add(cacheKey(jwkURL, key.Kid), key)
+		}
+
+		if key.Kid == kid {
+			res = key
+		}
+	}
+
+	if res == nil {
+		return nil, ErrPublicKeyNotFound
+	}
+
+	return res, nil
+}

internal/jwks/manager_test.go

@@ -0,0 +1,291 @@
+package jwks
+
+import (
+	"context"
+	"crypto/rand"
+	"crypto/rsa"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+
+	"github.com/rakutentech/jwk-go/jwk"
+	"github.com/stretchr/testify/require"
+)
+
+type testKey struct {
+	Kid string
+	Key any
+}
+
+func parseRSA(t *testing.T, k *JWK) *rsa.PublicKey {
+	t.Helper()
+	spec, err := k.ParseKeySpec()
+	require.NoError(t, err)
+	pub, ok := spec.Key.(*rsa.PublicKey)
+	require.True(t, ok)
+	return pub
+}
+
+func randomKeys() (*rsa.PrivateKey, *rsa.PublicKey, error) {
+	privateKey, err := rsa.GenerateKey(rand.Reader, 1024)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return privateKey, &privateKey.PublicKey, nil
+}
+
+func jwksHandler(keys ...testKey) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		specs := jwk.KeySpecSet{}
+
+		for _, key := range keys {
+			spec := jwk.NewSpecWithID(key.Kid, key.Key)
+			spec.Use = "sig"
+			specs.Keys = append(specs.Keys, *spec)
+		}
+
+		data, err := json.Marshal(specs)
+		if err != nil {
+			http.Error(w, "Server Error", http.StatusInternalServerError)
+			return
+		}
+
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write(data)
+	})
+}
+
+func TestManagerFetchKey_UnmarshalError(t *testing.T) {
+	mux := http.NewServeMux()
+	path := "/.well-known/jwks.json"
+	mux.HandleFunc(path, func(writer http.ResponseWriter, request *http.Request) {
+		writer.WriteHeader(http.StatusOK)
+		_, _ = writer.Write([]byte(`...`))
+	})
+	server := httptest.NewServer(mux)
+	defer server.Close()
+
+	manager, err := NewManager(server.URL + path)
+	require.NoError(t, err)
+
+	_, err = manager.FetchKey(context.Background(), "202101", nil)
+	require.ErrorIs(t, err, errUnmarshal)
+}
+
+func TestManagerFetchKey_KeyNotFound(t *testing.T) {
+	mux := http.NewServeMux()
+	path := "/.well-known/jwks.json"
+	mux.HandleFunc(path, func(writer http.ResponseWriter, request *http.Request) {
+		writer.WriteHeader(http.StatusOK)
+		_, _ = writer.Write([]byte(`{"keys": []}`))
+	})
+	server := httptest.NewServer(mux)
+	defer server.Close()
+
+	manager, err := NewManager(server.URL + path)
+	require.NoError(t, err)
+
+	_, err = manager.FetchKey(context.Background(), "202101", nil)
+	require.ErrorIs(t, err, ErrPublicKeyNotFound)
+}
+
+func TestManagerFetchKey_WrongStatusCode(t *testing.T) {
+	mux := http.NewServeMux()
+	path := "/.well-known/jwks.json"
+	mux.HandleFunc(path, func(writer http.ResponseWriter, request *http.Request) {
+		writer.WriteHeader(http.StatusInternalServerError)
+	})
+	server := httptest.NewServer(mux)
+	defer server.Close()
+
+	manager, err := NewManager(server.URL + path)
+	require.NoError(t, err)
+
+	_, err = manager.FetchKey(context.Background(), "202101", nil)
+	require.ErrorIs(t, err, errUnexpectedStatusCode)
+}
+
+func TestManagerInitialFetchKey(t *testing.T) {
+	_, pubKey, err := randomKeys()
+	require.NoError(t, err)
+
+	testCases := []struct {
+		Name    string
+		Handler http.Handler
+		Kid     string
+		Error   error
+	}{
+		{
+			Name:    "OK",
+			Handler: jwksHandler(testKey{"202101", pubKey}),
+			Kid:     "202101",
+		},
+		{
+			Name:    "NotFound",
+			Handler: jwksHandler(testKey{"202101", pubKey}),
+			Kid:     "202102",
+			Error:   ErrPublicKeyNotFound,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.Name, func(t *testing.T) {
+			r := require.New(t)
+
+			ts := httptest.NewServer(tc.Handler)
+			defer ts.Close()
+
+			manager, err := NewManager(ts.URL)
+			r.NoError(err)
+
+			key, err := manager.FetchKey(context.Background(), tc.Kid, nil)
+			if tc.Error != nil {
+				r.Error(err)
+				r.ErrorIs(err, tc.Error)
+			} else {
+				r.NoError(err)
+				r.Equal(tc.Kid, key.Kid)
+			}
+		})
+	}
+}
+
+func TestManagerFetchKey_PathTraversalRejected(t *testing.T) {
+	kid := "test-kid"
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		t.Fatal("request should have been rejected before reaching the server")
+	}))
+	defer ts.Close()
+
+	manager, err := NewManager(ts.URL + "/tenants/{{tenant}}/jwks.json")
+	require.NoError(t, err)
+
+	// Path traversal via ".." — must be rejected.
+	tokenVars := map[string]any{"tenant": "../../etc"}
+	_, err = manager.FetchKey(context.Background(), kid, tokenVars)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "traversal")
+
+	// Single ".." segment — also rejected.
+	tokenVars = map[string]any{"tenant": ".."}
+	_, err = manager.FetchKey(context.Background(), kid, tokenVars)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "traversal")
+
+	// Clean value — must succeed (uses real server).
+	_, pubKey, err := randomKeys()
+	require.NoError(t, err)
+
+	ts2 := httptest.NewServer(jwksHandler(testKey{kid, pubKey}))
+	defer ts2.Close()
+
+	manager, err = NewManager(ts2.URL + "/tenants/{{tenant}}/jwks.json")
+	require.NoError(t, err)
+
+	tokenVars = map[string]any{"tenant": "acme"}
+	key, err := manager.FetchKey(context.Background(), kid, tokenVars)
+	require.NoError(t, err)
+	require.Equal(t, kid, key.Kid)
+}
+
+func TestManagerCachedFetchKey(t *testing.T) {
+	testCases := []struct {
+		Name         string
+		Options      []Option
+		ExpectedSize int
+	}{
+		{
+			Name:         "Default",
+			ExpectedSize: 1,
+		},
+		{
+			Name:         "NoCacheLookup",
+			Options:      []Option{WithUseCache(false)},
+			ExpectedSize: 0,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.Name, func(t *testing.T) {
+			r := require.New(t)
+
+			ctx := context.Background()
+			kid := "202101"
+
+			_, pubKey, err := randomKeys()
+			r.NoError(err)
+
+			ts := httptest.NewServer(jwksHandler(testKey{kid, pubKey}))
+			defer ts.Close()
+
+			manager, err := NewManager(ts.URL, tc.Options...)
+			r.NoError(err)
+
+			key, err := manager.FetchKey(ctx, kid, nil)
+			r.NoError(err)
+			r.Equal(kid, key.Kid)
+
+			size, err := manager.cache.Len()
+			r.NoError(err)
+			r.Equal(tc.ExpectedSize, size)
+		})
+	}
+}
+
+// TestManagerFetchKey_CacheScopedByResolvedURL is a regression test for a
+// cross-trust-domain cache reuse bug: when JWKS URLs are templated from token
+// claims (e.g., per-tenant endpoints), a key cached from tenant A's endpoint
+// must NOT satisfy a lookup for tenant B's endpoint, even if both JWKS documents
+// advertise the same kid. kid values are not globally unique by spec and may
+// collide across issuers.
+func TestManagerFetchKey_CacheScopedByResolvedURL(t *testing.T) {
+	const sharedKid = "shared-kid"
+
+	_, tenantAPubKey, err := randomKeys()
+	require.NoError(t, err)
+	_, tenantBPubKey, err := randomKeys()
+	require.NoError(t, err)
+
+	var tenantARequests, tenantBRequests int32
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/tenant-a/jwks.json", func(w http.ResponseWriter, r *http.Request) {
+		atomic.AddInt32(&tenantARequests, 1)
+		jwksHandler(testKey{sharedKid, tenantAPubKey}).ServeHTTP(w, r)
+	})
+	mux.HandleFunc("/tenant-b/jwks.json", func(w http.ResponseWriter, r *http.Request) {
+		atomic.AddInt32(&tenantBRequests, 1)
+		jwksHandler(testKey{sharedKid, tenantBPubKey}).ServeHTTP(w, r)
+	})
+	ts := httptest.NewServer(mux)
+	defer ts.Close()
+
+	manager, err := NewManager(ts.URL + "/{{tenant}}/jwks.json")
+	require.NoError(t, err)
+
+	ctx := context.Background()
+
+	keyA, err := manager.FetchKey(ctx, sharedKid, map[string]any{"tenant": "tenant-a"})
+	require.NoError(t, err)
+	require.Equal(t, tenantAPubKey.N, parseRSA(t, keyA).N)
+	require.Equal(t, int32(1), atomic.LoadInt32(&tenantARequests))
+	require.Equal(t, int32(0), atomic.LoadInt32(&tenantBRequests))
+
+	// Fetching for tenant B with the same kid must consult tenant B's endpoint,
+	// not reuse the cached tenant A key.
+	keyB, err := manager.FetchKey(ctx, sharedKid, map[string]any{"tenant": "tenant-b"})
+	require.NoError(t, err)
+	require.Equal(t, tenantBPubKey.N, parseRSA(t, keyB).N)
+	require.Equal(t, int32(1), atomic.LoadInt32(&tenantARequests))
+	require.Equal(t, int32(1), atomic.LoadInt32(&tenantBRequests))
+
+	// Second fetch for tenant A should hit the cache — no new HTTP request.
+	_, err = manager.FetchKey(ctx, sharedKid, map[string]any{"tenant": "tenant-a"})
+	require.NoError(t, err)
+	require.Equal(t, int32(1), atomic.LoadInt32(&tenantARequests))
+	require.Equal(t, int32(1), atomic.LoadInt32(&tenantBRequests))
+}

internal/jwks/options.go

@@ -0,0 +1,28 @@
+package jwks
+
+import (
+	"net/http"
+)
+
+// Option is used for configuring key manager.
+type Option func(m *Manager)
+
+// WithCache sets custom cache. Default is TTLCache.
+func WithCache(c Cache) Option {
+	return func(m *Manager) { m.cache = c }
+}
+
+// WithHTTPClient sets custom http client. By default client with 1 sec timeout used.
+func WithHTTPClient(c *http.Client) Option {
+	return func(m *Manager) { m.client = c }
+}
+
+// WithUseCache defines useCache option. Default is true.
+func WithUseCache(flag bool) Option {
+	return func(m *Manager) { m.useCache = flag }
+}
+
+// WithMaxRetries defines max retries count if request has been failed. Default is 2.
+func WithMaxRetries(n uint) Option {
+	return func(m *Manager) { m.retries = n }
+}