Enrich Identity with upstream tokens and simplify upstreamswap middleware (RFC-0052 Phase 4)

[tgrunnagle]

Description

Wire upstream token enrichment into the auth middleware and simplify the upstreamswap middleware by removing its storage dependency. After this change, the TokenValidator eagerly loads all upstream access tokens into Identity.UpstreamTokens at JWT validation time, so upstreamswap only needs to read identity.UpstreamTokens[cfg.ProviderName] directly — no storage call, no context extraction, no ServiceGetter pattern.

Context

RFC-0052 introduces a sequential authorization chain that accumulates tokens from multiple upstream IDPs per session. The current upstreamswap middleware retrieves the correct upstream token by calling InProcessService.GetValidTokens(ctx, tsid) at request time, which couples the middleware to the UpstreamTokenStorage dependency via a ServiceGetter function. This phase decouples that by moving the storage read into the auth middleware layer — a single GetAllUpstreamTokens bulk read happens once per request during JWT validation and the results are stored in the Identity struct. All downstream middleware (including upstreamswap) then read from the pre-enriched Identity with no additional storage round-trips.

This is Phase 4 of the RFC-0052 implementation. It depends on Phase 1 (#4136), which introduced the GetAllUpstreamTokens bulk-read method on UpstreamTokenStorage. It may be developed in parallel with TASK-002 and TASK-003 after TASK-001 merges.

Dependencies: #4136
Blocks: None (terminal phase)

Acceptance Criteria

• Identity struct in pkg/auth/identity.go has a new UpstreamTokens map[string]string field (providerName → access token) with a doc comment stating it is populated by auth middleware after JWT validation, contains expired tokens as-is, and is empty for non-embedded-AS requests
• Identity.MarshalJSON() redacts UpstreamTokens values (replaces each token value with "REDACTED") while keeping provider name keys visible; provider names are safe to log but token values are not
• Identity.String() does not leak upstream token values (the existing format Identity{Subject:%q} is sufficient; no additional changes required unless token values would otherwise appear)
• TokenValidator struct in pkg/auth/token.go has a new optional field upstreamStorage storage.UpstreamTokenStorage (nil = no enrichment; used for non-embedded-AS deployments where no storage exists)
• Between claimsToIdentity (~line 1087) and WithIdentity (~line 1096) in TokenValidator.Middleware, when v.upstreamStorage is non-nil and the identity has a tsid claim, GetAllUpstreamTokens(r.Context(), tsid) is called and the result is used to populate identity.UpstreamTokens
• When GetAllUpstreamTokens returns an error, the middleware logs at WARN level and continues with an empty (or nil) UpstreamTokens map — the error is non-fatal and the request proceeds
• When the tsid claim is absent or v.upstreamStorage is nil, identity.UpstreamTokens remains nil/empty and the request proceeds without error
• GetAuthenticationMiddleware in pkg/auth/utils.go accepts an optional storage.UpstreamTokenStorage parameter and threads it through to NewTokenValidator
• CreateMiddleware in pkg/auth/middleware.go extracts the storage dependency from the runner and passes it to GetAuthenticationMiddleware
• upstreamswap.Config in pkg/auth/upstreamswap/middleware.go has a new ProviderName string field with a comment stating it is derived from Upstreams[0].Name at construction time and defaults to "default"
• ServiceGetter type and all references to serviceGetter/storageGetter are removed from pkg/auth/upstreamswap/middleware.go; createMiddlewareFunc no longer accepts a ServiceGetter parameter
• upstreamswap middleware reads identity.UpstreamTokens[cfg.ProviderName] directly from the pre-enriched Identity; when the token is absent or empty it returns HTTP 401 with a WWW-Authenticate: Bearer error="invalid_token" header
• addUpstreamSwapMiddleware in pkg/runner/middleware.go derives ProviderName from config.EmbeddedAuthServerConfig.Upstreams[0].Name, falling back to "default" when the name is empty; the derived ProviderName is set on upstreamswap.Config before serialization
• GetUpstreamTokenService() is removed from the MiddlewareRunner interface in pkg/transport/types/transport.go if no other middleware callers remain after this refactor; the mock at pkg/transport/types/mocks/mock_transport.go is regenerated
• pkg/auth/middleware_test.go is updated: tests verify GetAllUpstreamTokens is called when tsid claim is present, expired tokens are included in UpstreamTokens as-is, and UpstreamTokens is nil/empty when no tsid claim is present
• pkg/auth/upstreamswap/middleware_test.go is updated: tests verify ProviderName config field, middleware reads from identity.UpstreamTokens directly, middleware returns 401 when the named provider token is absent, and StorageGetter is no longer part of the test setup
• pkg/runner/middleware_test.go is updated: tests verify ProviderName is derived from Upstreams[0].Name, fallback to "default" when name is empty, and StorageGetter is no longer injected into middleware params
• 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)

