Make outbound DCR cache population race-safe across replicas

A reviewer found that two replicas racing on the same outbound DCR
(RFC 7591) cache-miss could each independently register a different
OAuth client with the upstream IdP — dynamic registration always
mints a fresh client_id/secret — then whichever replica's write
landed last in the shared Redis cache silently won. The losing
replica keeps the client it registered baked into its own config for
the rest of its process lifetime (DCR resolution runs once per
upstream at startup, never re-resolved), so it no longer agrees with
the durable cache about which client it holds credentials for.
dcrFlight (a singleflight.Group) only coalesces concurrent callers
within one process; it has no cross-replica reach.

Change the cache-population contract from upsert to create-if-absent,
returning the authoritative durable value either way: the caller's
own resolution on a successful claim, or the concurrent winner's
otherwise. CredentialStore.Put becomes PutIfAbsent, and
DCRCredentialStore.StoreDCRCredentials becomes
StoreDCRCredentialsIfAbsent; registerAndCache now returns whichever
resolution the store says is authoritative instead of trusting its
own local registration, and logs (at Debug, without ever including a
secret) when this replica lost the race. Callers MUST use the
returned value — RFC 7591 guarantees nothing about the two
registrations converging.

Redis claims the key with SET...NX (the same reservation-lock shape
already used twice in this file for ClientAssertionJWTValid and
ConsumeAssertionJWT), not WATCH/MULTI: unlike ReconcileConfiguredClient,
this write has no read-then-decide step to protect, so a plain atomic
NX claim is sufficient. On a lost claim it reads back the winner
through the existing GetDCRCredentials path rather than a second,
hand-rolled unmarshal, and retries the whole claim-or-read cycle
(bounded) if the winner's row evicts between the failed NX and the
read — its TTL can be as short as one second when the caller's
ClientSecretExpiresAt was already in the past, so this is a real,
reachable window, not a hypothetical one, and the alternative (a hard
error) would turn a retryable race into a permanent startup failure.

MemoryStorage's implementation treats an existing entry as absent
only when its ClientSecretExpiresAt is non-zero and already past —
otherwise it returns the existing entry unchanged rather than
overwriting it. A single process's dcrFlight already prevents a live
race there; this is contract symmetry with Redis, plus the correctness
case Redis gets from TTL eviction: without the expiry check, a
never-expiring entry can never be reclaimed, but a naive "any existing
entry blocks re-registration" check would also permanently pin an
already-expired one that should be re-registered.

Refs #6200

Signed-off-by: Jakub Hrozek <jakub@stacklok.com>

Author: jhrozek

pkg/auth/dcr/resolver.go

