diff --git a/src/auth/oidc-token.test.ts b/src/auth/oidc-token.test.ts index 7dfc6a2..900acad 100644 --- a/src/auth/oidc-token.test.ts +++ b/src/auth/oidc-token.test.ts @@ -6,6 +6,7 @@ import { jwt } from "better-auth/plugins/jwt"; import { createLocalJWKSet, jwtVerify } from "jose"; import { describe, expect, it } from "vite-plus/test"; import { jwtOptions } from "./jwt-options"; +import { serializeRefreshGrant } from "./refresh-grant"; import { standardIdTokenClaims } from "./policy"; const baseURL = "https://auth.polinetwork.test"; @@ -36,6 +37,7 @@ async function issueTokens(scopes: string[], customClaims = true) { loginPage: "/", consentPage: "/consent", scopes: ["openid", "profile", "email", "offline_access"], + refreshTokenReuseInterval: 10, customIdTokenClaims: customClaims ? ({ scopes }) => standardIdTokenClaims(appUser, scopes) : undefined, @@ -95,7 +97,11 @@ async function issueTokens(scopes: string[], customClaims = true) { }), ); expect(tokenResponse.status).toBe(200); - const tokens = (await tokenResponse.json()) as { id_token: string; access_token: string }; + const tokens = (await tokenResponse.json()) as { + id_token: string; + access_token: string; + refresh_token?: string; + }; const jwks = await auth.api.getJwks(); const verified = await jwtVerify(tokens.id_token, createLocalJWKSet(jwks), { issuer: `${baseURL}/api/auth`, @@ -109,7 +115,7 @@ async function issueTokens(scopes: string[], customClaims = true) { ); expect(userInfoResponse.status).toBe(200); const userInfo = (await userInfoResponse.json()) as Record; - return { payload: verified.payload, userInfo }; + return { payload: verified.payload, userInfo, auth, client, tokens }; } describe("OIDC standard claims in signed tokens", () => { @@ -142,3 +148,44 @@ describe("OIDC standard claims in signed tokens", () => { } }); }); + +describe("refresh token grants", () => { + it.each(["refresh_token", " refresh_token "])( + "answers concurrent refreshes of one token with the same rotated tokens (%j)", + async (grantType) => { + const { auth, client, tokens } = await issueTokens(["openid", "offline_access"]); + expect(tokens.refresh_token).toBeTruthy(); + const body = new URLSearchParams({ + grant_type: grantType, + refresh_token: tokens.refresh_token!, + client_id: client.client_id, + client_secret: client.client_secret!, + }).toString(); + const refresh = () => { + const request = new Request(`${baseURL}/api/auth/oauth2/token`, { + method: "POST", + headers: { + "content-type": "application/x-www-form-urlencoded", + "content-length": String(body.length), + }, + body, + }); + return serializeRefreshGrant(request, () => auth.handler(request)); + }; + const responses = await Promise.all([refresh(), refresh(), refresh()]); + const bodies = (await Promise.all(responses.map((response) => response.json()))) as { + refresh_token?: string; + id_token?: string; + }[]; + expect(responses.map((response) => response.status)).toEqual([200, 200, 200]); + const [first] = bodies; + expect(first.refresh_token).toBeTruthy(); + expect(first.refresh_token).not.toBe(tokens.refresh_token); + expect(first.id_token).toBeTruthy(); + for (const item of bodies) { + expect(item.refresh_token).toBe(first.refresh_token); + expect(item.id_token).toBe(first.id_token); + } + }, + ); +}); diff --git a/src/auth/refresh-grant.test.ts b/src/auth/refresh-grant.test.ts new file mode 100644 index 0000000..571ec51 --- /dev/null +++ b/src/auth/refresh-grant.test.ts @@ -0,0 +1,111 @@ +import { afterEach, expect, it, vi } from "vite-plus/test"; +import { serializeRefreshGrant } from "./refresh-grant"; + +afterEach(() => { + vi.useRealTimers(); +}); + +function refreshRequest(token: string) { + const body = new URLSearchParams({ + grant_type: "refresh_token", + refresh_token: token, + }).toString(); + return new Request("https://auth.polinetwork.test/api/auth/oauth2/token", { + method: "POST", + headers: { + "content-type": "application/x-www-form-urlencoded", + "content-length": String(body.length), + }, + body, + }); +} + +/** A handler that records its calls and only answers when `finish` is called. */ +function blockedHandler() { + const calls: string[] = []; + let finish = () => {}; + const finished = new Promise((resolve) => (finish = resolve)); + return { + calls, + finish, + handler: (name: string) => async () => { + calls.push(name); + await finished; + return new Response(name); + }, + }; +} + +it("runs a second refresh of the same token after the first", async () => { + const { calls, finish, handler } = blockedHandler(); + const first = serializeRefreshGrant(refreshRequest("token"), handler("first")); + const second = serializeRefreshGrant(refreshRequest("token"), handler("second")); + await vi.waitFor(() => expect(calls).toEqual(["first"])); + finish(); + await Promise.all([first, second]); + expect(calls).toEqual(["first", "second"]); +}); + +it("keeps the queue going when a refresh throws", async () => { + const failed = serializeRefreshGrant(refreshRequest("token"), async () => { + throw new Error("refresh failed"); + }); + const waiting = serializeRefreshGrant( + refreshRequest("token"), + async () => new Response("second"), + ); + await expect(failed).rejects.toThrow("refresh failed"); + expect(await (await waiting).text()).toBe("second"); + const later = await serializeRefreshGrant( + refreshRequest("token"), + async () => new Response("third"), + ); + expect(await later.text()).toBe("third"); +}); + +it("runs refreshes of different tokens at once", async () => { + const { calls, finish, handler } = blockedHandler(); + const first = serializeRefreshGrant(refreshRequest("one"), handler("one")); + const second = serializeRefreshGrant(refreshRequest("two"), handler("two")); + await vi.waitFor(() => expect(calls).toEqual(["one", "two"])); + finish(); + await Promise.all([first, second]); +}); + +it("stops waiting for a stuck request after five seconds", async () => { + vi.useFakeTimers(); + const { calls, finish, handler } = blockedHandler(); + const first = serializeRefreshGrant(refreshRequest("token"), handler("first")); + const second = serializeRefreshGrant(refreshRequest("token"), handler("second")); + await vi.advanceTimersByTimeAsync(4_999); + expect(calls).toEqual(["first"]); + await vi.advanceTimersByTimeAsync(1); + expect(calls).toEqual(["first", "second"]); + finish(); + await Promise.all([first, second]); +}); + +it("runs requests past the eighth for one token without queueing them", async () => { + const { calls, finish, handler } = blockedHandler(); + const requests = Array.from({ length: 9 }, (_, index) => + serializeRefreshGrant(refreshRequest("token"), handler(String(index))), + ); + // The running request and the ninth; which one is ninth depends on read order. + await vi.waitFor(() => expect(calls).toHaveLength(2)); + finish(); + await Promise.all(requests); + expect(calls).toHaveLength(9); +}); + +it("does not read requests without a known, small length", async () => { + const handle = vi.fn(async () => new Response("ok")); + const request = new Request("https://auth.polinetwork.test/api/auth/oauth2/token", { + method: "POST", + headers: { "content-type": "application/x-www-form-urlencoded" }, + body: "grant_type=refresh_token&refresh_token=token", + }); + const clone = vi.spyOn(request, "clone"); + await serializeRefreshGrant(request, handle); + expect(clone).not.toHaveBeenCalled(); + expect(handle).toHaveBeenCalledOnce(); +}); diff --git a/src/auth/refresh-grant.ts b/src/auth/refresh-grant.ts new file mode 100644 index 0000000..b10e784 --- /dev/null +++ b/src/auth/refresh-grant.ts @@ -0,0 +1,63 @@ +// A token request is a short form of known length; other bodies are not read and run unqueued. +const MAX_BODY_BYTES = 16 * 1024; +// Requests per token in the queue, the running one included; more run unqueued. +const MAX_QUEUED = 8; +// How long a request waits for the one ahead before running anyway. Below the provider's +// 10-second `refreshTokenReuseInterval`, so the replay is still there when it runs. +const MAX_WAIT_MS = 5_000; + +type Queue = { last: Promise; size: number }; +const queues = new Map(); + +/** The refresh token a token request presents, or null for any other request. */ +async function presentedRefreshToken(request: Request) { + if (request.method !== "POST" || !new URL(request.url).pathname.endsWith("/oauth2/token")) + return null; + if (!request.headers.get("content-type")?.includes("application/x-www-form-urlencoded")) + return null; + const length = Number(request.headers.get("content-length")); + if (!(length > 0 && length <= MAX_BODY_BYTES)) return null; + const body = new URLSearchParams(await request.clone().text()); + // The provider trims the grant type before matching it. + return body.get("grant_type")?.trim() === "refresh_token" + ? body.get("refresh_token") || null + : null; +} + +async function waitAtMost(previous: Promise, ms: number) { + let timer: ReturnType | undefined; + await Promise.race([previous, new Promise((resolve) => (timer = setTimeout(resolve, ms)))]); + clearTimeout(timer); +} + +/** + * Handles refresh grants that present the same refresh token one at a time. + * + * An app may refresh one token from several requests at once: Flux Web does when a page loads + * its data just after the ID token expired. The provider rotates the token for the first request + * and refuses the others with `invalid_grant`, which signs the person out. In turn, the later + * requests find the token just rotated and get the first response again, through the provider's + * `refreshTokenReuseInterval` replay. The queue is per process; production runs one replica. + */ +export async function serializeRefreshGrant( + request: Request, + handle: () => Promise, +): Promise { + const token = await presentedRefreshToken(request); + if (!token) return handle(); + const queue = queues.get(token) ?? { last: Promise.resolve(), size: 0 }; + if (queue.size >= MAX_QUEUED) return handle(); + const previous = queue.last; + let release = () => {}; + queue.last = new Promise((resolve) => (release = resolve)); + queue.size += 1; + queues.set(token, queue); + try { + await waitAtMost(previous, MAX_WAIT_MS); + return await handle(); + } finally { + release(); + queue.size -= 1; + if (queue.size === 0) queues.delete(token); + } +} diff --git a/src/routes/api/auth/$.ts b/src/routes/api/auth/$.ts index f4126e9..08ecbbb 100644 --- a/src/routes/api/auth/$.ts +++ b/src/routes/api/auth/$.ts @@ -2,9 +2,10 @@ import { createFileRoute } from "@tanstack/react-router"; import { auth } from "@/auth"; import { redirectLogoutPage } from "@/auth/logout-page"; import { rateLimitedResponse } from "@/auth/rate-limit"; +import { serializeRefreshGrant } from "@/auth/refresh-grant"; async function handle(request: Request) { - const response = await auth.handler(request); + const response = await serializeRefreshGrant(request, () => auth.handler(request)); if (response.status === 429 && response.headers.has("x-retry-after")) return rateLimitedResponse(request, response); return new URL(request.url).pathname.includes("/oauth2/end-session")