From 8b164c556ed7d0b5283cde5177ce43225f560844 Mon Sep 17 00:00:00 2001 From: Maxwell Date: Sun, 3 May 2026 17:47:46 -1000 Subject: [PATCH] fix: thread per-account cookie slot through OAuth flows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The multi-account refresh-token cookie slot wiring was half-implemented: every account's refresh token ended up on slot 0, so "+ Add Account" silently clobbered the previous account's `jmap_rt` cookie. On page refresh, only the most-recently-added account had a working refresh token; the others bounced to login. Three coordinated changes: 1. `app/[locale]/login/page.tsx` (handleOAuthLogin): write the next-free cookie slot to `sessionStorage['oauth_cookie_slot']` before redirecting to the IdP. `loginWithOAuth` already reads this key but it was never written, so it always defaulted to 0. 2. `stores/auth-store.ts` (loginWithOAuth): distinguish "no value set" (`rawSlot === null`) from "value is 0". Previously `parseInt(getItem(...) || '0')` collapsed both cases, making the `getNextCookieSlot()` fallback unreachable. 3. `stores/auth-store.ts` (loginWithServerSso) + `app/api/auth/sso/complete/route.ts`: pass the slot through the body of the POST and use it for `refreshTokenCookieName(slot)`. Same pattern as the existing `/api/auth/token POST` that already accepts a slot. The server defaults to 0 for back-compat with any caller that omits it. After the fix, signing in with multiple accounts produces distinct `jmap_rt`, `jmap_rt_1`, `jmap_rt_2`, ... cookies (matching the cookieSlot field in account-store) and all accounts survive a page refresh. Repro before the fix: - Sign in with one account, refresh — works. - Click "+ Add Account", sign in with a second account, refresh — second account vanishes from the dropdown; switching to the first account in the dropdown still shows the second account's identity in the From box. --- app/[locale]/login/page.tsx | 11 ++++++++++ app/api/auth/sso/complete/route.ts | 12 ++++++++--- stores/auth-store.ts | 32 +++++++++++++++++++++--------- 3 files changed, 43 insertions(+), 12 deletions(-) diff --git a/app/[locale]/login/page.tsx b/app/[locale]/login/page.tsx index 8e0872cb..5bdd1938 100644 --- a/app/[locale]/login/page.tsx +++ b/app/[locale]/login/page.tsx @@ -7,6 +7,7 @@ import { useTranslations } from "next-intl"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { useAuthStore } from "@/stores/auth-store"; +import { useAccountStore } from "@/stores/account-store"; import { useThemeStore } from "@/stores/theme-store"; import { useShallow } from "zustand/react/shallow"; import { useConfig } from "@/hooks/use-config"; @@ -460,6 +461,16 @@ export default function LoginPage() { sessionStorage.setItem("oauth_add_account_mode", "true"); } + // Persist the next-free cookie slot so loginWithOAuth (in stores/auth-store.ts) + // writes the refresh token to the correct per-account jmap_rt_ cookie. + // loginWithOAuth reads this key but it was previously never written, so every + // OAuth account collapsed onto slot 0 and clobbered earlier accounts' refresh + // tokens. getNextCookieSlot() returns 0 when no accounts exist (correct for + // first sign-in) and the lowest unused slot otherwise (correct for "+ Add + // Account"). + const nextSlot = useAccountStore.getState().getNextCookieSlot(); + sessionStorage.setItem("oauth_cookie_slot", nextSlot.toString()); + const authUrl = new URL(oauthMetadata.authorization_endpoint); authUrl.searchParams.set("response_type", "code"); authUrl.searchParams.set("client_id", oauthClientId); diff --git a/app/api/auth/sso/complete/route.ts b/app/api/auth/sso/complete/route.ts index fcbf6d33..eba9a7b7 100644 --- a/app/api/auth/sso/complete/route.ts +++ b/app/api/auth/sso/complete/route.ts @@ -13,12 +13,18 @@ export async function POST(request: NextRequest) { const cookieStore = await cookies(); try { - const { code, state } = await request.json(); + const { code, state, slot: bodySlot } = await request.json(); if (!code || !state) { return NextResponse.json({ error: 'Missing code or state' }, { status: 400 }); } + // Per-account refresh-token cookie slot. Without this the route hardcoded + // slot 0, so the "+ Add Account" flow overwrote the first account's + // refresh-token cookie. Default to 0 for back-compat with any caller that + // omits slot. Mirrors the validation in /api/auth/token POST. + const slot = typeof bodySlot === 'number' && bodySlot >= 0 && bodySlot <= 4 ? bodySlot : 0; + // Read and decrypt the pending SSO cookie const pendingCookie = cookieStore.get(SSO_PENDING_COOKIE)?.value; if (!pendingCookie) { @@ -58,9 +64,9 @@ export async function POST(request: NextRequest) { // Exchange code for tokens const tokens = await exchangeCodeForTokens(code, codeVerifier, redirectUri); - // Store refresh token + // Store refresh token in the per-account cookie slot. if (tokens.refresh_token) { - const cookieName = refreshTokenCookieName(0); + const cookieName = refreshTokenCookieName(slot); cookieStore.set(cookieName, tokens.refresh_token, getCookieOptions()); } diff --git a/stores/auth-store.ts b/stores/auth-store.ts index 008b020f..04955b59 100644 --- a/stores/auth-store.ts +++ b/stores/auth-store.ts @@ -601,12 +601,21 @@ export const useAuthStore = create()( set({ isLoading: true, error: null, isRateLimited: false, rateLimitUntil: null }); try { - // Determine slot for this account (use slot from sessionStorage if re-adding) + // Determine slot for this account (use slot from sessionStorage if re-adding). + // Note: `parseInt(getItem(...) || '0')` collapses "no value set" and + // "value is 0" into the same case, so the fallback to getNextCookieSlot() + // never fired for the common "+ Add Account" path — every OAuth account + // ended up on slot 0 and overwrote earlier accounts' refresh-token cookies. + // Distinguishing rawSlot === null from a parsed 0 fixes that. The page + // also writes oauth_cookie_slot before redirecting to the IdP. const accountStore = useAccountStore.getState(); - const pendingSlot = typeof window !== 'undefined' - ? parseInt(sessionStorage.getItem('oauth_cookie_slot') || '0', 10) - : 0; - const slot = pendingSlot >= 0 && pendingSlot <= 4 ? pendingSlot : accountStore.getNextCookieSlot(); + const rawSlot = typeof window !== 'undefined' + ? sessionStorage.getItem('oauth_cookie_slot') + : null; + const pendingSlot = rawSlot !== null ? parseInt(rawSlot, 10) : NaN; + const slot = !isNaN(pendingSlot) && pendingSlot >= 0 && pendingSlot <= 4 + ? pendingSlot + : accountStore.getNextCookieSlot(); const tokenRes = await apiFetch(`/api/auth/token?slot=${slot}`, { method: 'POST', @@ -718,12 +727,19 @@ export const useAuthStore = create()( set({ isLoading: true, error: null, isRateLimited: false, rateLimitUntil: null }); try { - // Server-side SSO: the server holds the PKCE verifier in an encrypted cookie + // Server-side SSO: the server holds the PKCE verifier in an encrypted cookie. + // Pass the next-free cookie slot so /api/auth/sso/complete writes the refresh + // token to the correct per-account jmap_rt_ cookie. Without this the + // route hardcoded slot 0, which broke "+ Add Account" by overwriting the + // first account's refresh-token cookie. + const accountStore = useAccountStore.getState(); + const slot = accountStore.getNextCookieSlot(); + const ssoRes = await apiFetch('/api/auth/sso/complete', { method: 'POST', headers: { 'Content-Type': 'application/json' }, credentials: 'include', - body: JSON.stringify({ code, state }), + body: JSON.stringify({ code, state, slot }), }); if (!ssoRes.ok) { @@ -741,8 +757,6 @@ export const useAuthStore = create()( throw new Error('Server URL not configured'); } - const accountStore = useAccountStore.getState(); - const refreshFn = get().refreshAccessToken; const client = JMAPClient.withBearer(ssoServerUrl, access_token, '', () => refreshFn()); await client.connect();