diff --git a/pkg/authserver/server/handlers/callback.go b/pkg/authserver/server/handlers/callback.go index 727cc36769..bd211c2c7b 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,8 +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. - _ = h.storage.DeleteUpstreamTokens(ctx, sessionID) + // 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 } @@ -398,19 +400,20 @@ 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 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. + // 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. 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, []string{pending.UpstreamProviderName}) } ar := h.buildAuthorizeRequesterFromPending(ctx, pending) if ar != nil { @@ -428,6 +431,53 @@ func (h *Handler) handleUpstreamError( http.Error(w, "upstream authentication failed", http.StatusBadGateway) } +// 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 (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 +// 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 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) { + for _, provider := range providers { + if provider == "" { + continue + } + // 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", //nolint:gosec // G706 - provider name from server-side chain state + "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 +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) - _ = h.storage.DeleteUpstreamTokens(ctx, sessionID) + // 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 } @@ -484,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.storage.DeleteUpstreamTokens(ctx, sessionID) + // 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 } @@ -492,8 +547,11 @@ 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. - _ = h.storage.DeleteUpstreamTokens(ctx, sessionID) + // 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 } @@ -501,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.storage.DeleteUpstreamTokens(ctx, sessionID) + // 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 @@ -539,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.storage.DeleteUpstreamTokens(ctx, sessionID) + // 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 } @@ -553,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.storage.DeleteUpstreamTokens(ctx, sessionID) + // 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 } @@ -561,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.storage.DeleteUpstreamTokens(ctx, sessionID) + // 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 add96cf56a..8bb6c50a1b 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,311 @@ 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) +} + +// 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"))) + + 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, + } + + 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 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") +} + +// 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") + + // 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 +// 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() @@ -1599,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 -// (DeleteUpstreamTokens 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. 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. 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) @@ -1638,6 +1950,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,13 +1962,19 @@ 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 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) 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 1f87ed1b39..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++ @@ -375,6 +422,28 @@ 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. + 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() + 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 @@ -382,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 { @@ -463,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", @@ -506,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 +}