From 424b6570b394b5352ae9afb331b7c8978533148e Mon Sep 17 00:00:00 2001
From: Trey <tgrunnagle@gmail.com>
Date: Sun, 5 Jul 2026 11:44:31 -0700
Subject: [PATCH 1/6] Filter upstream auth chain via callback hook

The embedded OAuth authorization server walks every configured upstream
on every authorization, prompting for and storing tokens that a given
authorization may not need. This widens the stored-token surface and
forces needless upstream prompts.

Add an optional UpstreamFilter hook (WithUpstreamFilter), consulted once
in the callback handler after the first upstream resolves, that narrows
the remaining chain to a subset of the configured upstreams:

- The first upstream is always required and is never passed to nor
  removable by the filter.
- The kept set is computed once and carried in PendingAuthorization
  (ChainUpstreams) so the filter is not re-run per leg.
- nextMissingUpstream and the cross-leg identity check now operate on
  the effective chain rather than the raw config.
- A filter error fails the authorization with a server error; it never
  silently falls back to walking every upstream.

With no filter configured, behaviour is unchanged: the handler walks all
configured upstreams as before. SingleLeg flows are unaffected and the
upstream-token storage key contracts are unchanged.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 pkg/authserver/server/handlers/callback.go    |  81 +++++---
 .../server/handlers/callback_test.go          | 183 ++++++++++++++++++
 pkg/authserver/server/handlers/handler.go     |  95 ++++++++-
 .../server/handlers/handler_chain_test.go     | 159 ++++++++++++++-
 pkg/authserver/storage/memory.go              |   2 +
 pkg/authserver/storage/memory_test.go         |   4 +-
 pkg/authserver/storage/redis.go               |   3 +
 pkg/authserver/storage/redis_test.go          |   4 +-
 pkg/authserver/storage/types.go               |  10 +
 9 files changed, 506 insertions(+), 35 deletions(-)

diff --git a/pkg/authserver/server/handlers/callback.go b/pkg/authserver/server/handlers/callback.go
index 55cf24d875..ab15dba91e 100644
--- a/pkg/authserver/server/handlers/callback.go
+++ b/pkg/authserver/server/handlers/callback.go
@@ -433,7 +433,24 @@ func (h *Handler) continueChainOrComplete(
 		return
 	}
 