Technical Approach

Recommended Implementation

Start with the data model in pkg/auth/identity.go — add the UpstreamTokens field and update MarshalJSON to redact token values. This is a self-contained change with no dependencies.

Next, add the upstreamStorage field to TokenValidator and insert the enrichment call in TokenValidator.Middleware. The insertion point is between the existing claimsToIdentity call and the WithIdentity call — this two-line gap (lines ~1087–1096) is where the new block belongs. Use an optional field pattern (nil = no enrichment) to preserve existing behavior for all non-embedded-AS deployments.

Then thread UpstreamTokenStorage from the runner through to TokenValidator by updating GetAuthenticationMiddleware in pkg/auth/utils.go (add an optional parameter or a functional option) and CreateMiddleware in pkg/auth/middleware.go (the runner already provides IDPTokenStorage() indirectly via embeddedAuthServer).

The threading path requires the runner to expose IDPTokenStorage() to the auth.CreateMiddleware factory. Currently the runner exposes GetUpstreamTokenService() on the MiddlewareRunner interface. After this phase, assess whether GetUpstreamTokenService() has any remaining callers — if upstreamswap no longer calls it, remove it from the interface and regenerate the mock.

Finally, simplify pkg/auth/upstreamswap/middleware.go: remove ServiceGetter, add ProviderName to Config, and rewrite createMiddlewareFunc to read identity.UpstreamTokens[cfg.ProviderName] instead of calling the service. Update addUpstreamSwapMiddleware in pkg/runner/middleware.go to derive ProviderName from config.EmbeddedAuthServerConfig.Upstreams[0].Name with fallback to "default".

Patterns & Frameworks

• Optional field pattern for storage dependency: Add upstreamStorage storage.UpstreamTokenStorage as an optional field on TokenValidator (nil = feature disabled). This is consistent with Go's convention of zero-value-safe types. Existing deployments without an embedded auth server pass nil and observe no behavior change.
• Enrichment is non-fatal: Per architecture doc core principle Implement secret store #4: if GetAllUpstreamTokens fails, log at WARN with slog.WarnContext and continue with empty UpstreamTokens. Do not fail the entire request — the downstream backend will produce its own auth error.
• No token values in logs: Per core principle Implement secret injection #5: identity.UpstreamTokens values (access tokens) must be redacted in MarshalJSON. Provider name keys are safe to log. The String() method currently returns only Identity{Subject:%q} — this format is already safe.
• slog structured logging: Use slog.WarnContext(r.Context(), "failed to load upstream tokens", "err", err) for the non-fatal error log. Do not use slog.Warn (use the context-aware variant).
• Immutable assignment in middleware: Use an immediately-invoked anonymous function for any conditional value derivation (e.g., when computing the fallback provider name), per ToolHive Go coding style.
• go.uber.org/mock (mockgen): If GetUpstreamTokenService() is removed from MiddlewareRunner, regenerate pkg/transport/types/mocks/mock_transport.go using go generate ./pkg/transport/types/... or the directive at the top of transport.go.
• SPDX license headers: All new or modified .go files must have // SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. and // SPDX-License-Identifier: Apache-2.0. Run task license-fix to add them automatically.
• Table-driven tests with require.NoError: Follow the existing patterns in middleware_test.go and upstreamswap/middleware_test.go. Use require.NoError(t, err) rather than t.Fatal.

Code Pointers

