Fix CIMD client renewal broken by create-only registration

Rebasing this stack onto main picked up CIMD's write-through client
persistence (main), which relies on RegisterClient acting as an
upsert to renew a resolved client's row on every document re-fetch.
This stack's own registration hardening made RegisterClient
create-only for DCR-issued clients, so every renewal after the first
fetch for a given CIMD client_id would silently fail, leaving stale
client data (and, absent the token-exchange-triggered RenewClientTTL
path, a stale TTL) in storage.

Add UpsertDCRIssuedClient, a narrow fourth ClientRegistry operation
distinct from both RegisterClient (unauthenticated DCR, stays
create-only) and ReconcileConfiguredClient (fingerprint-locked, the
wrong shape since a CIMD document can legitimately change between
fetches). It creates the row if absent, replaces and renews it only
when the existing row is itself DCR-issued, and refuses with
ErrAlreadyExists otherwise -- protecting a configured or SPIFFE
client from being clobbered. Wire CIMDStorageDecorator.fetch to call
it instead of RegisterClient, and give SPIFFEStorageDecorator the
same reserved-ID guard its other overrides already enforce.

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

Author: jhrozek

pkg/authserver/storage/cimd_decorator.go

@@ -276,10 +276,14 @@ func (d *CIMDStorageDecorator) fetch(ctx context.Context, id string) (fosite.Cli
 	// issue #6187). The client carries the DCR-issued marker (above) so the
 	// row gets the same anti-bloat TTL as DCR registrations: unauthenticated
 	// /oauth/authorize traffic can mint these rows, so they must never be
-	// permanent. Re-persisting on every fresh fetch keeps the snapshot
-	// current with the document. A persistence failure only degrades
-	// token-path rehydration, so it must not fail the resolution itself.
-	if err := d.RegisterClient(ctx, client); err != nil {
+	// permanent. UpsertDCRIssuedClient (not RegisterClient) is used because
+	// RegisterClient is strictly create-only and would fail every fetch after
+	// the first for the same client_id -- this method re-persists on every
+	// fresh fetch, keeping the stored snapshot current with the document,
+	// while still refusing to clobber a configured/SPIFFE-reconciled client
+	// at the same ID. A persistence failure only degrades token-path
+	// rehydration, so it must not fail the resolution itself.
+	if err := d.UpsertDCRIssuedClient(ctx, client); err != nil {
 		slog.WarnContext(ctx, "failed to persist resolved CIMD client",
 			"client_id", id, "error", err)
 	}

pkg/authserver/storage/cimd_decorator_test.go

@@ -950,14 +950,15 @@ func TestCIMDStorageDecorator_PersistsLoopbackResolvedClient(t *testing.T) {
 	require.NoError(t, err)
 }
 
-// registerFailingStorage wraps a Storage and fails every RegisterClient call,
-// for testing that a write-through persistence failure does not fail the
+// registerFailingStorage wraps a Storage and fails every UpsertDCRIssuedClient
+// call -- the write-through persistence path fetch() actually calls -- for
+// testing that a write-through persistence failure does not fail the
 // resolution itself.
 type registerFailingStorage struct {
 	Storage
 }
 
-func (*registerFailingStorage) RegisterClient(context.Context, fosite.Client) error {
+func (*registerFailingStorage) UpsertDCRIssuedClient(context.Context, fosite.Client) error {
 	return errors.New("register failed")
 }
 
@@ -977,3 +978,50 @@ func TestCIMDStorageDecorator_PersistFailureDoesNotFailResolution(t *testing.T)
 	require.NoError(t, err, "a write-through persistence failure must not fail the resolution")
 	assert.NotNil(t, client)
 }
+
+// TestCIMDStorageDecorator_RepeatFetchRenewsPersistedClient closes the gap
+// that let the RegisterClient/UpsertDCRIssuedClient bug go undetected:
+// nothing previously asserted the persisted row's state after a repeat
+// fetch of the same client_id. Two fetches of the same CIMD document, whose
+// served content changes between them, must both succeed, and the second
+// fetch's write-through must actually replace the stored row's data rather
+// than silently failing with ErrAlreadyExists (as it would have with the old
+// RegisterClient-only call, whose failure fetch() only logs).
+func TestCIMDStorageDecorator_RepeatFetchRenewsPersistedClient(t *testing.T) {
+	t.Parallel()
+
+	var callCount atomic.Int32
+	srv := serveCIMDDocWithFields(t, func(doc *cimd.ClientMetadataDocument) {
+		// Change the served document between the first and second fetch so a
+		// real re-persist is distinguishable from a no-op.
+		if callCount.Add(1) == 1 {
+			doc.RedirectURIs = []string{"https://example.com/callback-v1"}
+		} else {
+			doc.RedirectURIs = []string{"https://example.com/callback-v2"}
+		}
+	})
+	base := newTestBase(t)
+	dec := newEnabledDecorator(t, base, 10, time.Minute)
+	id := srv.URL + "/meta.json"
+	ctx := context.Background()
+
+	// Call fetch() directly (bypassing the in-process LRU cache) to force two
+	// real write-through persistence attempts, exactly as two independent
+	// process replicas resolving the same client_id would.
+	first, err := dec.fetch(ctx, id)
+	require.NoError(t, err)
+	assert.Equal(t, []string{"https://example.com/callback-v1"}, first.GetRedirectURIs())
+
+	stored, err := base.GetClient(ctx, id)
+	require.NoError(t, err, "first fetch must persist the row")
+	assert.Equal(t, []string{"https://example.com/callback-v1"}, stored.GetRedirectURIs())
+
+	second, err := dec.fetch(ctx, id)
+	require.NoError(t, err, "second fetch for the same client_id must not fail")
+	assert.Equal(t, []string{"https://example.com/callback-v2"}, second.GetRedirectURIs())
+
+	stored, err = base.GetClient(ctx, id)
+	require.NoError(t, err)
+	assert.Equal(t, []string{"https://example.com/callback-v2"}, stored.GetRedirectURIs(),
+		"the persisted row must be renewed with the newly-fetched document, not left stale")
+}

pkg/authserver/storage/memory.go

@@ -476,6 +476,38 @@ func (s *MemoryStorage) RegisterClient(_ context.Context, client fosite.Client)
 	return s.insertClientLocked(id, client)
 }
 
+// UpsertDCRIssuedClient creates or replaces a DCR-issued client at
+// client.GetID(). Unlike RegisterClient, an existing row is replaced (and its
+// eviction position refreshed) rather than rejected -- but only when the
+// existing row is itself DCR-issued; a configured/SPIFFE-reconciled row at
+// the same ID is protected and refuses with ErrAlreadyExists. See the
+// ClientRegistry interface doc for the full contract.
+func (s *MemoryStorage) UpsertDCRIssuedClient(_ context.Context, client fosite.Client) error {
+	if !registration.DCRIssued(client) {
+		return fmt.Errorf("client %q must carry the DCR-issued marker to use UpsertDCRIssuedClient", client.GetID())
+	}
+	id := client.GetID()
+	if err := ValidateRegisterableClientID(id); err != nil {
+		return err
+	}
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	existing, exists := s.clients[id]
+	if !exists {
+		return s.insertClientLocked(id, client)
+	}
+	if !registration.DCRIssued(existing) {
+		return fmt.Errorf("%w: client %q", ErrAlreadyExists, id)
+	}
+
+	s.clientOrder = slices.DeleteFunc(s.clientOrder, func(e clientOrderEntry) bool { return e.id == id })
+	s.clientOrder = append(s.clientOrder, clientOrderEntry{id: id, touchedAt: time.Now()})
+	s.clients[id] = client
+	return nil
+}
+
 // ReconcileConfiguredClient applies an operator-declared client: creates it
 // if absent, idempotently replaces a matching-fingerprint configured client
 // (the restart-with-unchanged-config case), or refuses with ErrAlreadyExists

pkg/authserver/storage/memory_test.go

@@ -380,6 +380,99 @@ func TestMemoryStorage_ReconcileConfiguredClient(t *testing.T) {
 	})
 }
 
