fix(auth): answer concurrent refreshes of one token instead of refusing them

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<string, unknown>;
-  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);
+      }
+    },
+  );
+});

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<void>((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();
+});

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<void>; size: number };
+const queues = new Map<string, Queue>();
+
+/** 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<void>, ms: number) {
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  await Promise.race([previous, new Promise<void>((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<Response>,
+): Promise<Response> {
+  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);
+  }
+}

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")