Implement multi-upstream sequential authorization chain in handler layer (RFC-0052 Phase 2)

[tgrunnagle]

Description

Replace the single-upstream Handler struct with a multi-upstream design that drives a server-side sequential authorization chain: the auth server transparently redirects users through every configured upstream IDP in order, accumulating tokens for each, before issuing a single authorization code to the OAuth client. This is Phase 2 of RFC-0052 and builds directly on the UpstreamTokenStorage interface changes landed in Phase 1 (#4136).

Context

RFC-0052 extends the embedded OAuth authorization server (pkg/authserver/) to support multiple upstream IDPs. Phase 1 (#4136) restructured the storage layer so that UpstreamTokenStorage keys on (sessionID, providerName) and added GetAllUpstreamTokens. Phase 2 wires that storage layer into the handler and server layer: the Handler struct is refactored to hold a map of upstream providers keyed by name, the authorize handler generates a server-side SessionID and kicks off the first leg of the chain, and the callback handler either continues the chain (redirect to next upstream) or issues the final authorization code (all upstreams satisfied).

This task also introduces a documented breaking change: ProviderIdentity records stored with providerID = "oauth2" (previously derived from upstream.Type()) will not match after this change switches to pending.UpstreamProviderName as the providerID. Affected users will receive new User records after upgrade. This must be documented in the PR description.

Dependencies: #4136 (Phase 1: Storage Layer)
Blocks: None (TASK-003 and TASK-004 are independent of this task after Phase 1)

Breaking Change — ProviderIdentity provider ID: ProviderIdentity records stored with providerID = "oauth2" or "oidc" (from upstream.Type()) will not match after this change switches to using pending.UpstreamProviderName (the configured name, e.g., "github", "corporate-idp"). Affected users will receive fresh User records. This must be documented in the PR description.

Acceptance Criteria

• Handler struct in pkg/authserver/server/handlers/handler.go replaces the single upstream upstream.OAuth2Provider field with upstreams map[string]upstream.OAuth2Provider (keyed by provider name) and upstreamOrder []string (preserves config order to define chain sequence)
• NewHandler constructor signature accepts upstreams map[string]upstream.OAuth2Provider and upstreamOrder []string in place of the single upstreamIDP upstream.OAuth2Provider parameter
• AuthorizeHandler generates SessionID server-side using rand.Text() (never read from any client-supplied query parameter), stores it in PendingAuthorization.SessionID, sets PendingAuthorization.UpstreamProviderName to upstreamOrder[0], and redirects to the first upstream
• nextMissingUpstream helper method is implemented on *Handler and returns the name of the first upstream in upstreamOrder whose tokens are not yet present in GetAllUpstreamTokens(ctx, sessionID), or empty string when all are satisfied; returns upstreamOrder[0] on GetAllUpstreamTokens error (fail-safe restart)
• CallbackHandler uses pending.SessionID (not a freshly generated ID) to key all StoreUpstreamTokens calls for the completed leg
• CallbackHandler deletes the used PendingAuthorization before creating the next leg's PendingAuthorization (CSRF protection: fresh InternalState per leg)
• CallbackHandler creates a new PendingAuthorization with a fresh InternalState, the same SessionID, and UpstreamProviderName set to the next upstream name, then redirects to that upstream when nextMissingUpstream returns a non-empty name
• CallbackHandler calls writeAuthorizationResponse (issues the final authorization code) when nextMissingUpstream returns an empty string (all upstreams satisfied), threading the accumulated SessionID into the fosite session
• CallbackHandler uses pending.UpstreamProviderName (not string(h.upstream.Type())) as the providerID argument to h.userResolver.ResolveUser and h.userResolver.UpdateLastAuthenticated — this is the breaking change
• CallbackHandler resolves the correct upstream.OAuth2Provider from h.upstreams[pending.UpstreamProviderName] rather than using a single h.upstream field
• server_impl.go (newServer) builds the upstream map and ordered slice from cfg.Upstreams and passes them to handlers.NewHandler; the single-upstream path (existing GetUpstream() call) is replaced with iteration over cfg.Upstreams
• handlers_test.go and helpers_test.go are updated to call NewHandler with the new signature; the testStorageState.upstreamTokens field and its mock wiring are updated to match Phase 1's StoreUpstreamTokens(ctx, sessionID, providerName, tokens) signature
• New tests cover: SessionID is server-generated and present in PendingAuthorization after AuthorizeHandler runs; AuthorizeHandler correctly selects upstreamOrder[0] as the first leg provider
• New tests cover: nextMissingUpstream returns the correct first unsatisfied upstream name; returns empty when all are satisfied; returns upstreamOrder[0] when GetAllUpstreamTokens errors
• New tests cover: CallbackHandler redirects to the second upstream when nextMissingUpstream returns a non-empty name (single-leg-done state); CallbackHandler issues an authorization code when all upstreams are satisfied (all-satisfied state)
• New tests cover: fresh InternalState is generated per chain leg (the old PendingAuthorization is deleted before the new one is stored); the same SessionID is threaded through both legs
• server_test.go is updated to pass the new NewHandler signature (via newServer factory path)
• A full sequential chain integration test is added in pkg/authserver/integration_test.go: two mock upstream IDPs, full authorize → leg-1-callback → leg-2-callback → auth code → token exchange flow, verifying both providers' tokens are stored under the correct SessionID and the issued JWT contains the expected tsid claim
• PR description documents the ProviderIdentity breaking change: records stored with providerID = "oauth2" or "oidc" will not match post-upgrade; affected users receive fresh User records
• task license-check passes (SPDX headers on all new/modified .go files)
• task lint-fix passes with no remaining lint errors
• task test passes (unit tests)

Technical Approach

Recommended Implementation

Start with handler.go: replace the upstream upstream.OAuth2Provider field with upstreams map[string]upstream.OAuth2Provider and upstreamOrder []string, and update NewHandler's signature. This immediately breaks authorize.go, callback.go, and server_impl.go at compile time — use the compiler to drive the remaining changes.

In authorize.go, generate the SessionID (via rand.Text()) and set both PendingAuthorization.SessionID and PendingAuthorization.UpstreamProviderName = h.upstreamOrder[0]. Resolve the first upstream provider from h.upstreams[h.upstreamOrder[0]] for the AuthorizationURL call.

In callback.go, the most significant change is the chain continuation logic: after StoreUpstreamTokens, call nextMissingUpstream(ctx, pending.SessionID). If a non-empty name is returned, delete the current PendingAuthorization, create a new one (fresh InternalState, same SessionID, new UpstreamProviderName), store it, resolve the next upstream from h.upstreams[nextProvider], build the authorization URL, and redirect. If empty, fall through to writeAuthorizationResponse. The providerID passed to ResolveUser changes from string(h.upstream.Type()) to pending.UpstreamProviderName.

In server_impl.go, replace the GetUpstream() → single-provider path with a loop over cfg.Upstreams that builds the map and ordered slice.

Add nextMissingUpstream as a private method on *Handler after the public methods, following the codebase convention of public methods first.

Patterns and Frameworks

• rand.Text() for SessionID: Already used in the current CallbackHandler for the single-upstream sessionID. Phase 2 moves generation to AuthorizeHandler so the same SessionID persists across all chain legs. Never read SessionID from client-supplied request parameters.
• Defensive fresh InternalState per leg: Call h.storage.DeletePendingAuthorization(ctx, internalState) for the completed leg before calling h.storage.StorePendingAuthorization(ctx, newSecrets.State, newPending). This is the CSRF protection invariant: no two legs share the same InternalState.
• Fail-safe restart in nextMissingUpstream: On GetAllUpstreamTokens error, return h.upstreamOrder[0] (restart from the beginning). This is conservative but correct: the user repeats all upstream authorizations rather than getting stuck in a broken partial state.
• Immutable variable assignment with anonymous functions: Use for any providerID or nextProvider derivation per the codebase's Go style conventions (see CLAUDE.md).
• slog structured logging: Follow the existing slog.Error/slog.Warn/slog.Debug calls in authorize.go and callback.go. Never log token values.
• go.uber.org/mock (gomock): Handler tests use mocks.NewMockStorage. After Phase 1 updates the StoreUpstreamTokens mock signature, update helpers_test.go to use (ctx, sessionID, providerName, tokens) call pattern.
• Table-driven tests with require.NoError: Follow the patterns in authorize_test.go and callback_test.go.

Code Pointers

• pkg/authserver/server/handlers/handler.go — Primary struct and constructor to modify: replace upstream upstream.OAuth2Provider with upstreams map[string]upstream.OAuth2Provider and upstreamOrder []string; update NewHandler signature. Add nextMissingUpstream as a private method here or in a new helper file.
• pkg/authserver/server/handlers/authorize.go — Authorize handler: add SessionID generation (rand.Text()); set PendingAuthorization.SessionID and PendingAuthorization.UpstreamProviderName; resolve first upstream from h.upstreams[h.upstreamOrder[0]].
• pkg/authserver/server/handlers/callback.go lines 83–147 — Callback handler: the existing sessionID := rand.Text() on line 113 moves to authorize.go; instead read sessionID from pending.SessionID. Replace providerID := string(h.upstream.Type()) (line 98) with providerID := pending.UpstreamProviderName. Add nextMissingUpstream call after StoreUpstreamTokens; branch on result.
• pkg/authserver/server_impl.go lines 127–138 — Server construction: replace cfg.GetUpstream() → options.upstreamFactory(ctx, upstreamCfg) → handlers.NewHandler(provider, authServerConfig, stor, upstreamIDP) with a loop over cfg.Upstreams that builds the map and ordered slice.
• pkg/authserver/server/handlers/helpers_test.go lines 193–205 — Update StoreUpstreamTokens mock expectation from (ctx, sessionID, tokens) to (ctx, sessionID, providerName, tokens) after Phase 1 lands. Also add GetAllUpstreamTokens mock expectation for the multi-upstream chain tests.
• pkg/authserver/integration_test.go — Add TestIntegration_MultiUpstreamSequentialChain following the existing TestIntegration_FullPKCEFlow pattern but with two mockoidc.MockOIDC instances representing two upstream providers.
• pkg/authserver/server_test.go — TestNewServer_Success and TestNew use the withUpstreamFactory option; confirm these tests compile correctly after newServer's internal NewHandler call is updated.

Component Interfaces

// pkg/authserver/server/handlers/handler.go

// Handler provides HTTP handlers for the OAuth authorization server endpoints.
type Handler struct {
    provider      fosite.OAuth2Provider
    config        *server.AuthorizationServerConfig
    storage       storage.Storage
    upstreams     map[string]upstream.OAuth2Provider // keyed by provider name
    upstreamOrder []string                           // config order → chain sequence
    userResolver  *UserResolver
}

// NewHandler creates a new Handler with the given dependencies.
func NewHandler(
    provider fosite.OAuth2Provider,
    config *server.AuthorizationServerConfig,
    stor storage.Storage,
    upstreams map[string]upstream.OAuth2Provider,
    upstreamOrder []string,
    userResolver *UserResolver,
) *Handler

// nextMissingUpstream returns the name of the first upstream in upstreamOrder
// whose tokens are not yet present in storage for the given sessionID.
// Returns "" when all upstreams are satisfied.
// Returns upstreamOrder[0] on storage error (fail-safe: restart chain from beginning).
func (h *Handler) nextMissingUpstream(ctx context.Context, sessionID string) string {
    stored, err := h.storage.GetAllUpstreamTokens(ctx, sessionID)
    if err != nil {
        return h.upstreamOrder[0] // restart from beginning on error
    }
    for _, name := range h.upstreamOrder {
        if _, exists := stored[name]; !exists {
            return name
        }
    }
    return "" // all satisfied
}

// pkg/authserver/server/handlers/authorize.go — key changes

// AuthorizeHandler — new fields set on PendingAuthorization:
pending := &storage.PendingAuthorization{
    // ... existing fields ...
    UpstreamProviderName: h.upstreamOrder[0],       // first upstream in chain
    SessionID:            rand.Text(),               // server-side generated TSID
}

// Resolve provider for first leg:
firstProvider := h.upstreams[h.upstreamOrder[0]]
upstreamURL, err := firstProvider.AuthorizationURL(secrets.State, secrets.PKCEChallenge, authOpts...)

// pkg/authserver/server/handlers/callback.go — key changes

// providerID switches from upstream.Type() to pending.UpstreamProviderName:
providerID := pending.UpstreamProviderName

// sessionID comes from pending (not freshly generated):
sessionID := pending.SessionID

// After StoreUpstreamTokens, determine next upstream:
nextProvider := h.nextMissingUpstream(ctx, sessionID)
if nextProvider != "" {
    // Delete current pending (CSRF: fresh InternalState per leg)
    _ = h.storage.DeletePendingAuthorization(ctx, internalState)

    // Build next leg's PendingAuthorization (same SessionID, fresh InternalState)
    nextSecrets := newUpstreamAuthSecrets()
    nextPending := &storage.PendingAuthorization{
        ClientID:             pending.ClientID,
        RedirectURI:          pending.RedirectURI,
        State:                pending.State,
        PKCEChallenge:        pending.PKCEChallenge,
        PKCEMethod:           pending.PKCEMethod,
        Scopes:               pending.Scopes,
        InternalState:        nextSecrets.State,
        UpstreamPKCEVerifier: nextSecrets.PKCEVerifier,
        UpstreamNonce:        nextSecrets.Nonce,
        CreatedAt:            time.Now(),
        UpstreamProviderName: nextProvider,
        SessionID:            sessionID,
    }
    if err := h.storage.StorePendingAuthorization(ctx, nextSecrets.State, nextPending); err != nil {
        // handle error
    }
    nextUpstream := h.upstreams[nextProvider]
    nextURL, err := nextUpstream.AuthorizationURL(nextSecrets.State, nextSecrets.PKCEChallenge, authOpts...)
    // ...
    http.Redirect(w, req, nextURL, http.StatusFound)
    return
}

// All satisfied — issue authorization code (existing writeAuthorizationResponse call)

// pkg/authserver/server_impl.go — replace GetUpstream() single-provider path

// Build upstream map and ordered slice from config
upstreams := make(map[string]upstream.OAuth2Provider, len(cfg.Upstreams))
upstreamOrder := make([]string, 0, len(cfg.Upstreams))
for i := range cfg.Upstreams {
    upCfg := &cfg.Upstreams[i]
    provider, err := options.upstreamFactory(ctx, upCfg)
    if err != nil {
        return nil, fmt.Errorf("failed to create upstream provider %q: %w", upCfg.Name, err)
    }
    upstreams[upCfg.Name] = provider
    upstreamOrder = append(upstreamOrder, upCfg.Name)
}

handlerInstance := handlers.NewHandler(provider, authServerConfig, stor, upstreams, upstreamOrder, handlers.NewUserResolver(stor))

Testing Strategy

Unit Tests — pkg/authserver/server/handlers/authorize_test.go

• TestAuthorizeHandler_MultiUpstream_SessionIDGenerated: verify that after AuthorizeHandler runs, the stored PendingAuthorization has a non-empty SessionID (generated server-side) and UpstreamProviderName == upstreamOrder[0]
• TestAuthorizeHandler_MultiUpstream_RedirectsToFirstUpstream: verify the 302 redirect URL corresponds to the first upstream provider's authorization URL
• TestAuthorizeHandler_NoUpstreams_ReturnsServerError: verify error handling when upstreamOrder is empty

Unit Tests — pkg/authserver/server/handlers/callback_test.go

• TestCallbackHandler_FirstLegComplete_RedirectsToSecondUpstream: configure two upstreams, simulate a callback completing the first leg (no tokens yet for the second upstream via GetAllUpstreamTokens), verify nextMissingUpstream returns the second provider name and the handler redirects to its authorization URL
• TestCallbackHandler_AllSatisfied_IssuesAuthCode: configure two upstreams with both already having tokens in GetAllUpstreamTokens, verify the handler calls writeAuthorizationResponse (issues auth code to client, 303 response)
• TestCallbackHandler_FreshInternalStatePerLeg: after completing leg 1, verify the used PendingAuthorization (keyed by internalState) is deleted and a new entry with a different InternalState is created for leg 2
• TestCallbackHandler_SessionIDThreadedAcrossLegs: verify the SessionID in the new leg-2 PendingAuthorization equals the SessionID from the leg-1 PendingAuthorization
• TestCallbackHandler_ProviderIDFromPendingName: verify ResolveUser and UpdateLastAuthenticated are called with pending.UpstreamProviderName (not upstream.Type())

Unit Tests — nextMissingUpstream helper

• Returns first upstream name when no tokens are stored (all missing)
• Returns second upstream name when only first upstream's tokens are stored
• Returns "" when all upstreams' tokens are stored
• Returns upstreamOrder[0] when GetAllUpstreamTokens returns an error

Integration Tests — pkg/authserver/integration_test.go

• TestIntegration_MultiUpstreamSequentialChain: full end-to-end test with two mockoidc.MockOIDC instances as upstream providers. Flow: GET /oauth/authorize → redirect to provider-1 → GET /oauth/callback (provider-1 code) → redirect to provider-2 → GET /oauth/callback (provider-2 code) → 303 to redirect_uri with auth code → POST /oauth/token → JWT access token. Verify: (1) both providers' tokens are stored under the correct SessionID; (2) the issued JWT contains the tsid claim equal to the SessionID; (3) the final auth code is issued only after both callbacks complete; (4) the client observes exactly two upstream redirects (one per provider) before receiving the auth code redirect.

Edge Cases

• Single-upstream deployment: verify backward compatibility — when upstreamOrder has exactly one entry, the handler behaves identically to the current single-upstream implementation (no second redirect, auth code issued after first callback)
• GetAllUpstreamTokens returns error during nextMissingUpstream: verify the chain restarts from upstreamOrder[0] (existing session data is preserved in storage; only the in-memory decision restarts)
• Replayed InternalState from a completed leg: verify LoadPendingAuthorization returns ErrNotFound (the entry was deleted after the leg completed), and the callback handler returns a 400 with an appropriate error

Out of Scope

• Storage layer changes (UpstreamTokenStorage interface, backends, PendingAuthorization struct): covered in TASK-001 (Implement multi-provider upstream token storage layer (RFC-0052 Phase 1) #4136) and must already be merged
• Config and operator validation changes (removing/moving len > 1 guard): covered in TASK-003
• Identity.UpstreamTokens enrichment and upstreamswap middleware refactor: covered in TASK-004
• UpstreamTokenRefresher multi-provider support: Phase 1 updated the storage call; full multi-provider refresh chain is deferred to a later RFC
• Distributed locking for concurrent token refresh races
• Dynamic upstream IDP registration at runtime

References

• RFC: docs/proposals/THV-0052-multi-upstream-idp-authserver.md
• Parent epic: Auth Server: multi-upstream provider support #3924
• Phase 1 (dependency): Implement multi-provider upstream token storage layer (RFC-0052 Phase 1) #4136
• Architecture doc for auth server storage: docs/arch/11-auth-server-storage.md
• Related Stacklok epic: https://github.com/stacklok/stacklok-epics/issues/251