-	nextProvider, err := h.nextMissingUpstream(ctx, sessionID)
+	// Resolve the effective chain of upstreams to walk. The first leg (which
+	// carries no computed chain) consults the optional filter once; every
+	// subsequent leg reuses the chain carried in the pending authorization so the
+	// filter is not re-run per leg. A legacy pending predating ChainUpstreams also
+	// lands here and is safely recomputed.
+	chain := pending.ChainUpstreams
+	if len(chain) == 0 {
+		computed, err := h.computeChain(ctx)
+		if err != nil {
+			slog.Error("failed to compute upstream chain", "error", err)
+			_ = h.storage.DeleteUpstreamTokens(ctx, sessionID)
+			h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to determine authorization chain"))
+			return
+		}
+		chain = computed
+	}
+
+	nextProvider, err := h.nextMissingUpstream(ctx, sessionID, chain)
 	if err != nil {
 		slog.Error("failed to determine next upstream", "error", err)
 		_ = h.storage.DeleteUpstreamTokens(ctx, sessionID)
@@ -442,28 +459,11 @@ func (h *Handler) continueChainOrComplete(
 	}
 
 	if nextProvider == "" {
-		// Defense-in-depth: verify identity consistency across chain legs.
-		// The subject was resolved from the first leg's upstream and carried through
-		// PendingAuthorization. Cross-check it against the stored upstream tokens.
-		if len(h.upstreams) > 1 {
-			allTokens, err := h.storage.GetAllUpstreamTokens(ctx, sessionID)
-			if err != nil {
-				slog.Error("failed to verify identity consistency", "error", err)
-				_ = h.storage.DeleteUpstreamTokens(ctx, sessionID)
-				h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to verify identity consistency"))
-				return
-			}
-			firstProvider := h.upstreams[0].Name
-			if firstTokens, ok := allTokens[firstProvider]; ok && firstTokens.UserID != subject {
-				slog.Error("identity mismatch between chain state and stored tokens",
-					"expected", subject,
-					"got", firstTokens.UserID,
-					"provider", firstProvider,
-				)
-				_ = h.storage.DeleteUpstreamTokens(ctx, sessionID)
-				h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("identity verification failed"))
-				return
-			}
+		if err := h.verifyChainIdentity(ctx, sessionID, chain, subject); err != nil {
+			slog.Error("chain identity verification failed", "error", err)
+			_ = h.storage.DeleteUpstreamTokens(ctx, sessionID)
+			h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("identity verification failed"))
+			return
 		}
 
 		// All upstreams satisfied — issue authorization code
@@ -495,6 +495,9 @@ func (h *Handler) continueChainOrComplete(
 		// Chain state
 		UpstreamProviderName: nextProvider,
 		SessionID:            sessionID,
+		// Carry the effective chain forward so the filter is computed once, on the
+		// first leg, and reused for every subsequent leg.
+		ChainUpstreams: chain,
 		// Carry resolved identity from first leg
 		ResolvedUserID:    subject,
 		ResolvedUserName:  name,
@@ -533,3 +536,35 @@ func (h *Handler) continueChainOrComplete(
 
 	http.Redirect(w, req, nextURL, http.StatusFound)
 }
+
+// verifyChainIdentity performs the defense-in-depth cross-leg identity check once
+// every leg of the effective chain is satisfied.
+//
+// subject MUST be the canonical ToolHive user ID resolved from the first leg's
+// upstream via that provider's configured subject-claim mapping (the OIDC
+// SubjectClaim, or the "sub" claim by default) and carried forward unchanged
+// through PendingAuthorization.ResolvedUserID. The caller resolves it exactly once,
+// on the first leg, from the claim-mapped upstream subject. This cross-checks it
+// against firstTokens.UserID — the same canonical ID persisted when the first leg
+// stored its tokens — so a chain whose legs disagree on the user is rejected.
+//
+// It gates on the effective chain rather than the raw config: chain[0] is always
+// the first (required) upstream, so a chain the filter narrowed to just that
+// upstream has no cross-legs to reconcile and the check is a no-op. Returns a
+// non-nil error when the storage lookup fails or an identity mismatch is detected;
+// the caller maps either to a server error.
+func (h *Handler) verifyChainIdentity(ctx context.Context, sessionID string, chain []string, subject string) error {
+	if len(chain) <= 1 {
+		return nil
+	}
+	allTokens, err := h.storage.GetAllUpstreamTokens(ctx, sessionID)
+	if err != nil {
+		return fmt.Errorf("failed to load upstream tokens for identity check: %w", err)
+	}
+	firstProvider := chain[0]
+	if firstTokens, ok := allTokens[firstProvider]; ok && firstTokens.UserID != subject {
+		return fmt.Errorf("identity mismatch on provider %q: chain subject %q != stored user %q",
+			firstProvider, subject, firstTokens.UserID)
+	}
+	return nil
+}
diff --git a/pkg/authserver/server/handlers/callback_test.go b/pkg/authserver/server/handlers/callback_test.go
index 484db3fded..2ad2627f2f 100644
--- a/pkg/authserver/server/handlers/callback_test.go
+++ b/pkg/authserver/server/handlers/callback_test.go
@@ -515,6 +515,189 @@ func TestCallbackHandler_TwoUpstreams_SecondLeg_IssuesCode(t *testing.T) {
 	assert.False(t, ok, "second leg pending should be consumed")
 }
 
+func TestCallbackHandler_TwoUpstreams_FilterDropsSecondLeg_IssuesCode(t *testing.T) {
+	t.Parallel()
+	// Filter keeps nothing, so the chain narrows to just the first upstream.
+	filter := &stubUpstreamFilter{keep: []string{}}
+	handler, storState, provider1, _ := multiUpstreamTestSetup(t, WithUpstreamFilter(filter))
+
+	sessionID := "chain-session-filter-drop"
+	firstLegState := "filter-drop-first-leg-state"
+	pending := &storage.PendingAuthorization{
+		ClientID:             testAuthClientID,
+		RedirectURI:          testAuthRedirectURI,
+		State:                "client-state-drop",
+		PKCEChallenge:        "client-challenge",
+		PKCEMethod:           "S256",
+		Scopes:               []string{"openid", "profile"},
+		InternalState:        firstLegState,
+		UpstreamPKCEVerifier: "filter-drop-verifier-1234567890123456789012",
+		UpstreamNonce:        "filter-drop-nonce",
+		UpstreamProviderName: "provider-1",
+		SessionID:            sessionID,
+		CreatedAt:            time.Now(),
+	}
+	storState.pendingAuths[firstLegState] = pending
+
+	req := httptest.NewRequest(http.MethodGet, "/oauth/callback?code=provider1-code&state="+firstLegState, nil)
+	rec := httptest.NewRecorder()
+	handler.CallbackHandler(rec, req)
+
+	// The filtered-out second leg means the chain completes at the first upstream:
+	// issue the authorization code (303) rather than redirect onward (302).
+	assert.Equal(t, http.StatusSeeOther, rec.Code, "filtered chain should complete at the first upstream")
+	location := rec.Header().Get("Location")
+	assert.Contains(t, location, testAuthRedirectURI, "redirect should be to the client's redirect_uri")
+	assert.Contains(t, location, "code=", "redirect should include an authorization code")
+	assert.NotContains(t, location, "error=", "redirect should not contain an error")
+
+	// The filter was consulted exactly once, with the non-first upstream names.
+	assert.Equal(t, 1, filter.calls, "filter should be consulted once on the first leg")
+	assert.Equal(t, []string{"provider-2"}, filter.capturedArgs, "filter receives non-first upstreams")
+
+	// provider-1's code was exchanged and no second-leg pending was created.
+	assert.Equal(t, "provider1-code", provider1.capturedCode)
+	for state, p := range storState.pendingAuths {
+		assert.NotEqualf(t, "provider-2", p.UpstreamProviderName,
+			"no second-leg pending should be created (state %q)", state)
+	}
+}
+
+func TestCallbackHandler_TwoUpstreams_FilterKeepsSecondLeg_CarriesChain(t *testing.T) {
+	t.Parallel()
+	filter := &stubUpstreamFilter{keep: []string{"provider-2"}}
+	handler, storState, _, _ := multiUpstreamTestSetup(t, WithUpstreamFilter(filter))
+
+	sessionID := "chain-session-filter-keep"
+	firstLegState := "filter-keep-first-leg-state"
+	pending := &storage.PendingAuthorization{
+		ClientID:             testAuthClientID,
+		RedirectURI:          testAuthRedirectURI,
+		State:                "client-state-keep",
+		PKCEChallenge:        "client-challenge",
+		PKCEMethod:           "S256",
+		Scopes:               []string{"openid", "profile"},
+		InternalState:        firstLegState,
+		UpstreamPKCEVerifier: "filter-keep-verifier-1234567890123456789012",
+		UpstreamNonce:        "filter-keep-nonce",
+		UpstreamProviderName: "provider-1",
+		SessionID:            sessionID,
+		CreatedAt:            time.Now(),
+	}
+	storState.pendingAuths[firstLegState] = pending
+
+	req := httptest.NewRequest(http.MethodGet, "/oauth/callback?code=provider1-code&state="+firstLegState, nil)
+	rec := httptest.NewRecorder()
+	handler.CallbackHandler(rec, req)
+
+	// Kept provider-2, so the chain continues onward.
+	assert.Equal(t, http.StatusFound, rec.Code, "kept second leg should redirect onward")
+	assert.Contains(t, rec.Header().Get("Location"), "https://idp2.example.com/authorize")
+	assert.Equal(t, 1, filter.calls, "filter should be consulted once on the first leg")
+
+	// The effective chain must be carried into the next leg's pending so the
+	// filter is not re-run per leg.
+	var nextPending *storage.PendingAuthorization
+	for state, p := range storState.pendingAuths {
+		if state != firstLegState && p.UpstreamProviderName == "provider-2" {
+			nextPending = p
+			break
+		}
+	}
+	require.NotNil(t, nextPending, "a second-leg pending should exist")
+	assert.Equal(t, []string{"provider-1", "provider-2"}, nextPending.ChainUpstreams,
+		"the computed chain should be carried forward")
+}
+
+func TestCallbackHandler_FilterError_FailsAuthorization(t *testing.T) {
+	t.Parallel()
+	filter := &stubUpstreamFilter{err: errors.New("filter unavailable")}
+	handler, storState, _, _ := multiUpstreamTestSetup(t, WithUpstreamFilter(filter))
+
+	sessionID := "chain-session-filter-error"
+	firstLegState := "filter-error-first-leg-state"
+	pending := &storage.PendingAuthorization{
+		ClientID:             testAuthClientID,
+		RedirectURI:          testAuthRedirectURI,
+		State:                "client-state-error",
+		PKCEChallenge:        "client-challenge",
+		PKCEMethod:           "S256",
+		Scopes:               []string{"openid", "profile"},
+		InternalState:        firstLegState,
+		UpstreamPKCEVerifier: "filter-error-verifier-123456789012345678901",
+		UpstreamNonce:        "filter-error-nonce",
+		UpstreamProviderName: "provider-1",
+		SessionID:            sessionID,
+		CreatedAt:            time.Now(),
+	}
+	storState.pendingAuths[firstLegState] = pending
+
+	req := httptest.NewRequest(http.MethodGet, "/oauth/callback?code=provider1-code&state="+firstLegState, nil)
+	rec := httptest.NewRecorder()
+	handler.CallbackHandler(rec, req)
+
+	// A filter error fails the authorization cleanly with a server error — no
+	// silent fallback to walking every upstream.
+	assert.Equal(t, http.StatusSeeOther, rec.Code, "filter error should produce a fosite error redirect")
+	location := rec.Header().Get("Location")
+	assert.Contains(t, location, "error=server_error", "should surface a server error")
+	assert.Contains(t, location, "state=client-state-error", "should preserve client state")
+	assert.Equal(t, 1, filter.calls)
+
+	// provider-1's already-stored tokens are cleaned up.
+	for key := range storState.upstreamTokens {
+		assert.Failf(t, "upstream tokens should be cleaned up", "found leftover token %q", key)
+	}
+}
+
+func TestCallbackHandler_SecondLeg_ReusesChain_DoesNotReRunFilter(t *testing.T) {
+	t.Parallel()
+	// A filter that errors if consulted — proving the second leg reuses the chain
+	// from the pending authorization rather than re-running the filter.
+	filter := &stubUpstreamFilter{err: errors.New("must not be called on a subsequent leg")}
+	handler, storState, _, provider2 := multiUpstreamTestSetup(t, WithUpstreamFilter(filter))
+
+	sessionID := "chain-session-reuse"
+	key1 := sessionID + ":provider-1"
+	storState.upstreamTokens[key1] = &storage.UpstreamTokens{
+		ProviderID:  "provider-1",
+		AccessToken: "provider1-access-token",
+		ExpiresAt:   time.Now().Add(time.Hour),
+		ClientID:    testAuthClientID,
+		UserID:      "resolved-user-id-from-leg1",
+	}
+
+	secondLegState := "reuse-second-leg-state"
+	pending := &storage.PendingAuthorization{
+		ClientID:             testAuthClientID,
+		RedirectURI:          testAuthRedirectURI,
+		State:                "client-state-reuse",
+		PKCEChallenge:        "client-challenge",
+		PKCEMethod:           "S256",
+		Scopes:               []string{"openid", "profile"},
+		InternalState:        secondLegState,
+		UpstreamPKCEVerifier: "reuse-verifier-98765432109876543210987654",
+		UpstreamNonce:        "reuse-nonce",
+		UpstreamProviderName: "provider-2",
+		SessionID:            sessionID,
+		ChainUpstreams:       []string{"provider-1", "provider-2"}, // computed on the first leg
+		ResolvedUserID:       "resolved-user-id-from-leg1",
+		ResolvedUserName:     "First Leg User",
+		ResolvedUserEmail:    "firstleg@example.com",
+		CreatedAt:            time.Now(),
+	}
+	storState.pendingAuths[secondLegState] = pending
+
+	req := httptest.NewRequest(http.MethodGet, "/oauth/callback?code=provider2-code&state="+secondLegState, nil)
+	rec := httptest.NewRecorder()
+	handler.CallbackHandler(rec, req)
+
+	assert.Equal(t, http.StatusSeeOther, rec.Code, "second leg should issue the authorization code")
+	assert.Contains(t, rec.Header().Get("Location"), "code=")
+	assert.Equal(t, 0, filter.calls, "filter must not be re-run on a subsequent leg")
+	assert.Equal(t, "provider2-code", provider2.capturedCode)
+}
+
 func TestCallbackHandler_SingleLeg_IssuesCodeWithoutChaining(t *testing.T) {
 	t.Parallel()
 	handler, storState, provider1, _ := multiUpstreamTestSetup(t)
diff --git a/pkg/authserver/server/handlers/handler.go b/pkg/authserver/server/handlers/handler.go
index b971554b02..ace69247cd 100644
--- a/pkg/authserver/server/handlers/handler.go
+++ b/pkg/authserver/server/handlers/handler.go
@@ -47,6 +47,26 @@ type Handler struct {
 	// expired upstream leg during chain evaluation instead of re-prompting. Nil
 	// when no refresher is configured; an expired leg is then treated as missing.
 	refresher storage.UpstreamTokenRefresher
+	// filter, when set, narrows the authorization chain to a subset of the
+	// configured upstreams once the first leg resolves. Nil when no filter is
+	// configured; the chain then walks all configured upstreams as before.
+	filter UpstreamFilter
+}
+
+// UpstreamFilter narrows the authorization chain to a subset of the configured
+// upstreams. It is consulted once in the callback handler, after the first
+// upstream (upstreams[0]) resolves. The first upstream is always required: it is
+// never passed to the filter and cannot be removed by it.
+//
+// FilterUpstreams receives the request context of the first leg's callback and
+// the names of the non-first configured upstreams, in configured order. It
+// returns the subset to keep. The handler preserves configured order and ignores
+// any returned name that is not one of the non-first configured upstreams, so the
+// filter cannot reorder the chain or introduce unknown providers. A returned
+// error fails the authorization with a server error — the handler never falls
+// back to walking every upstream on error.
+type UpstreamFilter interface {
+	FilterUpstreams(ctx context.Context, configured []string) ([]string, error)
 }
 
 // Option configures optional Handler behavior at construction time.
@@ -62,6 +82,16 @@ func WithUpstreamRefresher(r storage.UpstreamTokenRefresher) Option {
 	}
 }
 
+// WithUpstreamFilter injects a filter that narrows the authorization chain to a
+// subset of the configured upstreams once the first leg resolves. When unset, the
+// handler walks all configured upstreams — the behavior before this option
+// existed. See UpstreamFilter for the contract.
+func WithUpstreamFilter(f UpstreamFilter) Option {
+	return func(h *Handler) {
+		h.filter = f
+	}
+}
+
 // NewHandler creates a new Handler with the given dependencies.
 // upstreams defines the ordered sequence of upstream providers consulted
 // during multi-upstream authorization flows (e.g., sequential token acquisition).
@@ -141,7 +171,10 @@ func (h *Handler) WellKnownRoutes(r chi.Router) {
 
 // nextMissingUpstream returns the name of the next upstream provider in the
 // authorization chain that the user must (re-)authenticate with for this session.
-// Returns empty string if all upstreams are satisfied, or an error if the storage
+// It walks the provided chain — the effective, possibly filtered, ordered set of
+// upstream names for this authorization (see computeChain) — rather than the raw
+// configured list, so a leg the filter dropped is never prompted for. Returns
+// empty string if all legs in the chain are satisfied, or an error if the storage
 // lookup fails.
 //
 // A leg is satisfied when it has a stored token that is live (or asserts no
@@ -151,29 +184,77 @@ func (h *Handler) WellKnownRoutes(r chi.Router) {
 // succeeds. If refresh is impossible or fails, the leg is reported as missing so
 // the user is re-prompted up front, rather than the stale token surfacing as a
 // runtime auth error later at MCP-request token-swap time.
-func (h *Handler) nextMissingUpstream(ctx context.Context, sessionID string) (string, error) {
+func (h *Handler) nextMissingUpstream(ctx context.Context, sessionID string, chain []string) (string, error) {
 	stored, err := h.storage.GetAllUpstreamTokens(ctx, sessionID)
 	if err != nil {
 		return "", fmt.Errorf("failed to check upstream token state: %w", err)
 	}
-	for _, u := range h.upstreams {
-		tokens, ok := stored[u.Name]
+	for _, name := range chain {
+		tokens, ok := stored[name]
 		if !ok || tokens == nil {
 			// No token stored for this leg — prompt.
-			return u.Name, nil
+			return name, nil
 		}
 		// A live token (or one with no asserted expiry) satisfies the leg.
 		if tokens.ExpiresAt.IsZero() || !tokens.IsExpired(time.Now()) {
 			continue
 		}
 		// Expired — attempt a transparent refresh; prompt now if it can't be done.
-		if !h.refreshExpiredLeg(ctx, sessionID, u.Name, tokens) {
-			return u.Name, nil
+		if !h.refreshExpiredLeg(ctx, sessionID, name, tokens) {
+			return name, nil
 		}
 	}
 	return "", nil
 }
 
+// computeChain returns the ordered, effective set of upstream names this
+// authorization must walk. The first configured upstream always leads the chain
+// and is never filtered out. When no filter is configured, the chain is the full
+// configured list in order (the behavior before the filter hook existed). When a
+// filter is configured, it is consulted with the names of the non-first upstreams
+// (in configured order) and the chain becomes the first upstream plus the kept
+// subset. Configured order is preserved and any returned name that is not a
+// non-first configured upstream is ignored, so the filter can only narrow — never
+// reorder or extend — the chain.
+//
+// A filter error is returned to the caller so the authorization fails cleanly; it
+// never silently falls back to walking every upstream.
+func (h *Handler) computeChain(ctx context.Context) ([]string, error) {
+	chain := []string{h.upstreams[0].Name}
+	rest := h.upstreams[1:]
+	if len(rest) == 0 {
+		return chain, nil
+	}
+
+	restNames := make([]string, len(rest))
+	for i := range rest {
+		restNames[i] = rest[i].Name
+	}
+
+	if h.filter == nil {
+		return append(chain, restNames...), nil
+	}
+
+	keep, err := h.filter.FilterUpstreams(ctx, restNames)
+	if err != nil {
+		return nil, fmt.Errorf("upstream filter failed: %w", err)
+	}
+
+	keepSet := make(map[string]struct{}, len(keep))
+	for _, name := range keep {
+		keepSet[name] = struct{}{}
+	}
+	// Iterate configured order (not the filter's return order) so the chain
+	// preserves the operator-defined sequence and silently drops any name the
+	// filter returned that is not a non-first configured upstream.
+	for i := range rest {
+		if _, ok := keepSet[rest[i].Name]; ok {
+			chain = append(chain, rest[i].Name)
+		}
+	}
+	return chain, nil
+}
+
 // refreshExpiredLeg attempts to refresh an expired upstream token for one chain
 // leg. It returns true when the leg can be treated as authenticated (refresh
 // succeeded) and false when the user must be re-prompted: no refresher configured,
diff --git a/pkg/authserver/server/handlers/handler_chain_test.go b/pkg/authserver/server/handlers/handler_chain_test.go
index 114a8547ed..a573f8bb70 100644
--- a/pkg/authserver/server/handlers/handler_chain_test.go
+++ b/pkg/authserver/server/handlers/handler_chain_test.go
@@ -24,6 +24,159 @@ import (
 	"github.com/stacklok/toolhive/pkg/authserver/upstream"
 )
 
+// twoUpstreamChain is the effective chain for multiUpstreamTestSetup's two
+// configured providers, in configured order.
+var twoUpstreamChain = []string{"provider-1", "provider-2"}
+
+// stubUpstreamFilter is a hand-written test double for UpstreamFilter, mirroring
+// the mockIDPProvider pattern used elsewhere in this package. It records how many
+// times it was called and the names it was passed, and returns a canned keep set
+// or error.
+type stubUpstreamFilter struct {
+	keep         []string
+	err          error
+	calls        int
+	capturedArgs []string
+}
+
+func (f *stubUpstreamFilter) FilterUpstreams(_ context.Context, configured []string) ([]string, error) {
+	f.calls++
+	f.capturedArgs = configured
+	if f.err != nil {
+		return nil, f.err
+	}
+	return f.keep, nil
+}
+
+// newChainTestHandler builds a Handler over upstreams named after the given
+// slice (each backed by a throwaway OAuth2 mock provider), forwarding any
+// Options. It exists so chain-resolution tests can control the exact number and
+// order of configured upstreams.
+func newChainTestHandler(t *testing.T, names []string, opts ...Option) *Handler {
+	t.Helper()
+
+	provider, oauth2Config, stor, _ := baseTestSetup(t)
+	upstreams := make([]NamedUpstream, len(names))
+	for i, n := range names {
+		upstreams[i] = NamedUpstream{Name: n, Provider: &mockIDPProvider{providerType: upstream.ProviderTypeOAuth2}}
+	}
+	handler, err := NewHandler(provider, oauth2Config, stor, upstreams, opts...)
+	require.NoError(t, err)
+	return handler
+}
+
+func TestComputeChain(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name         string
+		upstreams    []string
+		filter       *stubUpstreamFilter
+		want         []string
+		wantErr      bool
+		wantCalls    int
+		wantFilterIn []string
+	}{
+		{
+			name:      "no filter walks all configured upstreams in order",
+			upstreams: []string{"provider-1", "provider-2", "provider-3"},
+			filter:    nil,
+			want:      []string{"provider-1", "provider-2", "provider-3"},
+		},
+		{
+			name:      "single upstream does not consult filter",
+			upstreams: []string{"provider-1"},
+			filter:    &stubUpstreamFilter{keep: []string{"anything"}},
+			want:      []string{"provider-1"},
+			wantCalls: 0,
+		},
+		{
+			name:         "filter keeps a subset, first upstream always leads",
+			upstreams:    []string{"provider-1", "provider-2", "provider-3"},
+			filter:       &stubUpstreamFilter{keep: []string{"provider-3"}},
+			want:         []string{"provider-1", "provider-3"},
+			wantCalls:    1,
+			wantFilterIn: []string{"provider-2", "provider-3"},
+		},
+		{
+			name:         "filter keeps none leaves only the first upstream",
+			upstreams:    []string{"provider-1", "provider-2", "provider-3"},
+			filter:       &stubUpstreamFilter{keep: []string{}},
+			want:         []string{"provider-1"},
+			wantCalls:    1,
+			wantFilterIn: []string{"provider-2", "provider-3"},
+		},
+		{
+			name:      "returned order is ignored, configured order is preserved",
+			upstreams: []string{"provider-1", "provider-2", "provider-3"},
+			filter:    &stubUpstreamFilter{keep: []string{"provider-3", "provider-2"}},
+			want:      []string{"provider-1", "provider-2", "provider-3"},
+			wantCalls: 1,
+		},
+		{
+			name:      "unknown and first-upstream names in keep set are ignored",
+			upstreams: []string{"provider-1", "provider-2"},
+			filter:    &stubUpstreamFilter{keep: []string{"provider-1", "does-not-exist", "provider-2"}},
+			want:      []string{"provider-1", "provider-2"},
+			wantCalls: 1,
+		},
+		{
+			name:      "filter error is propagated, no fallback to full walk",
+			upstreams: []string{"provider-1", "provider-2"},
+			filter:    &stubUpstreamFilter{err: errors.New("filter boom")},
+			wantErr:   true,
+			wantCalls: 1,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			var opts []Option
+			if tt.filter != nil {
+				opts = append(opts, WithUpstreamFilter(tt.filter))
+			}
+			handler := newChainTestHandler(t, tt.upstreams, opts...)
+
+			got, err := handler.computeChain(context.Background())
+			if tt.wantErr {
+				require.Error(t, err)
+				assert.Nil(t, got)
+			} else {
+				require.NoError(t, err)
+				assert.Equal(t, tt.want, got)
+			}
+			if tt.filter != nil {
+				assert.Equal(t, tt.wantCalls, tt.filter.calls, "filter call count")
+				if tt.wantFilterIn != nil {
+					assert.Equal(t, tt.wantFilterIn, tt.filter.capturedArgs,
+						"filter must receive non-first configured upstreams in order")
+				}
+			}
+		})
+	}
+}
+
+// TestNextMissingUpstream_RespectsChainSubset verifies that a leg absent from the
+// chain is never reported missing even when it has no stored token — the filter
+// dropped it, so the walk must not prompt for it.
+func TestNextMissingUpstream_RespectsChainSubset(t *testing.T) {
+	t.Parallel()
+
+	handler, storState, _, _ := multiUpstreamTestSetup(t)
+	// provider-1 satisfied; provider-2 has no token but is not in the chain.
+	storState.upstreamTokens["test-session:provider-1"] = &storage.UpstreamTokens{
+		ProviderID:  "provider-1",
+		AccessToken: "tok-1",
+		ExpiresAt:   time.Now().Add(time.Hour),
+	}
+
+	got, err := handler.nextMissingUpstream(context.Background(), "test-session", []string{"provider-1"})
+	require.NoError(t, err)
+	assert.Empty(t, got, "provider-2 is not in the chain, so it must not be prompted")
+}
+
 func TestNextMissingUpstream(t *testing.T) {
 	t.Parallel()
 
@@ -74,7 +227,7 @@ func TestNextMissingUpstream(t *testing.T) {
 			handler, storState, _, _ := multiUpstreamTestSetup(t)
 			tt.setupTokens(storState)
 
-			got, err := handler.nextMissingUpstream(context.Background(), "test-session")
+			got, err := handler.nextMissingUpstream(context.Background(), "test-session", twoUpstreamChain)
 			require.NoError(t, err)
 			assert.Equal(t, tt.want, got)
 		})
@@ -167,7 +320,7 @@ func TestNextMissingUpstream_RefreshExpired(t *testing.T) {
 			handler, storState, _, _ := multiUpstreamTestSetup(t, opts...)
 			tt.setupTokens(storState)
 
-			got, err := handler.nextMissingUpstream(context.Background(), "test-session")
+			got, err := handler.nextMissingUpstream(context.Background(), "test-session", twoUpstreamChain)
 			require.NoError(t, err)
 			assert.Equal(t, tt.want, got)
 		})
@@ -237,7 +390,7 @@ func TestNextMissingUpstream_StorageError(t *testing.T) {
 	)
 	require.NoError(t, err)
 
-	got, err := handler.nextMissingUpstream(context.Background(), "test-session")
+	got, err := handler.nextMissingUpstream(context.Background(), "test-session", twoUpstreamChain)
 	require.Error(t, err)
 	assert.ErrorContains(t, err, "failed to check upstream token state")
 	assert.ErrorIs(t, err, storageErr)
diff --git a/pkg/authserver/storage/memory.go b/pkg/authserver/storage/memory.go
index 1887b81a5b..639ad3dc23 100644
--- a/pkg/authserver/storage/memory.go
+++ b/pkg/authserver/storage/memory.go
@@ -965,6 +965,7 @@ func (s *MemoryStorage) StorePendingAuthorization(_ context.Context, state strin
 		ResolvedUserName:     pending.ResolvedUserName,
 		ResolvedUserEmail:    pending.ResolvedUserEmail,
 		SingleLeg:            pending.SingleLeg,
+		ChainUpstreams:       slices.Clone(pending.ChainUpstreams),
 		CreatedAt:            pending.CreatedAt,
 	}
 
@@ -1015,6 +1016,7 @@ func (s *MemoryStorage) LoadPendingAuthorization(_ context.Context, state string
 		ResolvedUserName:     pending.ResolvedUserName,
 		ResolvedUserEmail:    pending.ResolvedUserEmail,
 		SingleLeg:            pending.SingleLeg,
+		ChainUpstreams:       slices.Clone(pending.ChainUpstreams),
 		CreatedAt:            pending.CreatedAt,
 	}, nil
 }
diff --git a/pkg/authserver/storage/memory_test.go b/pkg/authserver/storage/memory_test.go
index d4997d9ae7..968238d342 100644
--- a/pkg/authserver/storage/memory_test.go
+++ b/pkg/authserver/storage/memory_test.go
@@ -871,7 +871,8 @@ func TestMemoryStorage_PendingAuthorization(t *testing.T) {
 			State: "client-state", PKCEChallenge: "challenge", PKCEMethod: "S256",
 			Scopes: []string{"openid", "profile"}, InternalState: state,
 			UpstreamPKCEVerifier: "verifier", UpstreamNonce: "nonce",
-			SingleLeg: true, CreatedAt: time.Now(),
+			SingleLeg: true, ChainUpstreams: []string{"provider-1", "provider-2"},
+			CreatedAt: time.Now(),
 		}
 	}
 
@@ -886,6 +887,7 @@ func TestMemoryStorage_PendingAuthorization(t *testing.T) {
 			assert.Equal(t, pending.PKCEChallenge, retrieved.PKCEChallenge)
 			assert.Equal(t, pending.Scopes, retrieved.Scopes)
 			assert.Equal(t, pending.SingleLeg, retrieved.SingleLeg)
+			assert.Equal(t, pending.ChainUpstreams, retrieved.ChainUpstreams)
 		})
 	})
 