• pkg/auth/identity.go — Add UpstreamTokens map[string]string field to Identity struct and update MarshalJSON to redact values. The SafeIdentity shadow struct pattern already handles redaction of Token; add a parallel map with redacted values for UpstreamTokens.
• pkg/auth/token.go lines 354–383 — TokenValidator struct definition. Add upstreamStorage storage.UpstreamTokenStorage as the last field. Lines 1086–1097 — the insertion point for the enrichment block between claimsToIdentity and WithIdentity.
• pkg/auth/token.go lines 557–655 — NewTokenValidator constructor. Add a WithUpstreamStorage(stor storage.UpstreamTokenStorage) TokenValidatorOption functional option (or equivalent) to wire the storage dependency without breaking the existing call signature.
• pkg/auth/utils.go line 64–76 — GetAuthenticationMiddleware. Thread UpstreamTokenStorage through here to NewTokenValidator. The function currently creates NewTokenValidator(ctx, *oidcConfig) — add an optional parameter or a functional option for the storage.
• pkg/auth/middleware.go lines 47–74 — CreateMiddleware factory. The runner exposes GetUpstreamTokenService() on types.MiddlewareRunner today; after this change a new method (or repurposed method) is needed to expose IDPTokenStorage(). Evaluate whether to add a GetUpstreamTokenStorage() storage.UpstreamTokenStorage method to MiddlewareRunner or reuse the existing interface extension point.
• pkg/auth/upstreamswap/middleware.go — Complete rewrite of createMiddlewareFunc. Remove ServiceGetter type (line 47) and the serviceGetter parameter. Add ProviderName to Config struct (line 32–38). In the middleware handler, replace the four-step tsid → serviceGetter → GetValidTokens → inject flow with a two-step UpstreamTokens[cfg.ProviderName] → inject read.
• pkg/runner/middleware.go lines 252–282 — addUpstreamSwapMiddleware. Derive ProviderName from config.EmbeddedAuthServerConfig.Upstreams[0].Name, defaulting to "default" when the name is empty. Set it on the upstreamswap.Config before creating the middleware config.
• pkg/transport/types/transport.go lines 65–87 — MiddlewareRunner interface. After removing the ServiceGetter dependency from upstreamswap, assess whether GetUpstreamTokenService() (line 86) has any remaining callers. If none remain, remove it and add a GetUpstreamTokenStorage() storage.UpstreamTokenStorage method if needed by the auth middleware enrichment path.
• pkg/transport/types/mocks/mock_transport.go — Regenerated file; do not edit by hand. Run go generate ./pkg/transport/types/... after any interface changes.
• pkg/auth/upstreamswap/middleware_test.go — Reference for the existing test patterns. After removing ServiceGetter, simplify requestWithIdentity to include UpstreamTokens in the Identity directly rather than using mock storage. Remove serviceGetterFromMocks and nilServiceGetter helpers that are no longer needed.
• pkg/runner/middleware_test.go lines 22–41 — createMinimalAuthServerConfig() helper with Upstreams[0].Name = "test-upstream". New test cases for ProviderName derivation should verify that the serialized params contain ProviderName: "test-upstream".

Component Interfaces

// pkg/auth/identity.go — updated Identity struct
type Identity struct {
    Subject   string
    Name      string
    Email     string
    Groups    []string
    Claims    map[string]any
    Token     string
    TokenType string
    Metadata  map[string]string

    // UpstreamTokens holds upstream access tokens keyed by provider name.
    // Populated by auth middleware after JWT validation when the embedded auth
    // server is in use (tsid claim present + upstreamStorage configured).
    // Expired tokens are included as-is; callers must not attempt refresh.
    // Empty for non-embedded-AS requests.
    UpstreamTokens map[string]string // providerName → access token
}

// MarshalJSON must redact UpstreamTokens values:
// SafeIdentity.UpstreamTokens: build map[string]string with same keys but "REDACTED" values

// pkg/auth/token.go — TokenValidator enrichment insertion point
// (between claimsToIdentity and WithIdentity)

if v.upstreamStorage != nil {
    if tsid, ok := identity.Claims[upstreamtoken.TokenSessionIDClaimKey].(string); ok && tsid != "" {
        allTokens, err := v.upstreamStorage.GetAllUpstreamTokens(r.Context(), tsid)
        if err != nil {
            slog.WarnContext(r.Context(), "failed to load upstream tokens", "err", err)
        } else {
            identity.UpstreamTokens = make(map[string]string, len(allTokens))
            for name, t := range allTokens {
                identity.UpstreamTokens[name] = t.AccessToken // include expired
            }
        }
    }
}

// pkg/auth/upstreamswap/middleware.go — updated Config and createMiddlewareFunc

type Config struct {
    HeaderStrategy   string `json:"header_strategy,omitempty" yaml:"header_strategy,omitempty"`
    CustomHeaderName string `json:"custom_header_name,omitempty" yaml:"custom_header_name,omitempty"`
    // ProviderName is the upstream provider whose token to inject.
    // Derived from Upstreams[0].Name at construction time; defaults to "default".
    ProviderName string `json:"provider_name" yaml:"provider_name"`
}
// ServiceGetter type and storageGetter field are removed entirely.

// createMiddlewareFunc: simplified handler reads from identity.UpstreamTokens
func createMiddlewareFunc(cfg *Config) types.MiddlewareFunction {
    // ... injectToken determined at startup time (unchanged) ...
    return func(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
            identity, ok := auth.IdentityFromContext(r.Context())
            if !ok {
                next.ServeHTTP(w, r)
                return
            }
            // No tsid check needed — if UpstreamTokens is empty, behave same as absent
            token, exists := identity.UpstreamTokens[cfg.ProviderName]
            if !exists || token == "" {
                writeUpstreamAuthRequired(w)
                return
            }
            injectToken(r, token)
            next.ServeHTTP(w, r)
        })
    }
}