+// TestMemoryStorage_UpsertDCRIssuedClient covers the create/replace/reject
+// matrix UpsertDCRIssuedClient must implement: create when absent, replace
+// (and renew the eviction position) when the existing record is itself
+// DCR-issued, refuse with ErrAlreadyExists when the existing record is NOT
+// DCR-issued (the critical protection: a configured/SPIFFE client must never
+// be clobbered by this path), and refuse when the incoming client itself does
+// not carry the DCR-issued marker (misuse guard).
+func TestMemoryStorage_UpsertDCRIssuedClient(t *testing.T) {
+	t.Parallel()
+
+	t.Run("creates when absent", func(t *testing.T) {
+		t.Parallel()
+		ctx := t.Context()
+		s := NewMemoryStorage()
+		defer s.Close()
+
+		client := dcrClient(t, "cimd-client")
+		require.NoError(t, s.UpsertDCRIssuedClient(ctx, client))
+
+		retrieved, err := s.GetClient(ctx, "cimd-client")
+		require.NoError(t, err)
+		assert.Equal(t, client, retrieved)
+	})
+
+	t.Run("replaces and renews when existing row is DCR-issued", func(t *testing.T) {
+		t.Parallel()
+		ctx := t.Context()
+		// MinClientAge(0) isolates this test to the replace/renew behaviour;
+		// the age-floor grace window has its own tests above.
+		s := NewMemoryStorage(WithMaxClients(2), WithMinClientAge(0))
+		defer s.Close()
+
+		first := dcrClient(t, "cimd-client")
+		require.NoError(t, s.UpsertDCRIssuedClient(ctx, first))
+		require.NoError(t, s.RegisterClient(ctx, dcrClient(t, "client-b")))
+
+		// A distinguishing field on the replacement proves the replace branch
+		// actually overwrote the stored data rather than being a silent no-op.
+		second, err := registration.New(registration.Config{
+			ID:                      "cimd-client",
+			TokenEndpointAuthMethod: oauthproto.TokenEndpointAuthMethodNone,
+			RedirectURIs:            []string{"https://app.example/cb-v2"},
+		})
+		require.NoError(t, err)
+		require.NoError(t, s.UpsertDCRIssuedClient(ctx, second))
+
+		retrieved, err := s.GetClient(ctx, "cimd-client")
+		require.NoError(t, err)
+		assert.Equal(t, second, retrieved, "second call must replace the stored row")
+
+		// Renewed eviction position: cimd-client was registered first (oldest)
+		// but the upsert must have moved it to the back, so overflow evicts
+		// client-b instead.
+		require.NoError(t, s.RegisterClient(ctx, dcrClient(t, "client-c")))
+		_, err = s.GetClient(ctx, "cimd-client")
+		require.NoError(t, err, "renewed client must survive eviction")
+		_, err = s.GetClient(ctx, "client-b")
+		requireNotFoundError(t, err)
+	})
+
+	t.Run("refuses to overwrite a non-DCR-issued client", func(t *testing.T) {
+		t.Parallel()
+		ctx := t.Context()
+		s := NewMemoryStorage()
+		defer s.Close()
+
+		configured := &mockClient{id: "configured", public: false}
+		require.NoError(t, s.ReconcileConfiguredClient(ctx, configured))
+
+		err := s.UpsertDCRIssuedClient(ctx, dcrClient(t, "configured"))
+		require.ErrorIs(t, err, ErrAlreadyExists)
+
+		// The original registration must be untouched.
+		retrieved, err := s.GetClient(ctx, "configured")
+		require.NoError(t, err)
+		assert.Equal(t, configured, retrieved)
+	})
+
+	t.Run("rejects a client not carrying the DCR-issued marker", func(t *testing.T) {
+		t.Parallel()
+		ctx := t.Context()
+		s := NewMemoryStorage()
+		defer s.Close()
+
+		err := s.UpsertDCRIssuedClient(ctx, &mockClient{id: "not-dcr"})
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "must carry the DCR-issued marker")
+
+		_, err = s.GetClient(ctx, "not-dcr")
+		require.ErrorIs(t, err, ErrNotFound)
+	})
+}
+
 // TestMemoryStorage_RegisterClient_Bounded pins the anti-DoS cap: the client
 // map is bounded by maxClients with oldest-first eviction among DCR-issued
 // clients only, and duplicate registrations fail without changing stored state.