diff --git a/pkg/authserver/storage/redis.go b/pkg/authserver/storage/redis.go
index bcddefa341..c5f53cf4a8 100644
--- a/pkg/authserver/storage/redis.go
+++ b/pkg/authserver/storage/redis.go
@@ -1420,6 +1420,7 @@ type storedPendingAuthorization struct {
 	ResolvedUserName     string   `json:"resolved_user_name,omitempty"`
 	ResolvedUserEmail    string   `json:"resolved_user_email,omitempty"`
 	SingleLeg            bool     `json:"single_leg,omitempty"`
+	ChainUpstreams       []string `json:"chain_upstreams,omitempty"`
 	CreatedAt            int64    `json:"created_at"`
 }
 
@@ -1450,6 +1451,7 @@ func (s *RedisStorage) StorePendingAuthorization(ctx context.Context, state stri
 		ResolvedUserName:     pending.ResolvedUserName,
 		ResolvedUserEmail:    pending.ResolvedUserEmail,
 		SingleLeg:            pending.SingleLeg,
+		ChainUpstreams:       slices.Clone(pending.ChainUpstreams),
 		CreatedAt:            pending.CreatedAt.Unix(),
 	}
 
@@ -1501,6 +1503,7 @@ func (s *RedisStorage) LoadPendingAuthorization(ctx context.Context, state strin
 		ResolvedUserName:     stored.ResolvedUserName,
 		ResolvedUserEmail:    stored.ResolvedUserEmail,
 		SingleLeg:            stored.SingleLeg,
+		ChainUpstreams:       slices.Clone(stored.ChainUpstreams),
 		CreatedAt:            createdAt,
 	}, nil
 }
diff --git a/pkg/authserver/storage/redis_test.go b/pkg/authserver/storage/redis_test.go
index e37394ce3d..50d3bbce63 100644
--- a/pkg/authserver/storage/redis_test.go
+++ b/pkg/authserver/storage/redis_test.go
@@ -1389,7 +1389,8 @@ func TestRedisStorage_PendingAuthorization(t *testing.T) {
 			State: "client-state", PKCEChallenge: "challenge", PKCEMethod: "S256",
 			Scopes: []string{"openid", "profile"}, InternalState: state,
 			UpstreamPKCEVerifier: "verifier", UpstreamNonce: "nonce",
-			SingleLeg: true, CreatedAt: time.Now(),
+			SingleLeg: true, ChainUpstreams: []string{"provider-1", "provider-2"},
+			CreatedAt: time.Now(),
 		}
 	}
 
@@ -1404,6 +1405,7 @@ func TestRedisStorage_PendingAuthorization(t *testing.T) {
 			assert.Equal(t, pending.PKCEChallenge, retrieved.PKCEChallenge)
 			assert.Equal(t, pending.Scopes, retrieved.Scopes)
 			assert.Equal(t, pending.SingleLeg, retrieved.SingleLeg)
+			assert.Equal(t, pending.ChainUpstreams, retrieved.ChainUpstreams)
 		})
 	})
 
diff --git a/pkg/authserver/storage/types.go b/pkg/authserver/storage/types.go
index 320f5920b2..268f9bd5fa 100644
--- a/pkg/authserver/storage/types.go
+++ b/pkg/authserver/storage/types.go
@@ -458,6 +458,16 @@ type PendingAuthorization struct {
 	// preserves the multi-upstream chaining behavior.
 	SingleLeg bool
 
+	// ChainUpstreams is the ordered, effective set of upstream provider names this
+	// authorization must walk, always led by the required first upstream. It is
+	// computed once when the first leg resolves — narrowed by the optional upstream
+	// filter when one is configured — and carried forward across subsequent legs so
+	// the filter is not re-run per leg. Empty on the first leg's inbound pending;
+	// populated on every subsequent leg. When empty (no chain has been computed yet,
+	// or a legacy pending predating this field), the callback recomputes it, which
+	// for the no-filter case yields all configured upstreams in order.
+	ChainUpstreams []string
+
 	// CreatedAt is when the pending authorization was created.
 	CreatedAt time.Time
 }

From e79acbf5a9588f8547b42967474f81c43c0aad0d Mon Sep 17 00:00:00 2001
From: Trey <tgrunnagle@gmail.com>
Date: Sun, 5 Jul 2026 15:36:22 -0700
Subject: [PATCH 2/6] Harden upstream chain resolution per review

Addresses stacklok/toolhive#5725 review comments:
- MEDIUM callback.go (3525733910): distinguish a true first leg
  (ResolvedUserID == "") from a subsequent leg; reject a subsequent-leg
  pending that lacks a computed chain instead of recomputing the filter
  against the wrong leg's request context.
- MEDIUM callback.go (3525733913): validate a chain loaded from storage
  (non-empty, leads with the first configured upstream, names only
  configured upstreams) so a tampered PendingAuthorization row cannot
  shrink the chain and disable the cross-leg identity check.
- MEDIUM callback.go (3525733916): restore discrete expected/got/provider
  slog fields on the identity-mismatch check after the verifyChainIdentity
  extraction, so log pipelines can still alert on it.
- MEDIUM redis_test.go (3525733919): add a backward-compatibility test
  that loads a legacy pending JSON lacking chain_upstreams and asserts an
  empty ChainUpstreams.

Also expands the UpstreamFilter doc comment (first-leg-context guarantee
across rolling upgrades, WithPlatformUser in ctx) and updates existing
chain tests to carry ChainUpstreams on subsequent legs.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 pkg/authserver/server/handlers/callback.go    | 41 +++++++-----
 .../server/handlers/callback_test.go          | 62 ++++++++++++++++++
 pkg/authserver/server/handlers/handler.go     | 65 ++++++++++++++++++-
 pkg/authserver/storage/redis_test.go          | 36 ++++++++++
 pkg/authserver/storage/types.go               |  6 +-
 5 files changed, 186 insertions(+), 24 deletions(-)

diff --git a/pkg/authserver/server/handlers/callback.go b/pkg/authserver/server/handlers/callback.go
index ab15dba91e..ea31221c16 100644
--- a/pkg/authserver/server/handlers/callback.go
+++ b/pkg/authserver/server/handlers/callback.go
@@ -433,21 +433,17 @@ func (h *Handler) continueChainOrComplete(
 		return
 	}
 