// pkg/runner/middleware.go — addUpstreamSwapMiddleware ProviderName derivation
func addUpstreamSwapMiddleware(
    middlewares []types.MiddlewareConfig,
    config *RunConfig,
) ([]types.MiddlewareConfig, error) {
    if config.EmbeddedAuthServerConfig == nil {
        return middlewares, nil
    }

    providerName := func() string {
        if len(config.EmbeddedAuthServerConfig.Upstreams) > 0 &&
            config.EmbeddedAuthServerConfig.Upstreams[0].Name != "" {
            return config.EmbeddedAuthServerConfig.Upstreams[0].Name
        }
        return "default"
    }()

    upstreamSwapConfig := config.UpstreamSwapConfig
    if upstreamSwapConfig == nil {
        upstreamSwapConfig = &upstreamswap.Config{}
    }
    upstreamSwapConfig.ProviderName = providerName
    // ... rest unchanged ...
}

Testing Strategy

Unit Tests — pkg/auth/middleware_test.go

• GetAllUpstreamTokens is called when tsid claim is present and upstreamStorage is non-nil; tokens are populated in identity.UpstreamTokens keyed by provider name
• Expired tokens are included in UpstreamTokens as-is (no filtering by expiry at enrichment time)
• UpstreamTokens is nil/empty when the tsid claim is absent from the JWT
• UpstreamTokens is nil/empty when upstreamStorage is nil (non-embedded-AS deployment)
• When GetAllUpstreamTokens returns an error, the middleware logs WARN and proceeds with empty UpstreamTokens (does not return HTTP error)
• MarshalJSON on Identity with populated UpstreamTokens redacts token values to "REDACTED" while preserving provider name keys

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

• Middleware reads identity.UpstreamTokens[cfg.ProviderName] and injects the token when present
• Middleware returns HTTP 401 with WWW-Authenticate: Bearer error="invalid_token" when identity.UpstreamTokens[cfg.ProviderName] is absent
• Middleware returns HTTP 401 when the token value is an empty string
• ProviderName config field is respected — using "provider-a" reads UpstreamTokens["provider-a"], not UpstreamTokens["provider-b"]
• Middleware proceeds without swap when identity has no UpstreamTokens map at all (nil map) — returns 401 since token is required for an embedded-AS-configured runner
• CreateMiddleware factory no longer calls GetUpstreamTokenService() on the runner mock
• TestCreateMiddleware test cases do not expect mockRunner.EXPECT().GetUpstreamTokenService() to be called

Unit Tests — pkg/runner/middleware_test.go

• addUpstreamSwapMiddleware with Upstreams[0].Name = "test-upstream" produces upstreamswap.Config{ProviderName: "test-upstream"}
• addUpstreamSwapMiddleware with Upstreams[0].Name = "" produces upstreamswap.Config{ProviderName: "default"}
• addUpstreamSwapMiddleware with empty Upstreams slice produces upstreamswap.Config{ProviderName: "default"}
• Serialized upstreamswap.MiddlewareParams.Config.ProviderName matches the derived name in all existing TestAddUpstreamSwapMiddleware cases

Edge Cases

• Identity.MarshalJSON() with an empty UpstreamTokens map (map[string]string{}) serializes to an empty object or omitted field, not "REDACTED"
• upstreamswap middleware with ProviderName = "" behaves deterministically (either returns 401 because UpstreamTokens[""] is absent, or validateConfig rejects an empty ProviderName — document the chosen behavior)
• TokenValidator with upstreamStorage = nil and a JWT containing a tsid claim does not panic and does not call any storage method

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
• Storage layer changes (UpstreamTokenStorage interface, backends): covered in TASK-001 (Implement multi-provider upstream token storage layer (RFC-0052 Phase 1) #4136)
• Transparent token refresh or re-authorization during middleware enrichment (deferred to a later RFC)
• Distributed locking for concurrent token refresh operations
• Multi-upstream IDP support in the CLI proxy runner or MCPServer/MCPRemoteProxy K8s workload types
• 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 (Storage Layer): Implement multi-provider upstream token storage layer (RFC-0052 Phase 1) #4136
• Architecture overview: docs/arch/11-auth-server-storage.md
• Related Stacklok epic: https://github.com/stacklok/stacklok-epics/issues/251