diff --git a/apps/apollo-vertex/templates/ai-chat/AiChatLoginGate.tsx b/apps/apollo-vertex/templates/ai-chat/AiChatLoginGate.tsx index f859527ce5..2865222dba 100644 --- a/apps/apollo-vertex/templates/ai-chat/AiChatLoginGate.tsx +++ b/apps/apollo-vertex/templates/ai-chat/AiChatLoginGate.tsx @@ -7,7 +7,7 @@ import { ChevronRight, LogIn, LogOut } from "lucide-react"; import type { ReactNode } from "react"; import { z } from "zod"; import { Button } from "@/components/ui/button"; -import { STORAGE_KEYS, toCodedAppFilePath } from "@/lib/auth"; +import { getStoredIdToken, STORAGE_KEYS, toCodedAppFilePath } from "@/lib/auth"; import { STALE_TIME_MS } from "@/lib/constants"; import { Select, @@ -18,10 +18,12 @@ import { } from "@/registry/select/select"; import { useAuth } from "@/registry/shell/shell-auth-provider"; import { + AICHAT_CLIENT_ID, AICHAT_CODED_APP_PATH, AICHAT_DIRECT_BASE_URL, AICHAT_IS_CODED_APP, AICHAT_PORTAL_API_BASE, + AICHAT_SCOPE, AICHAT_STORAGE_KEYS, } from "./ai-chat-example-utils"; @@ -169,23 +171,7 @@ export function AiChatLoginGate({ children }: AiChatLoginGateProps) { const orgTenant = resolveOrgTenant(orgInfo, selectedTenantId); function handleLogout() { - const tokenDataStr = localStorage.getItem(STORAGE_KEYS.TOKEN); - let idToken: string | null = null; - if (tokenDataStr) { - try { - const parsed: unknown = JSON.parse(tokenDataStr); - if ( - typeof parsed === "object" && - parsed !== null && - "id_token" in parsed && - typeof parsed.id_token === "string" - ) { - idToken = parsed.id_token; - } - } catch { - idToken = null; - } - } + const idToken = getStoredIdToken(AICHAT_CLIENT_ID, AICHAT_SCOPE); sessionStorage.setItem( STORAGE_KEYS.LOGOUT_RETURN_TO, diff --git a/packages/apollo-vertex/AGENTS.md b/packages/apollo-vertex/AGENTS.md index 5177151b02..1087ce20ca 100644 --- a/packages/apollo-vertex/AGENTS.md +++ b/packages/apollo-vertex/AGENTS.md @@ -57,8 +57,6 @@ pnpm --filter @uipath/apollo-react test -- --grep "ApButton" pnpm --filter @uipath/apollo-wind test:watch -- button.test.tsx ``` -Note: @uipath/apollo-vertex currently has no test files. Tests are in `packages/apollo-wind` and `packages/apollo-react`. - ## Code Style Guidelines ### Formatting (Biome) diff --git a/packages/apollo-vertex/entries/index.ts b/packages/apollo-vertex/entries/index.ts index e4e18e6310..b58a5e5f0c 100644 --- a/packages/apollo-vertex/entries/index.ts +++ b/packages/apollo-vertex/entries/index.ts @@ -168,6 +168,8 @@ export { TOKEN_QUERY_KEY, type TokenData, ensureValidToken, + getStoredIdToken, + getTokenQueryKey, login, logout, resolveReturnPath, diff --git a/packages/apollo-vertex/lib/auth.test.ts b/packages/apollo-vertex/lib/auth.test.ts new file mode 100644 index 0000000000..df9e0a9765 --- /dev/null +++ b/packages/apollo-vertex/lib/auth.test.ts @@ -0,0 +1,324 @@ +import { QueryClient, QueryObserver } from "@tanstack/react-query"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import { + ensureValidToken, + getStoredIdToken, + getTokenQueryKey, + login, + logout, + STORAGE_KEYS, + type TokenData, +} from "./auth"; + +const CLIENT_ID = "client-a"; +const OTHER_CLIENT_ID = "client-b"; +const BASE_URL = "https://example.test"; +const SCOPE = "openid offline_access OR.Jobs DataFabric.Data.Read"; +const NORMALIZED_SCOPE = "DataFabric.Data.Read OR.Jobs offline_access openid"; +const TOKEN_STORAGE_KEY = `${STORAGE_KEYS.TOKEN}:${CLIENT_ID}:${NORMALIZED_SCOPE}`; + +const tokenData = (overrides: Partial = {}): TokenData => ({ + access_token: "stored-access", + id_token: "id", + refresh_token: "refresh", + expires_in: 3600, + token_type: "Bearer", + expiresAt: Date.now() + 60 * 60 * 1000, + ...overrides, +}); + +const base64Url = (value: object) => + btoa(JSON.stringify(value)) + .replace(/=+$/, "") + .replaceAll("+", "-") + .replaceAll("/", "_"); + +const idTokenFor = (aud: string) => + `${base64Url({ alg: "none" })}.${base64Url({ aud })}.signature`; + +const storeToken = (key: string, overrides: Partial = {}) => { + localStorage.setItem(key, JSON.stringify(tokenData(overrides))); +}; + +const readToken = (key: string): unknown => { + const raw = localStorage.getItem(key); + return raw === null ? null : JSON.parse(raw); +}; + +const tokenResponse = (accessToken: string) => + new Response( + JSON.stringify({ + access_token: accessToken, + id_token: "id", + refresh_token: "refresh-2", + expires_in: 3600, + token_type: "Bearer", + }), + { status: 200 }, + ); + +const mockTokenEndpoint = (accessToken: string) => { + const fetchMock = vi.fn().mockResolvedValue(tokenResponse(accessToken)); + vi.stubGlobal("fetch", fetchMock); + return fetchMock; +}; + +const grantTypeOf = (init?: RequestInit) => + init?.body instanceof URLSearchParams ? init.body.get("grant_type") : null; + +const mockFailingTokenEndpoint = () => { + const fetchMock = vi + .fn() + .mockResolvedValue(new Response("invalid_grant", { status: 400 })); + vi.stubGlobal("fetch", fetchMock); + return fetchMock; +}; + +const ensureValidTokenFor = (scope: string, clientId = CLIENT_ID) => + ensureValidToken(clientId, scope, BASE_URL, "/"); + +const loginAndReturnWithCode = async (scope: string) => { + // jsdom reports the authorize redirect as an unimplemented navigation. + vi.spyOn(console, "error").mockImplementation(() => null); + await login(CLIENT_ID, scope, BASE_URL); + const state = sessionStorage.getItem(STORAGE_KEYS.STATE); + window.history.replaceState({}, "", `/?code=abc&state=${state}`); +}; + +const expectCallbackCleanedUp = () => { + expect(window.location.search).toBe(""); + expect(sessionStorage.getItem(STORAGE_KEYS.CODE_VERIFIER)).toBeNull(); + expect(sessionStorage.getItem(STORAGE_KEYS.STATE)).toBeNull(); + expect(sessionStorage.getItem(STORAGE_KEYS.REQUESTED_SCOPE)).toBeNull(); +}; + +beforeEach(() => { + localStorage.clear(); + sessionStorage.clear(); + window.history.replaceState({}, "", "/"); +}); + +afterEach(() => { + vi.unstubAllGlobals(); + vi.restoreAllMocks(); +}); + +describe("token storage key", () => { + it("finds the session whatever the scope order and spacing", async () => { + storeToken(TOKEN_STORAGE_KEY); + + await expect( + ensureValidTokenFor( + "DataFabric.Data.Read OR.Jobs openid offline_access", + ), + ).resolves.toBe("stored-access"); + }); + + it("starts without a session when the scope changes, and does not refresh the old one", async () => { + const fetchMock = mockTokenEndpoint("refreshed"); + storeToken(TOKEN_STORAGE_KEY, { expiresAt: Date.now() }); + + await expect( + ensureValidTokenFor(`${SCOPE} Traces.Api`), + ).resolves.toBeNull(); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("keeps separate sessions for apps on different clients", async () => { + storeToken(TOKEN_STORAGE_KEY, { access_token: "app-a" }); + storeToken(`${STORAGE_KEYS.TOKEN}:${OTHER_CLIENT_ID}:${NORMALIZED_SCOPE}`, { + access_token: "app-b", + }); + + await expect(ensureValidTokenFor(SCOPE)).resolves.toBe("app-a"); + await expect(ensureValidTokenFor(SCOPE, OTHER_CLIENT_ID)).resolves.toBe( + "app-b", + ); + }); + + it("ignores a session stored under the unscoped key", async () => { + storeToken(STORAGE_KEYS.TOKEN); + + await expect(ensureValidTokenFor(SCOPE)).resolves.toBeNull(); + }); + + it("writes a refreshed token back under the same key", async () => { + mockTokenEndpoint("refreshed"); + storeToken(TOKEN_STORAGE_KEY, { expiresAt: Date.now() }); + + await expect(ensureValidTokenFor(SCOPE)).resolves.toBe("refreshed"); + expect(readToken(TOKEN_STORAGE_KEY)).toMatchObject({ + access_token: "refreshed", + }); + }); + + it("clears the session when its refresh fails", async () => { + mockFailingTokenEndpoint(); + storeToken(TOKEN_STORAGE_KEY, { expiresAt: Date.now() }); + + await expect(ensureValidTokenFor(SCOPE)).resolves.toBeNull(); + expect(readToken(TOKEN_STORAGE_KEY)).toBeNull(); + }); +}); + +describe("getTokenQueryKey", () => { + it("normalizes the scope like the storage key", () => { + expect(getTokenQueryKey(CLIENT_ID, "OR.Jobs openid OR.Jobs")).toEqual([ + "auth", + "token", + CLIENT_ID, + "OR.Jobs openid", + ]); + }); +}); + +describe("login callback", () => { + it("stores the exchanged token under the client and scope key", async () => { + await loginAndReturnWithCode(SCOPE); + mockTokenEndpoint("fresh"); + + await expect(ensureValidTokenFor(SCOPE)).resolves.toBe("fresh"); + expect(readToken(TOKEN_STORAGE_KEY)).toMatchObject({ + access_token: "fresh", + }); + expectCallbackCleanedUp(); + }); + + it("does not redeem the code when the login was started with another scope", async () => { + await loginAndReturnWithCode("openid OR.Jobs"); + const fetchMock = mockTokenEndpoint("fresh"); + + await expect(ensureValidTokenFor(SCOPE)).resolves.toBeNull(); + expect(fetchMock).not.toHaveBeenCalled(); + expect(readToken(TOKEN_STORAGE_KEY)).toBeNull(); + expectCallbackCleanedUp(); + }); + + it("does not redeem the code when the login recorded no scope", async () => { + await loginAndReturnWithCode(SCOPE); + sessionStorage.removeItem(STORAGE_KEYS.REQUESTED_SCOPE); + const fetchMock = mockTokenEndpoint("fresh"); + + await expect(ensureValidTokenFor(SCOPE)).resolves.toBeNull(); + expect(fetchMock).not.toHaveBeenCalled(); + expectCallbackCleanedUp(); + }); + + it("cleans up when the state does not match", async () => { + await loginAndReturnWithCode(SCOPE); + window.history.replaceState({}, "", "/?code=abc&state=other"); + const fetchMock = mockTokenEndpoint("fresh"); + + await expect(ensureValidTokenFor(SCOPE)).resolves.toBeNull(); + expect(fetchMock).not.toHaveBeenCalled(); + expectCallbackCleanedUp(); + }); + + it("cleans up and redeems nothing when Identity returns an error", async () => { + await loginAndReturnWithCode(SCOPE); + const state = new URLSearchParams(window.location.search).get("state"); + window.history.replaceState({}, "", `/?error=access_denied&state=${state}`); + const fetchMock = mockTokenEndpoint("fresh"); + + await expect(ensureValidTokenFor(SCOPE)).resolves.toBeNull(); + expect(fetchMock).not.toHaveBeenCalled(); + expectCallbackCleanedUp(); + }); + + it("cleans up when the exchange fails", async () => { + await loginAndReturnWithCode(SCOPE); + mockFailingTokenEndpoint(); + + await expect(ensureValidTokenFor(SCOPE)).resolves.toBeNull(); + expectCallbackCleanedUp(); + }); + + it("refreshes a stored session only once when the callback fails", async () => { + storeToken(TOKEN_STORAGE_KEY, { expiresAt: Date.now() + 60 * 1000 }); + await loginAndReturnWithCode(SCOPE); + const fetchMock = vi.fn((_url: string, init?: RequestInit) => + Promise.resolve( + grantTypeOf(init) === "authorization_code" + ? new Response("invalid_grant", { status: 400 }) + : tokenResponse("refreshed"), + ), + ); + vi.stubGlobal("fetch", fetchMock); + const queryClient = new QueryClient(); + + // An observer makes the query active, like the shell's useQuery. + const unsubscribe = new QueryObserver(queryClient, { + queryKey: getTokenQueryKey(CLIENT_ID, SCOPE), + queryFn: () => ensureValidToken(CLIENT_ID, SCOPE, BASE_URL, "/"), + }).subscribe(() => null); + await vi.waitFor(() => expect(queryClient.isFetching()).toBe(0)); + unsubscribe(); + + const refreshCalls = fetchMock.mock.calls.filter( + ([, init]) => grantTypeOf(init) === "refresh_token", + ); + expect(refreshCalls).toHaveLength(1); + expect(readToken(TOKEN_STORAGE_KEY)).toMatchObject({ + access_token: "refreshed", + }); + }); +}); + +describe("getStoredIdToken", () => { + it("reads the id token of this app's session, not the unscoped one", () => { + storeToken(STORAGE_KEYS.TOKEN, { id_token: "legacy-id" }); + storeToken(TOKEN_STORAGE_KEY, { id_token: "app-id" }); + + expect(getStoredIdToken(CLIENT_ID, SCOPE)).toBe("app-id"); + }); + + it("returns null without a session", () => { + expect(getStoredIdToken(CLIENT_ID, SCOPE)).toBeNull(); + }); +}); + +describe("logout", () => { + it("ends every session of this client and keeps other clients' sessions", () => { + const earlierScopeKey = `${STORAGE_KEYS.TOKEN}:${CLIENT_ID}:OR.Jobs openid`; + const otherClientKey = `${STORAGE_KEYS.TOKEN}:${OTHER_CLIENT_ID}:${NORMALIZED_SCOPE}`; + storeToken(TOKEN_STORAGE_KEY); + storeToken(earlierScopeKey); + storeToken(STORAGE_KEYS.TOKEN, { id_token: idTokenFor(CLIENT_ID) }); + storeToken(otherClientKey); + + logout(new QueryClient(), CLIENT_ID); + + expect(readToken(TOKEN_STORAGE_KEY)).toBeNull(); + expect(readToken(earlierScopeKey)).toBeNull(); + expect(readToken(STORAGE_KEYS.TOKEN)).toBeNull(); + expect(readToken(otherClientKey)).not.toBeNull(); + }); + + it("resets this client's cached tokens and keeps other clients'", () => { + const queryClient = new QueryClient(); + const ownKeys = [ + getTokenQueryKey(CLIENT_ID, SCOPE), + getTokenQueryKey(CLIENT_ID, "openid"), + ]; + const otherClientKey = getTokenQueryKey(OTHER_CLIENT_ID, SCOPE); + for (const key of [...ownKeys, otherClientKey]) { + queryClient.setQueryData(key, "token"); + } + + logout(queryClient, CLIENT_ID); + + for (const key of ownKeys) { + expect(queryClient.getQueryData(key)).toBeUndefined(); + } + expect(queryClient.getQueryData(otherClientKey)).toBe("token"); + }); + + it("keeps the unscoped entry when it belongs to another client", () => { + storeToken(STORAGE_KEYS.TOKEN, { id_token: idTokenFor(OTHER_CLIENT_ID) }); + + logout(new QueryClient(), CLIENT_ID); + + expect(readToken(STORAGE_KEYS.TOKEN)).not.toBeNull(); + }); +}); diff --git a/packages/apollo-vertex/lib/auth.ts b/packages/apollo-vertex/lib/auth.ts index a51b878b0f..6211189372 100644 --- a/packages/apollo-vertex/lib/auth.ts +++ b/packages/apollo-vertex/lib/auth.ts @@ -1,4 +1,5 @@ import { useQueryClient } from "@tanstack/react-query"; +import { jwtDecode } from "jwt-decode"; import PKCEChallenge from "pkce-challenge"; import { toast } from "sonner"; import { z } from "zod"; @@ -19,6 +20,7 @@ export const STORAGE_KEYS = { TOKEN: "uipath_token", CODE_VERIFIER: "uipath_code_verifier", STATE: "uipath_state", + REQUESTED_SCOPE: "uipath_requested_scope", AUTH_RETURN_TO: "auth_return_to", LOGOUT_RETURN_TO: "logout_return_to", } as const; @@ -92,12 +94,49 @@ const shouldRefreshToken = (tokenData: TokenData): boolean => { return timeUntilExpiry < REFRESH_THRESHOLD_MS; }; -const saveTokenData = (tokenData: TokenData): void => { - localStorage.setItem(STORAGE_KEYS.TOKEN, JSON.stringify(tokenData)); +const normalizeScope = (scope: string): string => + [...new Set(scope.split(/\s+/).filter(Boolean))].toSorted().join(" "); + +// Keyed by client so apps on one host (one localStorage) keep their own sessions, +// and by scope so a scope change forces a new login: a refresh can't add scopes. +const getTokenStorageKey = (clientId: string, scope: string): string => + `${STORAGE_KEYS.TOKEN}:${clientId}:${normalizeScope(scope)}`; + +// Mirrors the storage key, so callers on another client or scope get their own +// cache entry. +export const getTokenQueryKey = (clientId: string, scope: string) => + [...TOKEN_QUERY_KEY, clientId, normalizeScope(scope)] as const; + +const readTokenData = (storageKey: string): TokenData | null => { + try { + const { data, success } = TokenDataSchema.safeParse( + JSON.parse(localStorage.getItem(storageKey) ?? "null"), + ); + return success ? data : null; + } catch { + return null; + } }; -const clearTokenData = (): void => { - localStorage.removeItem(STORAGE_KEYS.TOKEN); +export const getStoredIdToken = ( + clientId: string, + scope: string, +): string | null => + readTokenData(getTokenStorageKey(clientId, scope))?.id_token ?? null; + +const saveTokenData = (storageKey: string, tokenData: TokenData): void => { + localStorage.setItem(storageKey, JSON.stringify(tokenData)); +}; + +const clearTokenData = (storageKey: string): void => { + localStorage.removeItem(storageKey); +}; + +const clearPendingLogin = (): void => { + sessionStorage.removeItem(STORAGE_KEYS.CODE_VERIFIER); + sessionStorage.removeItem(STORAGE_KEYS.STATE); + sessionStorage.removeItem(STORAGE_KEYS.REQUESTED_SCOPE); + sessionStorage.removeItem(STORAGE_KEYS.AUTH_RETURN_TO); }; const fetchTokenData = async (baseUrl: string, body?: URLSearchParams) => { @@ -140,6 +179,7 @@ const refreshTokenIfNeeded = async ( tokenData: TokenData, clientId: string, baseUrl: string, + storageKey: string, ): Promise => { if (!tokenData.refresh_token) { return null; @@ -151,10 +191,10 @@ const refreshTokenIfNeeded = async ( clientId, baseUrl, ); - saveTokenData(newTokenData); + saveTokenData(storageKey, newTokenData); return newTokenData.access_token; } catch { - clearTokenData(); + clearTokenData(storageKey); return null; } }; @@ -179,6 +219,7 @@ const exchangeCodeForToken = ( const handleOAuthCallback = async ( clientId: string, + scope: string, baseUrl: string, redirectPath?: string, ): Promise => { @@ -187,25 +228,37 @@ const handleOAuthCallback = async ( const state = params.get("state"); const error = params.get("error"); - if (error) { - throw new Error(`OAuth error: ${error}`); - } - - if (!code || !state) { + if (!state || (!code && !error)) { return; } const storedState = sessionStorage.getItem(STORAGE_KEYS.STATE); - if (state !== storedState) { - throw new Error("State mismatch"); - } - const codeVerifier = sessionStorage.getItem(STORAGE_KEYS.CODE_VERIFIER); - if (!codeVerifier) { - throw new Error("Code verifier not found"); - } + const requestedScope = sessionStorage.getItem(STORAGE_KEYS.REQUESTED_SCOPE); + const returnTo = + sessionStorage.getItem(STORAGE_KEYS.AUTH_RETURN_TO) ?? + window.location.pathname; + clearPendingLogin(); try { + if (error !== null || code === null) { + throw new Error(`OAuth error: ${error ?? "no code returned"}`); + } + if (state !== storedState) { + throw new Error("State mismatch"); + } + if (!codeVerifier) { + throw new Error("Code verifier not found"); + } + // False when the login began under another scope, e.g. from a bundle + // deployed before this one; its token would lack the current scopes. + const loginUsedCurrentScope = + requestedScope !== null && + normalizeScope(requestedScope) === normalizeScope(scope); + if (!loginUsedCurrentScope) { + toast.info("Your sign-in settings changed. Please sign in again."); + return; + } const tokenData = await exchangeCodeForToken( code, codeVerifier, @@ -213,41 +266,29 @@ const handleOAuthCallback = async ( baseUrl, redirectPath, ); - sessionStorage.removeItem(STORAGE_KEYS.CODE_VERIFIER); - sessionStorage.removeItem(STORAGE_KEYS.STATE); - saveTokenData(tokenData); - const returnTo = sessionStorage.getItem(STORAGE_KEYS.AUTH_RETURN_TO); - sessionStorage.removeItem(STORAGE_KEYS.AUTH_RETURN_TO); - const targetPath = returnTo ?? window.location.pathname; - window.history.replaceState({}, document.title, targetPath); + saveTokenData(getTokenStorageKey(clientId, scope), tokenData); } catch (authError) { - clearTokenData(); - sessionStorage.removeItem(STORAGE_KEYS.AUTH_RETURN_TO); const message = authError instanceof Error ? authError.message : "Unknown error"; toast.error(`Authentication failed: ${message}`); - throw authError; + } finally { + // Strips ?code and ?state on every outcome, or each refetch retries the + // callback. It runs after the exchange, whose redirect_uri can be built + // from the current path. + window.history.replaceState({}, document.title, returnTo); } }; export const ensureValidToken = async ( - queryClient: ReturnType, clientId: string, + scope: string, baseUrl: string, redirectPath?: string, ): Promise => { - const params = new URLSearchParams(window.location.search); - const isInOAuthCallback = params.has("code") && params.has("state"); - - if (isInOAuthCallback) { - try { - await handleOAuthCallback(clientId, baseUrl, redirectPath); - } catch { - void queryClient.resetQueries({ queryKey: TOKEN_QUERY_KEY }); - } - } + await handleOAuthCallback(clientId, scope, baseUrl, redirectPath); - const tokenDataStr = localStorage.getItem(STORAGE_KEYS.TOKEN); + const storageKey = getTokenStorageKey(clientId, scope); + const tokenDataStr = localStorage.getItem(storageKey); if (!tokenDataStr) { return null; } @@ -256,19 +297,19 @@ export const ensureValidToken = async ( try { parsed = JSON.parse(tokenDataStr); } catch { - clearTokenData(); + clearTokenData(storageKey); return null; } const { data: tokenData, success } = TokenDataSchema.safeParse(parsed); if (!success) { - clearTokenData(); + clearTokenData(storageKey); return null; } if (isTokenExpired(tokenData) || shouldRefreshToken(tokenData)) { - return refreshTokenIfNeeded(tokenData, clientId, baseUrl); + return refreshTokenIfNeeded(tokenData, clientId, baseUrl, storageKey); } return tokenData.access_token; @@ -285,6 +326,7 @@ export const login = async ( sessionStorage.setItem(STORAGE_KEYS.CODE_VERIFIER, pkce.code_verifier); sessionStorage.setItem(STORAGE_KEYS.STATE, state); + sessionStorage.setItem(STORAGE_KEYS.REQUESTED_SCOPE, scope); if (redirectPath && window.location.pathname !== redirectPath) { sessionStorage.setItem( @@ -306,12 +348,38 @@ export const login = async ( window.location.href = `${getAuthorizationEndpoint(baseUrl)}?${params.toString()}`; }; +// The unscoped pre-upgrade entry may belong to another app on the host; its id +// token's audience names the client. An unreadable entry is no one's session. +const isUnscopedSessionOf = (clientId: string): boolean => { + const tokenData = readTokenData(STORAGE_KEYS.TOKEN); + if (tokenData === null) { + return true; + } + try { + const { aud } = jwtDecode(tokenData.id_token); + return Array.isArray(aud) ? aud.includes(clientId) : aud === clientId; + } catch { + return true; + } +}; + export const logout = ( queryClient: ReturnType, + clientId: string, ): void => { - clearTokenData(); - sessionStorage.removeItem(STORAGE_KEYS.CODE_VERIFIER); - sessionStorage.removeItem(STORAGE_KEYS.STATE); - sessionStorage.removeItem(STORAGE_KEYS.AUTH_RETURN_TO); - void queryClient.resetQueries({ queryKey: TOKEN_QUERY_KEY }); + // Signing out ends every session of this client, including ones from earlier + // scopes and its pre-upgrade entry, which still hold refresh tokens. + const clientPrefix = `${STORAGE_KEYS.TOKEN}:${clientId}:`; + const sessionKeys = Array.from({ length: localStorage.length }, (_, index) => + localStorage.key(index), + ).filter((key): key is string => + key === STORAGE_KEYS.TOKEN + ? isUnscopedSessionOf(clientId) + : key?.startsWith(clientPrefix) === true, + ); + for (const key of sessionKeys) { + clearTokenData(key); + } + clearPendingLogin(); + void queryClient.resetQueries({ queryKey: [...TOKEN_QUERY_KEY, clientId] }); }; diff --git a/packages/apollo-vertex/registry/shell/shell-auth-provider.tsx b/packages/apollo-vertex/registry/shell/shell-auth-provider.tsx index f11b929ca4..b4aa366f6f 100644 --- a/packages/apollo-vertex/registry/shell/shell-auth-provider.tsx +++ b/packages/apollo-vertex/registry/shell/shell-auth-provider.tsx @@ -9,7 +9,7 @@ import { useContext, } from "react"; import { z } from "zod"; -import { ensureValidToken, login, logout, TOKEN_QUERY_KEY } from "@/lib/auth"; +import { ensureValidToken, getTokenQueryKey, login, logout } from "@/lib/auth"; export interface AuthContextValue { user: UserInfo | null; @@ -75,19 +75,18 @@ interface UseAccessTokenProps { clientId: string; baseUrl: string; redirectPath?: string; + scope: string; } export const useAccessToken = ({ clientId, baseUrl, redirectPath, + scope, }: UseAccessTokenProps) => { - const queryClient = useQueryClient(); - const { data: token, isLoading } = useQuery({ - queryKey: TOKEN_QUERY_KEY, - queryFn: () => - ensureValidToken(queryClient, clientId, baseUrl, redirectPath), + queryKey: getTokenQueryKey(clientId, scope), + queryFn: () => ensureValidToken(clientId, scope, baseUrl, redirectPath), refetchInterval: 60 * 1000, enabled: typeof window !== "undefined", }); @@ -110,6 +109,7 @@ export const ShellAuthProvider: FC< clientId, baseUrl, redirectPath, + scope, }); const user = accessToken ? decodeJWT(accessToken) : null; const isAuthenticated = accessToken != null; @@ -119,7 +119,7 @@ export const ShellAuthProvider: FC< isAuthenticated, isLoading, login: () => login(clientId, scope, baseUrl, redirectPath), - logout: () => logout(queryClient), + logout: () => logout(queryClient, clientId), accessToken, }; diff --git a/packages/apollo-vertex/vitest.config.ts b/packages/apollo-vertex/vitest.config.ts index 979668f960..5f63b1ed24 100644 --- a/packages/apollo-vertex/vitest.config.ts +++ b/packages/apollo-vertex/vitest.config.ts @@ -7,7 +7,7 @@ export default defineConfig({ }, test: { environment: "jsdom", - include: ["registry/**/*.test.tsx"], + include: ["registry/**/*.test.tsx", "lib/**/*.test.ts"], coverage: { provider: "v8", reporter: ["text", "json", "json-summary", "html", "lcov"],