Implement multi-provider upstream token storage layer (RFC-0052 Phase 1)

[tgrunnagle]

Description

Add a providerName dimension to the UpstreamTokenStorage interface so that multiple upstream IDPs can have their tokens stored per-session without overwriting each other. This is the foundational storage layer change for RFC-0052 — all subsequent phases (handler chain, config/operator, identity enrichment) break at compile time without these new interface signatures landing first.

Context

RFC-0052 extends the embedded OAuth authorization server to support multiple upstream IDPs in a single instance. Today, UpstreamTokenStorage stores exactly one set of tokens per session (map[sessionID]*timedEntry), which makes it impossible to accumulate tokens from a second provider during a sequential authorization chain. This task restructures the storage layer to accommodate N providers per session by keying on (sessionID, providerName).

This is the root task of the RFC-0052 implementation plan. TASK-002 (Handler and Server Layer), TASK-003 (Config and Operator), and TASK-004 (Identity Enrichment) all depend on the new interface signatures produced here.

Dependencies: None (root task)
Blocks: TASK-002, TASK-003, TASK-004

Breaking Change — Redis key pattern: Upstream tokens stored under the old Redis key pattern {prefix}upstream:{sessionID} (a single key per session) will be unreachable after upgrade. There is no automatic migration. Existing sessions will require re-authentication after deploying this change. This must be documented in the PR description.

Acceptance Criteria

• UpstreamTokenStorage interface in pkg/authserver/storage/types.go has updated signatures: StoreUpstreamTokens(ctx, sessionID, providerName string, tokens *UpstreamTokens), GetUpstreamTokens(ctx, sessionID, providerName string), GetAllUpstreamTokens(ctx, sessionID string) (map[string]*UpstreamTokens, error), and DeleteUpstreamTokens(ctx, sessionID string) (session-scoped delete, signature unchanged)
• PendingAuthorization struct in pkg/authserver/storage/types.go has two new fields: UpstreamProviderName string and SessionID string with doc comments matching the architecture spec
• MemoryStorage.upstreamTokens field is restructured from map[string]*timedEntry[*UpstreamTokens] to map[string]map[string]*timedEntry[*UpstreamTokens] (outer key = sessionID, inner key = providerName)
• MemoryStorage.StoreUpstreamTokens and GetUpstreamTokens accept and use providerName; both validate that providerName is non-empty and return a clear error if it is
• MemoryStorage.GetAllUpstreamTokens is implemented: returns a map[string]*UpstreamTokens (defensive copies) for all providers under the given sessionID; returns an empty map (not an error) for an unknown sessionID
• MemoryStorage eviction TTL for upstream tokens uses max(accessToken.ExpiresAt, refreshToken.ExpiresAt) instead of tokens.ExpiresAt.Add(DefaultRefreshTokenTTL)
• Redis storage uses new key pattern {prefix}upstream:{sessionID}:{providerName} for each provider key (via new redisUpstreamKey helper in redis_keys.go)
• Redis storage maintains a session-scoped index SET at {prefix}upstream:idx:{sessionID} (using SADD) for efficient GetAllUpstreamTokens and session-level delete
• RedisStorage.StoreUpstreamTokens uses a Lua script for atomic write of the provider key + SADD to the session index SET
• RedisStorage.DeleteUpstreamTokens uses a Lua script for atomic SMEMBERS + DEL of all provider keys + DEL of the index SET
• RedisStorage.GetAllUpstreamTokens is implemented using SMEMBERS + MGET (two round-trips); returns empty map for unknown session
• RedisStorage.StoreUpstreamTokens and GetUpstreamTokens validate that providerName is non-empty and return a clear error if it is
• All call sites of StoreUpstreamTokens and GetUpstreamTokens within pkg/authserver/ are updated to pass a providerName argument: pkg/authserver/server/handlers/callback.go and pkg/authserver/refresher.go
• Storage mocks in pkg/authserver/storage/mocks/mock_storage.go are regenerated via mockgen to reflect the new interface signatures
• memory_test.go is updated: existing upstream token tests pass providerName; new test cases cover multi-provider store/get, session delete wipes all providers, GetAllUpstreamTokens returns all providers (including those with expired tokens), GetAllUpstreamTokens returns empty map for unknown session, and eviction TTL uses max(access.ExpiresAt, refresh.ExpiresAt)
• redis_test.go and redis_integration_test.go are updated with equivalent coverage for the Redis backend (miniredis and real Redis respectively)
• task license-check passes (all new/modified .go files have SPDX headers)
• task lint-fix passes with no remaining lint errors
• task test passes (unit tests)
• PR description documents the Redis breaking change (old key pattern {prefix}upstream:{sessionID} becomes unreachable; sessions require re-authentication)
• upstreamtoken.Service interface in pkg/auth/upstreamtoken/types.go has updated GetValidTokens signature: GetValidTokens(ctx context.Context, sessionID, providerName string) (*UpstreamCredential, error)
• InProcessService.GetValidTokens in pkg/auth/upstreamtoken/service.go accepts and uses providerName when calling storage.GetUpstreamTokens; refreshOrFail is updated to thread providerName through its call chain
• upstreamswap middleware Config struct in pkg/auth/upstreamswap/middleware.go gains a ProviderName string field; validateConfig returns an error if ProviderName is empty
• upstreamswap middleware passes cfg.ProviderName as the providerName argument to svc.GetValidTokens in createMiddlewareFunc
• The proxy runner derives ProviderName from the single upstream provider name in the embedded auth server config and stores it in the MiddlewareConfig parameters for the upstreamswap middleware during wiring (the proxy runner only ever has one upstream provider, so a single ProviderName in MiddlewareConfig is sufficient for Phase 1)
• Mocks for upstreamtoken.Service (if generated) are regenerated to reflect the new GetValidTokens signature
• pkg/auth/upstreamswap/middleware_test.go is updated: existing tests pass providerName to GetValidTokens; new test cases cover empty ProviderName config validation and correct forwarding of ProviderName to the service