@@ -470,19 +470,39 @@ func registerAndCache(
 	// failure leaves no in-memory state diverging from the cache: the
 	// next call simply re-resolves rather than reading a value the cache
 	// never saw.
-	if err := cache.Put(ctx, key, resolution); err != nil {
+	//
+	// authoritative may differ from resolution: RFC 7591 dynamic
+	// registration mints a brand-new, unique client_id/client_secret on
+	// every call, so if another replica raced this one to register against
+	// the same Key and won the durable claim first, cache.PutIfAbsent
+	// returns THAT replica's credentials — the only ones the shared cache
+	// (and hence any other replica or a future restart) will agree this
+	// caller holds. Returning resolution here instead would leave this
+	// process using a client_id the durable store does not recognize.
+	authoritative, err := cache.PutIfAbsent(ctx, key, resolution)
+	if err != nil {
 		return nil, newDCRStepError(dcrStepCacheWrite, req.Issuer, redirectURI,
 			fmt.Errorf("cache put: %w", err))
 	}
+	if authoritative.ClientID != resolution.ClientID {
+		//nolint:gosec // G706: client_id is public metadata per RFC 7591.
+		slog.Debug("dcr: registration superseded by concurrent winner",
+			"local_issuer", req.Issuer,
+			"upstream_id", key.UpstreamID,
+			"redirect_uri", redirectURI,
+			"registered_client_id", resolution.ClientID,
+			"authoritative_client_id", authoritative.ClientID,
+		)
+	}
 
 	//nolint:gosec // G706: client_id is public metadata per RFC 7591.
 	slog.Debug("dcr: registered new client",
 		"local_issuer", req.Issuer,
 		"upstream_id", key.UpstreamID,
 		"redirect_uri", redirectURI,
-		"client_id", resolution.ClientID,
+		"client_id", authoritative.ClientID,
 	)
-	return resolution, nil
+	return authoritative, nil
 }
 
 // LogStepError emits the single boundary slog.Error record for a DCR

pkg/auth/dcr/resolver_test.go

@@ -170,7 +170,8 @@ func TestResolveDCRCredentials_CacheHitShortCircuits(t *testing.T) {
 		AuthorizationEndpoint: "https://preloaded/authorize",
 		TokenEndpoint:         "https://preloaded/token",
 	}
-	require.NoError(t, cache.Put(context.Background(), key, preloaded))
+	_, err := cache.PutIfAbsent(context.Background(), key, preloaded)
+	require.NoError(t, err)
 
 	req := &Request{
 		Issuer:       issuer,
@@ -1510,8 +1511,8 @@ func (c *countingStore) Get(ctx context.Context, key Key) (*Resolution, bool, er
 	return res, ok, err
 }
 
-func (c *countingStore) Put(ctx context.Context, key Key, res *Resolution) error {
-	return c.inner.Put(ctx, key, res)
+func (c *countingStore) PutIfAbsent(ctx context.Context, key Key, res *Resolution) (*Resolution, error) {
+	return c.inner.PutIfAbsent(ctx, key, res)
 }
 
 // TestResolveDCRCredentials_SingleflightCoalescesConcurrentCallers pins the
@@ -1790,8 +1791,11 @@ func (f failingDCRStore) Get(_ context.Context, _ Key) (*Resolution, bool, error
 	return nil, false, nil
 }
 
-func (f failingDCRStore) Put(_ context.Context, _ Key, _ *Resolution) error {
-	return f.putErr
+func (f failingDCRStore) PutIfAbsent(_ context.Context, _ Key, res *Resolution) (*Resolution, error) {
+	if f.putErr != nil {
+		return nil, f.putErr
+	}
+	return res, nil
 }
 
 // TestResolveDCRCredentials_CacheGetFailureWrapped covers PR #5042 review
@@ -2017,7 +2021,7 @@ func (panickingPutDCRStore) Get(_ context.Context, _ Key) (*Resolution, bool, er
 	return nil, false, nil
 }
 
-func (s panickingPutDCRStore) Put(_ context.Context, _ Key, _ *Resolution) error {
+func (s panickingPutDCRStore) PutIfAbsent(_ context.Context, _ Key, _ *Resolution) (*Resolution, error) {
 	panic(s.panicValue)
 }
 

pkg/auth/dcr/store.go

@@ -50,11 +50,20 @@ type CredentialStore interface {
 	// key is not present. An error is returned only on backend failure.
 	Get(ctx context.Context, key Key) (*Resolution, bool, error)
 
-	// Put stores the resolution for key, overwriting any existing entry.
+	// PutIfAbsent claims key for resolution. Returns the authoritative
+	// durable value: the caller's own resolution on a successful claim, the
+	// concurrent winner's otherwise. Callers MUST use the returned value,
+	// not their input resolution — RFC 7591 dynamic registration mints a
+	// unique client_id/client_secret on every call, so a caller that lost
+	// the race and kept using its own resolution would hold credentials the
+	// durable store does not agree it owns.
+	//
 	// Implementations must reject a nil resolution with an error rather
 	// than silently succeeding — a no-op would leave callers with no
 	// debug trail for the subsequent Get miss.
-	Put(ctx context.Context, key Key, resolution *Resolution) error
+	//
+	// The returned *Resolution is always non-nil when err is nil.
+	PutIfAbsent(ctx context.Context, key Key, resolution *Resolution) (*Resolution, error)
 }
 
 // NewStorageBackedStore returns a CredentialStore that delegates to a
@@ -154,15 +163,19 @@ func (s *inMemoryStore) Get(ctx context.Context, key Key) (*Resolution, bool, er
 	return credentialsToResolution(creds), true, nil
 }
 
-// Put implements CredentialStore by delegating to the embedded
+// PutIfAbsent implements CredentialStore by delegating to the embedded
 // *storage.MemoryStorage. The nil-resolution rejection matches
-// storageBackedStore.Put; see that method for the rationale.
-func (s *inMemoryStore) Put(ctx context.Context, key Key, resolution *Resolution) error {
+// storageBackedStore.PutIfAbsent; see that method for the rationale.
+func (s *inMemoryStore) PutIfAbsent(ctx context.Context, key Key, resolution *Resolution) (*Resolution, error) {
 	if resolution == nil {
-		return fmt.Errorf("dcr: resolution must not be nil")
+		return nil, fmt.Errorf("dcr: resolution must not be nil")
 	}
 	creds := resolutionToCredentials(key, resolution)
-	return s.mem.StoreDCRCredentials(ctx, creds)
+	authoritative, err := s.mem.StoreDCRCredentialsIfAbsent(ctx, creds)
+	if err != nil {
+		return nil, err
+	}
+	return credentialsToResolution(authoritative), nil
 }
 
 // Close releases the embedded MemoryStorage cleanup goroutine. Safe to
@@ -201,18 +214,27 @@ func (s *storageBackedStore) Get(ctx context.Context, key Key) (*Resolution, boo
 	return credentialsToResolution(creds), true, nil
 }
 
-// Put implements CredentialStore.
+// PutIfAbsent implements CredentialStore.
 //
 // A nil resolution is rejected rather than silently no-oped: a caller
 // passing nil would otherwise get a successful return, observe a miss on
 // the next Get, and have no error trail to debug from. Failing loudly at
 // the boundary makes such bugs visible at the first call.
-func (s *storageBackedStore) Put(ctx context.Context, key Key, resolution *Resolution) error {
+//
+// The returned *Resolution is the authoritative durable value —
+// s.backend.StoreDCRCredentialsIfAbsent's own contract — so a caller whose
+// registration lost a concurrent claim on this key gets back the winner's
+// credentials, not its own.
+func (s *storageBackedStore) PutIfAbsent(ctx context.Context, key Key, resolution *Resolution) (*Resolution, error) {
 	if resolution == nil {
-		return fmt.Errorf("dcr: resolution must not be nil")
+		return nil, fmt.Errorf("dcr: resolution must not be nil")
 	}
 	creds := resolutionToCredentials(key, resolution)
-	return s.backend.StoreDCRCredentials(ctx, creds)
+	authoritative, err := s.backend.StoreDCRCredentialsIfAbsent(ctx, creds)
+	if err != nil {
+		return nil, err
+	}
+	return credentialsToResolution(authoritative), nil
 }
 
 // resolutionToCredentials converts a resolver-side *Resolution into the

pkg/auth/dcr/store_test.go

@@ -40,7 +40,10 @@ func TestStorageBackedStore_PutGet_RoundTrip(t *testing.T) {
 		CreatedAt:               time.Now(),
 	}
 
-	require.NoError(t, store.Put(ctx, key, resolution))
+	authoritative, err := store.PutIfAbsent(ctx, key, resolution)
+	require.NoError(t, err)
+	assert.Equal(t, resolution.ClientID, authoritative.ClientID,
+		"an uncontested claim must return the caller's own resolution")
 
 	got, ok, err := store.Get(ctx, key)
 	require.NoError(t, err)
@@ -119,11 +122,15 @@ func TestStorageBackedStore_DistinctKeysDoNotCollide(t *testing.T) {
 		}
 	}
 
-	require.NoError(t, store.Put(ctx, keyA, resolution("a")))
-	require.NoError(t, store.Put(ctx, keyB, resolution("b")))
-	require.NoError(t, store.Put(ctx, keyC, resolution("c")))
-	require.NoError(t, store.Put(ctx, keyD, resolution("d")))
-	require.NoError(t, store.Put(ctx, keyE, resolution("e")))
+	for _, put := range []struct {
+		key      Key
+		clientID string
+	}{
+		{keyA, "a"}, {keyB, "b"}, {keyC, "c"}, {keyD, "d"}, {keyE, "e"},
+	} {
+		_, err := store.PutIfAbsent(ctx, put.key, resolution(put.clientID))
+		require.NoError(t, err)
+	}
 
 	for _, tc := range []struct {
 		key      Key
@@ -142,7 +149,15 @@ func TestStorageBackedStore_DistinctKeysDoNotCollide(t *testing.T) {
 	}
 }
 
-func TestStorageBackedStore_Put_OverwritesExisting(t *testing.T) {
+// TestStorageBackedStore_PutIfAbsent_FirstClaimWins pins the create-if-absent
+// contract that replaced unconditional overwrite (see PutIfAbsent doc): a
+// second PutIfAbsent for a key that already holds a value must NOT overwrite
+// it. Instead it returns the existing (first) resolution as the authoritative
+// value, and the store keeps holding the first entry. This is the fix for the
+// concurrent-replica bug where two callers independently register distinct
+// RFC 7591 clients for the same key — only one registration may ever become
+// the durable, agreed-upon value.
+func TestStorageBackedStore_PutIfAbsent_FirstClaimWins(t *testing.T) {
 	t.Parallel()
 
 	store := newMemoryDCRStore(t)
@@ -161,21 +176,27 @@ func TestStorageBackedStore_Put_OverwritesExisting(t *testing.T) {
 		Authorization: "https://idp.example.com/authorize",
 		Token:         "https://idp.example.com/token",
 	}
-	require.NoError(t, store.Put(ctx, key, &Resolution{
+	first, err := store.PutIfAbsent(ctx, key, &Resolution{
 		ClientID:              "first",
 		AuthorizationEndpoint: endpoints.Authorization,
 		TokenEndpoint:         endpoints.Token,
-	}))
-	require.NoError(t, store.Put(ctx, key, &Resolution{
+	})
+	require.NoError(t, err)
+	assert.Equal(t, "first", first.ClientID, "the uncontested first claim returns its own resolution")
+
+	second, err := store.PutIfAbsent(ctx, key, &Resolution{
 		ClientID:              "second",
 		AuthorizationEndpoint: endpoints.Authorization,
 		TokenEndpoint:         endpoints.Token,
-	}))
+	})
+	require.NoError(t, err)
+	assert.Equal(t, "first", second.ClientID,
+		"the loser must get back the winner's resolution, not its own")
 
 	got, ok, err := store.Get(ctx, key)
 	require.NoError(t, err)
 	require.True(t, ok)
-	assert.Equal(t, "second", got.ClientID)
+	assert.Equal(t, "first", got.ClientID, "the store must keep the first-claimed entry, not overwrite it")
 }
 
 // TestStorageBackedStore_Put_RejectsNilResolution pins the
@@ -189,7 +210,7 @@ func TestStorageBackedStore_Put_RejectsNilResolution(t *testing.T) {
 	ctx := context.Background()
 	key := Key{Issuer: "https://idp.example.com", RedirectURI: "https://x.example.com/cb"}
 
-	err := store.Put(ctx, key, nil)
+	_, err := store.PutIfAbsent(ctx, key, nil)
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "must not be nil")
 
@@ -211,11 +232,12 @@ func TestStorageBackedStore_GetReturnsDefensiveCopy(t *testing.T) {
 		RedirectURI: "https://x.example.com/cb",
 		ScopesHash:  storage.ScopesHash([]string{"openid"}),
 	}
-	require.NoError(t, store.Put(ctx, key, &Resolution{
+	_, err := store.PutIfAbsent(ctx, key, &Resolution{
 		ClientID:              "orig",
 		AuthorizationEndpoint: "https://idp.example.com/authorize",
 		TokenEndpoint:         "https://idp.example.com/token",
-	}))
+	})
+	require.NoError(t, err)
 
 	got, ok, err := store.Get(ctx, key)
 	require.NoError(t, err)
@@ -289,14 +311,14 @@ func TestStorageBackedStore_ConcurrentAccess(t *testing.T) {
 					CreatedAt:             time.Now(),
 				}
 				if i%2 == 0 {
-					if err := store.Put(ctx, overlappingKey(i), resolution); err != nil {
+					if _, err := store.PutIfAbsent(ctx, overlappingKey(i), resolution); err != nil {
 						atomic.AddInt32(&errCount, 1)
 					}
 					if _, _, err := store.Get(ctx, overlappingKey(i)); err != nil {
 						atomic.AddInt32(&errCount, 1)
 					}
 				} else {
-					if err := store.Put(ctx, disjointKey(worker, i), resolution); err != nil {
+					if _, err := store.PutIfAbsent(ctx, disjointKey(worker, i), resolution); err != nil {
 						atomic.AddInt32(&errCount, 1)
 					}
 					if _, _, err := store.Get(ctx, disjointKey(worker, i)); err != nil {
@@ -508,18 +530,19 @@ func TestInMemoryStore_PutGetCloseShareBackend(t *testing.T) {
 		TokenEndpoint:         "https://idp.example.com/token",
 	}
 
-	require.NoError(t, store.Put(ctx, key, resolution))
+	_, err := store.PutIfAbsent(ctx, key, resolution)
+	require.NoError(t, err)
 
 	got, ok, err := store.Get(ctx, key)
 	require.NoError(t, err)
-	require.True(t, ok, "Get must see the value Put just wrote — confirms Put and Get share a backend")
+	require.True(t, ok, "Get must see the value PutIfAbsent just wrote — confirms Put and Get share a backend")
 	assert.Equal(t, "client-abc", got.ClientID)
 }
 
 // TestInMemoryStore_PutRejectsNilResolution mirrors the contract pinned
-// for storageBackedStore.Put: a nil resolution is rejected at the
+// for storageBackedStore.PutIfAbsent: a nil resolution is rejected at the
 // adapter boundary rather than silently no-oped, so the next Get miss
-// surfaces with a debug trail. inMemoryStore implements Put directly
+// surfaces with a debug trail. inMemoryStore implements PutIfAbsent directly
 // (not via embedding) — this test guards against a delegation
 // regression that omitted the nil check.
 func TestInMemoryStore_PutRejectsNilResolution(t *testing.T) {
@@ -528,7 +551,7 @@ func TestInMemoryStore_PutRejectsNilResolution(t *testing.T) {
 	store := NewInMemoryStore()
 	t.Cleanup(func() { _ = store.Close() })
 
-	err := store.Put(context.Background(), Key{Issuer: "https://idp.example.com"}, nil)
+	_, err := store.PutIfAbsent(context.Background(), Key{Issuer: "https://idp.example.com"}, nil)
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "resolution must not be nil")
 }

pkg/authserver/storage/memory.go

@@ -157,7 +157,7 @@ type MemoryStorage struct {
 	// dcrCredentials maps DCRKey -> DCRCredentials for RFC 7591 Dynamic Client
 	// Registration credentials. These entries come from OUTBOUND DCR: ToolHive
 	// acting as a DCR client to register itself against a configured upstream
-	// IdP (see pkg/auth/dcr/store.go, the only caller of StoreDCRCredentials).
+	// IdP (see pkg/auth/dcr/store.go, the only caller of StoreDCRCredentialsIfAbsent).
 	// This is not reachable from the inbound, unauthenticated /oauth/register
 	// handler, so unlike clients (bounded by maxClients/clientOrder because
 	// every inbound registration call mints an entry), growth here is bounded
@@ -1538,29 +1538,46 @@ func cloneDCRCredentials(c *DCRCredentials) *DCRCredentials {
 	return &cp
 }
 
-// StoreDCRCredentials persists DCR credentials for the given key.
-// The credentials are stored under their own Key field; callers must populate
-// it before calling. A defensive copy is made so subsequent caller mutations
-// do not affect persisted state.
+// StoreDCRCredentialsIfAbsent claims creds.Key for creds. The credentials
+// are stored under their own Key field; callers must populate it before
+// calling. A defensive copy is made so subsequent caller mutations do not
+// affect persisted state.
 //
-// Overwrites any existing entry for the same Key. The in-memory backend
-// applies no native TTL — DCR registrations are long-lived and bounded by
-// the operator-configured upstream count, and ClientSecretExpiresAt is
-// retained verbatim for callers to re-check on read (see the interface
-// docstring's "TTL handling" section).
+// Create-if-absent, not overwrite: a single process's dcrFlight singleflight
+// (see pkg/auth/dcr) already prevents concurrent same-key writers within
+// this process, so the check below is not closing a live race here — it is
+// contract symmetry with RedisStorage.StoreDCRCredentialsIfAbsent, which
+// DOES need it to prevent two replicas from independently registering RFC
+// 7591 clients for the same Key and racing on which write wins.
+//
+// The in-memory backend has no native TTL, so "absent" cannot be a plain
+// map-presence check: the Redis backend's rows self-evict via TTL derived
+// from ClientSecretExpiresAt, so a claim there naturally succeeds again once
+// the old row expires. To keep behaviour symmetric, an existing entry whose
+// ClientSecretExpiresAt is non-zero and already in the past is treated as
+// absent — the claim proceeds and overwrites it — so a secret's expiry does
+// not permanently block re-registration on this backend the way a naive
+// presence check would.
 //
 // Validation is delegated to validateDCRCredentialsForStore so the rejection
 // set stays in sync with sibling backends.
-func (s *MemoryStorage) StoreDCRCredentials(_ context.Context, creds *DCRCredentials) error {
+func (s *MemoryStorage) StoreDCRCredentialsIfAbsent(_ context.Context, creds *DCRCredentials) (*DCRCredentials, error) {
 	if err := validateDCRCredentialsForStore(creds); err != nil {
-		return err
+		return nil, err
 	}
 
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
+	if existing, ok := s.dcrCredentials[creds.Key]; ok {
+		expired := !existing.ClientSecretExpiresAt.IsZero() && time.Now().After(existing.ClientSecretExpiresAt)
+		if !expired {
+			return cloneDCRCredentials(existing), nil
+		}
+	}
+
 	s.dcrCredentials[creds.Key] = cloneDCRCredentials(creds)
-	return nil
+	return cloneDCRCredentials(creds), nil
 }
 
 // GetDCRCredentials retrieves DCR credentials by key.