pkg/authserver/storage/mocks/mock_storage.go

@@ -421,6 +421,81 @@ func (s *RedisStorage) RegisterClient(ctx context.Context, client fosite.Client)
 	return nil
 }
 
+// UpsertDCRIssuedClient creates or replaces a DCR-issued client at
+// client.GetID(). Unlike RegisterClient (create-only via SetNX), an existing
+// row is replaced -- and its TTL refreshed to DefaultDCRClientTTL -- when the
+// existing row is itself DCR-issued; it refuses with ErrAlreadyExists when the
+// existing row is NOT DCR-issued, protecting a configured/SPIFFE-reconciled
+// client from being clobbered. See the ClientRegistry interface doc for the
+// full contract.
+//
+// The read-check-write sequence for an existing key runs inside a Redis
+// WATCH/MULTI transaction, mirroring ReconcileConfiguredClient, so a
+// concurrent writer cannot interleave between the DCR-issued check and the
+// write; see maxConfiguredClientReconcileRetries for why this method retries
+// redis.TxFailedErr itself, up to the same bounded count. Unlike
+// ReconcileConfiguredClient (ttl=0, permanent), the write here always uses
+// DefaultDCRClientTTL: this row is TTL-bounded like any other DCR-issued
+// client.
+func (s *RedisStorage) UpsertDCRIssuedClient(ctx context.Context, client fosite.Client) error {
+	if !registration.DCRIssued(client) {
+		return fmt.Errorf("client %q must carry the DCR-issued marker to use UpsertDCRIssuedClient", client.GetID())
+	}
+	if err := ValidateRegisterableClientID(client.GetID()); err != nil {
+		return err
+	}
+
+	key := redisKey(s.keyPrefix, KeyTypeClient, client.GetID())
+	stored := buildStoredClient(client)
+	data, err := json.Marshal(stored) //nolint:gosec // G117 - internal Redis storage serialization, not exposed to users
+	if err != nil {
+		return fmt.Errorf("failed to marshal client: %w", err)
+	}
+
+	setPipelined := func(tx *redis.Tx) error {
+		_, err := tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
+			pipe.Set(ctx, key, data, DefaultDCRClientTTL)
+			return nil
+		})
+		return err
+	}
+
+	txFn := func(tx *redis.Tx) error {
+		existingData, getErr := tx.Get(ctx, key).Bytes()
+		if errors.Is(getErr, redis.Nil) {
+			return setPipelined(tx)
+		}
+		if getErr != nil {
+			return fmt.Errorf("failed to get existing client: %w", getErr)
+		}
+
+		ttl, ttlErr := tx.TTL(ctx, key).Result()
+		if ttlErr != nil {
+			return fmt.Errorf("failed to get existing client TTL: %w", ttlErr)
+		}
+
+		var existingStored storedClient
+		if unmarshalErr := json.Unmarshal(existingData, &existingStored); unmarshalErr != nil {
+			return fmt.Errorf("failed to unmarshal existing client: %w", unmarshalErr)
+		}
+
+		if !registration.DCRIssued(clientFromStored(existingStored, ttl >= 0)) {
+			return fmt.Errorf("%w: client %q", ErrAlreadyExists, client.GetID())
+		}
+
+		return setPipelined(tx)
+	}
+
+	var watchErr error
+	for attempt := 0; attempt < maxConfiguredClientReconcileRetries; attempt++ {
+		watchErr = s.client.Watch(ctx, txFn, key)
+		if !errors.Is(watchErr, redis.TxFailedErr) {
+			return watchErr
+		}
+	}
+	return watchErr
+}
+
 // storedClientFingerprintsEqual mirrors clientFingerprintsEqual but compares
 // the raw serialized fields directly rather than through the fosite.Client
 // interface returned by clientFromStored. This matters because