Technical Approach

Recommended Implementation

Start with pkg/authserver/storage/types.go — update the UpstreamTokenStorage interface and PendingAuthorization struct first so the compiler immediately surfaces every call site that needs updating. Then update memory.go (restructure the nested map and TTL logic), then redis_keys.go (add redisUpstreamKey), then redis.go (new key pattern, session index, Lua scripts, GetAllUpstreamTokens). Update the two call sites in callback.go and refresher.go last (they will need a providerName argument — derive this from pending.UpstreamProviderName in the callback handler, and from expired.ProviderID in the refresher). Regenerate mocks after all interface changes are complete.

Next, propagate providerName up through the token service layer: update upstreamtoken.Service.GetValidTokens in pkg/auth/upstreamtoken/types.go and InProcessService.GetValidTokens in pkg/auth/upstreamtoken/service.go. Finally, wire the provider name into the upstreamswap middleware: add ProviderName to Config, enforce it in validateConfig, pass it to GetValidTokens in createMiddlewareFunc, and update the proxy runner to derive the provider name from the embedded auth server config when constructing the middleware config.

Patterns and Frameworks

• Defensive copies on read and write: Follow the existing pattern in MemoryStorage.StoreUpstreamTokens and GetUpstreamTokens — make a full struct copy before storing and before returning, so external mutations do not affect stored state. Preserve this when restructuring to the nested map.
• Lua scripts for atomicity in Redis: Follow the existing storeUpstreamTokensScript pattern in redis.go. The new StoreUpstreamTokens Lua script must atomically write the provider key and SADD to the session index SET in one round-trip. The new DeleteUpstreamTokens Lua script must atomically SMEMBERS, DEL all member keys, and DEL the index SET.
• warnOnCleanupErr for best-effort secondary index cleanup: Use this existing helper when non-critical index cleanup operations fail (see redis.go).
• redisKey / redisProviderKey / redisSetKey key generation conventions: Follow the same fmt.Sprintf format in redis_keys.go for the new redisUpstreamKey helper.
• errors.Is() / fmt.Errorf with %w: Use for error wrapping and inspection throughout; never swallow errors silently.
• Table-driven tests with require.NoError: Follow patterns in memory_test.go and redis_test.go; use require.NoError(t, err) rather than t.Fatal.
• go.uber.org/mock (mockgen): Regenerate mocks with mockgen -destination=mocks/mock_storage.go -package=mocks -source=types.go Storage,PendingAuthorizationStorage,ClientRegistry,UpstreamTokenStorage,UpstreamTokenRefresher,UserStorage.
• SPDX license headers: All new or modified .go files must have SPDX headers; run task license-fix to add them automatically.

Code Pointers