-	// Resolve the effective chain of upstreams to walk. The first leg (which
-	// carries no computed chain) consults the optional filter once; every
-	// subsequent leg reuses the chain carried in the pending authorization so the
-	// filter is not re-run per leg. A legacy pending predating ChainUpstreams also
-	// lands here and is safely recomputed.
-	chain := pending.ChainUpstreams
-	if len(chain) == 0 {
-		computed, err := h.computeChain(ctx)
-		if err != nil {
-			slog.Error("failed to compute upstream chain", "error", err)
-			_ = h.storage.DeleteUpstreamTokens(ctx, sessionID)
-			h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to determine authorization chain"))
-			return
-		}
-		chain = computed
+	// Resolve the effective chain of upstreams to walk. The first leg computes it
+	// (consulting the optional filter with this leg's request context); every
+	// subsequent leg reuses the validated chain the first leg carried forward, so
+	// the filter is not re-run per leg. A subsequent leg whose pending predates the
+	// chain is rejected rather than recomputed against a later leg's context.
+	chain, err := h.resolveChain(ctx, pending)
+	if err != nil {
+		slog.Error("failed to resolve upstream chain", "error", err)
+		_ = h.storage.DeleteUpstreamTokens(ctx, sessionID)
+		h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to determine authorization chain"))
+		return
 	}
 
 	nextProvider, err := h.nextMissingUpstream(ctx, sessionID, chain)
@@ -460,7 +456,8 @@ func (h *Handler) continueChainOrComplete(
 
 	if nextProvider == "" {
 		if err := h.verifyChainIdentity(ctx, sessionID, chain, subject); err != nil {
-			slog.Error("chain identity verification failed", "error", err)
+			// verifyChainIdentity already logged the specific cause (with structured
+			// fields for a mismatch); here we just clean up and fail closed.
 			_ = h.storage.DeleteUpstreamTokens(ctx, sessionID)
 			h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("identity verification failed"))
 			return
@@ -559,12 +556,20 @@ func (h *Handler) verifyChainIdentity(ctx context.Context, sessionID string, cha
 	}
 	allTokens, err := h.storage.GetAllUpstreamTokens(ctx, sessionID)
 	if err != nil {
+		slog.Error("failed to load upstream tokens for chain identity check", "error", err)
 		return fmt.Errorf("failed to load upstream tokens for identity check: %w", err)
 	}
 	firstProvider := chain[0]
 	if firstTokens, ok := allTokens[firstProvider]; ok && firstTokens.UserID != subject {
-		return fmt.Errorf("identity mismatch on provider %q: chain subject %q != stored user %q",
-			firstProvider, subject, firstTokens.UserID)
+		// Emit the mismatch as discrete slog fields — not folded into the returned
+		// error string — so log pipelines can filter/alert on this defense-in-depth
+		// check, matching the logging from before this check was extracted here.
+		slog.Error("identity mismatch between chain state and stored tokens",
+			"expected", subject,
+			"got", firstTokens.UserID,
+			"provider", firstProvider,
+		)
+		return fmt.Errorf("identity mismatch on provider %q", firstProvider)
 	}
 	return nil
 }
diff --git a/pkg/authserver/server/handlers/callback_test.go b/pkg/authserver/server/handlers/callback_test.go
index 2ad2627f2f..c17108aecc 100644
--- a/pkg/authserver/server/handlers/callback_test.go
+++ b/pkg/authserver/server/handlers/callback_test.go
@@ -481,6 +481,7 @@ func TestCallbackHandler_TwoUpstreams_SecondLeg_IssuesCode(t *testing.T) {
 		UpstreamNonce:        "second-leg-nonce",
 		UpstreamProviderName: "provider-2",
 		SessionID:            sessionID,
+		ChainUpstreams:       []string{"provider-1", "provider-2"},
 		ResolvedUserID:       "resolved-user-id-from-leg1",
 		ResolvedUserName:     "First Leg User",
 		ResolvedUserEmail:    "firstleg@example.com",
@@ -698,6 +699,62 @@ func TestCallbackHandler_SecondLeg_ReusesChain_DoesNotReRunFilter(t *testing.T)
 	assert.Equal(t, "provider2-code", provider2.capturedCode)
 }
 
+func TestCallbackHandler_SubsequentLeg_MissingChain_FailsClosed(t *testing.T) {
+	t.Parallel()
+	handler, storState, _, _ := multiUpstreamTestSetup(t)
+
+	sessionID := "stale-pending-session"
+	const leg1User = "resolved-user-id-from-leg1"
+
+	// First leg already completed.
+	storState.upstreamTokens[sessionID+":provider-1"] = &storage.UpstreamTokens{
+		ProviderID:  "provider-1",
+		AccessToken: "p1-at",
+		ExpiresAt:   time.Now().Add(time.Hour),
+		ClientID:    testAuthClientID,
+		UserID:      leg1User,
+	}
+
+	// A subsequent-leg pending that predates ChainUpstreams: ResolvedUserID is set
+	// (not a first leg) but ChainUpstreams is empty, as an older build would have
+	// written it mid-rollout.
+	secondLegState := "stale-pending-state"
+	storState.pendingAuths[secondLegState] = &storage.PendingAuthorization{
+		ClientID:             testAuthClientID,
+		RedirectURI:          testAuthRedirectURI,
+		State:                "client-original-state",
+		PKCEChallenge:        "client-challenge",
+		PKCEMethod:           "S256",
+		Scopes:               []string{"openid"},
+		InternalState:        secondLegState,
+		UpstreamPKCEVerifier: "stale-verifier-1234567890123456789012345678",
+		UpstreamNonce:        "stale-nonce",
+		UpstreamProviderName: "provider-2",
+		SessionID:            sessionID,
+		ResolvedUserID:       leg1User,
+		ResolvedUserName:     "First Leg User",
+		ResolvedUserEmail:    "firstleg@example.com",
+		// ChainUpstreams intentionally empty (stale pending).
+		CreatedAt: time.Now(),
+	}
+
+	req := httptest.NewRequest(http.MethodGet, "/oauth/callback?code=provider2-code&state="+secondLegState, nil)
+	rec := httptest.NewRecorder()
+	handler.CallbackHandler(rec, req)
+
+	// The stale pending is rejected (fail closed) rather than recomputing the chain
+	// against this later leg's context.
+	assert.Equal(t, http.StatusSeeOther, rec.Code, "stale pending should produce a fosite error redirect")
+	location := rec.Header().Get("Location")
+	assert.Contains(t, location, "error=server_error")
+	assert.Contains(t, location, "state=client-original-state")
+
+	// Upstream tokens for the session are cleaned up.
+	for key := range storState.upstreamTokens {
+		assert.Failf(t, "upstream tokens should be cleaned up", "found leftover token %q", key)
+	}
+}
+
 func TestCallbackHandler_SingleLeg_IssuesCodeWithoutChaining(t *testing.T) {
 	t.Parallel()
 	handler, storState, provider1, _ := multiUpstreamTestSetup(t)
@@ -790,6 +847,7 @@ func TestCallbackHandler_TwoUpstreams_IdentityFromFirstLeg(t *testing.T) {
 		UpstreamNonce:        "identity-test-nonce",
 		UpstreamProviderName: "provider-2",
 		SessionID:            sessionID,
+		ChainUpstreams:       []string{"provider-1", "provider-2"},
 		ResolvedUserID:       firstLegUserID,
 		ResolvedUserName:     "First Leg Name",
 		ResolvedUserEmail:    "firstleg@example.com",
@@ -850,6 +908,7 @@ func TestCallbackHandler_TwoUpstreams_IdentityMismatch_RejectsChain(t *testing.T
 		UpstreamNonce:        "mismatch-nonce",
 		UpstreamProviderName: "provider-2",
 		SessionID:            sessionID,
+		ChainUpstreams:       []string{"provider-1", "provider-2"},
 		ResolvedUserID:       "correct-user-id", // does NOT match provider-1's UserID above
 		ResolvedUserName:     "Correct User",
 		ResolvedUserEmail:    "correct@example.com",
@@ -1224,9 +1283,11 @@ func TestCallbackHandler_RefreshTokenCarryForward(t *testing.T) {
 				// Synthetic carry-forward only applies on a subsequent leg, where the
 				// stable user identity is carried from the first leg (so the prior row
 				// is found by UserID) while the provider's own subject rotates per flow.
+				// A subsequent leg carries the effective chain computed on the first leg.
 				pendingAuth.ResolvedUserID = existingUserID
 				pendingAuth.ResolvedUserName = "First Leg User"
 				pendingAuth.ResolvedUserEmail = "firstleg@example.com"
+				pendingAuth.ChainUpstreams = []string{providerName}
 			}
 			storState.pendingAuths[internalState] = pendingAuth
 
@@ -1329,6 +1390,7 @@ func TestCallbackHandler_PlacesPlatformUserInContext_OnChainRead(t *testing.T) {
 		UpstreamNonce:        "second-leg-nonce",
 		UpstreamProviderName: "provider-2",
 		SessionID:            sessionID,
+		ChainUpstreams:       []string{"provider-1", "provider-2"},
 		ResolvedUserID:       leg1User,
 		ResolvedUserName:     "First Leg User",
 		ResolvedUserEmail:    "firstleg@example.com",
diff --git a/pkg/authserver/server/handlers/handler.go b/pkg/authserver/server/handlers/handler.go
index ace69247cd..5083681ddd 100644
--- a/pkg/authserver/server/handlers/handler.go
+++ b/pkg/authserver/server/handlers/handler.go
@@ -54,9 +54,9 @@ type Handler struct {
 }
 
 // UpstreamFilter narrows the authorization chain to a subset of the configured
-// upstreams. It is consulted once in the callback handler, after the first
-// upstream (upstreams[0]) resolves. The first upstream is always required: it is
-// never passed to the filter and cannot be removed by it.
+// upstreams. It is consulted exactly once per authorization, in the callback
+// handler, after the first upstream (upstreams[0]) resolves. The first upstream is
+// always required: it is never passed to the filter and cannot be removed by it.
 //
 // FilterUpstreams receives the request context of the first leg's callback and
 // the names of the non-first configured upstreams, in configured order. It
@@ -65,6 +65,16 @@ type Handler struct {
 // filter cannot reorder the chain or introduce unknown providers. A returned
 // error fails the authorization with a server error — the handler never falls
 // back to walking every upstream on error.
+//
+// The context carries the resolved canonical user via auth.WithPlatformUser (set
+// just before the filter runs), which is the most concrete signal a filter keyed
+// on the authenticated principal would use. It does not carry the requesting
+// client, scopes, or resource — those live in the pending authorization, not ctx.
+//
+// The "first-leg context" guarantee holds even across a rolling upgrade of a
+// persistent backend: a subsequent-leg pending that lacks a computed chain (e.g.
+// one written before this feature existed) is rejected and forces a fresh
+// authorization, rather than re-running the filter against a later leg's context.
 type UpstreamFilter interface {
 	FilterUpstreams(ctx context.Context, configured []string) ([]string, error)
 }
@@ -255,6 +265,55 @@ func (h *Handler) computeChain(ctx context.Context) ([]string, error) {
 	return chain, nil
 }
 
+// resolveChain returns the effective upstream chain for the leg that just
+// completed. A true first leg — identified by an as-yet-unresolved user
+// (pending.ResolvedUserID == "") — computes the chain now, consulting the optional
+// filter with this leg's request context. A subsequent leg reuses the chain the
+// first leg carried forward in the pending authorization, after validating it
+// against the configured upstreams.
+//
+// A subsequent leg whose pending carries no chain (e.g. one written before the
+// ChainUpstreams field existed, a possible rolling-upgrade window on a persistent
+// backend) is a hard error rather than a silent recompute: the filter's contract
+// is to run once, with the first leg's context, so recomputing here with a later
+// leg's context could narrow the chain against the wrong request. Failing closed
+// forces a fresh authorization that starts cleanly at the first leg.
+func (h *Handler) resolveChain(ctx context.Context, pending *storage.PendingAuthorization) ([]string, error) {
+	if pending.ResolvedUserID == "" {
+		// True first leg — compute with this (the first) leg's request context.
+		return h.computeChain(ctx)
+	}
+	if len(pending.ChainUpstreams) == 0 {
+		return nil, fmt.Errorf("subsequent chain leg is missing its computed upstream chain (stale pending authorization)")
+	}
+	if err := h.validateChain(pending.ChainUpstreams); err != nil {
+		return nil, fmt.Errorf("stored upstream chain is invalid: %w", err)
+	}
+	return pending.ChainUpstreams, nil
+}
+
+// validateChain rejects an upstream chain loaded from storage that does not agree
+// with the configured upstreams. A valid chain is non-empty, leads with the
+// required first upstream, and names only configured upstreams. This guards the
+// reuse path against a corrupt or tampered PendingAuthorization row shrinking an
+// in-flight chain to skip legs — which would also silently disable the cross-leg
+// identity check, since verifyChainIdentity is a no-op for a single-element chain.
+func (h *Handler) validateChain(chain []string) error {
+	if len(chain) == 0 {
+		return fmt.Errorf("chain is empty")
+	}
+	if chain[0] != h.upstreams[0].Name {
+		return fmt.Errorf("chain must lead with the first configured upstream %q, got %q",
+			h.upstreams[0].Name, chain[0])
+	}
+	for _, name := range chain {
+		if _, ok := h.upstreamByName(name); !ok {
+			return fmt.Errorf("chain contains unconfigured upstream %q", name)
+		}
+	}
+	return nil
+}
+
 // refreshExpiredLeg attempts to refresh an expired upstream token for one chain
 // leg. It returns true when the leg can be treated as authenticated (refresh
 // succeeded) and false when the user must be re-prompted: no refresher configured,
diff --git a/pkg/authserver/storage/redis_test.go b/pkg/authserver/storage/redis_test.go
index 50d3bbce63..917e8bef03 100644
--- a/pkg/authserver/storage/redis_test.go
+++ b/pkg/authserver/storage/redis_test.go
@@ -1409,6 +1409,42 @@ func TestRedisStorage_PendingAuthorization(t *testing.T) {
 		})
 	})
 
+	t.Run("legacy JSON without chain_upstreams decodes with empty ChainUpstreams", func(t *testing.T) {
+		// Pin the wire-shape contract: pre-feature Redis data that has no
+		// "chain_upstreams" key must deserialise to an empty ChainUpstreams, which
+		// the callback treats as "no chain computed yet". A JSON tag rename or a
+		// DisallowUnknownFields flip would break this without failing another test.
+		withRedisStorage(t, func(ctx context.Context, s *RedisStorage, mr *miniredis.Miniredis) {
+			legacyJSON := fmt.Sprintf(`{
+				"client_id": "legacy-client",
+				"redirect_uri": "https://example.com/callback",
+				"state": "client-state",
+				"pkce_challenge": "challenge",
+				"pkce_method": "S256",
+				"scopes": ["openid"],
+				"internal_state": "legacy-internal-state",
+				"upstream_pkce_verifier": "verifier",
+				"upstream_nonce": "nonce",
+				"session_id": "legacy-session",
+				"created_at": %d
+			}`, time.Now().Unix())
+
+			// Inject directly into miniredis, bypassing the Store path, to simulate a
+			// pre-feature row written without "chain_upstreams".
+			key := redisKey(s.keyPrefix, KeyTypePending, "legacy-internal-state")
+			require.NoError(t, mr.Set(key, legacyJSON))
+
+			retrieved, err := s.LoadPendingAuthorization(ctx, "legacy-internal-state")
+			require.NoError(t, err)
+			require.NotNil(t, retrieved)
+
+			assert.Empty(t, retrieved.ChainUpstreams, "legacy record must decode with empty ChainUpstreams")
+			// Sanity-check that unrelated fields still populate from the legacy blob.
+			assert.Equal(t, "legacy-client", retrieved.ClientID)
+			assert.Equal(t, "legacy-session", retrieved.SessionID)
+		})
+	})
+
 	t.Run("load non-existent", func(t *testing.T) {
 		withRedisStorage(t, func(ctx context.Context, s *RedisStorage, _ *miniredis.Miniredis) {
 			_, err := s.LoadPendingAuthorization(ctx, "non-existent")
diff --git a/pkg/authserver/storage/types.go b/pkg/authserver/storage/types.go
index 268f9bd5fa..91490cd0ab 100644
--- a/pkg/authserver/storage/types.go
+++ b/pkg/authserver/storage/types.go
@@ -463,9 +463,9 @@ type PendingAuthorization struct {
 	// computed once when the first leg resolves — narrowed by the optional upstream
 	// filter when one is configured — and carried forward across subsequent legs so
 	// the filter is not re-run per leg. Empty on the first leg's inbound pending;
-	// populated on every subsequent leg. When empty (no chain has been computed yet,
-	// or a legacy pending predating this field), the callback recomputes it, which
-	// for the no-filter case yields all configured upstreams in order.
+	// populated on every subsequent leg. A subsequent leg that arrives with this
+	// unset (e.g. a pending written before this field existed) is rejected rather
+	// than recomputed, so the filter is never re-run against a later leg's context.
 	ChainUpstreams []string
 
 	// CreatedAt is when the pending authorization was created.

From 0fcc8eca3d8dbaa06dc5616ef57c92c692118738 Mon Sep 17 00:00:00 2001
From: Trey <tgrunnagle@gmail.com>
Date: Sun, 5 Jul 2026 16:40:23 -0700
Subject: [PATCH 3/6] Expose first-upstream identity to chain filter
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A useful upstream filter needs to key its narrowing decision on who
authenticated — the subject and the authorization-relevant claims from
the first upstream — not just on request context. Widen the hook to
receive that identity explicitly.

- FilterUpstreams now takes the canonical platform user ID and the
  resolved auth.PrincipalInfo (claim-mapped Subject, Name, Email, and the
  Claims map used by authz policies) alongside the configured upstreams.
  The principal is passed as an explicit parameter, not via context: the
  callback deliberately avoids placing a bearer-less Identity in ctx.
- Capture the upstream claims that previously had nowhere to live: add
  Claims to upstream.Identity, populated from the validated ID token
  (OIDC) and the userinfo response (OAuth2). identityFromToken and
  synthetic OAuth2 resolve no structured claim set and leave it nil.
- The callback builds the principal from the first leg's resolved
  identity and threads it through continueChainOrComplete -> resolveChain
  -> computeChain to the filter.

Tests: an OIDC provider integration test (real provider + mock IdP)
asserting ID-token claims are captured into the resolved Identity, and a
callback integration test asserting the first upstream's subject and
claims reach FilterUpstreams.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 pkg/authserver/server/handlers/callback.go    | 28 ++++++--
 .../server/handlers/callback_test.go          | 62 +++++++++++++++++
 pkg/authserver/server/handlers/handler.go     | 67 +++++++++++++------
 .../server/handlers/handler_chain_test.go     | 40 ++++++++---
 pkg/authserver/upstream/oauth2.go             |  1 +
 pkg/authserver/upstream/oidc.go               | 11 +++
 pkg/authserver/upstream/oidc_test.go          | 48 +++++++++++++
 pkg/authserver/upstream/types.go              | 12 ++++
 8 files changed, 235 insertions(+), 34 deletions(-)

diff --git a/pkg/authserver/server/handlers/callback.go b/pkg/authserver/server/handlers/callback.go
index ea31221c16..bdb5f5ae6e 100644
--- a/pkg/authserver/server/handlers/callback.go
+++ b/pkg/authserver/server/handlers/callback.go
@@ -180,7 +180,20 @@ func (h *Handler) CallbackHandler(w http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	h.continueChainOrComplete(ctx, w, req, ar, pending, sessionID, subject, userName, userEmail)
+	// Build the credential-free principal for the optional upstream filter, keyed on
+	// the identity the first leg just established. providerSubject is the claim-mapped
+	// upstream subject; subject is the canonical ToolHive user ID. On subsequent legs
+	// the filter is not consulted, so this only drives filtering for upstreams[0].
+	// result.Claims carries the ID-token/userinfo claims (nil for synthetic upstreams).
+	principal := auth.PrincipalInfo{
+		Subject:        providerSubject,
+		PlatformUserID: subject,
+		Name:           userName,
+		Email:          userEmail,
+		Claims:         result.Claims,
+	}
+
+	h.continueChainOrComplete(ctx, w, req, ar, pending, sessionID, principal)
 }
 
 // maybeCarryForwardRefreshToken preserves a prior refresh token when the upstream IdP
@@ -409,10 +422,15 @@ func (h *Handler) continueChainOrComplete(
 	ar fosite.AuthorizeRequester,
 	pending *storage.PendingAuthorization,
 	sessionID string,
-	subject string,
-	name string,
-	email string,
+	principal auth.PrincipalInfo,
 ) {
+	// subject is the canonical ToolHive user ID used for chain state, token keying,
+	// and the cross-leg identity check. Note this is principal.PlatformUserID, NOT
+	// principal.Subject (which is the upstream provider's subject).
+	subject := principal.PlatformUserID
+	name := principal.Name
+	email := principal.Email
+
 	// SingleLeg authorizations intentionally bypass chain continuation: the caller
 	// scoped this flow to one specific upstream (e.g. a UI-initiated "connect one
 	// backend" request), so other configured-but-tokenless upstreams must not
@@ -438,7 +456,7 @@ func (h *Handler) continueChainOrComplete(
 	// subsequent leg reuses the validated chain the first leg carried forward, so
 	// the filter is not re-run per leg. A subsequent leg whose pending predates the
 	// chain is rejected rather than recomputed against a later leg's context.
-	chain, err := h.resolveChain(ctx, pending)
+	chain, err := h.resolveChain(ctx, pending, principal)
 	if err != nil {
 		slog.Error("failed to resolve upstream chain", "error", err)
 		_ = h.storage.DeleteUpstreamTokens(ctx, sessionID)
diff --git a/pkg/authserver/server/handlers/callback_test.go b/pkg/authserver/server/handlers/callback_test.go
index c17108aecc..b701d5bbfb 100644
--- a/pkg/authserver/server/handlers/callback_test.go
+++ b/pkg/authserver/server/handlers/callback_test.go
@@ -699,6 +699,68 @@ func TestCallbackHandler_SecondLeg_ReusesChain_DoesNotReRunFilter(t *testing.T)
 	assert.Equal(t, "provider2-code", provider2.capturedCode)
 }
 
+func TestCallbackHandler_Filter_ReceivesFirstUpstreamIdentity(t *testing.T) {
+	t.Parallel()
+	filter := &stubUpstreamFilter{keep: []string{"provider-2"}}
+	handler, storState, provider1, _ := multiUpstreamTestSetup(t, WithUpstreamFilter(filter))
+
+	// Model an OIDC first upstream that resolved ID-token claims an authz filter
+	// would key on. (The real OIDC provider's capture of these claims is covered by
+	// TestOIDCProviderImpl_ExchangeCodeForIdentity/"captures ID token claims…"; this
+	// test covers the callback -> filter propagation.)
+	provider1.providerType = upstream.ProviderTypeOIDC
+	provider1.exchangeResult.Claims = map[string]any{
+		"sub":    "user-from-provider1",
+		"email":  "firstleg@example.com",
+		"groups": []any{"engineering", "admins"},
+	}
+
+	sessionID := "chain-session-filter-identity"
+	firstLegState := "filter-identity-first-leg-state"
+	pending := &storage.PendingAuthorization{
+		ClientID:             testAuthClientID,
+		RedirectURI:          testAuthRedirectURI,
+		State:                "client-state-identity",
+		PKCEChallenge:        "client-challenge",
+		PKCEMethod:           "S256",
+		Scopes:               []string{"openid", "profile"},
+		InternalState:        firstLegState,
+		UpstreamPKCEVerifier: "filter-identity-verifier-123456789012345678",
+		UpstreamNonce:        "filter-identity-nonce",
+		UpstreamProviderName: "provider-1",
+		SessionID:            sessionID,
+		CreatedAt:            time.Now(),
+	}
+	storState.pendingAuths[firstLegState] = pending
+
+	req := httptest.NewRequest(http.MethodGet, "/oauth/callback?code=provider1-code&state="+firstLegState, nil)
+	rec := httptest.NewRecorder()
+	handler.CallbackHandler(rec, req)
+
+	// Kept provider-2, so the chain continues onward (proving the filter ran).
+	require.Equal(t, http.StatusFound, rec.Code, "kept second leg should redirect onward")
+	require.Equal(t, 1, filter.calls, "filter should be consulted once on the first leg")
+
+	// The filter received the first upstream's resolved identity.
+	assert.Equal(t, "user-from-provider1", filter.capturedPrincipal.Subject,
+		"principal.Subject must be the claim-mapped upstream subject")
+	assert.Equal(t, "firstleg@example.com", filter.capturedPrincipal.Email,
+		"principal.Email must come from the first upstream")
+	assert.Equal(t, map[string]any{
+		"sub":    "user-from-provider1",
+		"email":  "firstleg@example.com",
+		"groups": []any{"engineering", "admins"},
+	}, filter.capturedPrincipal.Claims, "principal.Claims must carry the first upstream's claims")
+
+	// The platform user ID is the canonical (resolved) ToolHive user: non-empty,
+	// equal to principal.PlatformUserID, and distinct from the raw upstream subject.
+	assert.NotEmpty(t, filter.capturedUser, "platform user ID must be populated")
+	assert.Equal(t, filter.capturedUser, filter.capturedPrincipal.PlatformUserID,
+		"the standalone platform user ID must mirror principal.PlatformUserID")
+	assert.NotEqual(t, "user-from-provider1", filter.capturedUser,
+		"platform user ID is the canonical ToolHive user, not the raw upstream subject")
+}
+
 func TestCallbackHandler_SubsequentLeg_MissingChain_FailsClosed(t *testing.T) {
 	t.Parallel()
 	handler, storState, _, _ := multiUpstreamTestSetup(t)
diff --git a/pkg/authserver/server/handlers/handler.go b/pkg/authserver/server/handlers/handler.go
index 5083681ddd..ba642b2f07 100644
--- a/pkg/authserver/server/handlers/handler.go
+++ b/pkg/authserver/server/handlers/handler.go
@@ -24,6 +24,7 @@ import (
 	"github.com/go-chi/chi/v5"
 	"github.com/ory/fosite"
 
+	"github.com/stacklok/toolhive/pkg/auth"
 	"github.com/stacklok/toolhive/pkg/authserver/server"
 	"github.com/stacklok/toolhive/pkg/authserver/storage"
 	"github.com/stacklok/toolhive/pkg/authserver/upstream"
@@ -54,29 +55,42 @@ type Handler struct {
 }
 
 // UpstreamFilter narrows the authorization chain to a subset of the configured
-// upstreams. It is consulted exactly once per authorization, in the callback
-// handler, after the first upstream (upstreams[0]) resolves. The first upstream is
-// always required: it is never passed to the filter and cannot be removed by it.
+// upstreams, keyed on the identity established by the first upstream. It is
+// consulted exactly once per authorization, in the callback handler, after the
+// first upstream (upstreams[0]) resolves. The first upstream is always required:
+// it is never passed to the filter and cannot be removed by it.
 //
-// FilterUpstreams receives the request context of the first leg's callback and
-// the names of the non-first configured upstreams, in configured order. It
-// returns the subset to keep. The handler preserves configured order and ignores
-// any returned name that is not one of the non-first configured upstreams, so the
-// filter cannot reorder the chain or introduce unknown providers. A returned
-// error fails the authorization with a server error — the handler never falls
-// back to walking every upstream on error.
+// FilterUpstreams receives, in order:
+//   - ctx: the request context of the first leg's callback.
+//   - platformUserID: the canonical ToolHive user ID resolved from the first
+//     upstream (the stable internal identifier; equals principal.PlatformUserID).
+//   - principal: the upstream-derived identity for authorization decisions. Its
+//     Subject is the claim-mapped upstream subject (OIDC SubjectClaim, or "sub");
+//     Claims carries the ID-token/userinfo claims (nil for OAuth2 identityFromToken
+//     or synthetic upstreams, which resolve no structured claim set); Name and
+//     Email are populated when the upstream provides them. It carries NO tokens —
+//     it is the credential-free auth.PrincipalInfo projection. The filter MUST
+//     treat principal (including its Claims map) as read-only.
+//   - configured: the names of the non-first configured upstreams, in configured
+//     order.
 //
-// The context carries the resolved canonical user via auth.WithPlatformUser (set
-// just before the filter runs), which is the most concrete signal a filter keyed
-// on the authenticated principal would use. It does not carry the requesting
-// client, scopes, or resource — those live in the pending authorization, not ctx.
+// It returns the subset to keep. The handler preserves configured order and
+// ignores any returned name that is not one of the non-first configured upstreams,
+// so the filter cannot reorder the chain or introduce unknown providers. A
+// returned error fails the authorization with a server error — the handler never
+// falls back to walking every upstream on error.
 //
-// The "first-leg context" guarantee holds even across a rolling upgrade of a
+// The "first-leg identity" guarantee holds even across a rolling upgrade of a
 // persistent backend: a subsequent-leg pending that lacks a computed chain (e.g.
 // one written before this feature existed) is rejected and forces a fresh
-// authorization, rather than re-running the filter against a later leg's context.
+// authorization, rather than re-running the filter against a later leg.
 type UpstreamFilter interface {
-	FilterUpstreams(ctx context.Context, configured []string) ([]string, error)
+	FilterUpstreams(
+		ctx context.Context,
+		platformUserID string,
+		principal auth.PrincipalInfo,
+		configured []string,
+	) ([]string, error)
 }
 
 // Option configures optional Handler behavior at construction time.
@@ -229,7 +243,10 @@ func (h *Handler) nextMissingUpstream(ctx context.Context, sessionID string, cha
 //
 // A filter error is returned to the caller so the authorization fails cleanly; it
 // never silently falls back to walking every upstream.
-func (h *Handler) computeChain(ctx context.Context) ([]string, error) {
+//
+// principal is the first upstream's resolved identity, passed through to the filter
+// (see UpstreamFilter); it is unused when no filter is configured.
+func (h *Handler) computeChain(ctx context.Context, principal auth.PrincipalInfo) ([]string, error) {
 	chain := []string{h.upstreams[0].Name}
 	rest := h.upstreams[1:]
 	if len(rest) == 0 {
@@ -245,7 +262,7 @@ func (h *Handler) computeChain(ctx context.Context) ([]string, error) {
 		return append(chain, restNames...), nil
 	}
 
-	keep, err := h.filter.FilterUpstreams(ctx, restNames)
+	keep, err := h.filter.FilterUpstreams(ctx, principal.PlatformUserID, principal, restNames)
 	if err != nil {
 		return nil, fmt.Errorf("upstream filter failed: %w", err)
 	}
@@ -278,10 +295,18 @@ func (h *Handler) computeChain(ctx context.Context) ([]string, error) {
 // is to run once, with the first leg's context, so recomputing here with a later
 // leg's context could narrow the chain against the wrong request. Failing closed
 // forces a fresh authorization that starts cleanly at the first leg.
-func (h *Handler) resolveChain(ctx context.Context, pending *storage.PendingAuthorization) ([]string, error) {
+//
+// principal is the identity resolved by the leg that just completed; it is passed
+// to computeChain (and thus the filter) only on the first leg, where it reflects
+// the identity provider. Subsequent legs reuse the stored chain and ignore it.
+func (h *Handler) resolveChain(
+	ctx context.Context,
+	pending *storage.PendingAuthorization,
+	principal auth.PrincipalInfo,
+) ([]string, error) {
 	if pending.ResolvedUserID == "" {
 		// True first leg — compute with this (the first) leg's request context.
-		return h.computeChain(ctx)
+		return h.computeChain(ctx, principal)
 	}
 	if len(pending.ChainUpstreams) == 0 {
 		return nil, fmt.Errorf("subsequent chain leg is missing its computed upstream chain (stale pending authorization)")
diff --git a/pkg/authserver/server/handlers/handler_chain_test.go b/pkg/authserver/server/handlers/handler_chain_test.go
index a573f8bb70..029f5dea3d 100644
--- a/pkg/authserver/server/handlers/handler_chain_test.go
+++ b/pkg/authserver/server/handlers/handler_chain_test.go
@@ -17,6 +17,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
 
+	"github.com/stacklok/toolhive/pkg/auth"
 	"github.com/stacklok/toolhive/pkg/authserver/server"
 	servercrypto "github.com/stacklok/toolhive/pkg/authserver/server/crypto"
 	"github.com/stacklok/toolhive/pkg/authserver/storage"
@@ -30,18 +31,28 @@ var twoUpstreamChain = []string{"provider-1", "provider-2"}
 
 // stubUpstreamFilter is a hand-written test double for UpstreamFilter, mirroring
 // the mockIDPProvider pattern used elsewhere in this package. It records how many
-// times it was called and the names it was passed, and returns a canned keep set
-// or error.
+// times it was called and the arguments it was passed (the non-first upstream
+// names, the platform user ID, and the resolved principal), and returns a canned
+// keep set or error.
 type stubUpstreamFilter struct {
-	keep         []string
-	err          error
-	calls        int
-	capturedArgs []string
+	keep              []string
+	err               error
+	calls             int
+	capturedArgs      []string
+	capturedUser      string
+	capturedPrincipal auth.PrincipalInfo
 }
 
-func (f *stubUpstreamFilter) FilterUpstreams(_ context.Context, configured []string) ([]string, error) {
+func (f *stubUpstreamFilter) FilterUpstreams(
+	_ context.Context,
+	platformUserID string,
+	principal auth.PrincipalInfo,
+	configured []string,
+) ([]string, error) {
 	f.calls++
 	f.capturedArgs = configured
+	f.capturedUser = platformUserID
+	f.capturedPrincipal = principal
 	if f.err != nil {
 		return nil, f.err
 	}
@@ -139,7 +150,14 @@ func TestComputeChain(t *testing.T) {
 			}
 			handler := newChainTestHandler(t, tt.upstreams, opts...)
 
-			got, err := handler.computeChain(context.Background())
+			principal := auth.PrincipalInfo{
+				PlatformUserID: "canonical-user-1",
+				Subject:        "upstream-sub-1",
+				Email:          "user@example.com",
+				Claims:         map[string]any{"groups": []any{"admins"}},
+			}
+
+			got, err := handler.computeChain(context.Background(), principal)
 			if tt.wantErr {
 				require.Error(t, err)
 				assert.Nil(t, got)
@@ -153,6 +171,12 @@ func TestComputeChain(t *testing.T) {
 					assert.Equal(t, tt.wantFilterIn, tt.filter.capturedArgs,
 						"filter must receive non-first configured upstreams in order")
 				}
+				if tt.wantCalls > 0 {
+					assert.Equal(t, principal.PlatformUserID, tt.filter.capturedUser,
+						"filter must receive the platform user ID")
+					assert.Equal(t, principal, tt.filter.capturedPrincipal,
+						"filter must receive the resolved principal (subject + claims) verbatim")
+				}
 			}
 		})
 	}
diff --git a/pkg/authserver/upstream/oauth2.go b/pkg/authserver/upstream/oauth2.go
index 22809b3ad0..11153f9dd9 100644
--- a/pkg/authserver/upstream/oauth2.go
+++ b/pkg/authserver/upstream/oauth2.go
@@ -496,6 +496,7 @@ func (p *BaseOAuth2Provider) ExchangeCodeForIdentity(ctx context.Context, code,
 			Subject: userInfo.Subject,
 			Name:    userInfo.Name,
 			Email:   userInfo.Email,
+			Claims:  userInfo.Claims,
 		}, nil
 	}
 
diff --git a/pkg/authserver/upstream/oidc.go b/pkg/authserver/upstream/oidc.go
index 36ed03572a..a9b1b992dc 100644
--- a/pkg/authserver/upstream/oidc.go
+++ b/pkg/authserver/upstream/oidc.go
@@ -344,11 +344,22 @@ func (p *OIDCProviderImpl) ExchangeCodeForIdentity(
 		return nil, fmt.Errorf("%w: %w", ErrIdentityResolutionFailed, err)
 	}
 
+	// Capture all ID-token claims for downstream authorization inputs. Best-effort:
+	// the subject/name/email above are already resolved, so a claims-extraction
+	// failure only loses the enrichment, it does not fail the login.
+	var allClaims map[string]any
+	if err := validatedToken.Claims(&allClaims); err != nil {
+		slog.Warn("failed to extract full claims from ID token for authorization inputs",
+			"error", err,
+		)
+	}
+
 	return &Identity{
 		Tokens:  exchanged.tokens,
 		Subject: subject,
 		Name:    idClaims.Name,
 		Email:   idClaims.Email,
+		Claims:  allClaims,
 	}, nil
 }
 
diff --git a/pkg/authserver/upstream/oidc_test.go b/pkg/authserver/upstream/oidc_test.go
index cd87281f37..83ecefbe39 100644
--- a/pkg/authserver/upstream/oidc_test.go
+++ b/pkg/authserver/upstream/oidc_test.go
@@ -605,6 +605,54 @@ func TestOIDCProviderImpl_ExchangeCodeForIdentity(t *testing.T) {
 		assert.NotEmpty(t, result.Tokens.IDToken)
 	})
 
+	t.Run("captures ID token claims for authorization inputs", func(t *testing.T) {
+		t.Parallel()
+
+		mock := newMockOIDCServer(t)
+		t.Cleanup(mock.Close)
+
+		// Issue an ID token carrying claims a downstream authorization filter would
+		// key on (email, groups) alongside the standard set.
+		mock.tokenHandler = func(w http.ResponseWriter, _ *http.Request) {
+			idToken := mock.signIDTokenWithClaims(
+				testClientID, "user-789", "", time.Now().Add(time.Hour),
+				map[string]any{
+					"email":  "dev@example.com",
+					"groups": []any{"engineering", "admins"},
+				},
+			)
+			resp := testTokenResponse{
+				AccessToken: "access-token",
+				TokenType:   "Bearer",
+				IDToken:     idToken,
+				ExpiresIn:   3600,
+			}
+			w.Header().Set("Content-Type", "application/json")
+			_ = json.NewEncoder(w).Encode(resp)
+		}
+
+		config := &OIDCConfig{
+			CommonOAuthConfig: CommonOAuthConfig{
+				ClientID:     testClientID,
+				ClientSecret: testClientSecret,
+				RedirectURI:  testRedirectURI,
+			},
+			Issuer: mock.issuer,
+		}
+
+		provider, err := NewOIDCProvider(ctx, config)
+		require.NoError(t, err)
+
+		result, err := provider.ExchangeCodeForIdentity(ctx, "test-code", "", "")
+		require.NoError(t, err)
+
+		require.NotNil(t, result.Claims, "ID token claims must be captured for authorization inputs")
+		assert.Equal(t, "user-789", result.Claims["sub"], "standard sub claim must be captured")
+		assert.Equal(t, "dev@example.com", result.Claims["email"], "custom email claim must be captured")
+		assert.Equal(t, []any{"engineering", "admins"}, result.Claims["groups"],
+			"multi-valued groups claim must be captured verbatim")
+	})
+
 	t.Run("successful exchange without nonce", func(t *testing.T) {
 		t.Parallel()
 
diff --git a/pkg/authserver/upstream/types.go b/pkg/authserver/upstream/types.go
index c7ca8b0553..f08a0a0534 100644
--- a/pkg/authserver/upstream/types.go
+++ b/pkg/authserver/upstream/types.go
@@ -53,6 +53,18 @@ type Identity struct {
 	// per-user keys (doing so grows `users` without bound). Use the
 	// synthesized Subject as an ephemeral session key only.
 	Synthetic bool
+
+	// Claims holds the raw claims resolved from the upstream IDP for this
+	// identity, for downstream authorization-policy inputs. Source by provider
+	// type:
+	//   - OIDC: all claims from the validated ID token.
+	//   - OAuth2 with userInfo: all claims from the userinfo response.
+	//   - OAuth2 with identityFromToken, or synthetic OAuth2: nil — those paths
+	//     resolve only scalar subject/name/email (or nothing), with no structured
+	//     claim set to carry.
+	// These are informational enrichment only: they are NOT used for user
+	// resolution or token storage. May be nil.
+	Claims map[string]any
 }
 
 // ErrIdentityResolutionFailed indicates identity could not be determined.

From 98a3834ff539254aad35e507ea4a71a857f9afe1 Mon Sep 17 00:00:00 2001
From: Trey <tgrunnagle@gmail.com>
Date: Sun, 5 Jul 2026 17:10:36 -0700
Subject: [PATCH 4/6] Add callback/chain integration tests for auth server

The embedded auth server integration suite stopped at the first-leg
redirect and never exercised the /oauth/callback or multi-upstream chain
traversal. Add coverage that drives the callback end to end against the
real embedded server and mock upstream IDPs:

- Single upstream: authorize -> callback -> the chain completes at the
  sole upstream and issues the authorization code to the client.
- Multi-upstream: the first callback continues the chain to the second
  upstream, and the second completes it and issues the code.

Adds a WithUpstreams helper option, an OAuth2 upstream builder, and a
Callback client method to support driving the flow.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 .../authserver/authserver_integration_test.go | 168 ++++++++++++++++++
 .../authserver/helpers/authserver.go          |  25 +++
 .../authserver/helpers/http_client.go         |   9 +
 3 files changed, 202 insertions(+)

diff --git a/test/integration/authserver/authserver_integration_test.go b/test/integration/authserver/authserver_integration_test.go
index 2fd8e22ef2..aa92eb52ae 100644
--- a/test/integration/authserver/authserver_integration_test.go
+++ b/test/integration/authserver/authserver_integration_test.go
@@ -8,7 +8,9 @@ import (
 	"crypto/ecdsa"
 	"crypto/elliptic"
 	"crypto/rand"
+	"crypto/sha256"
 	"crypto/x509"
+	"encoding/base64"
 	"encoding/json"
 	"encoding/pem"
 	"io"
@@ -201,6 +203,172 @@ func TestEmbeddedAuthServer_AuthorizationFlow(t *testing.T) {
 	})
 }
 
+// TestEmbeddedAuthServer_CallbackCompletesAuthorization drives the full
+// single-upstream authorization: authorize (redirect to the upstream) → callback
+// (upstream code exchanged, tokens stored) → the chain completes at the sole
+// upstream and the authorization code is issued back to the client. This
+// exercises the callback + chain-completion path that the authorize-only test
+// above does not reach.
+func TestEmbeddedAuthServer_CallbackCompletesAuthorization(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+
+	upstream := helpers.NewMockUpstreamIDP(t)
+	cfg := helpers.NewTestAuthServerConfig(t, upstream.URL())
+	authServer := helpers.NewEmbeddedAuthServer(ctx, t, cfg)
+
+	server := httptest.NewServer(authServer.Handler())
+	defer server.Close()
+
+	client := helpers.NewOAuthClient(server.URL)
+	clientID := registerAuthCodeClient(t, client)
+	challenge := pkceS256Challenge()
+
+	// Leg 1: authorize → redirect to the upstream; capture the internal state the
+	// server threaded to the upstream so we can drive the callback.
+	authParams := url.Values{
+		"response_type":         {"code"},
+		"client_id":             {clientID},
+		"redirect_uri":          {testClientRedirectURI},
+		"scope":                 {"openid"},
+		"state":                 {"client-state-single"},
+		"resource":              {cfg.AllowedAudiences[0]},
+		"code_challenge":        {challenge},
+		"code_challenge_method": {"S256"},
+	}
+	resp, err := client.StartAuthorization(authParams)
+	require.NoError(t, err)
+	location := resp.Header.Get("Location")
+	require.NoError(t, resp.Body.Close())
+	require.Equal(t, http.StatusFound, resp.StatusCode, "authorize should redirect to the upstream")
+	internalState := stateParam(t, location)
+	require.NotEmpty(t, internalState, "authorize must thread an internal state to the upstream")
+
+	// Callback: with only one upstream in the chain, the server exchanges the code
+	// and issues the authorization code back to the client.
+	cbResp, err := client.Callback("mock-auth-code", internalState)
+	require.NoError(t, err)
+	defer cbResp.Body.Close()
+
+	require.Equal(t, http.StatusSeeOther, cbResp.StatusCode,
+		"single-upstream chain should complete and issue the authorization code")
+	final := cbResp.Header.Get("Location")
+	assert.Contains(t, final, testClientRedirectURI, "should redirect back to the client")
+	assert.Contains(t, final, "code=", "should include an authorization code")
+	assert.Contains(t, final, "state=client-state-single", "should preserve the client state")
+	assert.NotContains(t, final, "error=", "should not be an error redirect")
+}
+
+// TestEmbeddedAuthServer_MultiUpstreamChain drives an authorization across two
+// configured upstreams: the first callback continues the chain to the second
+// upstream, and the second callback completes it and issues the authorization
+// code. This exercises the end-to-end chain traversal — the effective chain is
+// computed once, carried across legs, and walked to completion.
+func TestEmbeddedAuthServer_MultiUpstreamChain(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+
+	upstreamA := helpers.NewMockUpstreamIDP(t)
+	upstreamB := helpers.NewMockUpstreamIDP(t)
+	cfg := helpers.NewTestAuthServerConfig(t, upstreamA.URL(),
+		helpers.WithUpstreams([]authserver.UpstreamRunConfig{
+			helpers.NewOAuth2Upstream("provider-a", upstreamA.URL()),
+			helpers.NewOAuth2Upstream("provider-b", upstreamB.URL()),
+		}),
+	)
+	authServer := helpers.NewEmbeddedAuthServer(ctx, t, cfg)
+
+	server := httptest.NewServer(authServer.Handler())
+	defer server.Close()
+
+	client := helpers.NewOAuthClient(server.URL)
+	clientID := registerAuthCodeClient(t, client)
+	challenge := pkceS256Challenge()
+
+	authParams := url.Values{
+		"response_type":         {"code"},
+		"client_id":             {clientID},
+		"redirect_uri":          {testClientRedirectURI},
+		"scope":                 {"openid"},
+		"state":                 {"client-state-chain"},
+		"resource":              {cfg.AllowedAudiences[0]},
+		"code_challenge":        {challenge},
+		"code_challenge_method": {"S256"},
+	}
+
+	// Leg 1: authorize → redirect to the first upstream.
+	resp, err := client.StartAuthorization(authParams)
+	require.NoError(t, err)
+	locA := resp.Header.Get("Location")
+	require.NoError(t, resp.Body.Close())
+	require.Equal(t, http.StatusFound, resp.StatusCode)
+	assert.Contains(t, locA, upstreamA.URL(), "first leg targets provider-a")
+	stateA := stateParam(t, locA)
+	require.NotEmpty(t, stateA)
+
+	// Leg 1 callback → chain continues, redirect onward to the second upstream.
+	cbA, err := client.Callback("code-a", stateA)
+	require.NoError(t, err)
+	locB := cbA.Header.Get("Location")
+	require.NoError(t, cbA.Body.Close())
+	require.Equal(t, http.StatusFound, cbA.StatusCode, "chain should continue to the second upstream")
+	assert.Contains(t, locB, upstreamB.URL(), "second leg targets provider-b")
+	stateB := stateParam(t, locB)
+	require.NotEmpty(t, stateB)
+	require.NotEqual(t, stateA, stateB, "each leg uses a fresh internal state")
+
+	// Leg 2 callback → chain complete, authorization code issued to the client.
+	cbB, err := client.Callback("code-b", stateB)
+	require.NoError(t, err)
+	defer cbB.Body.Close()
+
+	require.Equal(t, http.StatusSeeOther, cbB.StatusCode, "completed chain should issue the code")
+	final := cbB.Header.Get("Location")
+	assert.Contains(t, final, testClientRedirectURI, "should redirect back to the client")
+	assert.Contains(t, final, "code=", "should include an authorization code")
+	assert.Contains(t, final, "state=client-state-chain", "should preserve the client state")
+	assert.NotContains(t, final, "error=", "should not be an error redirect")
+}
+
+// testClientRedirectURI is the client callback used by the authorization-flow tests.
+const testClientRedirectURI = "http://localhost:8080/callback"
+
+// registerAuthCodeClient performs DCR for an authorization-code client and returns
+// its client_id.
+func registerAuthCodeClient(t *testing.T, client *helpers.OAuthClient) string {
+	t.Helper()
+	regResult, statusCode, err := client.RegisterClient(map[string]interface{}{
+		"client_name":   "Integration Test Client",
+		"redirect_uris": []string{testClientRedirectURI},
+		"grant_types":   []string{"authorization_code", "refresh_token"},
+	})
+	require.NoError(t, err)
+	require.Equal(t, http.StatusCreated, statusCode, "client registration should succeed")
+	clientID, ok := regResult["client_id"].(string)
+	require.True(t, ok, "registration response must include a client_id")
+	return clientID
+}
+
+// stateParam extracts the `state` query parameter from a redirect Location.
+func stateParam(t *testing.T, location string) string {
+	t.Helper()
+	require.NotEmpty(t, location, "redirect must include a Location header")
+	u, err := url.Parse(location)
+	require.NoError(t, err)
+	return u.Query().Get("state")
+}
+
+// pkceS256Challenge returns a valid S256 code challenge for a fixed verifier. The
+// verifier itself is unused because these tests assert the authorization-code
+// redirect rather than completing the token exchange.
+func pkceS256Challenge() string {
+	const verifier = "integration-test-pkce-verifier-0123456789abcdef"
+	sum := sha256.Sum256([]byte(verifier))
+	return base64.RawURLEncoding.EncodeToString(sum[:])
+}
+
 // TestEmbeddedAuthServer_DynamicClientRegistration verifies DCR (RFC 7591) support.
 //
 //nolint:paralleltest,tparallel // Subtests share expensive test fixtures
diff --git a/test/integration/authserver/helpers/authserver.go b/test/integration/authserver/helpers/authserver.go
index f4939f445f..60fec7c992 100644
--- a/test/integration/authserver/helpers/authserver.go
+++ b/test/integration/authserver/helpers/authserver.go
@@ -53,6 +53,31 @@ func WithBaselineClientScopes(scopes []string) AuthServerOption {
 	}
 }
 
+// WithUpstreams overrides the default single test upstream with an explicit,
+// ordered list of upstream providers. Used to exercise multi-upstream
+// authorization chains.
+func WithUpstreams(upstreams []authserver.UpstreamRunConfig) AuthServerOption {
+	return func(c *authServerConfig) {
+		c.upstreams = upstreams
+	}
+}
+
+// NewOAuth2Upstream builds an OAuth2 upstream run config pointing at the given
+// mock upstream base URL. It has no userinfo endpoint, so the auth server resolves
+// a synthetic identity for it — sufficient for exercising chain traversal.
+func NewOAuth2Upstream(name, upstreamURL string) authserver.UpstreamRunConfig {
+	return authserver.UpstreamRunConfig{
+		Name: name,
+		Type: authserver.UpstreamProviderTypeOAuth2,
+		OAuth2Config: &authserver.OAuth2UpstreamRunConfig{
+			AuthorizationEndpoint: upstreamURL + "/authorize",
+			TokenEndpoint:         upstreamURL + "/token",
+			ClientID:              "test-client-id",
+			RedirectURI:           upstreamURL + "/callback",
+		},
+	}
+}
+
 // GetFreePort returns an available TCP port on localhost.
 func GetFreePort(tb testing.TB) int {
 	tb.Helper()
diff --git a/test/integration/authserver/helpers/http_client.go b/test/integration/authserver/helpers/http_client.go
index 550be402e1..31fefa53bf 100644
--- a/test/integration/authserver/helpers/http_client.go
+++ b/test/integration/authserver/helpers/http_client.go
@@ -116,6 +116,15 @@ func (c *OAuthClient) StartAuthorization(params url.Values) (*http.Response, err
 	return c.httpClient.Get(authURL)
 }
 
+// Callback drives the upstream callback leg: GET /oauth/callback with the given
+// code and internal state. The client does not follow redirects, so the returned
+// response exposes what the server issues — an onward redirect to the next
+// upstream, or a redirect back to the client with an authorization code.
+func (c *OAuthClient) Callback(code, state string) (*http.Response, error) {
+	params := url.Values{"code": {code}, "state": {state}}
+	return c.httpClient.Get(c.baseURL + "/oauth/callback?" + params.Encode())
+}
+
 // ExchangeToken performs a token exchange at the token endpoint.
 func (c *OAuthClient) ExchangeToken(params url.Values) (map[string]interface{}, int, error) {
 	resp, err := c.httpClient.PostForm(c.baseURL+"/oauth/token", params)

From 3dd8352c6a2bbf23538ab347d1769175c33277b0 Mon Sep 17 00:00:00 2001
From: Trey <tgrunnagle@gmail.com>
Date: Mon, 6 Jul 2026 07:33:02 -0700
Subject: [PATCH 5/6] Tighten chain validation and clarify filter docs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Addresses stacklok/toolhive#5725 review comments:
- MEDIUM handler.go (3528199173): validateChain now enforces an in-order,
  duplicate-free subsequence of the configured upstreams (config-cursor
  walk), so a tampered pending row cannot shrink/reorder the chain to skip
  legs (e.g. ["provider-1"] or ["provider-1","provider-1"]) and slip past
  the single-element identity-check no-op. Adds a TestValidateChain table
  covering wrong-first / unconfigured / out-of-order / duplicate.
- MEDIUM handler.go (body): NewHandler rejects duplicate upstream names —
  the name invariant is now load-bearing (chain keyed by name; tokens and
  upstreamByName key by name). Adds a constructor test.
- LOW handler.go (3528199199): document the Claims contract — values are
  untyped (assert defensively, never panic), nil is possible on a transient
  OIDC extraction failure (treat as fail-closed), and aud is the upstream's
  client_id, not this AS.
- LOW callback.go (3528199204): reword verifyChainIdentity to match scope —
  it reconciles only the first leg; intermediate legs are intentionally not
  identity-checked.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 .../server/handlers/authorize_test.go         | 18 ++++++
 pkg/authserver/server/handlers/callback.go    | 13 ++--
 pkg/authserver/server/handlers/handler.go     | 60 ++++++++++++++-----
 .../server/handlers/handler_chain_test.go     | 52 ++++++++++++++++
 4 files changed, 124 insertions(+), 19 deletions(-)

diff --git a/pkg/authserver/server/handlers/authorize_test.go b/pkg/authserver/server/handlers/authorize_test.go
index d8dab8df79..14415e6c81 100644
--- a/pkg/authserver/server/handlers/authorize_test.go
+++ b/pkg/authserver/server/handlers/authorize_test.go
@@ -185,6 +185,24 @@ func TestNewHandler_ErrorsOnNilConfig(t *testing.T) {
 		"NewHandler should error when AuthorizationServerConfig.Config is nil")
 }
 
+// TestNewHandler_ErrorsOnDuplicateUpstreamNames pins that the constructor rejects
+// duplicate upstream names. Names must be unique: upstreamByName returns the first
+// match, tokens are keyed by name, and the authorization chain is keyed by name, so
+// a duplicate would silently shadow a provider.
+func TestNewHandler_ErrorsOnDuplicateUpstreamNames(t *testing.T) {
+	t.Parallel()
+
+	_, oauth2Config, stor, _ := baseTestSetup(t)
+	upstreams := []NamedUpstream{
+		{Name: "dup", Provider: &mockIDPProvider{}},
+		{Name: "dup", Provider: &mockIDPProvider{}},
+	}
+
+	_, err := NewHandler(nil, oauth2Config, stor, upstreams)
+	require.Error(t, err, "NewHandler should reject duplicate upstream names")
+	assert.ErrorContains(t, err, "duplicate upstream name")
+}
+
 func TestAuthorizeHandler_RedirectsToUpstream(t *testing.T) {
 	t.Parallel()
 	handler, storState, mockUpstream := handlerTestSetup(t)
diff --git a/pkg/authserver/server/handlers/callback.go b/pkg/authserver/server/handlers/callback.go
index bdb5f5ae6e..7f7114c33b 100644
--- a/pkg/authserver/server/handlers/callback.go
+++ b/pkg/authserver/server/handlers/callback.go
@@ -552,8 +552,12 @@ func (h *Handler) continueChainOrComplete(
 	http.Redirect(w, req, nextURL, http.StatusFound)
 }
 
-// verifyChainIdentity performs the defense-in-depth cross-leg identity check once
-// every leg of the effective chain is satisfied.
+// verifyChainIdentity is a defense-in-depth check run once every leg of the
+// effective chain is satisfied. Despite the "chain" framing, it reconciles only
+// the first leg: it confirms the identity provider's stored token (chain[0]) still
+// belongs to the subject carried through the flow. Intermediate/later legs are
+// deliberately NOT identity-checked — those are connect-this-backend flows whose
+// upstream identity can legitimately differ from the first leg's user.
 //
 // subject MUST be the canonical ToolHive user ID resolved from the first leg's
 // upstream via that provider's configured subject-claim mapping (the OIDC
@@ -561,11 +565,12 @@ func (h *Handler) continueChainOrComplete(
 // through PendingAuthorization.ResolvedUserID. The caller resolves it exactly once,
 // on the first leg, from the claim-mapped upstream subject. This cross-checks it
 // against firstTokens.UserID — the same canonical ID persisted when the first leg
-// stored its tokens — so a chain whose legs disagree on the user is rejected.
+// stored its tokens — so a first leg whose stored user disagrees with the carried
+// subject is rejected.
 //
 // It gates on the effective chain rather than the raw config: chain[0] is always
 // the first (required) upstream, so a chain the filter narrowed to just that
-// upstream has no cross-legs to reconcile and the check is a no-op. Returns a
+// upstream has no first-leg cross-check to run and the check is a no-op. Returns a
 // non-nil error when the storage lookup fails or an identity mismatch is detected;
 // the caller maps either to a server error.
 func (h *Handler) verifyChainIdentity(ctx context.Context, sessionID string, chain []string, subject string) error {
diff --git a/pkg/authserver/server/handlers/handler.go b/pkg/authserver/server/handlers/handler.go
index ba642b2f07..cc68c83390 100644
--- a/pkg/authserver/server/handlers/handler.go
+++ b/pkg/authserver/server/handlers/handler.go
@@ -66,14 +66,22 @@ type Handler struct {
 //     upstream (the stable internal identifier; equals principal.PlatformUserID).
 //   - principal: the upstream-derived identity for authorization decisions. Its
 //     Subject is the claim-mapped upstream subject (OIDC SubjectClaim, or "sub");
-//     Claims carries the ID-token/userinfo claims (nil for OAuth2 identityFromToken
-//     or synthetic upstreams, which resolve no structured claim set); Name and
-//     Email are populated when the upstream provides them. It carries NO tokens —
-//     it is the credential-free auth.PrincipalInfo projection. The filter MUST
-//     treat principal (including its Claims map) as read-only.
+//     Claims carries the ID-token/userinfo claims; Name and Email are populated
+//     when the upstream provides them. It carries NO tokens — it is the
+//     credential-free auth.PrincipalInfo projection. The filter MUST treat
+//     principal (including its Claims map) as read-only.
 //   - configured: the names of the non-first configured upstreams, in configured
 //     order.
 //
+// Claims contract: values are untyped and come straight from the upstream IdP, so
+// implementations MUST assert types defensively (e.g. comma-ok on map/slice
+// access) — a malformed IdP response must not be able to panic the filter. Claims
+// is nil not only for OAuth2 identityFromToken / synthetic upstreams (which resolve
+// no structured claim set) but also on a transient OIDC claims-extraction failure,
+// so a security-relevant filter MUST treat nil/absent claims as fail-closed. Note
+// that the `aud` claim, when present, is the upstream IdP's client_id — not this
+// authorization server.
+//
 // It returns the subset to keep. The handler preserves configured order and
 // ignores any returned name that is not one of the non-first configured upstreams,
 // so the filter cannot reorder the chain or introduce unknown providers. A
@@ -121,9 +129,12 @@ func WithUpstreamFilter(f UpstreamFilter) Option {
 // during multi-upstream authorization flows (e.g., sequential token acquisition).
 //
 // Returns an error if config is nil, if config's embedded *fosite.Config is
-// nil, if upstreams is empty, or if any entry has an empty name or nil
-// provider. Catching misconfiguration here is far easier to diagnose than
-// a nil-deref panic deep inside an HTTP handler at request time.
+// nil, if upstreams is empty, or if any entry has an empty name, a nil provider,
+// or a duplicate name. Upstream names must be unique: upstreamByName returns the
+// first match, tokens are keyed by name, and the authorization chain is keyed by
+// name — a duplicate would silently shadow a provider. Catching misconfiguration
+// here is far easier to diagnose than a nil-deref panic or a shadowed provider
+// deep inside an HTTP handler at request time.
 func NewHandler(
 	provider fosite.OAuth2Provider,
 	config *server.AuthorizationServerConfig,
@@ -138,6 +149,7 @@ func NewHandler(
 	if len(upstreams) == 0 {
 		return nil, fmt.Errorf("handlers: upstreams must not be empty")
 	}
+	seen := make(map[string]struct{}, len(upstreams))
 	for _, u := range upstreams {
 		if u.Name == "" {
 			return nil, fmt.Errorf("handlers: upstream entry has empty name")
@@ -145,6 +157,10 @@ func NewHandler(
 		if u.Provider == nil {
 			return nil, fmt.Errorf("handlers: upstream %q has nil provider", u.Name)
 		}
+		if _, dup := seen[u.Name]; dup {
+			return nil, fmt.Errorf("handlers: duplicate upstream name %q", u.Name)
+		}
+		seen[u.Name] = struct{}{}
 	}
 	h := &Handler{
 		provider:     provider,
@@ -318,11 +334,17 @@ func (h *Handler) resolveChain(
 }
 
 // validateChain rejects an upstream chain loaded from storage that does not agree
-// with the configured upstreams. A valid chain is non-empty, leads with the
-// required first upstream, and names only configured upstreams. This guards the
-// reuse path against a corrupt or tampered PendingAuthorization row shrinking an
-// in-flight chain to skip legs — which would also silently disable the cross-leg
-// identity check, since verifyChainIdentity is a no-op for a single-element chain.
+// with the configured upstreams. A valid chain is a non-empty, in-order,
+// duplicate-free subsequence of the configured upstreams led by the required
+// first upstream. This guards the reuse path against a corrupt or tampered
+// PendingAuthorization row shrinking or reordering an in-flight chain to skip legs
+// — which would also silently disable the cross-leg identity check, since
+// verifyChainIdentity is a no-op for a single-element chain.
+//
+// Note: this validates that each name is currently configured, not that the name
+// still points at the same provider it did when the chain was computed. A name
+// rebound to a different IdP mid-flight (a rolling config change) still validates;
+// the frozen subject and the pending TTL bound that exposure.
 func (h *Handler) validateChain(chain []string) error {
 	if len(chain) == 0 {
 		return fmt.Errorf("chain is empty")
@@ -331,10 +353,18 @@ func (h *Handler) validateChain(chain []string) error {
 		return fmt.Errorf("chain must lead with the first configured upstream %q, got %q",
 			h.upstreams[0].Name, chain[0])
 	}
+	// Walk both, advancing the config cursor past each match. A name that never
+	// matches — unconfigured, out of order, or a duplicate of an already-consumed
+	// entry — is rejected, so the stored chain is a faithful narrowing of the config.
+	ci := 0
 	for _, name := range chain {
-		if _, ok := h.upstreamByName(name); !ok {
-			return fmt.Errorf("chain contains unconfigured upstream %q", name)
+		for ci < len(h.upstreams) && h.upstreams[ci].Name != name {
+			ci++
+		}
+		if ci >= len(h.upstreams) {
+			return fmt.Errorf("chain entry %q is unconfigured, out of order, or duplicated", name)
 		}
+		ci++ // consume the matched upstream so it cannot be reused
 	}
 	return nil
 }
diff --git a/pkg/authserver/server/handlers/handler_chain_test.go b/pkg/authserver/server/handlers/handler_chain_test.go
index 029f5dea3d..534048bdb3 100644
--- a/pkg/authserver/server/handlers/handler_chain_test.go
+++ b/pkg/authserver/server/handlers/handler_chain_test.go
@@ -182,6 +182,58 @@ func TestComputeChain(t *testing.T) {
 	}
 }
 
+func TestValidateChain(t *testing.T) {
+	t.Parallel()
+
+	// Configured order is provider-1, provider-2, provider-3.
+	handler := newChainTestHandler(t, []string{"provider-1", "provider-2", "provider-3"})
+
+	tests := []struct {
+		name    string
+		chain   []string
+		wantErr string // substring to match; empty means no error expected
+	}{
+		{name: "full configured chain", chain: []string{"provider-1", "provider-2", "provider-3"}},
+		{name: "in-order subsequence", chain: []string{"provider-1", "provider-3"}},
+		{name: "first upstream only", chain: []string{"provider-1"}},
+		{name: "empty chain", chain: nil, wantErr: "chain is empty"},
+		{
+			name: "wrong first upstream", chain: []string{"provider-2", "provider-3"},
+			wantErr: "must lead with the first configured upstream",
+		},
+		{
+			name: "unconfigured entry", chain: []string{"provider-1", "provider-x"},
+			wantErr: "unconfigured, out of order, or duplicated",
+		},
+		{
+			name: "out of order", chain: []string{"provider-1", "provider-3", "provider-2"},
+			wantErr: "unconfigured, out of order, or duplicated",
+		},
+		{
+			name: "duplicate first upstream", chain: []string{"provider-1", "provider-1"},
+			wantErr: "unconfigured, out of order, or duplicated",
+		},
+		{
+			name: "duplicate non-first upstream", chain: []string{"provider-1", "provider-2", "provider-2"},
+			wantErr: "unconfigured, out of order, or duplicated",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			err := handler.validateChain(tt.chain)
+			if tt.wantErr == "" {
+				assert.NoError(t, err)
+			} else {
+				require.Error(t, err)
+				assert.ErrorContains(t, err, tt.wantErr)
+			}
+		})
+	}
+}
+
 // TestNextMissingUpstream_RespectsChainSubset verifies that a leg absent from the
 // chain is never reported missing even when it has no stored token — the filter
 // dropped it, so the walk must not prompt for it.

From 4dffa432d33b88ec49eff4151171d25c0f2abf19 Mon Sep 17 00:00:00 2001
From: Trey <tgrunnagle@gmail.com>
Date: Mon, 6 Jul 2026 10:39:57 -0700
Subject: [PATCH 6/6] Drop redundant platformUserID param from UpstreamFilter

The canonical user is already carried as principal.PlatformUserID, so the
separate platformUserID argument to FilterUpstreams was redundant. Remove it
and have consumers read principal.PlatformUserID; update the doc, the stub
filter, and the tests accordingly.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 pkg/authserver/server/handlers/callback_test.go      | 10 ++++------
 pkg/authserver/server/handlers/handler.go            | 12 +++++-------
 pkg/authserver/server/handlers/handler_chain_test.go | 10 ++--------
 3 files changed, 11 insertions(+), 21 deletions(-)

diff --git a/pkg/authserver/server/handlers/callback_test.go b/pkg/authserver/server/handlers/callback_test.go
index b701d5bbfb..4d5173f3b6 100644
--- a/pkg/authserver/server/handlers/callback_test.go
+++ b/pkg/authserver/server/handlers/callback_test.go
@@ -752,12 +752,10 @@ func TestCallbackHandler_Filter_ReceivesFirstUpstreamIdentity(t *testing.T) {
 		"groups": []any{"engineering", "admins"},
 	}, filter.capturedPrincipal.Claims, "principal.Claims must carry the first upstream's claims")
 
-	// The platform user ID is the canonical (resolved) ToolHive user: non-empty,
-	// equal to principal.PlatformUserID, and distinct from the raw upstream subject.
-	assert.NotEmpty(t, filter.capturedUser, "platform user ID must be populated")
-	assert.Equal(t, filter.capturedUser, filter.capturedPrincipal.PlatformUserID,
-		"the standalone platform user ID must mirror principal.PlatformUserID")
-	assert.NotEqual(t, "user-from-provider1", filter.capturedUser,
+	// principal.PlatformUserID is the canonical (resolved) ToolHive user: non-empty
+	// and distinct from the raw upstream subject.
+	assert.NotEmpty(t, filter.capturedPrincipal.PlatformUserID, "platform user ID must be populated")
+	assert.NotEqual(t, "user-from-provider1", filter.capturedPrincipal.PlatformUserID,
 		"platform user ID is the canonical ToolHive user, not the raw upstream subject")
 }
 
diff --git a/pkg/authserver/server/handlers/handler.go b/pkg/authserver/server/handlers/handler.go
index cc68c83390..8a54204bb6 100644
--- a/pkg/authserver/server/handlers/handler.go
+++ b/pkg/authserver/server/handlers/handler.go
@@ -62,12 +62,11 @@ type Handler struct {
 //
 // FilterUpstreams receives, in order:
 //   - ctx: the request context of the first leg's callback.
-//   - platformUserID: the canonical ToolHive user ID resolved from the first
-//     upstream (the stable internal identifier; equals principal.PlatformUserID).
 //   - principal: the upstream-derived identity for authorization decisions. Its
-//     Subject is the claim-mapped upstream subject (OIDC SubjectClaim, or "sub");
-//     Claims carries the ID-token/userinfo claims; Name and Email are populated
-//     when the upstream provides them. It carries NO tokens — it is the
+//     PlatformUserID is the canonical ToolHive user ID (the stable internal
+//     identifier); Subject is the claim-mapped upstream subject (OIDC SubjectClaim,
+//     or "sub"); Claims carries the ID-token/userinfo claims; Name and Email are
+//     populated when the upstream provides them. It carries NO tokens — it is the
 //     credential-free auth.PrincipalInfo projection. The filter MUST treat
 //     principal (including its Claims map) as read-only.
 //   - configured: the names of the non-first configured upstreams, in configured
@@ -95,7 +94,6 @@ type Handler struct {
 type UpstreamFilter interface {
 	FilterUpstreams(
 		ctx context.Context,
-		platformUserID string,
 		principal auth.PrincipalInfo,
 		configured []string,
 	) ([]string, error)
@@ -278,7 +276,7 @@ func (h *Handler) computeChain(ctx context.Context, principal auth.PrincipalInfo
 		return append(chain, restNames...), nil
 	}
 
-	keep, err := h.filter.FilterUpstreams(ctx, principal.PlatformUserID, principal, restNames)
+	keep, err := h.filter.FilterUpstreams(ctx, principal, restNames)
 	if err != nil {
 		return nil, fmt.Errorf("upstream filter failed: %w", err)
 	}
diff --git a/pkg/authserver/server/handlers/handler_chain_test.go b/pkg/authserver/server/handlers/handler_chain_test.go
index 534048bdb3..ba8abf59ff 100644
--- a/pkg/authserver/server/handlers/handler_chain_test.go
+++ b/pkg/authserver/server/handlers/handler_chain_test.go
@@ -32,26 +32,22 @@ var twoUpstreamChain = []string{"provider-1", "provider-2"}
 // stubUpstreamFilter is a hand-written test double for UpstreamFilter, mirroring
 // the mockIDPProvider pattern used elsewhere in this package. It records how many
 // times it was called and the arguments it was passed (the non-first upstream
-// names, the platform user ID, and the resolved principal), and returns a canned
-// keep set or error.
+// names and the resolved principal), and returns a canned keep set or error.
 type stubUpstreamFilter struct {
 	keep              []string
 	err               error
 	calls             int
 	capturedArgs      []string
-	capturedUser      string
 	capturedPrincipal auth.PrincipalInfo
 }
 
 func (f *stubUpstreamFilter) FilterUpstreams(
 	_ context.Context,
-	platformUserID string,
 	principal auth.PrincipalInfo,
 	configured []string,
 ) ([]string, error) {
 	f.calls++
 	f.capturedArgs = configured
-	f.capturedUser = platformUserID
 	f.capturedPrincipal = principal
 	if f.err != nil {
 		return nil, f.err
@@ -172,10 +168,8 @@ func TestComputeChain(t *testing.T) {
 						"filter must receive non-first configured upstreams in order")
 				}
 				if tt.wantCalls > 0 {
-					assert.Equal(t, principal.PlatformUserID, tt.filter.capturedUser,
-						"filter must receive the platform user ID")
 					assert.Equal(t, principal, tt.filter.capturedPrincipal,
-						"filter must receive the resolved principal (subject + claims) verbatim")
+						"filter must receive the resolved principal (platform user, subject + claims) verbatim")
 				}
 			}
 		})