From 1a479f1715b4e86a61be288b4d4ee3d92872a099 Mon Sep 17 00:00:00 2001 From: Trey Date: Mon, 31 Aug 2026 09:05:39 -0700 Subject: [PATCH 1/4] fix(authserver): scope upstream-token callback cleanup to the chain's providers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The /oauth/callback error paths cleaned up "the tokens this authorization collected" by calling the session-wide DeleteUpstreamTokens(ctx, sessionID), which removes every provider row under that session. That is correct for the native (session, provider) storage scheme, where each authorization mints its own session. It is wrong under a storage decorator that re-keys upstream tokens by (gateway, user) so a session-independent caller (e.g. a user-facing UI) can find a user's tokens. Under that scheme every one of a user's connector tokens shares the session slot, so a session-wide delete on any failing login leg deletes the user's entire connector set — a transient Directory/storage blip on the most-traversed leg (chain resolution) wipes a user's whole historical token set, and correlated failures widen that to every user logging in during the window. It fails safe (never grants access) and is largely self-healing, but the availability hit is real. The callback is the only production caller of the session-wide delete, so: - Route all nine callback cleanup sites through a new cleanupUpstreamTokens helper that deletes per provider via DeleteUpstreamTokensForProvider (already on the storage interface, implemented by memory and redis), sourcing the provider list from the chain config — never a storage read, which a broadening decorator cannot scope. Sites with the resolved chain pass it; pre-resolve sites pass pending.ChainUpstreams plus the current leg. - A user-keyed decorator can now scope each delete to exactly this chain's providers, so a connector outside the chain is never touched. Behavior note: a stale subsequent-leg pending with no ChainUpstreams can only name the leg just processed; an earlier leg it cannot name is left to expire by TTL rather than risk a session-wide delete widening to the whole user. This matches the existing abandoned-flow orphan behavior. Adds a regression test asserting an out-of-chain provider's token survives a chain cleanup, and extends the handler mock with a DeleteUpstreamTokensForProvider capture. Signed-off-by: Trey --- pkg/authserver/server/handlers/callback.go | 75 ++++++++++--- .../server/handlers/callback_test.go | 100 +++++++++++++++--- .../server/handlers/helpers_test.go | 12 +++ 3 files changed, 160 insertions(+), 27 deletions(-) diff --git a/pkg/authserver/server/handlers/callback.go b/pkg/authserver/server/handlers/callback.go index 727cc36769..ac045f4ceb 100644 --- a/pkg/authserver/server/handlers/callback.go +++ b/pkg/authserver/server/handlers/callback.go @@ -182,8 +182,13 @@ func (h *Handler) CallbackHandler(w http.ResponseWriter, req *http.Request) { slog.Error("failed to store upstream tokens", "error", err, ) - // Clean up any tokens stored by earlier legs of a multi-upstream chain. - _ = h.storage.DeleteUpstreamTokens(ctx, sessionID) + // Clean up any tokens stored by earlier legs of a multi-upstream chain, + // scoped per provider so a user-keyed storage decorator that re-keys upstream + // tokens by (gateway, user) removes only this chain's providers and never the + // user's tokens for connectors outside it. providerID (this leg) is included + // even though its store just failed: any partial or carried-forward row is + // removed, and an absent row deletes as a no-op. + h.cleanupUpstreamTokens(ctx, sessionID, append([]string{providerID}, pending.ChainUpstreams...)) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to store session")) return } @@ -401,16 +406,18 @@ func (h *Handler) handleUpstreamError( // Clean up any upstream tokens stored by earlier legs of a multi-upstream chain. // On a subsequent leg the resolved user is carried in pending.ResolvedUserID; // place it in ctx via WithPlatformUser so a context-keyed storage decorator can - // resolve the canonical user for this delete (DeleteUpstreamTokens takes no tokens - // argument). WithPlatformUser, not WithIdentity: there is no authenticated identity - // at the callback, only the storage-scoped user. A first-leg error has no resolved - // user and no earlier-leg tokens to clean up, so the bare ctx is correct there. + // resolve the canonical user for the per-provider cleanup. WithPlatformUser, not + // WithIdentity: there is no authenticated identity at the callback, only the + // storage-scoped user. The cleanup is scoped to the chain's providers so a + // user-keyed decorator does not over-delete the user's other connectors; a + // first-leg error carries no ChainUpstreams and no earlier-leg tokens, so nothing + // is deleted there. if pending.SessionID != "" { cleanupCtx := ctx if pending.ResolvedUserID != "" { cleanupCtx = auth.WithPlatformUser(ctx, pending.ResolvedUserID) } - _ = h.storage.DeleteUpstreamTokens(cleanupCtx, pending.SessionID) + h.cleanupUpstreamTokens(cleanupCtx, pending.SessionID, pending.ChainUpstreams) } ar := h.buildAuthorizeRequesterFromPending(ctx, pending) if ar != nil { @@ -428,6 +435,42 @@ func (h *Handler) handleUpstreamError( http.Error(w, "upstream authentication failed", http.StatusBadGateway) } +// cleanupUpstreamTokens best-effort removes the upstream tokens an authorization +// chain has collected, one provider at a time. It is the failure-path counterpart +// to a successful chain: when a leg fails, the tokens earlier legs stored under this +// session must not be left mid-chain. +// +// It deletes per provider rather than calling the session-wide DeleteUpstreamTokens +// because that method carries no provider name. A storage decorator that re-keys +// upstream tokens by (gateway, user) — so a session-independent caller can find a +// user's tokens — cannot narrow a session-wide delete back to this chain, and would +// instead delete every connector token the user holds on the gateway. Naming each +// provider lets such a decorator scope the delete to exactly this chain's providers. +// +// providers is the effective chain (or the chain carried forward in the pending on a +// subsequent leg), optionally including the leg currently being processed. Passing a +// provider whose leg has not been walked, or whose row is already absent, is a no-op: +// the storage contract makes a missing-row delete non-fatal. Errors are logged and +// swallowed — cleanup is best-effort and must not mask the original failure. +func (h *Handler) cleanupUpstreamTokens(ctx context.Context, sessionID string, providers []string) { + seen := make(map[string]struct{}, len(providers)) + for _, provider := range providers { + if provider == "" { + continue + } + if _, dup := seen[provider]; dup { + continue + } + seen[provider] = struct{}{} + if err := h.storage.DeleteUpstreamTokensForProvider(ctx, sessionID, provider); err != nil { + slog.Warn("failed to clean up upstream token for provider", + "provider", provider, + "error", err, + ) + } + } +} + // continueChainOrComplete checks whether all upstream providers in the authorization // chain have been satisfied. If so, it issues the authorization code and redirects // to the client. If not, it redirects to the next upstream provider to continue @@ -476,7 +519,11 @@ func (h *Handler) continueChainOrComplete( chain, err := h.resolveChain(ctx, pending, principal) if err != nil { slog.Error("failed to resolve upstream chain", "error", err) - _ = h.storage.DeleteUpstreamTokens(ctx, sessionID) + // The chain could not be resolved, so scope the cleanup to what is known: this + // leg's provider plus any chain carried forward on a subsequent leg. An earlier + // leg not named here (e.g. a stale pending with no ChainUpstreams) is left to + // expire by TTL rather than risk a user-keyed decorator over-deleting. + h.cleanupUpstreamTokens(ctx, sessionID, append([]string{pending.UpstreamProviderName}, pending.ChainUpstreams...)) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to determine authorization chain")) return } @@ -484,7 +531,7 @@ func (h *Handler) continueChainOrComplete( nextProvider, err := h.nextMissingUpstream(ctx, sessionID, chain) if err != nil { slog.Error("failed to determine next upstream", "error", err) - _ = h.storage.DeleteUpstreamTokens(ctx, sessionID) + h.cleanupUpstreamTokens(ctx, sessionID, chain) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to check authorization chain state")) return } @@ -493,7 +540,7 @@ func (h *Handler) continueChainOrComplete( if err := h.verifyChainIdentity(ctx, sessionID, chain, subject); err != nil { // 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.cleanupUpstreamTokens(ctx, sessionID, chain) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("identity verification failed")) return } @@ -501,7 +548,7 @@ func (h *Handler) continueChainOrComplete( // All upstreams satisfied — issue authorization code if err := h.writeAuthorizationResponse(ctx, w, pending, sessionID, subject, name, email); err != nil { slog.Error("failed to create authorization response", "error", err) - _ = h.storage.DeleteUpstreamTokens(ctx, sessionID) + h.cleanupUpstreamTokens(ctx, sessionID, chain) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to create authorization code")) } return @@ -539,7 +586,7 @@ func (h *Handler) continueChainOrComplete( if err := h.storage.StorePendingAuthorization(ctx, secrets.State, nextPending); err != nil { slog.Error("failed to store next chain leg", "error", err) - _ = h.storage.DeleteUpstreamTokens(ctx, sessionID) + h.cleanupUpstreamTokens(ctx, sessionID, chain) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to continue authorization chain")) return } @@ -553,7 +600,7 @@ func (h *Handler) continueChainOrComplete( if !ok { slog.Error("next upstream provider not found", "provider", nextProvider) _ = h.storage.DeletePendingAuthorization(ctx, secrets.State) - _ = h.storage.DeleteUpstreamTokens(ctx, sessionID) + h.cleanupUpstreamTokens(ctx, sessionID, chain) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("upstream provider configuration error")) return } @@ -561,7 +608,7 @@ func (h *Handler) continueChainOrComplete( if err != nil { slog.Error("failed to build next upstream authorization URL", "error", err) _ = h.storage.DeletePendingAuthorization(ctx, secrets.State) - _ = h.storage.DeleteUpstreamTokens(ctx, sessionID) + h.cleanupUpstreamTokens(ctx, sessionID, chain) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to build authorization URL")) return } diff --git a/pkg/authserver/server/handlers/callback_test.go b/pkg/authserver/server/handlers/callback_test.go index add96cf56a..9f25ac9d33 100644 --- a/pkg/authserver/server/handlers/callback_test.go +++ b/pkg/authserver/server/handlers/callback_test.go @@ -931,10 +931,15 @@ func TestCallbackHandler_SubsequentLeg_MissingChain_FailsClosed(t *testing.T) { 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) - } + // Fail-closed cleanup is scoped per provider so a user-keyed storage decorator does + // not over-delete the user's other connectors. This stale pending carries no + // ChainUpstreams, so only the leg just processed (provider-2) can be named and + // removed; the earlier leg (provider-1) is not nameable here and is left to expire + // by its own TTL rather than risk a session-wide delete widening to the whole user. + assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-2", + "the leg just processed should be cleaned up") + assert.Contains(t, storState.upstreamTokens, sessionID+":provider-1", + "an earlier leg with no ChainUpstreams to name it is left to TTL, not session-wide deleted") } func TestCallbackHandler_SingleLeg_IssuesCodeWithoutChaining(t *testing.T) { @@ -1234,6 +1239,74 @@ func TestCallbackHandler_TwoUpstreams_AuthorizationURLError_CleansUp(t *testing. } } +// TestCallbackHandler_Cleanup_LeavesOutOfChainProviderTokens is the regression guard for +// the failure that motivated per-provider cleanup. The callback's chain-cleanup must +// delete only the providers in THIS authorization's chain, never a token the storage +// holds for the same session under some other provider. +// +// It matters because a storage decorator can re-key upstream tokens from toolhive's +// native (session, provider) scheme to a per-(gateway, user) scheme so a +// session-independent caller can find a user's tokens. Under that decorator every one of +// a user's connector tokens shares the session slot, so the old session-wide +// DeleteUpstreamTokens on any failing login leg deleted the user's entire connector set. +// Naming each chain provider lets such a decorator scope the delete to exactly this +// chain, so a provider outside the chain survives — asserted here with a raw +// (session, provider)-keyed backend that makes the scoping directly observable. +func TestCallbackHandler_Cleanup_LeavesOutOfChainProviderTokens(t *testing.T) { + t.Parallel() + handler, storState, _, mockProvider2 := multiUpstreamTestSetup(t) + + // provider-2 fails to build its authorization URL, driving the chain-cleanup path + // with the effective chain [provider-1, provider-2]. + mockProvider2.authURLErr = errors.New("authorization URL error") + + sessionID := "chain-session-out-of-chain" + + // A token for a provider that is NOT part of this login's chain, stored under the + // same session slot — this is what a (gateway, user)-keyed decorator looks like from + // the backend's side: every connector the user holds shares one session key. + outOfChainKey := sessionID + ":provider-out-of-chain" + storState.upstreamTokens[outOfChainKey] = &storage.UpstreamTokens{ + ProviderID: "provider-out-of-chain", + AccessToken: "keep-me", + ExpiresAt: time.Now().Add(time.Hour), + ClientID: testAuthClientID, + UserID: "user-1", + } + + firstLegState := "out-of-chain-first-leg-state" + storState.pendingAuths[firstLegState] = &storage.PendingAuthorization{ + ClientID: testAuthClientID, + RedirectURI: testAuthRedirectURI, + State: "client-state-out-of-chain", + PKCEChallenge: "challenge-out-of-chain", + PKCEMethod: "S256", + Scopes: []string{"openid"}, + InternalState: firstLegState, + UpstreamPKCEVerifier: "out-of-chain-verifier-1234567890123456789012", + UpstreamNonce: "out-of-chain-nonce", + UpstreamProviderName: "provider-1", + SessionID: sessionID, + CreatedAt: time.Now(), + } + + req := httptest.NewRequest(http.MethodGet, "/oauth/callback?code=p1-code&state="+firstLegState, nil) + rec := httptest.NewRecorder() + handler.CallbackHandler(rec, req) + + assert.Equal(t, http.StatusSeeOther, rec.Code, "should return fosite error redirect") + + // The failing chain's own provider-1 token is cleaned up... + assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-1", + "the failing chain's own provider token should be cleaned up") + // ...but the out-of-chain provider's token is left untouched. This is the property a + // user-keyed decorator relies on so a failing login never wipes the user's other + // connectors. + require.Contains(t, storState.upstreamTokens, outOfChainKey, + "a token for a provider outside this chain must survive the cleanup") + assert.Equal(t, "keep-me", storState.upstreamTokens[outOfChainKey].AccessToken) +} + func TestCallbackHandler_TwoUpstreams_StorePendingError_CleansUp(t *testing.T) { t.Parallel() @@ -1600,15 +1673,15 @@ func TestCallbackHandler_PlacesPlatformUserInContext_OnChainRead(t *testing.T) { // TestCallbackHandler_PlacesPlatformUserInContext_OnUpstreamErrorCleanup verifies that // when a subsequent-leg callback returns an upstream error, the earlier-leg cleanup -// (DeleteUpstreamTokens in handleUpstreamError) runs with the resolved canonical user -// already in the request context. +// (per-provider DeleteUpstreamTokensForProvider in handleUpstreamError) runs with the +// resolved canonical user already in the request context. // // This is the error-path sibling of the chain-read test above. handleUpstreamError // runs from an early return at the top of CallbackHandler, before the happy-path -// user injection, so it must place the user itself. DeleteUpstreamTokens takes only -// (ctx, sessionID) — no tokens argument — so a context-keyed storage decorator can -// resolve the canonical user only from ctx. The resolved user is carried forward from -// leg 1 via pending.ResolvedUserID. +// user injection, so it must place the user itself. DeleteUpstreamTokensForProvider +// carries no tokens argument — so a context-keyed storage decorator can resolve the +// canonical user only from ctx. The resolved user is carried forward from leg 1 via +// pending.ResolvedUserID, and pending.ChainUpstreams names the earlier leg to clean. func TestCallbackHandler_PlacesPlatformUserInContext_OnUpstreamErrorCleanup(t *testing.T) { t.Parallel() handler, storState, _, _ := multiUpstreamTestSetup(t) @@ -1638,6 +1711,7 @@ func TestCallbackHandler_PlacesPlatformUserInContext_OnUpstreamErrorCleanup(t *t InternalState: secondLegState, UpstreamProviderName: "provider-2", SessionID: sessionID, + ChainUpstreams: []string{"provider-1", "provider-2"}, ResolvedUserID: leg1User, ResolvedUserName: "First Leg User", ResolvedUserEmail: "firstleg@example.com", @@ -1649,11 +1723,11 @@ func TestCallbackHandler_PlacesPlatformUserInContext_OnUpstreamErrorCleanup(t *t rec := httptest.NewRecorder() handler.CallbackHandler(rec, req) - // handleUpstreamError cleans up earlier-leg tokens via DeleteUpstreamTokens; the - // harness captured the ctx it ran with. Assert the callback placed the resolved + // handleUpstreamError cleans up earlier-leg tokens via DeleteUpstreamTokensForProvider; + // the harness captured the ctx it ran with. Assert the callback placed the resolved // canonical user into that ctx so a context-keyed decorator can delete the row. uid, ok := auth.PlatformUserFromContext(storState.deleteUpstreamCtx) - require.True(t, ok, "handleUpstreamError must place platform user in ctx before the cleanup DeleteUpstreamTokens") + require.True(t, ok, "handleUpstreamError must place platform user in ctx before the cleanup delete") require.Equal(t, leg1User, uid) // The error path must NOT place a stub Identity under the identity key. diff --git a/pkg/authserver/server/handlers/helpers_test.go b/pkg/authserver/server/handlers/helpers_test.go index 1f87ed1b39..7086dbb0c0 100644 --- a/pkg/authserver/server/handlers/helpers_test.go +++ b/pkg/authserver/server/handlers/helpers_test.go @@ -375,6 +375,18 @@ func baseTestSetup(t *testing.T, opts ...baseTestSetupOption) (fosite.OAuth2Prov return nil }).AnyTimes() + stor.EXPECT().DeleteUpstreamTokensForProvider(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, sessionID, providerName string) error { + // DeleteUpstreamTokensForProvider takes (ctx, sessionID, providerName) but + // still carries no tokens argument, so a context-keyed storage decorator + // resolves the user from ctx exactly as DeleteUpstreamTokens does. Capture + // the ctx here too so the callback's per-provider cleanup path is covered by + // the same context-placement assertions. Deleting an absent row is a no-op. + storState.deleteUpstreamCtx = ctx + delete(storState.upstreamTokens, sessionID+":"+providerName) + return nil + }).AnyTimes() + stor.EXPECT().GetAllUpstreamTokens(gomock.Any(), gomock.Any()).DoAndReturn( func(ctx context.Context, sessionID string) (map[string]*storage.UpstreamTokens, error) { // GetAllUpstreamTokens takes only (ctx, sessionID) — no tokens argument to From a5d4e6ca53db47741b6593af592c894c9300d0ee Mon Sep 17 00:00:00 2001 From: Trey Date: Mon, 31 Aug 2026 11:20:47 -0700 Subject: [PATCH 2/4] fix(authserver): guard per-provider cleanup and polish its docs Addresses stacklok/toolhive#6476 review comments: - MEDIUM callback.go (3897108470): skip provider names that are not a configured upstream inside cleanupUpstreamTokens, so the three pre-resolveChain sites that pass an unvalidated stored chain can no longer name an arbitrary delete target; a no-op for the post-resolveChain sites, which already pass a validateChain-checked chain. - MEDIUM callback.go (3897108480): the WithPlatformUser comment now names DeleteUpstreamTokensForProvider instead of the removed session-wide delete. - LOW callback.go (3897108524): trim the (gateway, user) rationale duplicated at the two call sites down to their call-site nuance; the shared rationale lives on the cleanupUpstreamTokens doc comment. - LOW callback.go (3897108537): annotate the cleanup WARN with //nolint:gosec // G706, matching the other server-storage-derived logs in this file. - LOW callback.go (3897108558): note the per-provider round-trip cost (bounded to the chain length, error-path only) on the doc comment. - LOW callback.go (3897108562): state the providers parameter's contract rather than what today's callers happen to pass. Signed-off-by: Trey --- pkg/authserver/server/handlers/callback.go | 57 +++++++++++++--------- 1 file changed, 34 insertions(+), 23 deletions(-) diff --git a/pkg/authserver/server/handlers/callback.go b/pkg/authserver/server/handlers/callback.go index ac045f4ceb..bdd971d36f 100644 --- a/pkg/authserver/server/handlers/callback.go +++ b/pkg/authserver/server/handlers/callback.go @@ -150,8 +150,8 @@ func (h *Handler) CallbackHandler(w http.ResponseWriter, req *http.Request) { // Place the resolved canonical user in the request context so callback storage // calls that carry no tokens argument — GetAllUpstreamTokens during chain - // consistency, DeleteUpstreamTokens on cleanup — can resolve the user from - // context. We use WithPlatformUser, not WithIdentity: no ToolHive bearer has + // consistency, DeleteUpstreamTokensForProvider on cleanup — can resolve the user + // from context. We use WithPlatformUser, not WithIdentity: no ToolHive bearer has // been issued at the callback, so there is no authenticated identity to assert — // only the canonical user for storage keying. (StoreUpstreamTokens does not need // this; it keys off tokens.UserID below.) @@ -182,12 +182,10 @@ func (h *Handler) CallbackHandler(w http.ResponseWriter, req *http.Request) { slog.Error("failed to store upstream tokens", "error", err, ) - // Clean up any tokens stored by earlier legs of a multi-upstream chain, - // scoped per provider so a user-keyed storage decorator that re-keys upstream - // tokens by (gateway, user) removes only this chain's providers and never the - // user's tokens for connectors outside it. providerID (this leg) is included - // even though its store just failed: any partial or carried-forward row is - // removed, and an absent row deletes as a no-op. + // Clean up tokens stored by earlier legs of a multi-upstream chain (see + // cleanupUpstreamTokens for the per-provider scoping rationale). providerID + // (this leg) is included even though its store just failed: any partial or + // carried-forward row is removed, and an absent row deletes as a no-op. h.cleanupUpstreamTokens(ctx, sessionID, append([]string{providerID}, pending.ChainUpstreams...)) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to store session")) return @@ -403,15 +401,14 @@ func (h *Handler) handleUpstreamError( pending, err := h.storage.LoadPendingAuthorization(ctx, internalState) if err == nil { _ = h.storage.DeletePendingAuthorization(ctx, internalState) - // Clean up any upstream tokens stored by earlier legs of a multi-upstream chain. - // On a subsequent leg the resolved user is carried in pending.ResolvedUserID; - // place it in ctx via WithPlatformUser so a context-keyed storage decorator can - // resolve the canonical user for the per-provider cleanup. WithPlatformUser, not - // WithIdentity: there is no authenticated identity at the callback, only the - // storage-scoped user. The cleanup is scoped to the chain's providers so a - // user-keyed decorator does not over-delete the user's other connectors; a - // first-leg error carries no ChainUpstreams and no earlier-leg tokens, so nothing - // is deleted there. + // Clean up upstream tokens stored by earlier legs of a multi-upstream chain + // (see cleanupUpstreamTokens for the per-provider scoping rationale). On a + // subsequent leg the resolved user is carried in pending.ResolvedUserID; place + // it in ctx via WithPlatformUser so a context-keyed storage decorator resolves + // the canonical user for the cleanup. WithPlatformUser, not WithIdentity: there + // is no authenticated identity at the callback, only the storage-scoped user. A + // first-leg error carries no ChainUpstreams and no earlier-leg tokens, so + // nothing is deleted there. if pending.SessionID != "" { cleanupCtx := ctx if pending.ResolvedUserID != "" { @@ -447,11 +444,18 @@ func (h *Handler) handleUpstreamError( // instead delete every connector token the user holds on the gateway. Naming each // provider lets such a decorator scope the delete to exactly this chain's providers. // -// providers is the effective chain (or the chain carried forward in the pending on a -// subsequent leg), optionally including the leg currently being processed. Passing a -// provider whose leg has not been walked, or whose row is already absent, is a no-op: -// the storage contract makes a missing-row delete non-fatal. Errors are logged and -// swallowed — cleanup is best-effort and must not mask the original failure. +// Each name is checked against the configured upstream set before it is deleted, so a +// name that is not a configured upstream — e.g. from a corrupted or tampered pending +// row reaching a pre-resolveChain cleanup site, which unlike the post-resolve sites has +// not been through validateChain — is skipped rather than turned into a delete. +// +// providers is the set of provider names whose upstream-token rows should be removed; +// empty entries, duplicates, and names outside the configured upstream set are ignored, +// and deleting an already-absent row is a no-op (the storage contract makes a +// missing-row delete non-fatal). Errors are logged and swallowed — cleanup is +// best-effort and must not mask the original failure. Cost is one storage round-trip +// per distinct provider; a chain is the corporate primary plus a handful of connectors, +// so the sequential deletes are bounded and only run on error paths. func (h *Handler) cleanupUpstreamTokens(ctx context.Context, sessionID string, providers []string) { seen := make(map[string]struct{}, len(providers)) for _, provider := range providers { @@ -462,8 +466,15 @@ func (h *Handler) cleanupUpstreamTokens(ctx context.Context, sessionID string, p continue } seen[provider] = struct{}{} + // Only a configured upstream can legitimately have a stored token, so skip + // anything else: it keeps a tampered pending row reaching a pre-resolveChain + // cleanup site from naming an arbitrary delete target. The post-resolveChain + // sites pass a validateChain-checked chain, so this is a no-op for them. + if _, ok := h.upstreamByName(provider); !ok { + continue + } if err := h.storage.DeleteUpstreamTokensForProvider(ctx, sessionID, provider); err != nil { - slog.Warn("failed to clean up upstream token for provider", + slog.Warn("failed to clean up upstream token for provider", //nolint:gosec // G706 - provider name from server-side chain state "provider", provider, "error", err, ) From 9c4e1a701cbdf44ef03827c5687147158c423c0f Mon Sep 17 00:00:00 2001 From: Trey Date: Mon, 31 Aug 2026 11:28:46 -0700 Subject: [PATCH 3/4] test(authserver): cover the per-provider cleanup paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses stacklok/toolhive#6476 review comments: - MEDIUM callback.go (3897108490): add withStoreUpstreamTokensError and a subsequent-leg store-failure test, exercising the one cleanup site where the dedup in cleanupUpstreamTokens is load-bearing (ChainUpstreams already contains the current provider). - MEDIUM callback.go (3897108499): add tests for the nextMissingUpstream-failure site (GetAllUpstreamTokens error) and the writeAuthorizationResponse-failure site on a fully satisfied chain (GetClient failing only on its second call). The third branch it names — "next upstream provider not found" — is unreachable because resolveChain runs validateChain over the chain before it is walked, so every nextProvider is a configured upstream; covered by reply, not a test. - MEDIUM callback_test.go (3897108506): the OnUpstreamErrorCleanup test now asserts the earlier-leg row was actually deleted, not merely that the ctx was captured. - LOW callback.go (3897108545): add withDeleteUpstreamTokensForProviderError and a test proving cleanup continues past a per-provider delete failure to remove sibling providers. - LOW helpers_test.go (3897108551): the DeleteUpstreamTokensForProvider mock now mirrors the real backends' rejection of an empty session or provider name. Extracts multiUpstreamProviders and adds multiUpstreamTestSetupWithStorage so the two-upstream setups can thread storage error injection without duplicating the mock providers. Signed-off-by: Trey --- .../server/handlers/callback_test.go | 212 ++++++++++++++++++ .../server/handlers/helpers_test.go | 103 ++++++++- 2 files changed, 304 insertions(+), 11 deletions(-) diff --git a/pkg/authserver/server/handlers/callback_test.go b/pkg/authserver/server/handlers/callback_test.go index 9f25ac9d33..08de280341 100644 --- a/pkg/authserver/server/handlers/callback_test.go +++ b/pkg/authserver/server/handlers/callback_test.go @@ -1307,6 +1307,212 @@ func TestCallbackHandler_Cleanup_LeavesOutOfChainProviderTokens(t *testing.T) { assert.Equal(t, "keep-me", storState.upstreamTokens[outOfChainKey].AccessToken) } +// TestCallbackHandler_StoreUpstreamTokensError_CleansUpChain covers the +// StoreUpstreamTokens-failure cleanup site on a subsequent leg. It is the one site +// where cleanupUpstreamTokens' dedup is load-bearing: the provider list is +// append([]string{providerID}, pending.ChainUpstreams...), and ChainUpstreams already +// contains providerID on a non-first leg, so the same provider is named twice. +func TestCallbackHandler_StoreUpstreamTokensError_CleansUpChain(t *testing.T) { + t.Parallel() + handler, storState, _, _ := multiUpstreamTestSetupWithStorage(t, + withStoreUpstreamTokensError(errors.New("storage unavailable"))) + + sessionID := "store-err-session" + const leg1User = "resolved-user-id-from-leg1" + + // First leg already completed: provider-1's token exists. + storState.upstreamTokens[sessionID+":provider-1"] = &storage.UpstreamTokens{ + ProviderID: "provider-1", + AccessToken: "p1-at", + ExpiresAt: time.Now().Add(time.Hour), + ClientID: testAuthClientID, + UserID: leg1User, + } + + // Subsequent-leg pending whose ChainUpstreams already contains provider-2 (this + // leg), so the cleanup list names provider-2 twice — exercising the dedup. + secondLegState := "store-err-second-leg-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: "store-err-verifier-123456789012345678901234", + UpstreamNonce: "store-err-nonce", + UpstreamProviderName: "provider-2", + SessionID: sessionID, + ChainUpstreams: []string{"provider-1", "provider-2"}, + ResolvedUserID: leg1User, + ResolvedUserName: "First Leg User", + ResolvedUserEmail: "firstleg@example.com", + CreatedAt: time.Now(), + } + + 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, "a store failure should produce a fosite error redirect") + assert.Contains(t, rec.Header().Get("Location"), "error=server_error") + + // The earlier leg's token is cleaned up; provider-2's store never landed. + assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-1", + "the earlier chain leg's token should be cleaned up on a store failure") + assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-2", + "provider-2's store failed, so no row should exist") +} + +// TestCallbackHandler_NextMissingUpstreamError_CleansUpChain covers the +// nextMissingUpstream-failure cleanup site: a chain-state read (GetAllUpstreamTokens) +// error after the first leg has stored its token. +func TestCallbackHandler_NextMissingUpstreamError_CleansUpChain(t *testing.T) { + t.Parallel() + handler, storState, _, _ := multiUpstreamTestSetupWithStorage(t, + withGetAllUpstreamTokensError(errors.New("storage unavailable"))) + + sessionID := "next-missing-err-session" + firstLegState := "next-missing-err-first-leg-state" + storState.pendingAuths[firstLegState] = &storage.PendingAuthorization{ + ClientID: testAuthClientID, + RedirectURI: testAuthRedirectURI, + State: "client-state-next-missing", + PKCEChallenge: "challenge-next-missing", + PKCEMethod: "S256", + Scopes: []string{"openid"}, + InternalState: firstLegState, + UpstreamPKCEVerifier: "next-missing-verifier-12345678901234567890", + UpstreamNonce: "next-missing-nonce", + UpstreamProviderName: "provider-1", + SessionID: sessionID, + CreatedAt: time.Now(), + } + + req := httptest.NewRequest(http.MethodGet, "/oauth/callback?code=p1-code&state="+firstLegState, nil) + rec := httptest.NewRecorder() + handler.CallbackHandler(rec, req) + + assert.Equal(t, http.StatusSeeOther, rec.Code, "a chain-state read failure should produce a fosite error redirect") + assert.Contains(t, rec.Header().Get("Location"), "error=server_error") + + // The first leg's just-stored token is cleaned up when the chain-state read fails. + assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-1", + "the stored first-leg token should be cleaned up on a nextMissingUpstream failure") +} + +// TestCallbackHandler_WriteAuthorizationResponseError_OnSatisfiedChain_CleansUp covers +// the cleanup site reached when writeAuthorizationResponse fails after a multi-upstream +// chain is fully satisfied and the identity check has passed. GetClient is wired to fail +// only on its second call: the first (which builds the error-redirect requester) +// succeeds so the failure still produces a redirect, the second (inside +// writeAuthorizationResponse) fails to drive the cleanup. +func TestCallbackHandler_WriteAuthorizationResponseError_OnSatisfiedChain_CleansUp(t *testing.T) { + t.Parallel() + handler, storState, _, _ := multiUpstreamTestSetupWithStorage(t, + withGetClientErrorAfterCalls(1, errors.New("get client unavailable"))) + + sessionID := "satisfied-writeresp-err-session" + const leg1User = "resolved-user-id-from-leg1" + + // First leg already completed, keyed to leg1User so verifyChainIdentity passes. + storState.upstreamTokens[sessionID+":provider-1"] = &storage.UpstreamTokens{ + ProviderID: "provider-1", + AccessToken: "p1-at", + ExpiresAt: time.Now().Add(time.Hour), + ClientID: testAuthClientID, + UserID: leg1User, + } + + secondLegState := "satisfied-writeresp-err-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: "satisfied-verifier-123456789012345678901234", + UpstreamNonce: "satisfied-nonce", + UpstreamProviderName: "provider-2", + SessionID: sessionID, + ChainUpstreams: []string{"provider-1", "provider-2"}, + ResolvedUserID: leg1User, + ResolvedUserName: "First Leg User", + ResolvedUserEmail: "firstleg@example.com", + CreatedAt: time.Now(), + } + + 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, "a response-write failure should produce a fosite error redirect") + + // Both legs of the satisfied chain are cleaned up when the final response write fails. + assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-1", + "the first leg's token should be cleaned up") + assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-2", + "the final leg's token should be cleaned up") +} + +// TestCallbackHandler_Cleanup_ContinuesWhenAProviderDeleteFails covers +// cleanupUpstreamTokens' warn-and-continue path: one provider's delete errors, and the +// cleanup must still remove the sibling providers rather than aborting on the first +// failure. +func TestCallbackHandler_Cleanup_ContinuesWhenAProviderDeleteFails(t *testing.T) { + t.Parallel() + handler, storState, _, _ := multiUpstreamTestSetupWithStorage(t, + withDeleteUpstreamTokensForProviderError("provider-1", errors.New("delete failed"))) + + sessionID := "cleanup-partial-fail-session" + + // provider-1 has a tampered UserID so verifyChainIdentity fails and the chain is + // cleaned up; provider-1's delete is wired to fail. + storState.upstreamTokens[sessionID+":provider-1"] = &storage.UpstreamTokens{ + ProviderID: "provider-1", + AccessToken: "p1-at", + ExpiresAt: time.Now().Add(time.Hour), + ClientID: testAuthClientID, + UserID: "tampered-user-id", + } + + secondLegState := "cleanup-partial-fail-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: "partial-fail-verifier-12345678901234567890", + UpstreamNonce: "partial-fail-nonce", + UpstreamProviderName: "provider-2", + SessionID: sessionID, + ChainUpstreams: []string{"provider-1", "provider-2"}, + ResolvedUserID: "correct-user-id", + ResolvedUserName: "Correct User", + ResolvedUserEmail: "correct@example.com", + CreatedAt: time.Now(), + } + + 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) + + // provider-1's delete failed, so its row is still present... + assert.Contains(t, storState.upstreamTokens, sessionID+":provider-1", + "the provider whose delete failed should still be present") + // ...but the cleanup continued and removed provider-2 anyway. + assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-2", + "cleanup must continue past a per-provider delete failure and remove sibling providers") +} + func TestCallbackHandler_TwoUpstreams_StorePendingError_CleansUp(t *testing.T) { t.Parallel() @@ -1730,6 +1936,12 @@ func TestCallbackHandler_PlacesPlatformUserInContext_OnUpstreamErrorCleanup(t *t require.True(t, ok, "handleUpstreamError must place platform user in ctx before the cleanup delete") require.Equal(t, leg1User, uid) + // The earlier leg's row must actually be gone — not merely have had its ctx + // captured. A regression passing an empty or wrong provider list would still + // populate deleteUpstreamCtx (the mock runs regardless) but leave the row. + assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-1", + "the earlier-leg token should be deleted by the cleanup, not just context-propagated") + // The error path must NOT place a stub Identity under the identity key. _, hasIdentity := auth.IdentityFromContext(storState.deleteUpstreamCtx) require.False(t, hasIdentity, "handleUpstreamError must not place an Identity (only a platform user) in ctx") diff --git a/pkg/authserver/server/handlers/helpers_test.go b/pkg/authserver/server/handlers/helpers_test.go index 7086dbb0c0..bb9bfebf4f 100644 --- a/pkg/authserver/server/handlers/helpers_test.go +++ b/pkg/authserver/server/handlers/helpers_test.go @@ -105,9 +105,14 @@ type testStorageState struct { type baseTestSetupOption func(*baseTestSetupConfig) type baseTestSetupConfig struct { - storePendingErr error // if non-nil, StorePendingAuthorization always returns this error - getLatestUpstreamTokensErr error // if non-nil, GetLatestUpstreamTokensForUser always returns this error - createUserErr error // if non-nil, CreateUser always returns this error + storePendingErr error // if non-nil, StorePendingAuthorization always returns this error + getLatestUpstreamTokensErr error // if non-nil, GetLatestUpstreamTokensForUser always returns this error + createUserErr error // if non-nil, CreateUser always returns this error + storeUpstreamTokensErr error // if non-nil, StoreUpstreamTokens always returns this error + getAllUpstreamTokensErr error // if non-nil, GetAllUpstreamTokens always returns this error + deleteForProviderErrs map[string]error // per-provider error for DeleteUpstreamTokensForProvider + getClientErr error // if non-nil, GetClient returns this after getClientErrAfter successful calls + getClientErrAfter int // number of GetClient calls that succeed before getClientErr kicks in } func withStorePendingError(err error) baseTestSetupOption { @@ -116,6 +121,40 @@ func withStorePendingError(err error) baseTestSetupOption { } } +func withStoreUpstreamTokensError(err error) baseTestSetupOption { + return func(c *baseTestSetupConfig) { + c.storeUpstreamTokensErr = err + } +} + +func withGetAllUpstreamTokensError(err error) baseTestSetupOption { + return func(c *baseTestSetupConfig) { + c.getAllUpstreamTokensErr = err + } +} + +// withDeleteUpstreamTokensForProviderError makes DeleteUpstreamTokensForProvider return +// err for the named provider only, so a test can exercise cleanupUpstreamTokens' +// warn-and-continue path while sibling providers still delete. +func withDeleteUpstreamTokensForProviderError(provider string, err error) baseTestSetupOption { + return func(c *baseTestSetupConfig) { + if c.deleteForProviderErrs == nil { + c.deleteForProviderErrs = make(map[string]error) + } + c.deleteForProviderErrs[provider] = err + } +} + +// withGetClientErrorAfterCalls makes GetClient return err starting with the +// (after+1)-th call, so a test can fail the GetClient inside writeAuthorizationResponse +// while the earlier GetClient that builds the error-redirect requester still succeeds. +func withGetClientErrorAfterCalls(after int, err error) baseTestSetupOption { + return func(c *baseTestSetupConfig) { + c.getClientErr = err + c.getClientErrAfter = after + } +} + func withGetLatestUpstreamTokensError(err error) baseTestSetupOption { return func(c *baseTestSetupConfig) { c.getLatestUpstreamTokensErr = err @@ -194,7 +233,12 @@ func baseTestSetup(t *testing.T, opts ...baseTestSetupOption) (fosite.OAuth2Prov // Setup mock expectations for GetClient. Looks up storState.clients so tests // can register additional clients (e.g. a loopback client under its own ID) // after baseTestSetup returns. + getClientCalls := 0 stor.EXPECT().GetClient(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, id string) (fosite.Client, error) { + getClientCalls++ + if setupCfg.getClientErr != nil && getClientCalls > setupCfg.getClientErrAfter { + return nil, setupCfg.getClientErr + } if c, ok := storState.clients[id]; ok { return c, nil } @@ -354,6 +398,9 @@ func baseTestSetup(t *testing.T, opts ...baseTestSetupOption) (fosite.OAuth2Prov // Keyed by "sessionID:providerName" to support multiple providers per session. stor.EXPECT().StoreUpstreamTokens(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( func(_ context.Context, sessionID, providerName string, tokens *storage.UpstreamTokens) error { + if setupCfg.storeUpstreamTokensErr != nil { + return setupCfg.storeUpstreamTokensErr + } key := sessionID + ":" + providerName storState.upstreamTokens[key] = tokens storState.idpTokenCount++ @@ -381,8 +428,18 @@ func baseTestSetup(t *testing.T, opts ...baseTestSetupOption) (fosite.OAuth2Prov // still carries no tokens argument, so a context-keyed storage decorator // resolves the user from ctx exactly as DeleteUpstreamTokens does. Capture // the ctx here too so the callback's per-provider cleanup path is covered by - // the same context-placement assertions. Deleting an absent row is a no-op. + // the same context-placement assertions. storState.deleteUpstreamCtx = ctx + // Mirror the real backends, which reject an empty session or provider rather + // than deleting; cleanupUpstreamTokens filters empty names, so this only + // fires if a future caller bug lets one through. + if sessionID == "" || providerName == "" { + return fosite.ErrInvalidRequest + } + if err := setupCfg.deleteForProviderErrs[providerName]; err != nil { + return err + } + // Deleting an absent row is a no-op. delete(storState.upstreamTokens, sessionID+":"+providerName) return nil }).AnyTimes() @@ -394,6 +451,9 @@ func baseTestSetup(t *testing.T, opts ...baseTestSetupOption) (fosite.OAuth2Prov // only from ctx. Capture the ctx here so a test can assert the callback // placed the identity into it before this read runs. storState.getAllUpstreamCtx = ctx + if setupCfg.getAllUpstreamTokensErr != nil { + return nil, setupCfg.getAllUpstreamTokensErr + } result := make(map[string]*storage.UpstreamTokens) prefix := sessionID + ":" for key, tokens := range storState.upstreamTokens { @@ -475,13 +535,9 @@ func handlerTestSetup(t *testing.T, opts ...baseTestSetupOption) (*Handler, *tes return handler, storState, mockUpstream } -// multiUpstreamTestSetup creates a test setup with two upstream providers ("provider-1" and "provider-2") -// for testing multi-upstream authorization chain logic. Any Option values are forwarded to NewHandler. -func multiUpstreamTestSetup(t *testing.T, opts ...Option) (*Handler, *testStorageState, *mockIDPProvider, *mockIDPProvider) { - t.Helper() - - provider, oauth2Config, stor, storState := baseTestSetup(t) - +// multiUpstreamProviders builds the two mock IdP providers ("provider-1" and +// "provider-2") and the NamedUpstream slice shared by the multi-upstream setups. +func multiUpstreamProviders() (*mockIDPProvider, *mockIDPProvider, []NamedUpstream) { mockProvider1 := &mockIDPProvider{ providerType: upstream.ProviderTypeOAuth2, authorizationURL: "https://idp1.example.com/authorize", @@ -518,8 +574,33 @@ func multiUpstreamTestSetup(t *testing.T, opts ...Option) (*Handler, *testStorag {Name: "provider-1", Provider: mockProvider1}, {Name: "provider-2", Provider: mockProvider2}, } + return mockProvider1, mockProvider2, upstreams +} + +// multiUpstreamTestSetup creates a test setup with two upstream providers ("provider-1" and "provider-2") +// for testing multi-upstream authorization chain logic. Any Option values are forwarded to NewHandler. +func multiUpstreamTestSetup(t *testing.T, opts ...Option) (*Handler, *testStorageState, *mockIDPProvider, *mockIDPProvider) { + t.Helper() + + provider, oauth2Config, stor, storState := baseTestSetup(t) + mockProvider1, mockProvider2, upstreams := multiUpstreamProviders() handler, err := NewHandler(provider, oauth2Config, stor, upstreams, opts...) require.NoError(t, err) return handler, storState, mockProvider1, mockProvider2 } + +// multiUpstreamTestSetupWithStorage is multiUpstreamTestSetup with storage-behavior +// overrides (error injection) threaded into baseTestSetup, so a test can drive the +// callback's cleanup paths (store failure, chain-read failure, per-provider delete +// failure) on a real two-upstream chain. +func multiUpstreamTestSetupWithStorage(t *testing.T, storageOpts ...baseTestSetupOption) (*Handler, *testStorageState, *mockIDPProvider, *mockIDPProvider) { + t.Helper() + + provider, oauth2Config, stor, storState := baseTestSetup(t, storageOpts...) + mockProvider1, mockProvider2, upstreams := multiUpstreamProviders() + handler, err := NewHandler(provider, oauth2Config, stor, upstreams) + require.NoError(t, err) + + return handler, storState, mockProvider1, mockProvider2 +} From ac6b820f2ecff71b6f770667df50df35b514f9bd Mon Sep 17 00:00:00 2001 From: Trey Date: Mon, 31 Aug 2026 11:46:04 -0700 Subject: [PATCH 4/4] fix(authserver): scope callback cleanup to the current leg, not the whole chain MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A failed authorization leg now cleans up only the leg being processed and leaves the user's earlier legs intact, rather than deleting every provider in the chain. Under a storage decorator that re-keys upstream tokens by (gateway, user), each earlier leg is the user's own live credential for another connector, so a transient failure on one login leg must not take the others down with it — most importantly at the chain-resolution site, the issue's most-reachable trigger (a Directory/filter blip on leg 1). The one exception is the identity-mismatch fail-closed (verifyChainIdentity): a detected mismatch means the session's stored state is no longer trustworthy, so that site still tears down the whole chain. - 8 of the 9 cleanup sites now pass just the current leg's provider; the verifyChainIdentity site keeps the whole-chain delete. - Drop the now-dead dedup in cleanupUpstreamTokens: every caller passes a single provider or a validateChain-checked, duplicate-free chain, and a repeated name is a harmless idempotent no-op delete. - Update the affected tests to assert earlier legs survive an ordinary error, and add a direct cleanupUpstreamTokens test covering the skip-unconfigured / skip-empty guard. Signed-off-by: Trey --- pkg/authserver/server/handlers/callback.go | 90 ++++++++++--------- .../server/handlers/callback_test.go | 79 +++++++++++----- 2 files changed, 104 insertions(+), 65 deletions(-) diff --git a/pkg/authserver/server/handlers/callback.go b/pkg/authserver/server/handlers/callback.go index bdd971d36f..bd211c2c7b 100644 --- a/pkg/authserver/server/handlers/callback.go +++ b/pkg/authserver/server/handlers/callback.go @@ -182,11 +182,10 @@ func (h *Handler) CallbackHandler(w http.ResponseWriter, req *http.Request) { slog.Error("failed to store upstream tokens", "error", err, ) - // Clean up tokens stored by earlier legs of a multi-upstream chain (see - // cleanupUpstreamTokens for the per-provider scoping rationale). providerID - // (this leg) is included even though its store just failed: any partial or - // carried-forward row is removed, and an absent row deletes as a no-op. - h.cleanupUpstreamTokens(ctx, sessionID, append([]string{providerID}, pending.ChainUpstreams...)) + // Clean up only this leg's token (see cleanupUpstreamTokens): its store just + // failed, so this removes any partial or carried-forward row and is otherwise a + // no-op. Earlier legs are the user's own valid tokens and are left intact. + h.cleanupUpstreamTokens(ctx, sessionID, []string{providerID}) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to store session")) return } @@ -401,20 +400,20 @@ func (h *Handler) handleUpstreamError( pending, err := h.storage.LoadPendingAuthorization(ctx, internalState) if err == nil { _ = h.storage.DeletePendingAuthorization(ctx, internalState) - // Clean up upstream tokens stored by earlier legs of a multi-upstream chain - // (see cleanupUpstreamTokens for the per-provider scoping rationale). On a + // Clean up only this leg's token (see cleanupUpstreamTokens); earlier legs are + // the user's own valid tokens and are left intact. This leg failed at the + // upstream IdP before a token was stored, so the delete is normally a no-op, + // but it also removes a row left by a prior attempt on the same leg. On a // subsequent leg the resolved user is carried in pending.ResolvedUserID; place // it in ctx via WithPlatformUser so a context-keyed storage decorator resolves // the canonical user for the cleanup. WithPlatformUser, not WithIdentity: there - // is no authenticated identity at the callback, only the storage-scoped user. A - // first-leg error carries no ChainUpstreams and no earlier-leg tokens, so - // nothing is deleted there. + // is no authenticated identity at the callback, only the storage-scoped user. if pending.SessionID != "" { cleanupCtx := ctx if pending.ResolvedUserID != "" { cleanupCtx = auth.WithPlatformUser(ctx, pending.ResolvedUserID) } - h.cleanupUpstreamTokens(cleanupCtx, pending.SessionID, pending.ChainUpstreams) + h.cleanupUpstreamTokens(cleanupCtx, pending.SessionID, []string{pending.UpstreamProviderName}) } ar := h.buildAuthorizeRequesterFromPending(ctx, pending) if ar != nil { @@ -432,17 +431,21 @@ func (h *Handler) handleUpstreamError( http.Error(w, "upstream authentication failed", http.StatusBadGateway) } -// cleanupUpstreamTokens best-effort removes the upstream tokens an authorization -// chain has collected, one provider at a time. It is the failure-path counterpart -// to a successful chain: when a leg fails, the tokens earlier legs stored under this -// session must not be left mid-chain. +// cleanupUpstreamTokens best-effort removes upstream tokens on a failed authorization +// leg, one provider at a time. Callers pass the leg currently being processed, so an +// ordinary error — a transient store/read failure, or a chain that cannot be resolved — +// takes down only that leg's token and leaves the user's earlier legs intact. That +// matters under a storage decorator that re-keys upstream tokens by (gateway, user), +// where each earlier leg is the user's own live credential for another connector rather +// than throwaway session state, so a failed connector login must not delete the others. +// The one exception is the identity-mismatch fail-closed, which passes the whole chain +// to tear down a session whose stored identity is no longer trustworthy. // // It deletes per provider rather than calling the session-wide DeleteUpstreamTokens -// because that method carries no provider name. A storage decorator that re-keys -// upstream tokens by (gateway, user) — so a session-independent caller can find a -// user's tokens — cannot narrow a session-wide delete back to this chain, and would -// instead delete every connector token the user holds on the gateway. Naming each -// provider lets such a decorator scope the delete to exactly this chain's providers. +// because that method carries no provider name: a (gateway, user)-keyed decorator cannot +// narrow a session-wide delete back to one leg and would delete every connector token +// the user holds on the gateway. Naming the provider lets such a decorator scope the +// delete to exactly the named leg(s). // // Each name is checked against the configured upstream set before it is deleted, so a // name that is not a configured upstream — e.g. from a corrupted or tampered pending @@ -450,22 +453,15 @@ func (h *Handler) handleUpstreamError( // not been through validateChain — is skipped rather than turned into a delete. // // providers is the set of provider names whose upstream-token rows should be removed; -// empty entries, duplicates, and names outside the configured upstream set are ignored, -// and deleting an already-absent row is a no-op (the storage contract makes a -// missing-row delete non-fatal). Errors are logged and swallowed — cleanup is -// best-effort and must not mask the original failure. Cost is one storage round-trip -// per distinct provider; a chain is the corporate primary plus a handful of connectors, -// so the sequential deletes are bounded and only run on error paths. +// empty entries and names outside the configured upstream set are skipped, and deleting +// an already-absent row is a no-op (the storage contract makes a missing-row delete +// non-fatal), so a repeated name is harmless. Errors are logged and swallowed — cleanup +// is best-effort and must not mask the original failure. func (h *Handler) cleanupUpstreamTokens(ctx context.Context, sessionID string, providers []string) { - seen := make(map[string]struct{}, len(providers)) for _, provider := range providers { if provider == "" { continue } - if _, dup := seen[provider]; dup { - continue - } - seen[provider] = struct{}{} // Only a configured upstream can legitimately have a stored token, so skip // anything else: it keeps a tampered pending row reaching a pre-resolveChain // cleanup site from naming an arbitrary delete target. The post-resolveChain @@ -530,11 +526,11 @@ func (h *Handler) continueChainOrComplete( chain, err := h.resolveChain(ctx, pending, principal) if err != nil { slog.Error("failed to resolve upstream chain", "error", err) - // The chain could not be resolved, so scope the cleanup to what is known: this - // leg's provider plus any chain carried forward on a subsequent leg. An earlier - // leg not named here (e.g. a stale pending with no ChainUpstreams) is left to - // expire by TTL rather than risk a user-keyed decorator over-deleting. - h.cleanupUpstreamTokens(ctx, sessionID, append([]string{pending.UpstreamProviderName}, pending.ChainUpstreams...)) + // Clean up only this leg's token; earlier legs are the user's own valid tokens + // and are left intact. This is the issue's most-reachable trigger — a transient + // filter/Directory failure on leg 1 — so scoping it to the one leg is what keeps + // a blip from taking the user's other connectors down with it. + h.cleanupUpstreamTokens(ctx, sessionID, []string{pending.UpstreamProviderName}) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to determine authorization chain")) return } @@ -542,7 +538,8 @@ func (h *Handler) continueChainOrComplete( nextProvider, err := h.nextMissingUpstream(ctx, sessionID, chain) if err != nil { slog.Error("failed to determine next upstream", "error", err) - h.cleanupUpstreamTokens(ctx, sessionID, chain) + // Clean up only this leg's token; earlier legs are left intact. + h.cleanupUpstreamTokens(ctx, sessionID, []string{pending.UpstreamProviderName}) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to check authorization chain state")) return } @@ -550,7 +547,10 @@ func (h *Handler) continueChainOrComplete( if nextProvider == "" { if err := h.verifyChainIdentity(ctx, sessionID, chain, subject); err != nil { // verifyChainIdentity already logged the specific cause (with structured - // fields for a mismatch); here we just clean up and fail closed. + // fields for a mismatch). This is the one cleanup site that tears down the + // WHOLE chain rather than just this leg: a detected identity mismatch means + // the session's stored state is no longer trustworthy, so we fail closed and + // remove every leg's token rather than leave a possibly cross-keyed row. h.cleanupUpstreamTokens(ctx, sessionID, chain) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("identity verification failed")) return @@ -559,7 +559,10 @@ func (h *Handler) continueChainOrComplete( // All upstreams satisfied — issue authorization code if err := h.writeAuthorizationResponse(ctx, w, pending, sessionID, subject, name, email); err != nil { slog.Error("failed to create authorization response", "error", err) - h.cleanupUpstreamTokens(ctx, sessionID, chain) + // The chain was fully satisfied and identity-verified; only the response + // write failed (a retryable server error). Clean up just this final leg and + // leave the earlier legs, so a retry re-collects only this one. + h.cleanupUpstreamTokens(ctx, sessionID, []string{pending.UpstreamProviderName}) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to create authorization code")) } return @@ -597,7 +600,8 @@ func (h *Handler) continueChainOrComplete( if err := h.storage.StorePendingAuthorization(ctx, secrets.State, nextPending); err != nil { slog.Error("failed to store next chain leg", "error", err) - h.cleanupUpstreamTokens(ctx, sessionID, chain) + // Clean up only this leg's token; earlier legs are left intact. + h.cleanupUpstreamTokens(ctx, sessionID, []string{pending.UpstreamProviderName}) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to continue authorization chain")) return } @@ -611,7 +615,8 @@ func (h *Handler) continueChainOrComplete( if !ok { slog.Error("next upstream provider not found", "provider", nextProvider) _ = h.storage.DeletePendingAuthorization(ctx, secrets.State) - h.cleanupUpstreamTokens(ctx, sessionID, chain) + // Clean up only this leg's token; earlier legs are left intact. + h.cleanupUpstreamTokens(ctx, sessionID, []string{pending.UpstreamProviderName}) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("upstream provider configuration error")) return } @@ -619,7 +624,8 @@ func (h *Handler) continueChainOrComplete( if err != nil { slog.Error("failed to build next upstream authorization URL", "error", err) _ = h.storage.DeletePendingAuthorization(ctx, secrets.State) - h.cleanupUpstreamTokens(ctx, sessionID, chain) + // Clean up only this leg's token; earlier legs are left intact. + h.cleanupUpstreamTokens(ctx, sessionID, []string{pending.UpstreamProviderName}) h.provider.WriteAuthorizeError(ctx, w, ar, fosite.ErrServerError.WithHint("failed to build authorization URL")) return } diff --git a/pkg/authserver/server/handlers/callback_test.go b/pkg/authserver/server/handlers/callback_test.go index 08de280341..8bb6c50a1b 100644 --- a/pkg/authserver/server/handlers/callback_test.go +++ b/pkg/authserver/server/handlers/callback_test.go @@ -1307,12 +1307,43 @@ func TestCallbackHandler_Cleanup_LeavesOutOfChainProviderTokens(t *testing.T) { assert.Equal(t, "keep-me", storState.upstreamTokens[outOfChainKey].AccessToken) } -// TestCallbackHandler_StoreUpstreamTokensError_CleansUpChain covers the -// StoreUpstreamTokens-failure cleanup site on a subsequent leg. It is the one site -// where cleanupUpstreamTokens' dedup is load-bearing: the provider list is -// append([]string{providerID}, pending.ChainUpstreams...), and ChainUpstreams already -// contains providerID on a non-first leg, so the same provider is named twice. -func TestCallbackHandler_StoreUpstreamTokensError_CleansUpChain(t *testing.T) { +// TestCleanupUpstreamTokens_SkipsUnconfiguredAndEmpty exercises cleanupUpstreamTokens +// directly: it skips empty names and any provider that is not a configured upstream — the +// guard that stops a corrupted or tampered pending row reaching a pre-resolveChain +// cleanup site from naming an arbitrary delete target — while removing the configured +// ones. +func TestCleanupUpstreamTokens_SkipsUnconfiguredAndEmpty(t *testing.T) { + t.Parallel() + handler, storState, _, _ := multiUpstreamTestSetup(t) + + sessionID := "direct-cleanup-session" + for _, p := range []string{"provider-1", "provider-2", "not-configured"} { + storState.upstreamTokens[sessionID+":"+p] = &storage.UpstreamTokens{ + ProviderID: p, + AccessToken: "at", + ExpiresAt: time.Now().Add(time.Hour), + ClientID: testAuthClientID, + UserID: "user-1", + } + } + + handler.cleanupUpstreamTokens(t.Context(), sessionID, + []string{"provider-1", "", "not-configured", "provider-2"}) + + // Configured providers are removed... + assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-1") + assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-2") + // ...while the unconfigured name is skipped (its row survives) and the empty name is + // a no-op rather than a panic. + assert.Contains(t, storState.upstreamTokens, sessionID+":not-configured", + "a name that is not a configured upstream must not be deleted") +} + +// TestCallbackHandler_StoreUpstreamTokensError_LeavesEarlierLegs covers the +// StoreUpstreamTokens-failure cleanup site on a subsequent leg: the cleanup is scoped to +// the leg being processed (provider-2, whose store just failed), so the user's earlier +// leg (provider-1) is left intact rather than wiped by the failure. +func TestCallbackHandler_StoreUpstreamTokensError_LeavesEarlierLegs(t *testing.T) { t.Parallel() handler, storState, _, _ := multiUpstreamTestSetupWithStorage(t, withStoreUpstreamTokensError(errors.New("storage unavailable"))) @@ -1329,8 +1360,6 @@ func TestCallbackHandler_StoreUpstreamTokensError_CleansUpChain(t *testing.T) { UserID: leg1User, } - // Subsequent-leg pending whose ChainUpstreams already contains provider-2 (this - // leg), so the cleanup list names provider-2 twice — exercising the dedup. secondLegState := "store-err-second-leg-state" storState.pendingAuths[secondLegState] = &storage.PendingAuthorization{ ClientID: testAuthClientID, @@ -1358,9 +1387,10 @@ func TestCallbackHandler_StoreUpstreamTokensError_CleansUpChain(t *testing.T) { assert.Equal(t, http.StatusSeeOther, rec.Code, "a store failure should produce a fosite error redirect") assert.Contains(t, rec.Header().Get("Location"), "error=server_error") - // The earlier leg's token is cleaned up; provider-2's store never landed. - assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-1", - "the earlier chain leg's token should be cleaned up on a store failure") + // The earlier leg's token survives — cleanup is scoped to the failed leg only. + assert.Contains(t, storState.upstreamTokens, sessionID+":provider-1", + "an earlier leg's token must survive a later leg's store failure") + // provider-2's store failed, so it never had a row. assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-2", "provider-2's store failed, so no row should exist") } @@ -1451,11 +1481,12 @@ func TestCallbackHandler_WriteAuthorizationResponseError_OnSatisfiedChain_Cleans assert.Equal(t, http.StatusSeeOther, rec.Code, "a response-write failure should produce a fosite error redirect") - // Both legs of the satisfied chain are cleaned up when the final response write fails. - assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-1", - "the first leg's token should be cleaned up") + // Only the final leg is cleaned up; the response-write failure is retryable and the + // earlier leg is left intact so a retry re-collects only the last one. assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-2", "the final leg's token should be cleaned up") + assert.Contains(t, storState.upstreamTokens, sessionID+":provider-1", + "the earlier leg's token must survive a retryable response-write failure") } // TestCallbackHandler_Cleanup_ContinuesWhenAProviderDeleteFails covers @@ -1878,16 +1909,18 @@ func TestCallbackHandler_PlacesPlatformUserInContext_OnChainRead(t *testing.T) { } // TestCallbackHandler_PlacesPlatformUserInContext_OnUpstreamErrorCleanup verifies that -// when a subsequent-leg callback returns an upstream error, the earlier-leg cleanup -// (per-provider DeleteUpstreamTokensForProvider in handleUpstreamError) runs with the -// resolved canonical user already in the request context. +// when a subsequent-leg callback returns an upstream error, the per-provider cleanup +// (DeleteUpstreamTokensForProvider in handleUpstreamError) runs with the resolved +// canonical user already in the request context, and that it leaves the earlier leg. // // This is the error-path sibling of the chain-read test above. handleUpstreamError // runs from an early return at the top of CallbackHandler, before the happy-path // user injection, so it must place the user itself. DeleteUpstreamTokensForProvider // carries no tokens argument — so a context-keyed storage decorator can resolve the // canonical user only from ctx. The resolved user is carried forward from leg 1 via -// pending.ResolvedUserID, and pending.ChainUpstreams names the earlier leg to clean. +// pending.ResolvedUserID. The cleanup is scoped to the leg being processed (provider-2, +// which errored at the upstream so it has no row), so the earlier leg (provider-1) +// survives. func TestCallbackHandler_PlacesPlatformUserInContext_OnUpstreamErrorCleanup(t *testing.T) { t.Parallel() handler, storState, _, _ := multiUpstreamTestSetup(t) @@ -1936,11 +1969,11 @@ func TestCallbackHandler_PlacesPlatformUserInContext_OnUpstreamErrorCleanup(t *t require.True(t, ok, "handleUpstreamError must place platform user in ctx before the cleanup delete") require.Equal(t, leg1User, uid) - // The earlier leg's row must actually be gone — not merely have had its ctx - // captured. A regression passing an empty or wrong provider list would still - // populate deleteUpstreamCtx (the mock runs regardless) but leave the row. - assert.NotContains(t, storState.upstreamTokens, sessionID+":provider-1", - "the earlier-leg token should be deleted by the cleanup, not just context-propagated") + // The earlier leg's token survives: cleanup is scoped to the leg being processed + // (provider-2), not the whole chain, so a failed connector login never wipes the + // user's other connectors. + assert.Contains(t, storState.upstreamTokens, sessionID+":provider-1", + "an earlier leg's token must survive a later leg's upstream error") // The error path must NOT place a stub Identity under the identity key. _, hasIdentity := auth.IdentityFromContext(storState.deleteUpstreamCtx)