• pkg/authserver/storage/types.go — Primary file to modify: UpstreamTokenStorage interface signatures (add providerName param to StoreUpstreamTokens/GetUpstreamTokens, add GetAllUpstreamTokens), and PendingAuthorization struct (add UpstreamProviderName and SessionID fields). The //go:generate mockgen directive at the top of this file controls mock regeneration.
• pkg/authserver/storage/memory.go lines 48–80 — MemoryStorage struct definition: change upstreamTokens map[string]*timedEntry[*UpstreamTokens] to map[string]map[string]*timedEntry[*UpstreamTokens].
• pkg/authserver/storage/memory.go lines 685–777 — Current StoreUpstreamTokens, GetUpstreamTokens, DeleteUpstreamTokens: shows the defensive copy pattern, timedEntry wrapper, eviction TTL calculation, and error wrapping to replicate when restructuring. TTL calculation on line 699 uses tokens.ExpiresAt.Add(DefaultRefreshTokenTTL) — this must change to max(accessToken.ExpiresAt, refreshToken.ExpiresAt) (note: UpstreamTokens currently has only one ExpiresAt field; confirm whether the RFC intent is to track RefreshToken expiry separately or to use existing ExpiresAt as the access token expiry).
• pkg/authserver/storage/redis_keys.go — Add redisUpstreamKey(prefix, sessionID, providerName string) string returning "{prefix}upstream:{sessionID}:{providerName}". The session index SET key should follow the redisSetKey pattern: redisSetKey(prefix, "upstream:idx", sessionID) → "{prefix}upstream:idx:{sessionID}".
• pkg/authserver/storage/redis.go lines 889–1004 — Current Redis StoreUpstreamTokens, GetUpstreamTokens, DeleteUpstreamTokens and storeUpstreamTokensScript Lua script. The new implementation must replace the single-key pattern with the two-key pattern (provider key + session index SET). The marshalUpstreamTokensWithTTL helper (line ~885) is reusable.
• pkg/authserver/storage/mocks/mock_storage.go — Generated file; do not edit by hand. Regenerate after all interface changes.
• pkg/authserver/server/handlers/callback.go line 127 — Call site for StoreUpstreamTokens; must add providerName argument. At this call site, providerName should come from pending.UpstreamProviderName (set by the authorize handler when it creates the PendingAuthorization). Note: in the current single-upstream handler, pending.UpstreamProviderName does not exist yet — adding it to PendingAuthorization here is sufficient for Phase 1 correctness; Phase 2 will drive multi-upstream chain logic that populates it per-leg.
• pkg/authserver/server/handlers/callback.go line 141 — DeleteUpstreamTokens call: signature is unchanged (session-scoped).
• pkg/authserver/refresher.go line 70 — Call site for StoreUpstreamTokens; must add providerName argument. Use expired.ProviderID as the providerName (the refresher already has access to the expired token struct which carries the provider ID).
• pkg/authserver/storage/memory_test.go line 429 — TestMemoryStorage_UpstreamTokens test suite: update all test cases with the new providerName parameter and add new multi-provider and GetAllUpstreamTokens cases.
• pkg/authserver/storage/redis_test.go and pkg/authserver/storage/redis_integration_test.go — Redis equivalents; follow the same test pattern updates.
• pkg/auth/upstreamtoken/types.go — Service interface: update GetValidTokens signature to add providerName string as the second parameter (after sessionID).
• pkg/auth/upstreamtoken/service.go line 50 — InProcessService.GetValidTokens: add providerName string parameter; pass it to storage.GetUpstreamTokens and thread it through refreshOrFail (which calls RefreshAndStore — the refresher signature itself is not changing in Phase 1, but providerName must be available so StoreUpstreamTokens inside the refresher can use it).
• pkg/auth/upstreamswap/middleware.go line 32 — Config struct: add ProviderName string \json:"provider_name" yaml:"provider_name"`` field representing the single upstream provider this proxy runner instance serves.
• pkg/auth/upstreamswap/middleware.go line 96 — validateConfig: add a check that cfg.ProviderName != ""; return a descriptive error if empty.
• pkg/auth/upstreamswap/middleware.go line 190 — svc.GetValidTokens(r.Context(), tsid) call site: update to svc.GetValidTokens(r.Context(), tsid, cfg.ProviderName).
• Proxy runner wiring code (search for upstreamswap or MiddlewareConfig construction in pkg/runner/ or cmd/thv/) — locate where the upstreamswap middleware config is built and add ProviderName derived from the single upstream IDP provider name in the embedded auth server config.

Component Interfaces

// pkg/authserver/storage/types.go

// UpstreamTokenStorage — updated interface
type UpstreamTokenStorage interface {
    // StoreUpstreamTokens stores upstream IDP tokens for (sessionID, providerName).
    // Returns an error if sessionID or providerName is empty.
    StoreUpstreamTokens(ctx context.Context, sessionID, providerName string, tokens *UpstreamTokens) error

    // GetUpstreamTokens retrieves upstream IDP tokens for (sessionID, providerName).
    // Returns ErrNotFound if the entry does not exist.
    // Returns ErrExpired (with token data) if the access token has expired — caller
    // may use the refresh token for transparent renewal.
    // Returns ErrInvalidBinding if binding validation fails.
    GetUpstreamTokens(ctx context.Context, sessionID, providerName string) (*UpstreamTokens, error)

    // GetAllUpstreamTokens returns all upstream tokens for the session, keyed by
    // provider name. Returns an empty map (not an error) for an unknown sessionID.
    // Expired tokens are included; callers that need freshness checks should use
    // GetUpstreamTokens per-provider.
    GetAllUpstreamTokens(ctx context.Context, sessionID string) (map[string]*UpstreamTokens, error)

    // DeleteUpstreamTokens removes all upstream tokens for the session (all providers).
    // Returns ErrNotFound if the session does not exist.
    DeleteUpstreamTokens(ctx context.Context, sessionID string) error
}

// PendingAuthorization — new fields added to existing struct
type PendingAuthorization struct {
    // ... existing fields unchanged ...

    // UpstreamProviderName is the provider being authenticated in this chain leg.
    // Set by the authorize handler when creating the PendingAuthorization;
    // read by the callback handler to key token storage.
    UpstreamProviderName string

    // SessionID is the TSID for the session being accumulated across all chain legs.
    // Generated server-side only using rand.Text(); never accepted from client input.
    SessionID string
}

// pkg/authserver/storage/redis_keys.go — new helper

// redisUpstreamKey generates a Redis key for a per-provider upstream token entry.
// Format: "{prefix}upstream:{sessionID}:{providerName}"
// Note: sessionID and providerName must not contain colons in practice
// (sessionID is rand.Text() output; providerName is a config-supplied name).
func redisUpstreamKey(prefix, sessionID, providerName string) string {
    return fmt.Sprintf("%s%s:%s:%s", prefix, KeyTypeUpstream, sessionID, providerName)
}

// Session index SET key (use existing redisSetKey):
// redisSetKey(prefix, "upstream:idx", sessionID) → "{prefix}upstream:idx:{sessionID}"

// pkg/authserver/storage/memory.go — MemoryStorage struct field change

// upstreamTokens maps sessionID -> providerName -> timedEntry.
// Outer key is the session TSID; inner key is the provider name.
upstreamTokens map[string]map[string]*timedEntry[*UpstreamTokens]

// pkg/auth/upstreamtoken/types.go — updated Service interface

// Service owns the upstream token lifecycle: read, refresh, error handling.
type Service interface {
    // GetValidTokens returns a valid upstream credential for (sessionID, providerName).
    // It transparently refreshes expired access tokens using the refresh token.
    //
    // Returns:
    //   - *UpstreamCredential on success
    //   - ErrSessionNotFound if no upstream tokens exist for the session
    //   - ErrNoRefreshToken if the access token is expired and no refresh token is available
    //   - ErrRefreshFailed if the refresh attempt fails (e.g., revoked refresh token)
    GetValidTokens(ctx context.Context, sessionID, providerName string) (*UpstreamCredential, error)
}

// pkg/auth/upstreamswap/middleware.go — Config struct update

// Config holds configuration for upstream swap middleware.
type Config struct {
    // ProviderName is the upstream IDP provider name used to key token storage.
    // For the proxy runner this is derived from the single upstream provider in
    // the embedded auth server config. Required — validation returns an error if empty.
    ProviderName string `json:"provider_name" yaml:"provider_name"`

    // HeaderStrategy determines how to inject the token: "replace" (default) or "custom".
    HeaderStrategy string `json:"header_strategy,omitempty" yaml:"header_strategy,omitempty"`

    // CustomHeaderName is the header name when HeaderStrategy is "custom".
    CustomHeaderName string `json:"custom_header_name,omitempty" yaml:"custom_header_name,omitempty"`
}

Testing Strategy

Unit Tests — pkg/authserver/storage/memory_test.go

• TestMemoryStorage_UpstreamTokens: update all existing subtests to pass providerName (e.g., "provider-a")
• Multi-provider store/get: store tokens for two providers under the same sessionID; verify each GetUpstreamTokens call returns the correct tokens for each provider independently
• Cross-provider isolation: storing tokens for provider A does not overwrite or affect tokens for provider B
• GetAllUpstreamTokens — two providers: store two providers, call GetAllUpstreamTokens, verify map has exactly two entries with correct defensive copies
• GetAllUpstreamTokens — unknown session: call GetAllUpstreamTokens for a sessionID that has never been stored; verify returns empty map and no error
• GetAllUpstreamTokens — includes expired tokens: store a token with an expired ExpiresAt; verify GetAllUpstreamTokens includes it (no expiry filtering at bulk-read level)
• Session delete wipes all providers: store two providers, call DeleteUpstreamTokens, verify both providers are gone (subsequent GetUpstreamTokens for each returns ErrNotFound)
• Eviction TTL uses max(access.ExpiresAt, refresh.ExpiresAt): confirm the cleanup loop evicts entries at the correct deadline (test using manipulated time.Now or very short TTLs)
• StoreUpstreamTokens with empty providerName returns a non-nil error
• GetUpstreamTokens with empty providerName returns a non-nil error

Unit Tests — pkg/authserver/storage/redis_test.go (miniredis)

• Same multi-provider store/get coverage as memory tests above
• GetAllUpstreamTokens returns empty map for unknown session (SMEMBERS on non-existent key)
• Session index SET ({prefix}upstream:idx:{sessionID}) is populated on store and cleaned up atomically on delete
• StoreUpstreamTokens with empty providerName returns a non-nil error
• GetUpstreamTokens with empty providerName returns a non-nil error

Integration Tests — pkg/authserver/storage/redis_integration_test.go

• Multi-provider round-trip: store two providers, retrieve individually and via GetAllUpstreamTokens, verify correct values (requires //go:build integration tag and real Redis via testcontainers)
• Session delete wipes all providers atomically in Redis
• TTL expiry: keys expire from Redis after the configured TTL (verify via PTTL or waiting)

Unit Tests — pkg/auth/upstreamswap/middleware_test.go

• Update existing tests: mock svc.GetValidTokens now expects (ctx, sessionID, providerName) — update expectations to pass the configured ProviderName
• Empty ProviderName in config: CreateMiddleware (or validateConfig) returns a non-nil error
• ProviderName forwarding: verify the middleware passes cfg.ProviderName to GetValidTokens (not a hardcoded value)

Edge Cases

• GetAllUpstreamTokens with a sessionID that has one expired and one valid provider: verify both are returned (expiry is per-provider at the GetUpstreamTokens level, not filtered at bulk-read)
• Concurrent StoreUpstreamTokens for different providers under the same sessionID (memory backend): verify no data corruption under the sync.RWMutex
• Redis Lua script atomicity: verify that a StoreUpstreamTokens failure partway through does not leave the session index SET in an inconsistent state (e.g., SADD succeeds but key write fails — Lua rollback)

Out of Scope

• Handler layer changes (sequential chain logic, nextMissingUpstream, multi-upstream authorize/callback): covered in TASK-002
• Config and operator validation changes (removing/moving len > 1 guard): covered in TASK-003
• Identity.UpstreamTokens enrichment: covered in TASK-004
• Full Redis migration tooling for tokens stored under the old {prefix}upstream:{sessionID} key pattern
• UpstreamTokenRefresher.RefreshAndStore multi-provider support (the refresher in pkg/authserver/refresher.go currently takes a single sessionID; Phase 1 only updates the StoreUpstreamTokens call within it to pass providerName — no signature change to RefreshAndStore itself)
• Transparent token refresh or distributed locking for concurrent refresh races
• Dynamic upstream IDP registration at runtime
• Multi-provider upstreamswap middleware support (TASK-004 will handle the case where the proxy runner needs to serve multiple upstream providers; Phase 1 configures a single ProviderName per middleware instance)

References

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